tenant.openagents/omega

1use anyhow::{Result, anyhow};
2use collections::{BTreeMap, HashMap};
3use credentials_provider::CredentialsProvider;
4
5use futures::{FutureExt, Stream, StreamExt, future::BoxFuture, stream::BoxStream};
6use gpui::{App, AppContext, AsyncApp, Context, Entity, Global, SharedString, Task};
7use http_client::{CustomHeaders, HttpClient};
8use language_model::{
9    ApiKeyConfiguration, ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel,
10    LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
11    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
12    LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
13    LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, ProviderSettingsView,
14    RateLimiter, Role, StopReason, TokenUsage, env_var,
15};
16pub use mistral::{MISTRAL_API_URL, StreamResponse};
17pub use settings::MistralAvailableModel as AvailableModel;
18use settings::{Settings, SettingsStore};
19use std::pin::Pin;
20use std::sync::{Arc, LazyLock};
21use strum::IntoEnumIterator;
22use ui::IconName;
23
24use language_model::util::{fix_streamed_json, parse_tool_arguments};
25
26const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("mistral");
27const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("Mistral");
28
29const API_KEY_ENV_VAR_NAME: &str = "MISTRAL_API_KEY";
30static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME);
31pub(crate) const RESERVED_HEADER_NAMES: &[&str] = &["x-affinity"];
32
33#[derive(Default, Clone, Debug, PartialEq)]
34pub struct MistralSettings {
35    pub api_url: String,
36    pub available_models: Vec<AvailableModel>,
37    pub custom_headers: CustomHeaders,
38}
39
40pub struct MistralLanguageModelProvider {
41    http_client: Arc<dyn HttpClient>,
42    pub state: Entity<State>,
43}
44
45pub struct State {
46    api_key_state: ApiKeyState,
47    credentials_provider: Arc<dyn CredentialsProvider>,
48}
49
50impl State {
51    fn is_authenticated(&self) -> bool {
52        self.api_key_state.has_key()
53    }
54
55    fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {
56        let credentials_provider = self.credentials_provider.clone();
57        let api_url = MistralLanguageModelProvider::api_url(cx);
58        self.api_key_state.store(
59            api_url,
60            api_key,
61            |this| &mut this.api_key_state,
62            credentials_provider,
63            cx,
64        )
65    }
66
67    fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
68        let credentials_provider = self.credentials_provider.clone();
69        let api_url = MistralLanguageModelProvider::api_url(cx);
70        self.api_key_state.load_if_needed(
71            api_url,
72            |this| &mut this.api_key_state,
73            credentials_provider,
74            cx,
75        )
76    }
77}
78
79struct GlobalMistralLanguageModelProvider(Arc<MistralLanguageModelProvider>);
80
81impl Global for GlobalMistralLanguageModelProvider {}
82
83impl MistralLanguageModelProvider {
84    pub fn try_global(cx: &App) -> Option<&Arc<MistralLanguageModelProvider>> {
85        cx.try_global::<GlobalMistralLanguageModelProvider>()
86            .map(|this| &this.0)
87    }
88
89    pub fn global(
90        http_client: Arc<dyn HttpClient>,
91        credentials_provider: Arc<dyn CredentialsProvider>,
92        cx: &mut App,
93    ) -> Arc<Self> {
94        if let Some(this) = cx.try_global::<GlobalMistralLanguageModelProvider>() {
95            return this.0.clone();
96        }
97        let state = cx.new(|cx| {
98            cx.observe_global::<SettingsStore>(|this: &mut State, cx| {
99                let credentials_provider = this.credentials_provider.clone();
100                let api_url = Self::api_url(cx);
101                this.api_key_state.handle_url_change(
102                    api_url,
103                    |this| &mut this.api_key_state,
104                    credentials_provider,
105                    cx,
106                );
107                cx.notify();
108            })
109            .detach();
110            State {
111                api_key_state: ApiKeyState::new(Self::api_url(cx), (*API_KEY_ENV_VAR).clone()),
112                credentials_provider,
113            }
114        });
115
116        let this = Arc::new(Self { http_client, state });
117        cx.set_global(GlobalMistralLanguageModelProvider(this));
118        cx.global::<GlobalMistralLanguageModelProvider>().0.clone()
119    }
120
121    fn create_language_model(&self, model: mistral::Model) -> Arc<dyn LanguageModel> {
122        Arc::new(MistralLanguageModel {
123            id: LanguageModelId::from(model.id().to_string()),
124            model,
125            state: self.state.clone(),
126            http_client: self.http_client.clone(),
127            request_limiter: RateLimiter::new(4),
128        })
129    }
130
131    fn settings(cx: &App) -> &MistralSettings {
132        &crate::AllLanguageModelSettings::get_global(cx).mistral
133    }
134
135    pub fn api_url(cx: &App) -> SharedString {
136        let api_url = &Self::settings(cx).api_url;
137        if api_url.is_empty() {
138            mistral::MISTRAL_API_URL.into()
139        } else {
140            SharedString::new(api_url.as_str())
141        }
142    }
143}
144
145impl LanguageModelProviderState for MistralLanguageModelProvider {
146    type ObservableEntity = State;
147
148    fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
149        Some(self.state.clone())
150    }
151}
152
153impl LanguageModelProvider for MistralLanguageModelProvider {
154    fn id(&self) -> LanguageModelProviderId {
155        PROVIDER_ID
156    }
157
158    fn name(&self) -> LanguageModelProviderName {
159        PROVIDER_NAME
160    }
161
162    fn icon(&self) -> IconOrSvg {
163        IconOrSvg::Icon(IconName::AiMistral)
164    }
165
166    fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
167        Some(self.create_language_model(mistral::Model::default()))
168    }
169
170    fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
171        Some(self.create_language_model(mistral::Model::default_fast()))
172    }
173
174    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
175        let mut models = BTreeMap::default();
176
177        // Add base models from mistral::Model::iter()
178        for model in mistral::Model::iter() {
179            if !matches!(model, mistral::Model::Custom { .. }) {
180                models.insert(model.id().to_string(), model);
181            }
182        }
183
184        // Override with available models from settings
185        for model in &Self::settings(cx).available_models {
186            models.insert(
187                model.name.clone(),
188                mistral::Model::Custom {
189                    name: model.name.clone(),
190                    display_name: model.display_name.clone(),
191                    max_tokens: model.max_tokens,
192                    max_output_tokens: model.max_output_tokens,
193                    max_completion_tokens: model.max_completion_tokens,
194                    supports_tools: model.supports_tools,
195                    supports_images: model.supports_images,
196                    supports_thinking: model.supports_thinking,
197                },
198            );
199        }
200
201        models
202            .into_values()
203            .map(|model| {
204                Arc::new(MistralLanguageModel {
205                    id: LanguageModelId::from(model.id().to_string()),
206                    model,
207                    state: self.state.clone(),
208                    http_client: self.http_client.clone(),
209                    request_limiter: RateLimiter::new(4),
210                }) as Arc<dyn LanguageModel>
211            })
212            .collect()
213    }
214
215    fn is_authenticated(&self, cx: &App) -> bool {
216        self.state.read(cx).is_authenticated()
217    }
218
219    fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
220        self.state.update(cx, |state, cx| state.authenticate(cx))
221    }
222
223    fn settings_view(&self, cx: &mut App) -> Option<ProviderSettingsView> {
224        let state = self.state.read(cx);
225        Some(ProviderSettingsView::ApiKey(ApiKeyConfiguration::new(
226            state.api_key_state.has_key(),
227            state.api_key_state.is_from_env_var(),
228            state.api_key_state.env_var_name().clone(),
229            "https://console.mistral.ai/api-keys".into(),
230        )))
231    }
232
233    fn set_api_key(&self, api_key: Option<String>, cx: &mut App) -> Task<Result<()>> {
234        self.state
235            .update(cx, |state, cx| state.set_api_key(api_key, cx))
236    }
237}
238
239pub struct MistralLanguageModel {
240    id: LanguageModelId,
241    model: mistral::Model,
242    state: Entity<State>,
243    http_client: Arc<dyn HttpClient>,
244    request_limiter: RateLimiter,
245}
246
247impl MistralLanguageModel {
248    fn stream_completion(
249        &self,
250        request: mistral::Request,
251        affinity: Option<String>,
252        cx: &AsyncApp,
253    ) -> BoxFuture<
254        'static,
255        Result<futures::stream::BoxStream<'static, Result<mistral::StreamResponse>>>,
256    > {
257        let http_client = self.http_client.clone();
258
259        let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| {
260            let api_url = MistralLanguageModelProvider::api_url(cx);
261            let extra_headers = MistralLanguageModelProvider::settings(cx)
262                .custom_headers
263                .clone();
264            (state.api_key_state.key(&api_url), api_url, extra_headers)
265        });
266
267        let future = self.request_limiter.stream(async move {
268            let Some(api_key) = api_key else {
269                return Err(LanguageModelCompletionError::NoApiKey {
270                    provider: PROVIDER_NAME,
271                });
272            };
273            let request = mistral::stream_completion(
274                http_client.as_ref(),
275                &api_url,
276                &api_key,
277                request,
278                affinity,
279                &extra_headers,
280            );
281            let response = request.await?;
282            Ok(response)
283        });
284
285        async move { Ok(future.await?.boxed()) }.boxed()
286    }
287}
288
289impl LanguageModel for MistralLanguageModel {
290    fn id(&self) -> LanguageModelId {
291        self.id.clone()
292    }
293
294    fn name(&self) -> LanguageModelName {
295        LanguageModelName::from(self.model.display_name().to_string())
296    }
297
298    fn provider_id(&self) -> LanguageModelProviderId {
299        PROVIDER_ID
300    }
301
302    fn provider_name(&self) -> LanguageModelProviderName {
303        PROVIDER_NAME
304    }
305
306    fn supports_tools(&self) -> bool {
307        self.model.supports_tools()
308    }
309
310    fn supports_streaming_tools(&self) -> bool {
311        true
312    }
313
314    fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool {
315        self.model.supports_tools()
316    }
317
318    fn supports_images(&self) -> bool {
319        self.model.supports_images()
320    }
321
322    fn telemetry_id(&self) -> String {
323        format!("mistral/{}", self.model.id())
324    }
325
326    fn max_token_count(&self) -> u64 {
327        self.model.max_token_count()
328    }
329
330    fn max_output_tokens(&self) -> Option<u64> {
331        self.model.max_output_tokens()
332    }
333
334    fn stream_completion(
335        &self,
336        request: LanguageModelRequest,
337        cx: &AsyncApp,
338    ) -> BoxFuture<
339        'static,
340        Result<
341            BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
342            LanguageModelCompletionError,
343        >,
344    > {
345        let (request, affinity) =
346            match into_mistral(request, self.model.clone(), self.max_output_tokens()) {
347                Ok(request) => request,
348                Err(error) => return async move { Err(error.into()) }.boxed(),
349            };
350        let stream = self.stream_completion(request, affinity, cx);
351
352        async move {
353            let stream = stream.await?;
354            let mapper = MistralEventMapper::new();
355            Ok(mapper.map_stream(stream).boxed())
356        }
357        .boxed()
358    }
359}
360
361pub fn into_mistral(
362    request: LanguageModelRequest,
363    model: mistral::Model,
364    max_output_tokens: Option<u64>,
365) -> Result<(mistral::Request, Option<String>)> {
366    if request.contains_custom_tool_input() {
367        anyhow::bail!("Mistral does not support custom tools");
368    }
369
370    let stream = true;
371
372    let mut messages = Vec::new();
373    for message in &request.messages {
374        match message.role {
375            Role::User => {
376                let mut message_content = mistral::MessageContent::empty();
377                for content in &message.content {
378                    match content {
379                        MessageContent::Text(text) => {
380                            message_content
381                                .push_part(mistral::MessagePart::Text { text: text.clone() });
382                        }
383                        MessageContent::Image(image_content) => {
384                            if model.supports_images() {
385                                message_content.push_part(mistral::MessagePart::ImageUrl {
386                                    image_url: image_content.to_base64_url(),
387                                });
388                            }
389                        }
390                        MessageContent::Thinking { text, .. } => {
391                            if model.supports_thinking() {
392                                message_content.push_part(mistral::MessagePart::Thinking {
393                                    thinking: vec![mistral::ThinkingPart::Text {
394                                        text: text.clone(),
395                                    }],
396                                });
397                            }
398                        }
399                        MessageContent::RedactedThinking(_) => {}
400                        MessageContent::Compaction(_) => {}
401                        MessageContent::ToolUse(_) => {
402                            // Tool use is not supported in User messages for Mistral
403                        }
404                        MessageContent::ToolResult(tool_result) => {
405                            let mut text_parts: Vec<String> = Vec::new();
406                            for part in &tool_result.content {
407                                match part {
408                                    LanguageModelToolResultContent::Text(text) => {
409                                        text_parts.push(text.to_string());
410                                    }
411                                    LanguageModelToolResultContent::Image(_) => {
412                                        text_parts.push("[Tool responded with an image, but Omega doesn't support these in Mistral models yet]".to_string());
413                                    }
414                                }
415                            }
416                            messages.push(mistral::RequestMessage::Tool {
417                                content: text_parts.join("\n"),
418                                tool_call_id: tool_result.tool_use_id.to_string(),
419                            });
420                        }
421                    }
422                }
423                if !matches!(message_content, mistral::MessageContent::Plain { ref content } if content.is_empty())
424                {
425                    messages.push(mistral::RequestMessage::User {
426                        content: message_content,
427                    });
428                }
429            }
430            Role::Assistant => {
431                for content in &message.content {
432                    match content {
433                        MessageContent::Text(text) if text.is_empty() => {
434                            // Mistral API returns a 400 if there's neither content nor tool_calls
435                        }
436                        MessageContent::Text(text) => {
437                            messages.push(mistral::RequestMessage::Assistant {
438                                content: Some(mistral::MessageContent::Plain {
439                                    content: text.clone(),
440                                }),
441                                tool_calls: Vec::new(),
442                            });
443                        }
444                        MessageContent::Thinking { text, .. } => {
445                            if model.supports_thinking() {
446                                messages.push(mistral::RequestMessage::Assistant {
447                                    content: Some(mistral::MessageContent::Multipart {
448                                        content: vec![mistral::MessagePart::Thinking {
449                                            thinking: vec![mistral::ThinkingPart::Text {
450                                                text: text.clone(),
451                                            }],
452                                        }],
453                                    }),
454                                    tool_calls: Vec::new(),
455                                });
456                            }
457                        }
458                        MessageContent::RedactedThinking(_) => {}
459                        MessageContent::Image(_) => {}
460                        MessageContent::Compaction(_) => {}
461                        MessageContent::ToolUse(tool_use) => {
462                            let input = tool_use.input.as_json().ok_or_else(|| {
463                                anyhow!("Mistral does not support custom tool calls")
464                            })?;
465                            let tool_call = mistral::ToolCall {
466                                id: tool_use.id.to_string(),
467                                content: mistral::ToolCallContent::Function {
468                                    function: mistral::FunctionContent {
469                                        name: tool_use.name.to_string(),
470                                        arguments: serde_json::to_string(input).unwrap_or_default(),
471                                    },
472                                },
473                            };
474
475                            if let Some(mistral::RequestMessage::Assistant { tool_calls, .. }) =
476                                messages.last_mut()
477                            {
478                                tool_calls.push(tool_call);
479                            } else {
480                                messages.push(mistral::RequestMessage::Assistant {
481                                    content: None,
482                                    tool_calls: vec![tool_call],
483                                });
484                            }
485                        }
486                        MessageContent::ToolResult(_) => {
487                            // Tool results are not supported in Assistant messages
488                        }
489                    }
490                }
491            }
492            Role::System => {
493                for content in &message.content {
494                    match content {
495                        MessageContent::Text(text) => {
496                            messages.push(mistral::RequestMessage::System {
497                                content: mistral::MessageContent::Plain {
498                                    content: text.clone(),
499                                },
500                            });
501                        }
502                        MessageContent::Thinking { text, .. } => {
503                            if model.supports_thinking() {
504                                messages.push(mistral::RequestMessage::System {
505                                    content: mistral::MessageContent::Multipart {
506                                        content: vec![mistral::MessagePart::Thinking {
507                                            thinking: vec![mistral::ThinkingPart::Text {
508                                                text: text.clone(),
509                                            }],
510                                        }],
511                                    },
512                                });
513                            }
514                        }
515                        MessageContent::RedactedThinking(_) => {}
516                        MessageContent::Compaction(_) => {}
517                        MessageContent::Image(_)
518                        | MessageContent::ToolUse(_)
519                        | MessageContent::ToolResult(_) => {
520                            // Images and tools are not supported in System messages
521                        }
522                    }
523                }
524            }
525        }
526    }
527
528    Ok((
529        mistral::Request {
530            model: model.id().to_string(),
531            messages,
532            stream,
533            stream_options: if stream {
534                Some(mistral::StreamOptions {
535                    stream_tool_calls: Some(true),
536                })
537            } else {
538                None
539            },
540            max_tokens: max_output_tokens,
541            temperature: request.temperature,
542            response_format: None,
543            tool_choice: match request.tool_choice {
544                Some(LanguageModelToolChoice::Auto) if !request.tools.is_empty() => {
545                    Some(mistral::ToolChoice::Auto)
546                }
547                Some(LanguageModelToolChoice::Any) if !request.tools.is_empty() => {
548                    Some(mistral::ToolChoice::Any)
549                }
550                Some(LanguageModelToolChoice::None) => Some(mistral::ToolChoice::None),
551                _ if !request.tools.is_empty() => Some(mistral::ToolChoice::Auto),
552                _ => None,
553            },
554            parallel_tool_calls: if !request.tools.is_empty() {
555                Some(false)
556            } else {
557                None
558            },
559            tools: request
560                .tools
561                .into_iter()
562                .map(|tool| {
563                    let input_schema = match tool.input {
564                        language_model::LanguageModelRequestToolInput::Function {
565                            input_schema,
566                            ..
567                        } => input_schema,
568                        language_model::LanguageModelRequestToolInput::Custom { .. } => {
569                            return Err(anyhow::anyhow!("Mistral does not support custom tools"));
570                        }
571                    };
572                    Ok(mistral::ToolDefinition::Function {
573                        function: mistral::FunctionDefinition {
574                            name: tool.name,
575                            description: Some(tool.description),
576                            parameters: Some(input_schema),
577                        },
578                    })
579                })
580                .collect::<Result<_>>()?,
581            reasoning_effort: if model.supports_thinking() && request.thinking_allowed {
582                Some(mistral::ReasoningEffort::High)
583            } else {
584                None
585            },
586        },
587        request.thread_id,
588    ))
589}
590
591pub struct MistralEventMapper {
592    tool_calls_by_index: HashMap<usize, RawToolCall>,
593}
594
595impl MistralEventMapper {
596    pub fn new() -> Self {
597        Self {
598            tool_calls_by_index: HashMap::default(),
599        }
600    }
601
602    pub fn map_stream(
603        mut self,
604        events: Pin<Box<dyn Send + Stream<Item = Result<StreamResponse>>>>,
605    ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
606    {
607        events.flat_map(move |event| {
608            futures::stream::iter(match event {
609                Ok(event) => self.map_event(event),
610                Err(error) => vec![Err(LanguageModelCompletionError::from(error))],
611            })
612        })
613    }
614
615    pub fn map_event(
616        &mut self,
617        event: mistral::StreamResponse,
618    ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
619        let Some(choice) = event.choices.first() else {
620            return vec![Err(LanguageModelCompletionError::from(anyhow!(
621                "Response contained no choices"
622            )))];
623        };
624
625        let mut events = Vec::new();
626        if let Some(content) = choice.delta.content.as_ref() {
627            match content {
628                mistral::MessageContentDelta::Text(text) => {
629                    events.push(Ok(LanguageModelCompletionEvent::Text(text.clone())));
630                }
631                mistral::MessageContentDelta::Parts(parts) => {
632                    for part in parts {
633                        match part {
634                            mistral::MessagePart::Text { text } => {
635                                events.push(Ok(LanguageModelCompletionEvent::Text(text.clone())));
636                            }
637                            mistral::MessagePart::Thinking { thinking } => {
638                                for tp in thinking.iter().cloned() {
639                                    match tp {
640                                        mistral::ThinkingPart::Text { text } => {
641                                            events.push(Ok(
642                                                LanguageModelCompletionEvent::Thinking {
643                                                    text,
644                                                    signature: None,
645                                                },
646                                            ));
647                                        }
648                                    }
649                                }
650                            }
651                            mistral::MessagePart::ImageUrl { .. } => {
652                                // We currently don't emit a separate event for images in responses.
653                            }
654                        }
655                    }
656                }
657            }
658        }
659
660        if let Some(tool_calls) = choice.delta.tool_calls.as_ref() {
661            for tool_call in tool_calls {
662                let entry = self.tool_calls_by_index.entry(tool_call.index).or_default();
663
664                if let Some(tool_id) = tool_call.id.clone()
665                    && !tool_id.is_empty()
666                    && tool_id != "null"
667                {
668                    entry.id = tool_id;
669                }
670
671                if let Some(function) = tool_call.function.as_ref() {
672                    if let Some(name) = function.name.clone()
673                        && !name.is_empty()
674                    {
675                        entry.name = name;
676                    }
677
678                    if let Some(arguments) = function.arguments.clone() {
679                        entry.arguments.push_str(&arguments);
680                    }
681                }
682
683                if !entry.id.is_empty() && !entry.name.is_empty() {
684                    if let Ok(input) = serde_json::from_str::<serde_json::Value>(
685                        &fix_streamed_json(&entry.arguments),
686                    ) {
687                        events.push(Ok(LanguageModelCompletionEvent::ToolUse(
688                            LanguageModelToolUse {
689                                id: entry.id.clone().into(),
690                                name: entry.name.as_str().into(),
691                                is_input_complete: false,
692                                input: language_model::LanguageModelToolUseInput::Json(input),
693                                raw_input: entry.arguments.clone(),
694                                thought_signature: None,
695                            },
696                        )));
697                    }
698                }
699            }
700        }
701
702        if let Some(usage) = event.usage {
703            events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
704                input_tokens: usage.prompt_tokens,
705                output_tokens: usage.completion_tokens,
706                cache_creation_input_tokens: 0,
707                cache_read_input_tokens: 0,
708            })));
709        }
710
711        if let Some(finish_reason) = choice.finish_reason.as_deref() {
712            match finish_reason {
713                "stop" => {
714                    events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));
715                }
716                "tool_calls" => {
717                    events.extend(self.process_tool_calls());
718                    events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)));
719                }
720                unexpected => {
721                    log::error!("Unexpected Mistral stop_reason: {unexpected:?}");
722                    events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));
723                }
724            }
725        }
726
727        events
728    }
729
730    fn process_tool_calls(
731        &mut self,
732    ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
733        let mut results = Vec::new();
734
735        for (_, tool_call) in self.tool_calls_by_index.drain() {
736            if tool_call.id.is_empty() || tool_call.name.is_empty() {
737                results.push(Err(LanguageModelCompletionError::from(anyhow!(
738                    "Received incomplete tool call: missing id or name"
739                ))));
740                continue;
741            }
742
743            match parse_tool_arguments(&tool_call.arguments) {
744                Ok(input) => results.push(Ok(LanguageModelCompletionEvent::ToolUse(
745                    LanguageModelToolUse {
746                        id: tool_call.id.into(),
747                        name: tool_call.name.into(),
748                        is_input_complete: true,
749                        input: language_model::LanguageModelToolUseInput::Json(input),
750                        raw_input: tool_call.arguments,
751                        thought_signature: None,
752                    },
753                ))),
754                Err(error) => {
755                    results.push(Ok(LanguageModelCompletionEvent::ToolUseJsonParseError {
756                        id: tool_call.id.into(),
757                        tool_name: tool_call.name.into(),
758                        raw_input: tool_call.arguments.into(),
759                        json_parse_error: error.to_string(),
760                    }))
761                }
762            }
763        }
764
765        results
766    }
767}
768
769#[derive(Default)]
770struct RawToolCall {
771    id: String,
772    name: String,
773    arguments: String,
774}
775
776#[cfg(test)]
777mod tests {
778    use super::*;
779    use language_model::{LanguageModelImage, LanguageModelRequestMessage, MessageContent};
780
781    fn tool_call_chunk(
782        id: Option<&str>,
783        name: Option<&str>,
784        arguments: Option<&str>,
785        finish_reason: Option<&str>,
786    ) -> mistral::StreamResponse {
787        mistral::StreamResponse {
788            id: "resp".into(),
789            object: "chat.completion.chunk".into(),
790            created: 0,
791            model: "test".into(),
792            choices: vec![mistral::StreamChoice {
793                index: 0,
794                delta: mistral::StreamDelta {
795                    role: None,
796                    content: None,
797                    tool_calls: if finish_reason.is_some() {
798                        None
799                    } else {
800                        Some(vec![mistral::ToolCallChunk {
801                            index: 0,
802                            id: id.map(Into::into),
803                            function: Some(mistral::FunctionChunk {
804                                name: name.map(Into::into),
805                                arguments: arguments.map(Into::into),
806                            }),
807                        }])
808                    },
809                },
810                finish_reason: finish_reason.map(Into::into),
811            }],
812            usage: None,
813        }
814    }
815
816    #[test]
817    fn test_streaming_tool_call_ignores_null_id() {
818        // Mistral's streaming API sometimes sends `"id": "null"` in continuation chunks.
819        let mut mapper = MistralEventMapper::new();
820
821        mapper.map_event(tool_call_chunk(
822            Some("real_id_123"),
823            Some("read_file"),
824            Some("{\"path\":"),
825            None,
826        ));
827        mapper.map_event(tool_call_chunk(
828            Some("null"),
829            None,
830            Some("\"a.txt\"}"),
831            None,
832        ));
833        let events = mapper.map_event(tool_call_chunk(None, None, None, Some("tool_calls")));
834
835        let Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) = &events[0] else {
836            panic!("Expected first event to be ToolUse, got: {:?}", events[0]);
837        };
838
839        assert_eq!(tool_use.id.to_string(), "real_id_123");
840        assert_eq!(tool_use.name.as_ref(), "read_file");
841        assert_eq!(
842            tool_use.input,
843            language_model::LanguageModelToolUseInput::Json(serde_json::json!({"path": "a.txt"}))
844        );
845    }
846
847    #[test]
848    fn test_into_mistral_basic_conversion() {
849        let request = LanguageModelRequest {
850            messages: vec![
851                LanguageModelRequestMessage {
852                    role: Role::System,
853                    content: vec![MessageContent::Text("System prompt".into())],
854                    cache: false,
855                    reasoning_details: None,
856                },
857                LanguageModelRequestMessage {
858                    role: Role::User,
859                    content: vec![MessageContent::Text("Hello".into())],
860                    cache: false,
861                    reasoning_details: None,
862                },
863                // should skip empty assistant messages
864                LanguageModelRequestMessage {
865                    role: Role::Assistant,
866                    content: vec![MessageContent::Text("".into())],
867                    cache: false,
868                    reasoning_details: None,
869                },
870            ],
871            temperature: Some(0.5),
872            tools: vec![],
873            tool_choice: None,
874            thread_id: Some("abcdef".into()),
875            prompt_id: None,
876            intent: None,
877            stop: vec![],
878            thinking_allowed: true,
879            thinking_effort: None,
880            speed: Default::default(),
881            compact_at_tokens: None,
882        };
883
884        let (mistral_request, affinity) =
885            into_mistral(request, mistral::Model::MistralSmallLatest, None).unwrap();
886
887        assert_eq!(mistral_request.model, "mistral-small-latest");
888        assert_eq!(mistral_request.temperature, Some(0.5));
889        assert_eq!(mistral_request.messages.len(), 2);
890        assert!(mistral_request.stream);
891        assert_eq!(affinity, Some("abcdef".into()));
892    }
893
894    #[test]
895    fn test_into_mistral_reasoning_effort() {
896        let request = |thinking_allowed| LanguageModelRequest {
897            messages: vec![LanguageModelRequestMessage {
898                role: Role::User,
899                content: vec![MessageContent::Text("Hello".into())],
900                cache: false,
901                reasoning_details: None,
902            }],
903            temperature: None,
904            tools: vec![],
905            tool_choice: None,
906            thread_id: None,
907            prompt_id: None,
908            intent: None,
909            stop: vec![],
910            thinking_allowed,
911            thinking_effort: None,
912            speed: Default::default(),
913            compact_at_tokens: None,
914        };
915
916        let (mistral_request, _) =
917            into_mistral(request(true), mistral::Model::MistralMediumLatest, None).unwrap();
918        assert_eq!(
919            mistral_request.reasoning_effort,
920            Some(mistral::ReasoningEffort::High)
921        );
922
923        let (mistral_request, _) =
924            into_mistral(request(false), mistral::Model::MistralMediumLatest, None).unwrap();
925        assert_eq!(mistral_request.reasoning_effort, None);
926
927        let (mistral_request, _) =
928            into_mistral(request(true), mistral::Model::CodestralLatest, None).unwrap();
929        assert_eq!(mistral_request.reasoning_effort, None);
930    }
931
932    #[test]
933    fn test_into_mistral_with_image() {
934        let request = LanguageModelRequest {
935            messages: vec![LanguageModelRequestMessage {
936                role: Role::User,
937                content: vec![
938                    MessageContent::Text("What's in this image?".into()),
939                    MessageContent::Image(LanguageModelImage {
940                        source: "base64data".into(),
941                    }),
942                ],
943                cache: false,
944                reasoning_details: None,
945            }],
946            tools: vec![],
947            tool_choice: None,
948            temperature: None,
949            thread_id: None,
950            prompt_id: None,
951            intent: None,
952            stop: vec![],
953            thinking_allowed: true,
954            thinking_effort: None,
955            speed: None,
956            compact_at_tokens: None,
957        };
958
959        let (mistral_request, _) =
960            into_mistral(request, mistral::Model::MistralSmallLatest, None).unwrap();
961
962        assert_eq!(mistral_request.messages.len(), 1);
963        assert!(matches!(
964            &mistral_request.messages[0],
965            mistral::RequestMessage::User {
966                content: mistral::MessageContent::Multipart { .. }
967            }
968        ));
969
970        if let mistral::RequestMessage::User {
971            content: mistral::MessageContent::Multipart { content },
972        } = &mistral_request.messages[0]
973        {
974            assert_eq!(content.len(), 2);
975            assert!(matches!(
976                &content[0],
977                mistral::MessagePart::Text { text } if text == "What's in this image?"
978            ));
979            assert!(matches!(
980                &content[1],
981                mistral::MessagePart::ImageUrl { image_url } if image_url.starts_with("data:image/png;base64,")
982            ));
983        }
984    }
985}
986