tenant.openagents/omega

1use anyhow::{Result, anyhow};
2use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
3use http_client::{
4    AsyncBody, CustomHeaders, HttpClient, HttpRequestExt, Method, Request as HttpRequest,
5    RequestBuilderExt,
6};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::convert::TryFrom;
10use strum::EnumIter;
11
12pub const MISTRAL_API_URL: &str = "https://api.mistral.ai/v1";
13
14#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
15#[serde(rename_all = "lowercase")]
16pub enum Role {
17    User,
18    Assistant,
19    System,
20    Tool,
21}
22
23impl TryFrom<String> for Role {
24    type Error = anyhow::Error;
25
26    fn try_from(value: String) -> Result<Self> {
27        match value.as_str() {
28            "user" => Ok(Self::User),
29            "assistant" => Ok(Self::Assistant),
30            "system" => Ok(Self::System),
31            "tool" => Ok(Self::Tool),
32            _ => anyhow::bail!("invalid role '{value}'"),
33        }
34    }
35}
36
37impl From<Role> for String {
38    fn from(val: Role) -> Self {
39        match val {
40            Role::User => "user".to_owned(),
41            Role::Assistant => "assistant".to_owned(),
42            Role::System => "system".to_owned(),
43            Role::Tool => "tool".to_owned(),
44        }
45    }
46}
47
48#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
49#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)]
50pub enum Model {
51    #[serde(rename = "codestral-latest", alias = "codestral-latest")]
52    #[default]
53    CodestralLatest,
54
55    #[serde(rename = "mistral-large-latest", alias = "mistral-large-latest")]
56    MistralLargeLatest,
57    #[serde(rename = "mistral-medium-latest", alias = "mistral-medium-latest")]
58    MistralMediumLatest,
59    #[serde(rename = "mistral-small-latest", alias = "mistral-small-latest")]
60    MistralSmallLatest,
61
62    #[serde(rename = "ministral-3b-latest", alias = "ministral-3b-latest")]
63    Ministral3bLatest,
64    #[serde(rename = "ministral-8b-latest", alias = "ministral-8b-latest")]
65    Ministral8bLatest,
66    #[serde(rename = "ministral-14b-latest", alias = "ministral-14b-latest")]
67    Ministral14bLatest,
68
69    #[serde(rename = "custom")]
70    Custom {
71        name: String,
72        /// The name displayed in the UI, such as in the agent panel model dropdown menu.
73        display_name: Option<String>,
74        max_tokens: u64,
75        max_output_tokens: Option<u64>,
76        max_completion_tokens: Option<u64>,
77        supports_tools: Option<bool>,
78        supports_images: Option<bool>,
79        supports_thinking: Option<bool>,
80    },
81}
82
83impl Model {
84    pub fn default_fast() -> Self {
85        Model::MistralSmallLatest
86    }
87
88    pub fn from_id(id: &str) -> Result<Self> {
89        match id {
90            "codestral-latest" => Ok(Self::CodestralLatest),
91            "mistral-large-latest" => Ok(Self::MistralLargeLatest),
92            "mistral-medium-latest" => Ok(Self::MistralMediumLatest),
93            "mistral-small-latest" => Ok(Self::MistralSmallLatest),
94            "ministral-3b-latest" => Ok(Self::Ministral3bLatest),
95            "ministral-8b-latest" => Ok(Self::Ministral8bLatest),
96            "ministral-14b-latest" => Ok(Self::Ministral14bLatest),
97            invalid_id => anyhow::bail!("invalid model id '{invalid_id}'"),
98        }
99    }
100
101    pub fn id(&self) -> &str {
102        match self {
103            Self::CodestralLatest => "codestral-latest",
104            Self::MistralLargeLatest => "mistral-large-latest",
105            Self::MistralMediumLatest => "mistral-medium-latest",
106            Self::MistralSmallLatest => "mistral-small-latest",
107            Self::Ministral3bLatest => "ministral-3b-latest",
108            Self::Ministral8bLatest => "ministral-8b-latest",
109            Self::Ministral14bLatest => "ministral-14b-latest",
110            Self::Custom { name, .. } => name,
111        }
112    }
113
114    pub fn display_name(&self) -> &str {
115        match self {
116            Self::CodestralLatest => "codestral-latest",
117            Self::MistralLargeLatest => "mistral-large-latest",
118            Self::MistralMediumLatest => "mistral-medium-latest",
119            Self::MistralSmallLatest => "mistral-small-latest",
120            Self::Ministral3bLatest => "ministral-3b-latest",
121            Self::Ministral8bLatest => "ministral-8b-latest",
122            Self::Ministral14bLatest => "ministral-14b-latest",
123            Self::Custom {
124                name, display_name, ..
125            } => display_name.as_ref().unwrap_or(name),
126        }
127    }
128
129    pub fn max_token_count(&self) -> u64 {
130        match self {
131            Self::CodestralLatest => 128000,
132            Self::MistralLargeLatest => 256000,
133            Self::MistralMediumLatest => 256000,
134            Self::MistralSmallLatest => 256000,
135            Self::Ministral3bLatest => 256000,
136            Self::Ministral8bLatest => 256000,
137            Self::Ministral14bLatest => 256000,
138            Self::Custom { max_tokens, .. } => *max_tokens,
139        }
140    }
141
142    pub fn max_output_tokens(&self) -> Option<u64> {
143        match self {
144            Self::Custom {
145                max_output_tokens, ..
146            } => *max_output_tokens,
147            _ => None,
148        }
149    }
150
151    pub fn supports_tools(&self) -> bool {
152        match self {
153            Self::CodestralLatest
154            | Self::MistralLargeLatest
155            | Self::MistralMediumLatest
156            | Self::MistralSmallLatest
157            | Self::Ministral3bLatest
158            | Self::Ministral8bLatest
159            | Self::Ministral14bLatest => true,
160            Self::Custom { supports_tools, .. } => supports_tools.unwrap_or(false),
161        }
162    }
163
164    pub fn supports_images(&self) -> bool {
165        match self {
166            Self::MistralLargeLatest
167            | Self::MistralMediumLatest
168            | Self::MistralSmallLatest
169            | Self::Ministral3bLatest
170            | Self::Ministral8bLatest
171            | Self::Ministral14bLatest => true,
172            Self::CodestralLatest => false,
173            Self::Custom {
174                supports_images, ..
175            } => supports_images.unwrap_or(false),
176        }
177    }
178
179    pub fn supports_thinking(&self) -> bool {
180        match self {
181            Self::MistralMediumLatest | Self::MistralSmallLatest => true,
182            Self::Custom {
183                supports_thinking, ..
184            } => supports_thinking.unwrap_or(false),
185            _ => false,
186        }
187    }
188}
189
190#[derive(Debug, Serialize, Deserialize)]
191pub struct Request {
192    pub model: String,
193    pub messages: Vec<RequestMessage>,
194    pub stream: bool,
195    #[serde(default, skip_serializing_if = "Option::is_none")]
196    pub stream_options: Option<StreamOptions>,
197    #[serde(default, skip_serializing_if = "Option::is_none")]
198    pub max_tokens: Option<u64>,
199    #[serde(default, skip_serializing_if = "Option::is_none")]
200    pub temperature: Option<f32>,
201    #[serde(default, skip_serializing_if = "Option::is_none")]
202    pub response_format: Option<ResponseFormat>,
203    #[serde(default, skip_serializing_if = "Option::is_none")]
204    pub tool_choice: Option<ToolChoice>,
205    #[serde(default, skip_serializing_if = "Option::is_none")]
206    pub parallel_tool_calls: Option<bool>,
207    #[serde(default, skip_serializing_if = "Vec::is_empty")]
208    pub tools: Vec<ToolDefinition>,
209    #[serde(default, skip_serializing_if = "Option::is_none")]
210    pub reasoning_effort: Option<ReasoningEffort>,
211}
212
213#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
214#[serde(rename_all = "lowercase")]
215pub enum ReasoningEffort {
216    None,
217    High,
218}
219
220#[derive(Debug, Serialize, Deserialize)]
221pub struct StreamOptions {
222    #[serde(default, skip_serializing_if = "Option::is_none")]
223    pub stream_tool_calls: Option<bool>,
224}
225
226#[derive(Debug, Serialize, Deserialize)]
227#[serde(rename_all = "snake_case")]
228pub enum ResponseFormat {
229    Text,
230    #[serde(rename = "json_object")]
231    JsonObject,
232}
233
234#[derive(Debug, Serialize, Deserialize)]
235#[serde(tag = "type", rename_all = "snake_case")]
236pub enum ToolDefinition {
237    Function { function: FunctionDefinition },
238}
239
240#[derive(Debug, Serialize, Deserialize)]
241pub struct FunctionDefinition {
242    pub name: String,
243    pub description: Option<String>,
244    pub parameters: Option<Value>,
245}
246
247#[derive(Debug, Serialize, Deserialize)]
248#[serde(rename_all = "lowercase")]
249pub enum ToolChoice {
250    Auto,
251    Required,
252    None,
253    Any,
254    #[serde(untagged)]
255    Function(ToolDefinition),
256}
257
258#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
259#[serde(tag = "role", rename_all = "lowercase")]
260pub enum RequestMessage {
261    Assistant {
262        #[serde(flatten)]
263        #[serde(default, skip_serializing_if = "Option::is_none")]
264        content: Option<MessageContent>,
265        #[serde(default, skip_serializing_if = "Vec::is_empty")]
266        tool_calls: Vec<ToolCall>,
267    },
268    User {
269        #[serde(flatten)]
270        content: MessageContent,
271    },
272    System {
273        #[serde(flatten)]
274        content: MessageContent,
275    },
276    Tool {
277        content: String,
278        tool_call_id: String,
279    },
280}
281
282#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
283#[serde(untagged)]
284pub enum MessageContent {
285    #[serde(rename = "content")]
286    Plain { content: String },
287    #[serde(rename = "content")]
288    Multipart { content: Vec<MessagePart> },
289}
290
291impl MessageContent {
292    pub fn empty() -> Self {
293        Self::Plain {
294            content: String::new(),
295        }
296    }
297
298    pub fn push_part(&mut self, part: MessagePart) {
299        match self {
300            Self::Plain { content } => match part {
301                MessagePart::Text { text } => {
302                    content.push_str(&text);
303                }
304                part => {
305                    let mut parts = if content.is_empty() {
306                        Vec::new()
307                    } else {
308                        vec![MessagePart::Text {
309                            text: content.clone(),
310                        }]
311                    };
312                    parts.push(part);
313                    *self = Self::Multipart { content: parts };
314                }
315            },
316            Self::Multipart { content } => {
317                content.push(part);
318            }
319        }
320    }
321}
322
323#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
324#[serde(tag = "type", rename_all = "snake_case")]
325pub enum MessagePart {
326    Text { text: String },
327    ImageUrl { image_url: String },
328    Thinking { thinking: Vec<ThinkingPart> },
329}
330
331// Backwards-compatibility alias for provider code that refers to ContentPart
332pub type ContentPart = MessagePart;
333
334#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
335#[serde(tag = "type", rename_all = "snake_case")]
336pub enum ThinkingPart {
337    Text { text: String },
338}
339
340#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
341pub struct ToolCall {
342    pub id: String,
343    #[serde(flatten)]
344    pub content: ToolCallContent,
345}
346
347#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
348#[serde(tag = "type", rename_all = "lowercase")]
349pub enum ToolCallContent {
350    Function { function: FunctionContent },
351}
352
353#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
354pub struct FunctionContent {
355    pub name: String,
356    pub arguments: String,
357}
358
359#[derive(Serialize, Deserialize, Debug)]
360pub struct Usage {
361    pub prompt_tokens: u64,
362    pub completion_tokens: u64,
363    pub total_tokens: u64,
364}
365
366#[derive(Serialize, Deserialize, Debug)]
367pub struct StreamResponse {
368    pub id: String,
369    pub object: String,
370    pub created: u64,
371    pub model: String,
372    pub choices: Vec<StreamChoice>,
373    pub usage: Option<Usage>,
374}
375
376#[derive(Serialize, Deserialize, Debug)]
377pub struct StreamChoice {
378    pub index: u32,
379    pub delta: StreamDelta,
380    pub finish_reason: Option<String>,
381}
382
383#[derive(Serialize, Deserialize, Debug, Clone)]
384pub struct StreamDelta {
385    pub role: Option<Role>,
386    #[serde(default, skip_serializing_if = "Option::is_none")]
387    pub content: Option<MessageContentDelta>,
388    #[serde(default, skip_serializing_if = "Option::is_none")]
389    pub tool_calls: Option<Vec<ToolCallChunk>>,
390}
391
392#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
393#[serde(untagged)]
394pub enum MessageContentDelta {
395    Text(String),
396    Parts(Vec<MessagePart>),
397}
398
399#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
400pub struct ToolCallChunk {
401    pub index: usize,
402    pub id: Option<String>,
403    pub function: Option<FunctionChunk>,
404}
405
406#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
407pub struct FunctionChunk {
408    pub name: Option<String>,
409    pub arguments: Option<String>,
410}
411
412pub async fn stream_completion(
413    client: &dyn HttpClient,
414    api_url: &str,
415    api_key: &str,
416    request: Request,
417    affinity: Option<String>,
418    extra_headers: &CustomHeaders,
419) -> Result<BoxStream<'static, Result<StreamResponse>>> {
420    let uri = format!("{api_url}/chat/completions");
421    let request_builder = HttpRequest::builder()
422        .method(Method::POST)
423        .uri(uri)
424        .header("Content-Type", "application/json")
425        .header("Authorization", format!("Bearer {}", api_key.trim()))
426        .when_some(affinity, |this, affinity| {
427            this.header("x-affinity", affinity)
428        })
429        .extra_headers(extra_headers);
430
431    let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;
432    let mut response = client.send(request).await?;
433
434    if response.status().is_success() {
435        let reader = BufReader::new(response.into_body());
436        Ok(reader
437            .lines()
438            .filter_map(|line| async move {
439                match line {
440                    Ok(line) => {
441                        let line = line.strip_prefix("data: ")?;
442                        if line == "[DONE]" {
443                            None
444                        } else {
445                            match serde_json::from_str(line) {
446                                Ok(response) => Some(Ok(response)),
447                                Err(error) => Some(Err(anyhow!(error))),
448                            }
449                        }
450                    }
451                    Err(error) => Some(Err(anyhow!(error))),
452                }
453            })
454            .boxed())
455    } else {
456        let mut body = String::new();
457        response.body_mut().read_to_string(&mut body).await?;
458        anyhow::bail!(
459            "Failed to connect to Mistral API: {} {}",
460            response.status(),
461            body,
462        );
463    }
464}
465