tenant.openagents/omega

1use anyhow::{Result, anyhow};
2use futures::{
3    AsyncBufReadExt, AsyncReadExt,
4    io::BufReader,
5    stream::{BoxStream, StreamExt},
6};
7use http_client::{
8    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt,
9};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use std::convert::TryFrom;
13
14pub const DEEPSEEK_API_URL: &str = "https://api.deepseek.com/v1";
15
16#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
17#[serde(rename_all = "lowercase")]
18pub enum Role {
19    User,
20    Assistant,
21    System,
22    Tool,
23}
24
25impl TryFrom<String> for Role {
26    type Error = anyhow::Error;
27
28    fn try_from(value: String) -> Result<Self> {
29        match value.as_str() {
30            "user" => Ok(Self::User),
31            "assistant" => Ok(Self::Assistant),
32            "system" => Ok(Self::System),
33            "tool" => Ok(Self::Tool),
34            _ => anyhow::bail!("invalid role '{value}'"),
35        }
36    }
37}
38
39impl From<Role> for String {
40    fn from(val: Role) -> Self {
41        match val {
42            Role::User => "user".to_owned(),
43            Role::Assistant => "assistant".to_owned(),
44            Role::System => "system".to_owned(),
45            Role::Tool => "tool".to_owned(),
46        }
47    }
48}
49
50#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
51#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
52pub enum Model {
53    #[serde(rename = "deepseek-v4-flash")]
54    V4Flash,
55    #[serde(rename = "deepseek-v4-pro")]
56    #[default]
57    V4Pro,
58    #[serde(rename = "custom")]
59    Custom {
60        name: String,
61        /// The name displayed in the UI, such as in the agent panel model dropdown menu.
62        display_name: Option<String>,
63        max_tokens: u64,
64        max_output_tokens: Option<u64>,
65    },
66}
67
68impl Model {
69    pub fn default_fast() -> Self {
70        Model::V4Flash
71    }
72
73    pub fn from_id(id: &str) -> Result<Self> {
74        match id {
75            "deepseek-v4-flash" => Ok(Self::V4Flash),
76            "deepseek-v4-pro" => Ok(Self::V4Pro),
77            _ => anyhow::bail!("invalid model id {id}"),
78        }
79    }
80
81    pub fn id(&self) -> &str {
82        match self {
83            Self::V4Flash => "deepseek-v4-flash",
84            Self::V4Pro => "deepseek-v4-pro",
85            Self::Custom { name, .. } => name,
86        }
87    }
88
89    pub fn display_name(&self) -> &str {
90        match self {
91            Self::V4Flash => "DeepSeek V4 Flash",
92            Self::V4Pro => "DeepSeek V4 Pro",
93            Self::Custom {
94                name, display_name, ..
95            } => display_name.as_ref().unwrap_or(name).as_str(),
96        }
97    }
98
99    pub fn max_token_count(&self) -> u64 {
100        match self {
101            Self::V4Flash | Self::V4Pro => 1_000_000,
102            Self::Custom { max_tokens, .. } => *max_tokens,
103        }
104    }
105
106    pub fn max_output_tokens(&self) -> Option<u64> {
107        match self {
108            Self::V4Flash | Self::V4Pro => Some(384_000),
109            Self::Custom {
110                max_output_tokens, ..
111            } => *max_output_tokens,
112        }
113    }
114}
115
116#[derive(Debug, Serialize, Deserialize)]
117pub struct Request {
118    pub model: String,
119    pub messages: Vec<RequestMessage>,
120    pub stream: bool,
121    #[serde(default, skip_serializing_if = "Option::is_none")]
122    pub max_tokens: Option<u64>,
123    #[serde(default, skip_serializing_if = "Option::is_none")]
124    pub temperature: Option<f32>,
125    #[serde(default, skip_serializing_if = "Option::is_none")]
126    pub thinking: Option<Thinking>,
127    #[serde(default, skip_serializing_if = "Option::is_none")]
128    pub reasoning_effort: Option<ReasoningEffort>,
129    #[serde(default, skip_serializing_if = "Option::is_none")]
130    pub response_format: Option<ResponseFormat>,
131    #[serde(default, skip_serializing_if = "Option::is_none")]
132    pub tool_choice: Option<ToolChoice>,
133    #[serde(default, skip_serializing_if = "Vec::is_empty")]
134    pub tools: Vec<ToolDefinition>,
135}
136
137#[derive(Debug, Serialize, Deserialize)]
138pub struct Thinking {
139    #[serde(rename = "type")]
140    pub kind: ThinkingType,
141}
142
143#[derive(Debug, Serialize, Deserialize, Clone, Copy, Eq, PartialEq)]
144#[serde(rename_all = "lowercase")]
145pub enum ThinkingType {
146    Enabled,
147    Disabled,
148}
149
150#[derive(Debug, Serialize, Deserialize, Clone, Copy, Eq, PartialEq)]
151#[serde(rename_all = "lowercase")]
152pub enum ReasoningEffort {
153    High,
154    Max,
155}
156
157#[derive(Debug, Serialize, Deserialize)]
158#[serde(rename_all = "snake_case")]
159pub enum ResponseFormat {
160    Text,
161    #[serde(rename = "json_object")]
162    JsonObject,
163}
164
165#[derive(Debug, Serialize, Deserialize, Clone, Copy, Eq, PartialEq)]
166#[serde(rename_all = "lowercase")]
167pub enum ToolChoice {
168    None,
169    Auto,
170    Required,
171}
172
173#[derive(Debug, Serialize, Deserialize)]
174#[serde(tag = "type", rename_all = "snake_case")]
175pub enum ToolDefinition {
176    Function { function: FunctionDefinition },
177}
178
179#[derive(Debug, Serialize, Deserialize)]
180pub struct FunctionDefinition {
181    pub name: String,
182    pub description: Option<String>,
183    pub parameters: Option<Value>,
184}
185
186#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
187#[serde(tag = "role", rename_all = "lowercase")]
188pub enum RequestMessage {
189    Assistant {
190        content: Option<String>,
191        #[serde(default, skip_serializing_if = "Vec::is_empty")]
192        tool_calls: Vec<ToolCall>,
193        #[serde(default, skip_serializing_if = "Option::is_none")]
194        reasoning_content: Option<String>,
195    },
196    User {
197        content: String,
198    },
199    System {
200        content: String,
201    },
202    Tool {
203        content: String,
204        tool_call_id: String,
205    },
206}
207
208#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
209pub struct ToolCall {
210    pub id: String,
211    #[serde(flatten)]
212    pub content: ToolCallContent,
213}
214
215#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
216#[serde(tag = "type", rename_all = "lowercase")]
217pub enum ToolCallContent {
218    Function { function: FunctionContent },
219}
220
221#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
222pub struct FunctionContent {
223    pub name: String,
224    pub arguments: String,
225}
226
227#[derive(Serialize, Deserialize, Debug)]
228pub struct Response {
229    pub id: String,
230    pub object: String,
231    pub created: u64,
232    pub model: String,
233    pub choices: Vec<Choice>,
234    pub usage: Usage,
235    #[serde(default, skip_serializing_if = "Option::is_none")]
236    pub reasoning_content: Option<String>,
237}
238
239#[derive(Serialize, Deserialize, Debug)]
240pub struct Usage {
241    pub prompt_tokens: u64,
242    pub completion_tokens: u64,
243    pub total_tokens: u64,
244    #[serde(default)]
245    pub prompt_cache_hit_tokens: u64,
246    #[serde(default)]
247    pub prompt_cache_miss_tokens: u64,
248}
249
250#[derive(Serialize, Deserialize, Debug)]
251pub struct Choice {
252    pub index: u32,
253    pub message: RequestMessage,
254    pub finish_reason: Option<String>,
255}
256
257#[derive(Serialize, Deserialize, Debug)]
258pub struct StreamResponse {
259    pub id: String,
260    pub object: String,
261    pub created: u64,
262    pub model: String,
263    pub choices: Vec<StreamChoice>,
264    pub usage: Option<Usage>,
265}
266
267#[derive(Serialize, Deserialize, Debug)]
268pub struct StreamChoice {
269    pub index: u32,
270    pub delta: StreamDelta,
271    pub finish_reason: Option<String>,
272}
273
274#[derive(Serialize, Deserialize, Debug)]
275pub struct StreamDelta {
276    pub role: Option<Role>,
277    pub content: Option<String>,
278    #[serde(default, skip_serializing_if = "Option::is_none")]
279    pub tool_calls: Option<Vec<ToolCallChunk>>,
280    #[serde(default, skip_serializing_if = "Option::is_none")]
281    pub reasoning_content: Option<String>,
282}
283
284#[derive(Serialize, Deserialize, Debug)]
285pub struct ToolCallChunk {
286    pub index: usize,
287    pub id: Option<String>,
288    pub function: Option<FunctionChunk>,
289}
290
291#[derive(Serialize, Deserialize, Debug)]
292pub struct FunctionChunk {
293    pub name: Option<String>,
294    pub arguments: Option<String>,
295}
296
297pub async fn stream_completion(
298    client: &dyn HttpClient,
299    api_url: &str,
300    api_key: &str,
301    request: Request,
302    extra_headers: &CustomHeaders,
303) -> Result<BoxStream<'static, Result<StreamResponse>>> {
304    let uri = format!("{api_url}/chat/completions");
305    let request = HttpRequest::builder()
306        .method(Method::POST)
307        .uri(uri)
308        .header("Content-Type", "application/json")
309        .header("Authorization", format!("Bearer {}", api_key.trim()))
310        .extra_headers(extra_headers)
311        .body(AsyncBody::from(serde_json::to_string(&request)?))?;
312    let mut response = client.send(request).await?;
313
314    if response.status().is_success() {
315        let reader = BufReader::new(response.into_body());
316        Ok(reader
317            .lines()
318            .filter_map(|line| async move {
319                match line {
320                    Ok(line) => {
321                        let line = line.strip_prefix("data: ")?;
322                        if line == "[DONE]" {
323                            None
324                        } else {
325                            Some(serde_json::from_str(line).map_err(Into::into))
326                        }
327                    }
328                    Err(error) => Some(Err(anyhow!(error))),
329                }
330            })
331            .boxed())
332    } else {
333        let mut body = String::new();
334        response.body_mut().read_to_string(&mut body).await?;
335        anyhow::bail!(
336            "Failed to connect to DeepSeek API: {} {}",
337            response.status(),
338            body,
339        );
340    }
341}
342