1use crate::client::{
42 self, ApiKey, Capabilities, Capable, DebugExt, ModelLister, Nothing, Provider, ProviderBuilder,
43 ProviderClient,
44};
45use crate::completion::{GetTokenUsage, Usage};
46use crate::http_client::{self, HttpClientExt};
47use crate::message::DocumentSourceKind;
48use crate::model::{Model, ModelList, ModelListingError};
49use crate::streaming::RawStreamingChoice;
50use crate::telemetry::{CompletionOperation, CompletionSpanBuilder};
51use crate::{
52 OneOrMany,
53 completion::{self, CompletionError, CompletionRequest},
54 embeddings::{self, EmbeddingError},
55 json_utils, message,
56 message::{ImageDetail, Text},
57 streaming,
58 wasm_compat::{WasmCompatSend, WasmCompatSync},
59};
60use async_stream::try_stream;
61use bytes::Bytes;
62use futures::StreamExt;
63use serde::{Deserialize, Serialize};
64use serde_json::{Value, json};
65use std::{convert::TryFrom, str::FromStr};
66use tracing_futures::Instrument;
67const OLLAMA_API_BASE_URL: &str = "http://localhost:11434";
70
71#[derive(Debug, Default, Clone)]
74pub struct OllamaApiKey(Option<String>);
75
76impl ApiKey for OllamaApiKey {
77 fn into_header(
78 self,
79 ) -> Option<http_client::Result<(http::header::HeaderName, http::header::HeaderValue)>> {
80 self.0.map(http_client::make_auth_header)
81 }
82}
83
84impl From<Nothing> for OllamaApiKey {
85 fn from(_: Nothing) -> Self {
86 Self(None)
87 }
88}
89
90impl From<String> for OllamaApiKey {
91 fn from(key: String) -> Self {
92 if key.is_empty() {
93 Self(None)
94 } else {
95 Self(Some(key))
96 }
97 }
98}
99
100impl From<&str> for OllamaApiKey {
101 fn from(key: &str) -> Self {
102 if key.is_empty() {
103 Self(None)
104 } else {
105 Self(Some(key.to_owned()))
106 }
107 }
108}
109
110#[derive(Debug, Default, Clone, Copy)]
111pub struct OllamaExt;
112
113#[derive(Debug, Default, Clone, Copy)]
114pub struct OllamaBuilder;
115
116impl Provider for OllamaExt {
117 type Builder = OllamaBuilder;
118 const VERIFY_PATH: &'static str = "api/tags";
119}
120
121impl<H> Capabilities<H> for OllamaExt {
122 type Completion = Capable<CompletionModel<H>>;
123 type Transcription = Nothing;
124 type Embeddings = Capable<EmbeddingModel<H>>;
125 type ModelListing = Capable<OllamaModelLister<H>>;
126 #[cfg(feature = "image")]
127 type ImageGeneration = Nothing;
128
129 #[cfg(feature = "audio")]
130 type AudioGeneration = Nothing;
131 type Rerank = Nothing;
132}
133
134impl DebugExt for OllamaExt {}
135
136impl ProviderBuilder for OllamaBuilder {
137 type Extension<H>
138 = OllamaExt
139 where
140 H: HttpClientExt;
141 type ApiKey = OllamaApiKey;
142
143 const BASE_URL: &'static str = OLLAMA_API_BASE_URL;
144
145 fn build<H>(
146 _builder: &client::ClientBuilder<Self, Self::ApiKey, H>,
147 ) -> http_client::Result<Self::Extension<H>>
148 where
149 H: HttpClientExt,
150 {
151 Ok(OllamaExt)
152 }
153}
154
155pub type Client<H = reqwest::Client> = client::Client<OllamaExt, H>;
156pub type ClientBuilder<H = crate::markers::Missing> =
157 client::ClientBuilder<OllamaBuilder, OllamaApiKey, H>;
158
159impl ProviderClient for Client {
160 type Input = OllamaApiKey;
161 type Error = crate::client::ProviderClientError;
162
163 fn from_env() -> Result<Self, Self::Error> {
164 let api_base = crate::client::optional_env_var("OLLAMA_API_BASE_URL")?
165 .unwrap_or_else(|| OLLAMA_API_BASE_URL.to_string());
166
167 let api_key = crate::client::optional_env_var("OLLAMA_API_KEY")?
168 .map(OllamaApiKey::from)
169 .unwrap_or_default();
170
171 Self::builder()
172 .api_key(api_key)
173 .base_url(&api_base)
174 .build()
175 .map_err(Into::into)
176 }
177
178 fn from_val(api_key: Self::Input) -> Result<Self, Self::Error> {
179 Self::builder().api_key(api_key).build().map_err(Into::into)
180 }
181}
182
183pub const ALL_MINILM: &str = "all-minilm";
186pub const NOMIC_EMBED_TEXT: &str = "nomic-embed-text";
187
188fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
189 match identifier {
190 ALL_MINILM => Some(384),
191 NOMIC_EMBED_TEXT => Some(768),
192 _ => None,
193 }
194}
195
196#[derive(Debug, Serialize, Deserialize)]
197pub struct EmbeddingResponse {
198 pub model: String,
199 pub embeddings: Vec<Vec<f64>>,
200 #[serde(default)]
201 pub total_duration: Option<u64>,
202 #[serde(default)]
203 pub load_duration: Option<u64>,
204 #[serde(default)]
205 pub prompt_eval_count: Option<u64>,
206}
207
208#[derive(Clone)]
211pub struct EmbeddingModel<T = reqwest::Client> {
212 client: Client<T>,
213 pub model: String,
214 ndims: usize,
215}
216
217impl<T> EmbeddingModel<T> {
218 pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
219 Self {
220 client,
221 model: model.into(),
222 ndims,
223 }
224 }
225
226 pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
227 Self {
228 client,
229 model: model.into(),
230 ndims,
231 }
232 }
233}
234
235impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
236where
237 T: HttpClientExt + Clone + 'static,
238{
239 type Client = Client<T>;
240
241 fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
242 let model = model.into();
243 let dims = dims
244 .or(model_dimensions_from_identifier(&model))
245 .unwrap_or_default();
246 Self::new(client.clone(), model, dims)
247 }
248
249 const MAX_DOCUMENTS: usize = 1024;
250 fn ndims(&self) -> usize {
251 self.ndims
252 }
253
254 async fn embed_texts(
255 &self,
256 documents: impl IntoIterator<Item = String>,
257 ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
258 let docs: Vec<String> = documents.into_iter().collect();
259
260 let body = serde_json::to_vec(&json!({
261 "model": self.model,
262 "input": docs
263 }))?;
264
265 let req = self
266 .client
267 .post("api/embed")?
268 .body(body)
269 .map_err(|e| EmbeddingError::HttpError(e.into()))?;
270
271 let response = self.client.send::<_, Vec<u8>>(req).await?;
272
273 let status = response.status();
274 if !status.is_success() {
275 let text = http_client::text(response).await?;
276 return Err(EmbeddingError::from_http_response(status, text));
277 }
278
279 let bytes: Vec<u8> = response.into_body().await?;
280
281 let api_resp: EmbeddingResponse = serde_json::from_slice(&bytes)?;
282
283 if api_resp.embeddings.len() != docs.len() {
284 return Err(EmbeddingError::ResponseError(
285 "Number of returned embeddings does not match input".into(),
286 ));
287 }
288 Ok(api_resp
289 .embeddings
290 .into_iter()
291 .zip(docs.into_iter())
292 .map(|(vec, document)| embeddings::Embedding { document, vec })
293 .collect())
294 }
295}
296
297pub const LLAMA3_2: &str = "llama3.2";
300pub const LLAVA: &str = "llava";
301pub const MISTRAL: &str = "mistral";
302
303#[derive(Debug, Serialize, Deserialize)]
304pub struct CompletionResponse {
305 pub model: String,
306 pub created_at: String,
307 pub message: Message,
308 pub done: bool,
309 #[serde(default)]
310 pub done_reason: Option<String>,
311 #[serde(default)]
312 pub total_duration: Option<u64>,
313 #[serde(default)]
314 pub load_duration: Option<u64>,
315 #[serde(default)]
316 pub prompt_eval_count: Option<u64>,
317 #[serde(default)]
318 pub prompt_eval_duration: Option<u64>,
319 #[serde(default)]
320 pub eval_count: Option<u64>,
321 #[serde(default)]
322 pub eval_duration: Option<u64>,
323}
324impl TryFrom<CompletionResponse> for completion::CompletionResponse<CompletionResponse> {
325 type Error = CompletionError;
326 fn try_from(resp: CompletionResponse) -> Result<Self, Self::Error> {
327 match resp.message {
328 Message::Assistant {
330 content,
331 thinking,
332 tool_calls,
333 ..
334 } => {
335 let mut assistant_contents = Vec::new();
336 let permits_omitted_think_start = resp.model.to_ascii_lowercase().contains("qwen3");
337 let (legacy_thinking, visible_content) =
338 if matches!(thinking.as_deref(), None | Some("")) {
339 split_legacy_thinking(&content, permits_omitted_think_start)
340 } else {
341 (None, content.as_str())
342 };
343 if let Some(thinking) = thinking.as_deref().filter(|t| !t.is_empty()) {
349 assistant_contents.push(completion::AssistantContent::reasoning(thinking));
350 }
351 if let Some(legacy_thinking) = legacy_thinking {
352 assistant_contents
353 .push(completion::AssistantContent::reasoning(legacy_thinking));
354 }
355 if !visible_content.is_empty() {
357 assistant_contents.push(completion::AssistantContent::text(visible_content));
358 }
359 for tc in tool_calls.iter() {
362 assistant_contents.push(completion::AssistantContent::tool_call(
363 tc.function.name.clone(),
364 tc.function.name.clone(),
365 tc.function.arguments.clone(),
366 ));
367 }
368 let choice = OneOrMany::many(assistant_contents).map_err(|_| {
369 CompletionError::ResponseError("No content provided".to_owned())
370 })?;
371 let prompt_tokens = resp.prompt_eval_count.unwrap_or(0);
372 let completion_tokens = resp.eval_count.unwrap_or(0);
373
374 let raw_response = CompletionResponse {
375 model: resp.model,
376 created_at: resp.created_at,
377 done: resp.done,
378 done_reason: resp.done_reason,
379 total_duration: resp.total_duration,
380 load_duration: resp.load_duration,
381 prompt_eval_count: resp.prompt_eval_count,
382 prompt_eval_duration: resp.prompt_eval_duration,
383 eval_count: resp.eval_count,
384 eval_duration: resp.eval_duration,
385 message: Message::Assistant {
386 content,
387 thinking,
388 images: None,
389 name: None,
390 tool_calls,
391 },
392 };
393
394 Ok(completion::CompletionResponse {
395 choice,
396 usage: Usage {
397 input_tokens: prompt_tokens,
398 output_tokens: completion_tokens,
399 total_tokens: prompt_tokens + completion_tokens,
400 cached_input_tokens: 0,
401 cache_creation_input_tokens: 0,
402 tool_use_prompt_tokens: 0,
403 reasoning_tokens: 0,
404 },
405 raw_response,
406 message_id: None,
407 })
408 }
409 _ => Err(CompletionError::ResponseError(
410 "Chat response does not include an assistant message".into(),
411 )),
412 }
413 }
414}
415
416fn split_legacy_thinking(content: &str, permits_omitted_start: bool) -> (Option<&str>, &str) {
421 let trimmed = content.trim_start();
422 let split = if let Some(reasoning_start) = trimmed.strip_prefix("<think>") {
423 reasoning_start.split_once("</think>")
424 } else if permits_omitted_start {
425 trimmed.split_once("\n</think>\n\n")
429 } else {
430 None
431 };
432 let Some((reasoning, visible)) = split else {
433 return (None, content);
434 };
435
436 let reasoning = reasoning.trim();
437 if reasoning.is_empty() {
438 return (None, visible.trim_start());
439 }
440
441 (Some(reasoning), visible.trim_start())
442}
443
444#[derive(Debug, Serialize, Deserialize)]
445pub(super) struct OllamaCompletionRequest {
446 model: String,
447 pub messages: Vec<Message>,
448 #[serde(skip_serializing_if = "Vec::is_empty")]
449 tools: Vec<ToolDefinition>,
450 pub stream: bool,
451 #[serde(skip_serializing_if = "Option::is_none")]
452 think: Option<Think>,
453 #[serde(skip_serializing_if = "Option::is_none")]
454 keep_alive: Option<String>,
455 #[serde(skip_serializing_if = "Option::is_none")]
456 format: Option<schemars::Schema>,
457 options: serde_json::Value,
458}
459
460impl TryFrom<(&str, CompletionRequest)> for OllamaCompletionRequest {
461 type Error = CompletionError;
462
463 fn try_from((model, req): (&str, CompletionRequest)) -> Result<Self, Self::Error> {
464 let chat_history = req.chat_history_with_documents();
465 let model = req.model.clone().unwrap_or_else(|| model.to_string());
466 if req.tool_choice.is_some() {
467 tracing::warn!("WARNING: `tool_choice` not supported for Ollama");
468 }
469 let mut partial_history = vec![];
471 partial_history.extend(chat_history);
472
473 let mut full_history: Vec<Message> = match &req.preamble {
475 Some(preamble) => vec![Message::system(preamble)],
476 None => vec![],
477 };
478
479 full_history.extend(
481 partial_history
482 .into_iter()
483 .map(message::Message::try_into)
484 .collect::<Result<Vec<Vec<Message>>, _>>()?
485 .into_iter()
486 .flatten()
487 .collect::<Vec<_>>(),
488 );
489
490 let mut think: Option<Think> = None;
491 let mut keep_alive: Option<String> = None;
492
493 let mut base_options = serde_json::Map::new();
497 if let Some(temperature) = req.temperature {
498 base_options.insert("temperature".to_string(), json!(temperature));
499 }
500 if let Some(max_tokens) = req.max_tokens {
501 base_options.insert("num_predict".to_string(), json!(max_tokens));
502 }
503 let base_options = Value::Object(base_options);
504
505 let options = if let Some(mut extra) = req.additional_params {
506 if let Some(obj) = extra.as_object_mut() {
508 if let Some(think_val) = obj.remove("think") {
510 think = Some(match think_val {
511 Value::Bool(think) => Think::Bool(think),
512 Value::String(think) => Think::Level(match think.to_lowercase().as_str() {
513 "low" => Level::Low,
514 "medium" => Level::Medium,
515 "high" => Level::High,
516 "max" => Level::Max,
517 _ => {
518 return Err(CompletionError::RequestError(
519 "`think` must be a 'low', 'medium', 'high', 'max' or bool"
520 .into(),
521 ));
522 }
523 }),
524 _ => {
525 return Err(CompletionError::RequestError(
526 "`think` must be a 'low', 'medium', 'high', 'max' or bool".into(),
527 ));
528 }
529 });
530 }
531
532 if let Some(keep_alive_val) = obj.remove("keep_alive") {
534 keep_alive = Some(
535 keep_alive_val
536 .as_str()
537 .ok_or_else(|| {
538 CompletionError::RequestError(
539 "`keep_alive` must be a string".into(),
540 )
541 })?
542 .to_string(),
543 );
544 }
545 }
546
547 json_utils::merge(base_options, extra)
548 } else {
549 base_options
550 };
551
552 Ok(Self {
553 model: model.to_string(),
554 messages: full_history,
555 stream: false,
556 think,
557 keep_alive,
558 format: req.output_schema,
559 tools: req
560 .tools
561 .clone()
562 .into_iter()
563 .map(ToolDefinition::from)
564 .collect::<Vec<_>>(),
565 options,
566 })
567 }
568}
569
570#[derive(Clone)]
571pub struct CompletionModel<T = reqwest::Client> {
572 client: Client<T>,
573 pub model: String,
574}
575
576impl<T> CompletionModel<T> {
577 pub fn new(client: Client<T>, model: &str) -> Self {
578 Self {
579 client,
580 model: model.to_owned(),
581 }
582 }
583}
584
585#[derive(Debug, Clone, Serialize, Deserialize)]
586#[serde(untagged)]
587enum Think {
588 Bool(bool),
589 Level(Level),
590}
591
592#[derive(Debug, Clone, Serialize, Deserialize)]
593#[serde(rename_all = "lowercase")]
594enum Level {
595 Low,
596 Medium,
597 High,
598 Max,
599}
600
601#[derive(Clone, Serialize, Deserialize, Debug)]
604pub struct StreamingCompletionResponse {
605 pub done_reason: Option<String>,
606 pub total_duration: Option<u64>,
607 pub load_duration: Option<u64>,
608 pub prompt_eval_count: Option<u64>,
609 pub prompt_eval_duration: Option<u64>,
610 pub eval_count: Option<u64>,
611 pub eval_duration: Option<u64>,
612}
613
614impl GetTokenUsage for StreamingCompletionResponse {
615 fn token_usage(&self) -> crate::completion::Usage {
616 let mut usage = crate::completion::Usage::new();
617 let input_tokens = self.prompt_eval_count.unwrap_or_default();
618 let output_tokens = self.eval_count.unwrap_or_default();
619 usage.input_tokens = input_tokens;
620 usage.output_tokens = output_tokens;
621 usage.total_tokens = input_tokens + output_tokens;
622
623 usage
624 }
625}
626
627#[derive(Default)]
633struct NdjsonBuffer {
634 buf: Vec<u8>,
635}
636
637impl NdjsonBuffer {
638 fn new() -> Self {
639 Self::default()
640 }
641
642 fn decode(&mut self, chunk: &[u8]) -> Vec<Vec<u8>> {
645 self.buf.extend_from_slice(chunk);
646
647 let mut lines = Vec::new();
648 while let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
649 let mut line: Vec<u8> = self.buf.drain(..=pos).collect();
650 line.pop();
651 if !line.is_empty() {
652 lines.push(line);
653 }
654 }
655 lines
656 }
657}
658
659impl<T> completion::CompletionModel for CompletionModel<T>
660where
661 T: HttpClientExt + Clone + Default + std::fmt::Debug + Send + 'static,
662{
663 type Response = CompletionResponse;
664 type StreamingResponse = StreamingCompletionResponse;
665
666 type Client = Client<T>;
667
668 fn make(client: &Self::Client, model: impl Into<String>) -> Self {
669 Self::new(client.clone(), model.into().as_str())
670 }
671
672 async fn completion(
673 &self,
674 completion_request: CompletionRequest,
675 ) -> Result<completion::CompletionResponse<Self::Response>, CompletionError> {
676 let system_instructions = completion_request.preamble.clone();
677 let record_telemetry_content = completion_request.record_telemetry_content;
678 let request = OllamaCompletionRequest::try_from((self.model.as_ref(), completion_request))?;
679 let span = CompletionSpanBuilder::new("ollama", &request.model, CompletionOperation::Chat)
680 .system_instructions(system_instructions.as_deref(), record_telemetry_content)
681 .build();
682
683 if tracing::enabled!(tracing::Level::TRACE) {
684 tracing::trace!(target: "rig::completions",
685 "Ollama completion request: {}",
686 serde_json::to_string_pretty(&request)?
687 );
688 }
689
690 let body = serde_json::to_vec(&request)?;
691
692 let req = self
693 .client
694 .post("api/chat")?
695 .body(body)
696 .map_err(http_client::Error::from)?;
697
698 let async_block = async move {
699 let response = self.client.send::<_, Bytes>(req).await?;
700 let status = response.status();
701 let response_body = response.into_body().into_future().await?.to_vec();
702
703 if !status.is_success() {
704 return Err(CompletionError::from_http_response(
705 status,
706 String::from_utf8_lossy(&response_body),
707 ));
708 }
709
710 let response: CompletionResponse = serde_json::from_slice(&response_body)?;
711 let span = tracing::Span::current();
712 span.record("gen_ai.response.model", &response.model);
713 span.record(
714 "gen_ai.usage.input_tokens",
715 response.prompt_eval_count.unwrap_or_default(),
716 );
717 span.record(
718 "gen_ai.usage.output_tokens",
719 response.eval_count.unwrap_or_default(),
720 );
721
722 if tracing::enabled!(tracing::Level::TRACE) {
723 tracing::trace!(target: "rig::completions",
724 "Ollama completion response: {}",
725 serde_json::to_string_pretty(&response)?
726 );
727 }
728
729 let response: completion::CompletionResponse<CompletionResponse> =
730 response.try_into()?;
731
732 Ok(response)
733 };
734
735 tracing::Instrument::instrument(async_block, span).await
736 }
737
738 async fn stream(
739 &self,
740 request: CompletionRequest,
741 ) -> Result<streaming::StreamingCompletionResponse<Self::StreamingResponse>, CompletionError>
742 {
743 let system_instructions = request.preamble.clone();
744 let record_telemetry_content = request.record_telemetry_content;
745 let mut request = OllamaCompletionRequest::try_from((self.model.as_ref(), request))?;
746 let span = CompletionSpanBuilder::new(
747 "ollama",
748 &request.model,
749 CompletionOperation::ChatStreaming,
750 )
751 .system_instructions(system_instructions.as_deref(), record_telemetry_content)
752 .build();
753 request.stream = true;
754
755 if tracing::enabled!(tracing::Level::TRACE) {
756 tracing::trace!(target: "rig::completions",
757 "Ollama streaming completion request: {}",
758 serde_json::to_string_pretty(&request)?
759 );
760 }
761
762 let body = serde_json::to_vec(&request)?;
763
764 let req = self
765 .client
766 .post("api/chat")?
767 .body(body)
768 .map_err(http_client::Error::from)?;
769
770 let response = self
771 .client
772 .send_streaming(req)
773 .instrument(span.clone())
774 .await?;
775 let status = response.status();
776 let mut byte_stream = response.into_body();
777
778 if !status.is_success() {
779 let mut body = Vec::new();
780 while let Some(chunk) = byte_stream.next().await {
781 match chunk {
782 Ok(bytes) => body.extend_from_slice(&bytes),
783 Err(e) => {
784 tracing::warn!(error = %e, "failed reading Ollama error-response body; preserving partial body");
785 break;
786 }
787 }
788 }
789 return Err(CompletionError::from_http_response(
790 status,
791 String::from_utf8_lossy(&body),
792 ));
793 }
794
795 let stream = try_stream! {
796 let span = tracing::Span::current();
797 let mut line_buf = NdjsonBuffer::new();
798
799 while let Some(chunk) = byte_stream.next().await {
800 let bytes = chunk.map_err(|e| http_client::Error::Instance(e.into()))?;
801
802 for line in line_buf.decode(&bytes) {
803 tracing::debug!(target: "rig", "Received NDJSON line from Ollama: {}", String::from_utf8_lossy(&line));
804
805 let response: CompletionResponse = serde_json::from_slice(&line)?;
806
807 if response.done {
808 span.record("gen_ai.response.model", &response.model);
809 }
810
811 if let Message::Assistant { content, thinking, tool_calls, .. } = response.message {
812 if let Some(thinking_content) = thinking && !thinking_content.is_empty() {
813 yield RawStreamingChoice::ReasoningDelta {
814 id: None,
815 reasoning: thinking_content,
816 };
817 }
818
819 if !content.is_empty() {
820 yield RawStreamingChoice::Message(content);
821 }
822
823 for tool_call in tool_calls {
824 yield RawStreamingChoice::ToolCall(
825 crate::streaming::RawStreamingToolCall::new(String::new(), tool_call.function.name, tool_call.function.arguments)
826 );
827 }
828 }
829
830 if response.done {
831 span.record("gen_ai.usage.input_tokens", response.prompt_eval_count);
832 span.record("gen_ai.usage.output_tokens", response.eval_count);
833 yield RawStreamingChoice::FinalResponse(
834 StreamingCompletionResponse {
835 total_duration: response.total_duration,
836 load_duration: response.load_duration,
837 prompt_eval_count: response.prompt_eval_count,
838 prompt_eval_duration: response.prompt_eval_duration,
839 eval_count: response.eval_count,
840 eval_duration: response.eval_duration,
841 done_reason: response.done_reason,
842 }
843 );
844 break;
845 }
846 }
847 }
848 }.instrument(span);
849
850 Ok(streaming::StreamingCompletionResponse::stream(Box::pin(
851 stream,
852 )))
853 }
854}
855
856#[derive(Debug, Deserialize)]
859struct ListModelsResponse {
860 models: Vec<ListModelEntry>,
861}
862
863#[derive(Debug, Deserialize)]
864struct ListModelEntry {
865 name: String,
866 model: String,
867}
868
869impl From<ListModelEntry> for Model {
870 fn from(value: ListModelEntry) -> Self {
871 Model::new(value.model, value.name)
872 }
873}
874
875#[derive(Clone)]
877pub struct OllamaModelLister<H = reqwest::Client> {
878 client: Client<H>,
879}
880
881impl<H> ModelLister<H> for OllamaModelLister<H>
882where
883 H: HttpClientExt + WasmCompatSend + WasmCompatSync + 'static,
884{
885 type Client = Client<H>;
886
887 fn new(client: Self::Client) -> Self {
888 Self { client }
889 }
890
891 async fn list_all(&self) -> Result<ModelList, ModelListingError> {
892 let path = "/api/tags";
893 let req = self.client.get(path)?.body(http_client::NoBody)?;
894 let response = self.client.send::<_, Vec<u8>>(req).await?;
895
896 if !response.status().is_success() {
897 let status_code = response.status().as_u16();
898 let body = response.into_body().await?;
899 return Err(ModelListingError::api_error_with_context(
900 "Ollama",
901 path,
902 status_code,
903 &body,
904 ));
905 }
906
907 let body = response.into_body().await?;
908 let api_resp: ListModelsResponse = serde_json::from_slice(&body).map_err(|error| {
909 ModelListingError::parse_error_with_context("Ollama", path, &error, &body)
910 })?;
911 let models = api_resp.models.into_iter().map(Model::from).collect();
912
913 Ok(ModelList::new(models))
914 }
915}
916
917#[derive(Clone, Debug, Deserialize, Serialize)]
921pub struct ToolDefinition {
922 #[serde(rename = "type")]
923 pub type_field: String, pub function: completion::ToolDefinition,
925}
926
927impl From<crate::completion::ToolDefinition> for ToolDefinition {
929 fn from(tool: crate::completion::ToolDefinition) -> Self {
930 ToolDefinition {
931 type_field: "function".to_owned(),
932 function: completion::ToolDefinition {
933 name: tool.name,
934 description: tool.description,
935 parameters: tool.parameters,
936 },
937 }
938 }
939}
940
941#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
942pub struct ToolCall {
943 #[serde(default, rename = "type")]
944 pub r#type: ToolType,
945 pub function: Function,
946}
947#[derive(Default, Debug, Serialize, Deserialize, PartialEq, Clone)]
948#[serde(rename_all = "lowercase")]
949pub enum ToolType {
950 #[default]
951 Function,
952}
953#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
954pub struct Function {
955 pub name: String,
956 pub arguments: Value,
957}
958
959#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
962#[serde(tag = "role", rename_all = "lowercase")]
963pub enum Message {
964 User {
965 content: String,
966 #[serde(skip_serializing_if = "Option::is_none")]
967 images: Option<Vec<String>>,
968 #[serde(skip_serializing_if = "Option::is_none")]
969 name: Option<String>,
970 },
971 Assistant {
972 #[serde(default)]
973 content: String,
974 #[serde(skip_serializing_if = "Option::is_none")]
975 thinking: Option<String>,
976 #[serde(skip_serializing_if = "Option::is_none")]
977 images: Option<Vec<String>>,
978 #[serde(skip_serializing_if = "Option::is_none")]
979 name: Option<String>,
980 #[serde(default, deserialize_with = "json_utils::null_or_vec")]
981 tool_calls: Vec<ToolCall>,
982 },
983 System {
984 content: String,
985 #[serde(skip_serializing_if = "Option::is_none")]
986 images: Option<Vec<String>>,
987 #[serde(skip_serializing_if = "Option::is_none")]
988 name: Option<String>,
989 },
990 #[serde(rename = "tool")]
991 ToolResult {
992 #[serde(rename = "tool_name")]
993 name: String,
994 content: String,
995 },
996}
997
998fn user_message_from_content(
1002 content: Vec<crate::message::UserContent>,
1003) -> Result<Message, crate::message::MessageError> {
1004 let mut texts = Vec::new();
1005 let mut images = Vec::new();
1006
1007 for content in content {
1008 match content {
1009 crate::message::UserContent::Text(crate::message::Text { text, .. }) => {
1010 texts.push(text);
1011 }
1012 crate::message::UserContent::Image(crate::message::Image {
1013 data: DocumentSourceKind::Base64(data),
1014 ..
1015 }) => images.push(data),
1016 crate::message::UserContent::Image(_) => {
1017 return Err(crate::message::MessageError::ConversionError(
1018 "Ollama images must be base64 encoded data".into(),
1019 ));
1020 }
1021 crate::message::UserContent::Document(crate::message::Document {
1022 data: DocumentSourceKind::Base64(data) | DocumentSourceKind::String(data),
1023 ..
1024 }) => texts.push(data),
1025 crate::message::UserContent::Document(_) => {
1026 return Err(crate::message::MessageError::ConversionError(
1027 "Ollama documents must be string or base64 encoded data".into(),
1028 ));
1029 }
1030 crate::message::UserContent::Audio(_) => {
1031 return Err(crate::message::MessageError::ConversionError(
1032 "Ollama does not support audio user content".into(),
1033 ));
1034 }
1035 crate::message::UserContent::Video(_) => {
1036 return Err(crate::message::MessageError::ConversionError(
1037 "Ollama does not support video user content".into(),
1038 ));
1039 }
1040 crate::message::UserContent::ToolResult(_) => {
1041 return Err(crate::message::MessageError::ConversionError(
1042 "tool results must be converted to a separate Ollama message".into(),
1043 ));
1044 }
1045 }
1046 }
1047
1048 Ok(Message::User {
1049 content: texts.join(" "),
1050 images: (!images.is_empty()).then_some(images),
1051 name: None,
1052 })
1053}
1054
1055impl TryFrom<crate::message::Message> for Vec<Message> {
1058 type Error = crate::message::MessageError;
1059 fn try_from(internal_msg: crate::message::Message) -> Result<Self, Self::Error> {
1060 use crate::message::Message as InternalMessage;
1061 match internal_msg {
1062 InternalMessage::System { content } => Ok(vec![Message::System {
1063 content,
1064 images: None,
1065 name: None,
1066 }]),
1067 InternalMessage::User { content, .. } => {
1068 let mut messages = Vec::new();
1069 let mut pending_user_content = Vec::new();
1070
1071 for content in content {
1072 match content {
1073 crate::message::UserContent::ToolResult(crate::message::ToolResult {
1074 id,
1075 content,
1076 ..
1077 }) => {
1078 if !pending_user_content.is_empty() {
1079 messages.push(user_message_from_content(std::mem::take(
1080 &mut pending_user_content,
1081 ))?);
1082 }
1083
1084 let content = content
1085 .into_iter()
1086 .map(|content| match content {
1087 crate::message::ToolResultContent::Text(text) => Ok(text.text),
1088 crate::message::ToolResultContent::Json { value } => {
1089 Ok(value.to_string())
1090 }
1091 crate::message::ToolResultContent::Image(_) => {
1092 Err(crate::message::MessageError::ConversionError(
1093 "Ollama does not support images in tool results".into(),
1094 ))
1095 }
1096 })
1097 .collect::<Result<Vec<_>, _>>()?
1098 .join("\n");
1099 messages.push(Message::ToolResult { name: id, content });
1100 }
1101 content => pending_user_content.push(content),
1102 }
1103 }
1104
1105 if !pending_user_content.is_empty() {
1106 messages.push(user_message_from_content(pending_user_content)?);
1107 }
1108
1109 Ok(messages)
1110 }
1111 InternalMessage::Assistant { content, .. } => {
1112 let mut thinking: Option<String> = None;
1113 let mut text_content = Vec::new();
1114 let mut tool_calls = Vec::new();
1115
1116 for content in content.into_iter() {
1117 match content {
1118 crate::message::AssistantContent::Text(text) => {
1119 text_content.push(text.text)
1120 }
1121 crate::message::AssistantContent::ToolCall(tool_call) => {
1122 tool_calls.push(tool_call)
1123 }
1124 crate::message::AssistantContent::Reasoning(reasoning) => {
1125 let display = reasoning.display_text();
1126 if !display.is_empty() {
1127 thinking = Some(display);
1128 }
1129 }
1130 crate::message::AssistantContent::Image(_) => {
1131 return Err(crate::message::MessageError::ConversionError(
1132 "Ollama currently doesn't support images.".into(),
1133 ));
1134 }
1135 }
1136 }
1137
1138 Ok(vec![Message::Assistant {
1141 content: text_content.join(" "),
1142 thinking,
1143 images: None,
1144 name: None,
1145 tool_calls: tool_calls
1146 .into_iter()
1147 .map(|tool_call| tool_call.into())
1148 .collect::<Vec<_>>(),
1149 }])
1150 }
1151 }
1152 }
1153}
1154
1155impl From<Message> for crate::completion::Message {
1158 fn from(msg: Message) -> Self {
1159 match msg {
1160 Message::User { content, .. } => crate::completion::Message::User {
1161 content: OneOrMany::one(crate::completion::message::UserContent::Text(Text::new(
1162 content,
1163 ))),
1164 },
1165 Message::Assistant {
1166 content,
1167 thinking,
1168 tool_calls,
1169 ..
1170 } => {
1171 let mut assistant_contents = Vec::new();
1172 if let Some(thinking) = thinking.filter(|t| !t.is_empty()) {
1174 assistant_contents.push(
1175 crate::completion::message::AssistantContent::reasoning(thinking),
1176 );
1177 }
1178 assistant_contents.push(crate::completion::message::AssistantContent::Text(
1179 Text::new(content),
1180 ));
1181 for tc in tool_calls {
1182 assistant_contents.push(
1183 crate::completion::message::AssistantContent::tool_call(
1184 tc.function.name.clone(),
1185 tc.function.name,
1186 tc.function.arguments,
1187 ),
1188 );
1189 }
1190 let content =
1191 OneOrMany::from_iter_optional(assistant_contents).unwrap_or_else(|| {
1192 OneOrMany::one(crate::completion::message::AssistantContent::Text(
1193 Text::new(String::new()),
1194 ))
1195 });
1196
1197 crate::completion::Message::Assistant { id: None, content }
1198 }
1199 Message::System { content, .. } => crate::completion::Message::User {
1201 content: OneOrMany::one(crate::completion::message::UserContent::Text(Text::new(
1202 content,
1203 ))),
1204 },
1205 Message::ToolResult { name, content } => crate::completion::Message::User {
1206 content: OneOrMany::one(message::UserContent::tool_result(
1207 name,
1208 OneOrMany::one(message::ToolResultContent::text(content)),
1209 )),
1210 },
1211 }
1212 }
1213}
1214
1215impl Message {
1216 pub fn system(content: &str) -> Self {
1218 Message::System {
1219 content: content.to_owned(),
1220 images: None,
1221 name: None,
1222 }
1223 }
1224}
1225
1226impl From<crate::message::ToolCall> for ToolCall {
1229 fn from(tool_call: crate::message::ToolCall) -> Self {
1230 Self {
1231 r#type: ToolType::Function,
1232 function: Function {
1233 name: tool_call.function.name,
1234 arguments: tool_call.function.arguments,
1235 },
1236 }
1237 }
1238}
1239
1240#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
1241pub struct SystemContent {
1242 #[serde(default)]
1243 r#type: SystemContentType,
1244 text: String,
1245}
1246
1247#[derive(Default, Debug, Serialize, Deserialize, PartialEq, Clone)]
1248#[serde(rename_all = "lowercase")]
1249pub enum SystemContentType {
1250 #[default]
1251 Text,
1252}
1253
1254impl From<String> for SystemContent {
1255 fn from(s: String) -> Self {
1256 SystemContent {
1257 r#type: SystemContentType::default(),
1258 text: s,
1259 }
1260 }
1261}
1262
1263impl FromStr for SystemContent {
1264 type Err = std::convert::Infallible;
1265 fn from_str(s: &str) -> Result<Self, Self::Err> {
1266 Ok(SystemContent {
1267 r#type: SystemContentType::default(),
1268 text: s.to_string(),
1269 })
1270 }
1271}
1272
1273#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
1274pub struct AssistantContent {
1275 pub text: String,
1276}
1277
1278impl FromStr for AssistantContent {
1279 type Err = std::convert::Infallible;
1280 fn from_str(s: &str) -> Result<Self, Self::Err> {
1281 Ok(AssistantContent { text: s.to_owned() })
1282 }
1283}
1284
1285#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
1286#[serde(tag = "type", rename_all = "lowercase")]
1287pub enum UserContent {
1288 Text { text: String },
1289 Image { image_url: ImageUrl },
1290 }
1292
1293impl FromStr for UserContent {
1294 type Err = std::convert::Infallible;
1295 fn from_str(s: &str) -> Result<Self, Self::Err> {
1296 Ok(UserContent::Text { text: s.to_owned() })
1297 }
1298}
1299
1300#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
1301pub struct ImageUrl {
1302 pub url: String,
1303 #[serde(default)]
1304 pub detail: ImageDetail,
1305}
1306
1307#[cfg(test)]
1312mod tests {
1313 use super::*;
1314 use serde_json::json;
1315
1316 #[test]
1317 fn splits_legacy_reasoning_with_or_without_opening_marker() {
1318 assert_eq!(
1319 split_legacy_thinking("<think>private reasoning</think>\n\nvisible answer", false),
1320 (Some("private reasoning"), "visible answer")
1321 );
1322 assert_eq!(
1323 split_legacy_thinking("private reasoning\n</think>\n\nvisible answer", true),
1324 (Some("private reasoning"), "visible answer")
1325 );
1326 }
1327
1328 #[test]
1329 fn leaves_unterminated_or_inline_reasoning_markers_visible() {
1330 assert_eq!(
1331 split_legacy_thinking("<think>unterminated", true),
1332 (None, "<think>unterminated")
1333 );
1334 assert_eq!(
1335 split_legacy_thinking("The literal marker is <think>.", true),
1336 (None, "The literal marker is <think>.")
1337 );
1338 assert_eq!(
1339 split_legacy_thinking(" visible indentation", true),
1340 (None, " visible indentation")
1341 );
1342 assert_eq!(
1343 split_legacy_thinking("The closing token </think> is XML-like.", true),
1344 (None, "The closing token </think> is XML-like.")
1345 );
1346 assert_eq!(
1347 split_legacy_thinking("Example:\n</think>\nis a closing tag.", true),
1348 (None, "Example:\n</think>\nis a closing tag.")
1349 );
1350 }
1351
1352 #[tokio::test]
1354 async fn test_chat_completion() {
1355 let sample_chat_response = json!({
1357 "model": "llama3.2",
1358 "created_at": "2023-08-04T19:22:45.499127Z",
1359 "message": {
1360 "role": "assistant",
1361 "content": "The sky is blue because of Rayleigh scattering.",
1362 "images": null,
1363 "tool_calls": [
1364 {
1365 "type": "function",
1366 "function": {
1367 "name": "get_current_weather",
1368 "arguments": {
1369 "location": "San Francisco, CA",
1370 "format": "celsius"
1371 }
1372 }
1373 }
1374 ]
1375 },
1376 "done": true,
1377 "total_duration": 8000000000u64,
1378 "load_duration": 6000000u64,
1379 "prompt_eval_count": 61u64,
1380 "prompt_eval_duration": 400000000u64,
1381 "eval_count": 468u64,
1382 "eval_duration": 7700000000u64
1383 });
1384 let sample_text = sample_chat_response.to_string();
1385
1386 let chat_resp: CompletionResponse =
1387 serde_json::from_str(&sample_text).expect("Invalid JSON structure");
1388 let conv: completion::CompletionResponse<CompletionResponse> =
1389 chat_resp.try_into().unwrap();
1390 assert!(
1391 !conv.choice.is_empty(),
1392 "Expected non-empty choice in chat response"
1393 );
1394 }
1395
1396 #[test]
1398 fn test_message_conversion() {
1399 let provider_msg = Message::User {
1401 content: "Test message".to_owned(),
1402 images: None,
1403 name: None,
1404 };
1405 let comp_msg: crate::completion::Message = provider_msg.into();
1407 match comp_msg {
1408 crate::completion::Message::User { content } => {
1409 let first_content = content.first();
1411 match first_content {
1413 crate::completion::message::UserContent::Text(text_struct) => {
1414 assert_eq!(text_struct.text, "Test message");
1415 }
1416 _ => panic!("Expected text content in conversion"),
1417 }
1418 }
1419 _ => panic!("Conversion from provider Message to completion Message failed"),
1420 }
1421 }
1422
1423 #[test]
1424 fn mixed_user_content_preserves_message_order() {
1425 use crate::OneOrMany;
1426 use crate::message::{Message as RigMessage, ToolResultContent, UserContent};
1427
1428 let message = RigMessage::User {
1429 content: OneOrMany::many(vec![
1430 UserContent::text("before"),
1431 UserContent::tool_result(
1432 "lookup",
1433 OneOrMany::one(ToolResultContent::json(json!({ "ok": true }))),
1434 ),
1435 UserContent::text("after"),
1436 ])
1437 .expect("mixed content is non-empty"),
1438 };
1439
1440 let messages = Vec::<Message>::try_from(message).expect("mixed content should convert");
1441 assert_eq!(messages.len(), 3);
1442 assert!(matches!(
1443 &messages[0],
1444 Message::User { content, .. } if content == "before"
1445 ));
1446 assert!(matches!(
1447 &messages[1],
1448 Message::ToolResult { name, content }
1449 if name == "lookup" && content == r#"{"ok":true}"#
1450 ));
1451 assert!(matches!(
1452 &messages[2],
1453 Message::User { content, .. } if content == "after"
1454 ));
1455 }
1456
1457 #[test]
1458 fn unsupported_user_content_returns_a_conversion_error() {
1459 use crate::OneOrMany;
1460 use crate::message::{ImageMediaType, Message as RigMessage, UserContent};
1461
1462 let message = RigMessage::User {
1463 content: OneOrMany::one(UserContent::image_url(
1464 "https://example.com/image.png",
1465 Some(ImageMediaType::PNG),
1466 None,
1467 )),
1468 };
1469
1470 let error = Vec::<Message>::try_from(message).expect_err("URL image should be rejected");
1471 assert!(error.to_string().contains("base64"));
1472 }
1473
1474 #[test]
1476 fn test_tool_definition_conversion() {
1477 let internal_tool = crate::completion::ToolDefinition {
1479 name: "get_current_weather".to_owned(),
1480 description: "Get the current weather for a location".to_owned(),
1481 parameters: json!({
1482 "type": "object",
1483 "properties": {
1484 "location": {
1485 "type": "string",
1486 "description": "The location to get the weather for, e.g. San Francisco, CA"
1487 },
1488 "format": {
1489 "type": "string",
1490 "description": "The format to return the weather in, e.g. 'celsius' or 'fahrenheit'",
1491 "enum": ["celsius", "fahrenheit"]
1492 }
1493 },
1494 "required": ["location", "format"]
1495 }),
1496 };
1497 let ollama_tool: ToolDefinition = internal_tool.into();
1499 assert_eq!(ollama_tool.type_field, "function");
1500 assert_eq!(ollama_tool.function.name, "get_current_weather");
1501 assert_eq!(
1502 ollama_tool.function.description,
1503 "Get the current weather for a location"
1504 );
1505 let params = &ollama_tool.function.parameters;
1507 assert_eq!(params["properties"]["location"]["type"], "string");
1508 }
1509
1510 #[tokio::test]
1512 async fn test_chat_completion_with_thinking() {
1513 let sample_response = json!({
1514 "model": "qwen-thinking",
1515 "created_at": "2023-08-04T19:22:45.499127Z",
1516 "message": {
1517 "role": "assistant",
1518 "content": "The answer is 42.",
1519 "thinking": "Let me think about this carefully. The question asks for the meaning of life...",
1520 "images": null,
1521 "tool_calls": []
1522 },
1523 "done": true,
1524 "total_duration": 8000000000u64,
1525 "load_duration": 6000000u64,
1526 "prompt_eval_count": 61u64,
1527 "prompt_eval_duration": 400000000u64,
1528 "eval_count": 468u64,
1529 "eval_duration": 7700000000u64
1530 });
1531
1532 let chat_resp: CompletionResponse =
1533 serde_json::from_value(sample_response).expect("Failed to deserialize");
1534
1535 if let Message::Assistant {
1537 thinking, content, ..
1538 } = &chat_resp.message
1539 {
1540 assert_eq!(
1541 thinking.as_ref().unwrap(),
1542 "Let me think about this carefully. The question asks for the meaning of life..."
1543 );
1544 assert_eq!(content, "The answer is 42.");
1545 } else {
1546 panic!("Expected Assistant message");
1547 }
1548 }
1549
1550 #[tokio::test]
1552 async fn test_chat_completion_without_thinking() {
1553 let sample_response = json!({
1554 "model": "llama3.2",
1555 "created_at": "2023-08-04T19:22:45.499127Z",
1556 "message": {
1557 "role": "assistant",
1558 "content": "Hello!",
1559 "images": null,
1560 "tool_calls": []
1561 },
1562 "done": true,
1563 "total_duration": 8000000000u64,
1564 "load_duration": 6000000u64,
1565 "prompt_eval_count": 10u64,
1566 "prompt_eval_duration": 400000000u64,
1567 "eval_count": 5u64,
1568 "eval_duration": 7700000000u64
1569 });
1570
1571 let chat_resp: CompletionResponse =
1572 serde_json::from_value(sample_response).expect("Failed to deserialize");
1573
1574 if let Message::Assistant {
1576 thinking, content, ..
1577 } = &chat_resp.message
1578 {
1579 assert!(thinking.is_none());
1580 assert_eq!(content, "Hello!");
1581 } else {
1582 panic!("Expected Assistant message");
1583 }
1584 }
1585
1586 #[test]
1588 fn test_streaming_response_with_thinking() {
1589 let sample_chunk = json!({
1590 "model": "qwen-thinking",
1591 "created_at": "2023-08-04T19:22:45.499127Z",
1592 "message": {
1593 "role": "assistant",
1594 "content": "",
1595 "thinking": "Analyzing the problem...",
1596 "images": null,
1597 "tool_calls": []
1598 },
1599 "done": false
1600 });
1601
1602 let chunk: CompletionResponse =
1603 serde_json::from_value(sample_chunk).expect("Failed to deserialize");
1604
1605 if let Message::Assistant {
1606 thinking, content, ..
1607 } = &chunk.message
1608 {
1609 assert_eq!(thinking.as_ref().unwrap(), "Analyzing the problem...");
1610 assert_eq!(content, "");
1611 } else {
1612 panic!("Expected Assistant message");
1613 }
1614 }
1615
1616 #[test]
1618 fn test_message_conversion_with_thinking() {
1619 let reasoning_content = crate::message::Reasoning::new("Step 1: Consider the problem");
1621
1622 let internal_msg = crate::message::Message::Assistant {
1623 id: None,
1624 content: crate::OneOrMany::many(vec![
1625 crate::message::AssistantContent::Reasoning(reasoning_content),
1626 crate::message::AssistantContent::Text(crate::message::Text::new(
1627 "The answer is X".to_string(),
1628 )),
1629 ])
1630 .unwrap(),
1631 };
1632
1633 let provider_msgs: Vec<Message> = internal_msg.try_into().unwrap();
1635 assert_eq!(provider_msgs.len(), 1);
1636
1637 if let Message::Assistant {
1638 thinking, content, ..
1639 } = &provider_msgs[0]
1640 {
1641 assert_eq!(thinking.as_ref().unwrap(), "Step 1: Consider the problem");
1642 assert_eq!(content, "The answer is X");
1643 } else {
1644 panic!("Expected Assistant message with thinking");
1645 }
1646 }
1647
1648 #[tokio::test]
1655 async fn nonstreaming_response_preserves_thinking_as_reasoning() {
1656 let sample_response = json!({
1657 "model": "qwen3:4b",
1658 "created_at": "2023-08-04T19:22:45.499127Z",
1659 "message": {
1660 "role": "assistant",
1661 "content": "",
1662 "thinking": "The user asked for the weather in Berlin. I should call get_weather with location=Berlin.",
1663 "images": null,
1664 "tool_calls": [
1665 { "type": "function", "function": { "name": "get_weather", "arguments": { "location": "Berlin" } } }
1666 ]
1667 },
1668 "done": true,
1669 "done_reason": "stop",
1670 "total_duration": 8000000000u64,
1671 "load_duration": 6000000u64,
1672 "prompt_eval_count": 61u64,
1673 "prompt_eval_duration": 400000000u64,
1674 "eval_count": 468u64,
1675 "eval_duration": 7700000000u64
1676 });
1677
1678 let raw: CompletionResponse =
1679 serde_json::from_value(sample_response).expect("deserialize ollama response");
1680 let completed: completion::CompletionResponse<CompletionResponse> =
1681 raw.try_into().expect("convert to completion response");
1682
1683 let reasoning = completed.choice.iter().find_map(|c| match c {
1684 completion::AssistantContent::Reasoning(r) => Some(r.clone()),
1685 _ => None,
1686 });
1687 let has_tool_call = completed
1688 .choice
1689 .iter()
1690 .any(|c| matches!(c, completion::AssistantContent::ToolCall(_)));
1691
1692 assert!(has_tool_call, "tool call should survive the conversion");
1693 let reasoning = reasoning.expect(
1694 "non-streaming response must surface `thinking` as AssistantContent::Reasoning (issue #1926)",
1695 );
1696 assert_eq!(
1697 reasoning.display_text(),
1698 "The user asked for the weather in Berlin. I should call get_weather with location=Berlin.",
1699 );
1700 }
1701
1702 #[test]
1704 fn test_empty_thinking_content() {
1705 let sample_response = json!({
1706 "model": "llama3.2",
1707 "created_at": "2023-08-04T19:22:45.499127Z",
1708 "message": {
1709 "role": "assistant",
1710 "content": "Response",
1711 "thinking": "",
1712 "images": null,
1713 "tool_calls": []
1714 },
1715 "done": true,
1716 "total_duration": 8000000000u64,
1717 "load_duration": 6000000u64,
1718 "prompt_eval_count": 10u64,
1719 "prompt_eval_duration": 400000000u64,
1720 "eval_count": 5u64,
1721 "eval_duration": 7700000000u64
1722 });
1723
1724 let chat_resp: CompletionResponse =
1725 serde_json::from_value(sample_response).expect("Failed to deserialize");
1726
1727 if let Message::Assistant {
1728 thinking, content, ..
1729 } = &chat_resp.message
1730 {
1731 assert_eq!(thinking.as_ref().unwrap(), "");
1733 assert_eq!(content, "Response");
1734 } else {
1735 panic!("Expected Assistant message");
1736 }
1737 }
1738
1739 #[test]
1741 fn test_thinking_with_tool_calls() {
1742 let sample_response = json!({
1743 "model": "qwen-thinking",
1744 "created_at": "2023-08-04T19:22:45.499127Z",
1745 "message": {
1746 "role": "assistant",
1747 "content": "Let me check the weather.",
1748 "thinking": "User wants weather info, I should use the weather tool",
1749 "images": null,
1750 "tool_calls": [
1751 {
1752 "type": "function",
1753 "function": {
1754 "name": "get_weather",
1755 "arguments": {
1756 "location": "San Francisco"
1757 }
1758 }
1759 }
1760 ]
1761 },
1762 "done": true,
1763 "total_duration": 8000000000u64,
1764 "load_duration": 6000000u64,
1765 "prompt_eval_count": 30u64,
1766 "prompt_eval_duration": 400000000u64,
1767 "eval_count": 50u64,
1768 "eval_duration": 7700000000u64
1769 });
1770
1771 let chat_resp: CompletionResponse =
1772 serde_json::from_value(sample_response).expect("Failed to deserialize");
1773
1774 if let Message::Assistant {
1775 thinking,
1776 content,
1777 tool_calls,
1778 ..
1779 } = &chat_resp.message
1780 {
1781 assert_eq!(
1782 thinking.as_ref().unwrap(),
1783 "User wants weather info, I should use the weather tool"
1784 );
1785 assert_eq!(content, "Let me check the weather.");
1786 assert_eq!(tool_calls.len(), 1);
1787 assert_eq!(tool_calls[0].function.name, "get_weather");
1788 } else {
1789 panic!("Expected Assistant message with thinking and tool calls");
1790 }
1791 }
1792
1793 #[test]
1795 fn test_completion_request_with_think_param() {
1796 use crate::OneOrMany;
1797 use crate::completion::Message as CompletionMessage;
1798 use crate::message::{Text, UserContent};
1799
1800 let completion_request = CompletionRequest {
1802 model: None,
1803 preamble: Some("You are a helpful assistant.".to_string()),
1804 chat_history: OneOrMany::one(CompletionMessage::User {
1805 content: OneOrMany::one(UserContent::Text(Text::new("What is 2 + 2?".to_string()))),
1806 }),
1807 documents: vec![],
1808 tools: vec![],
1809 temperature: Some(0.7),
1810 max_tokens: Some(1024),
1811 tool_choice: None,
1812 additional_params: Some(json!({
1813 "think": true,
1814 "keep_alive": "-1m",
1815 "num_ctx": 4096
1816 })),
1817 output_schema: None,
1818 record_telemetry_content: false,
1819 };
1820
1821 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
1823 .expect("Failed to create Ollama request");
1824
1825 let serialized =
1827 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
1828
1829 let expected = json!({
1835 "model": "qwen3:8b",
1836 "messages": [
1837 {
1838 "role": "system",
1839 "content": "You are a helpful assistant."
1840 },
1841 {
1842 "role": "user",
1843 "content": "What is 2 + 2?"
1844 }
1845 ],
1846 "stream": false,
1847 "think": true,
1848 "keep_alive": "-1m",
1849 "options": {
1850 "temperature": 0.7,
1851 "num_predict": 1024,
1852 "num_ctx": 4096
1853 }
1854 });
1855
1856 assert_eq!(serialized, expected);
1857 }
1858
1859 #[test]
1861 fn test_completion_request_with_level_low_think_param() {
1862 use crate::OneOrMany;
1863 use crate::completion::Message as CompletionMessage;
1864 use crate::message::{Text, UserContent};
1865
1866 let completion_request = CompletionRequest {
1868 model: None,
1869 preamble: Some("You are a helpful assistant.".to_string()),
1870 chat_history: OneOrMany::one(CompletionMessage::User {
1871 content: OneOrMany::one(UserContent::Text(Text::new("What is 2 + 2?".to_string()))),
1872 }),
1873 documents: vec![],
1874 tools: vec![],
1875 temperature: Some(0.7),
1876 max_tokens: Some(1024),
1877 tool_choice: None,
1878 additional_params: Some(json!({
1879 "think": "low",
1880 "keep_alive": "-1m",
1881 "num_ctx": 4096
1882 })),
1883 output_schema: None,
1884 record_telemetry_content: false,
1885 };
1886
1887 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
1889 .expect("Failed to create Ollama request");
1890
1891 let serialized =
1893 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
1894
1895 let expected = json!({
1901 "model": "qwen3:8b",
1902 "messages": [
1903 {
1904 "role": "system",
1905 "content": "You are a helpful assistant."
1906 },
1907 {
1908 "role": "user",
1909 "content": "What is 2 + 2?"
1910 }
1911 ],
1912 "stream": false,
1913 "think": "low",
1914 "keep_alive": "-1m",
1915 "options": {
1916 "temperature": 0.7,
1917 "num_predict": 1024,
1918 "num_ctx": 4096
1919 }
1920 });
1921
1922 assert_eq!(serialized, expected);
1923 }
1924
1925 #[test]
1927 fn test_completion_request_with_level_medium_think_param() {
1928 use crate::OneOrMany;
1929 use crate::completion::Message as CompletionMessage;
1930 use crate::message::{Text, UserContent};
1931
1932 let completion_request = CompletionRequest {
1934 model: None,
1935 preamble: Some("You are a helpful assistant.".to_string()),
1936 chat_history: OneOrMany::one(CompletionMessage::User {
1937 content: OneOrMany::one(UserContent::Text(Text::new("What is 2 + 2?".to_string()))),
1938 }),
1939 documents: vec![],
1940 tools: vec![],
1941 temperature: Some(0.7),
1942 max_tokens: Some(1024),
1943 tool_choice: None,
1944 additional_params: Some(json!({
1945 "think": "medium",
1946 "keep_alive": "-1m",
1947 "num_ctx": 4096
1948 })),
1949 output_schema: None,
1950 record_telemetry_content: false,
1951 };
1952
1953 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
1955 .expect("Failed to create Ollama request");
1956
1957 let serialized =
1959 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
1960
1961 let expected = json!({
1967 "model": "qwen3:8b",
1968 "messages": [
1969 {
1970 "role": "system",
1971 "content": "You are a helpful assistant."
1972 },
1973 {
1974 "role": "user",
1975 "content": "What is 2 + 2?"
1976 }
1977 ],
1978 "stream": false,
1979 "think": "medium",
1980 "keep_alive": "-1m",
1981 "options": {
1982 "temperature": 0.7,
1983 "num_predict": 1024,
1984 "num_ctx": 4096
1985 }
1986 });
1987
1988 assert_eq!(serialized, expected);
1989 }
1990
1991 #[test]
1993 fn test_completion_request_with_level_high_think_param() {
1994 use crate::OneOrMany;
1995 use crate::completion::Message as CompletionMessage;
1996 use crate::message::{Text, UserContent};
1997
1998 let completion_request = CompletionRequest {
2000 model: None,
2001 preamble: Some("You are a helpful assistant.".to_string()),
2002 chat_history: OneOrMany::one(CompletionMessage::User {
2003 content: OneOrMany::one(UserContent::Text(Text::new("What is 2 + 2?".to_string()))),
2004 }),
2005 documents: vec![],
2006 tools: vec![],
2007 temperature: Some(0.7),
2008 max_tokens: Some(1024),
2009 tool_choice: None,
2010 additional_params: Some(json!({
2011 "think": "high",
2012 "keep_alive": "-1m",
2013 "num_ctx": 4096
2014 })),
2015 output_schema: None,
2016 record_telemetry_content: false,
2017 };
2018
2019 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
2021 .expect("Failed to create Ollama request");
2022
2023 let serialized =
2025 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2026
2027 let expected = json!({
2033 "model": "qwen3:8b",
2034 "messages": [
2035 {
2036 "role": "system",
2037 "content": "You are a helpful assistant."
2038 },
2039 {
2040 "role": "user",
2041 "content": "What is 2 + 2?"
2042 }
2043 ],
2044 "stream": false,
2045 "think": "high",
2046 "keep_alive": "-1m",
2047 "options": {
2048 "temperature": 0.7,
2049 "num_predict": 1024,
2050 "num_ctx": 4096
2051 }
2052 });
2053
2054 assert_eq!(serialized, expected);
2055 }
2056
2057 #[test]
2059 fn test_completion_request_with_level_invalid_think_param() {
2060 use crate::OneOrMany;
2061 use crate::completion::Message as CompletionMessage;
2062 use crate::message::{Text, UserContent};
2063
2064 let completion_request = CompletionRequest {
2066 model: None,
2067 preamble: Some("You are a helpful assistant.".to_string()),
2068 chat_history: OneOrMany::one(CompletionMessage::User {
2069 content: OneOrMany::one(UserContent::Text(Text::new("What is 2 + 2?".to_string()))),
2070 }),
2071 documents: vec![],
2072 tools: vec![],
2073 temperature: Some(0.7),
2074 max_tokens: Some(1024),
2075 tool_choice: None,
2076 additional_params: Some(json!({
2077 "think": "invalid",
2078 "keep_alive": "-1m",
2079 "num_ctx": 4096
2080 })),
2081 output_schema: None,
2082 record_telemetry_content: false,
2083 };
2084
2085 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request));
2087
2088 assert!(ollama_request.is_err())
2089 }
2090
2091 #[test]
2094 fn test_completion_request_with_think_omitted_by_default() {
2095 use crate::OneOrMany;
2096 use crate::completion::Message as CompletionMessage;
2097 use crate::message::{Text, UserContent};
2098
2099 let completion_request = CompletionRequest {
2101 model: None,
2102 preamble: Some("You are a helpful assistant.".to_string()),
2103 chat_history: OneOrMany::one(CompletionMessage::User {
2104 content: OneOrMany::one(UserContent::Text(Text::new("Hello!".to_string()))),
2105 }),
2106 documents: vec![],
2107 tools: vec![],
2108 temperature: Some(0.5),
2109 max_tokens: None,
2110 tool_choice: None,
2111 additional_params: None,
2112 output_schema: None,
2113 record_telemetry_content: false,
2114 };
2115
2116 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2118 .expect("Failed to create Ollama request");
2119
2120 let serialized =
2122 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2123
2124 let expected = json!({
2127 "model": "llama3.2",
2128 "messages": [
2129 {
2130 "role": "system",
2131 "content": "You are a helpful assistant."
2132 },
2133 {
2134 "role": "user",
2135 "content": "Hello!"
2136 }
2137 ],
2138 "stream": false,
2139 "options": {
2140 "temperature": 0.5
2141 }
2142 });
2143
2144 assert_eq!(serialized, expected);
2145 }
2146
2147 #[test]
2151 fn test_completion_request_num_predict_from_additional_params_wins() {
2152 use crate::OneOrMany;
2153 use crate::completion::Message as CompletionMessage;
2154 use crate::message::{Text, UserContent};
2155
2156 let completion_request = CompletionRequest {
2157 model: None,
2158 preamble: None,
2159 chat_history: OneOrMany::one(CompletionMessage::User {
2160 content: OneOrMany::one(UserContent::Text(Text::new("Hello!".to_string()))),
2161 }),
2162 documents: vec![],
2163 tools: vec![],
2164 temperature: None,
2165 max_tokens: Some(1024),
2166 tool_choice: None,
2167 additional_params: Some(json!({ "num_predict": 42 })),
2168 output_schema: None,
2169 record_telemetry_content: false,
2170 };
2171
2172 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2173 .expect("Failed to create Ollama request");
2174 let serialized =
2175 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2176
2177 assert_eq!(serialized["options"], json!({ "num_predict": 42 }));
2178 assert_eq!(serialized.get("max_tokens"), None);
2179 }
2180
2181 #[test]
2186 fn test_completion_request_num_predict_without_additional_params() {
2187 use crate::OneOrMany;
2188 use crate::completion::Message as CompletionMessage;
2189 use crate::message::{Text, UserContent};
2190
2191 let completion_request = CompletionRequest {
2192 model: None,
2193 preamble: None,
2194 chat_history: OneOrMany::one(CompletionMessage::User {
2195 content: OneOrMany::one(UserContent::Text(Text::new("Hello!".to_string()))),
2196 }),
2197 documents: vec![],
2198 tools: vec![],
2199 temperature: Some(0.7),
2200 max_tokens: Some(1024),
2201 tool_choice: None,
2202 additional_params: None,
2203 output_schema: None,
2204 record_telemetry_content: false,
2205 };
2206
2207 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2208 .expect("Failed to create Ollama request");
2209 let serialized =
2210 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2211
2212 assert_eq!(
2213 serialized["options"],
2214 json!({ "temperature": 0.7, "num_predict": 1024 })
2215 );
2216 assert_eq!(serialized.get("max_tokens"), None);
2218 assert_eq!(serialized.get("temperature"), None);
2219 }
2220
2221 #[test]
2225 fn test_completion_request_options_omit_unset_parameters() {
2226 use crate::OneOrMany;
2227 use crate::completion::Message as CompletionMessage;
2228 use crate::message::{Text, UserContent};
2229
2230 let completion_request = CompletionRequest {
2231 model: None,
2232 preamble: None,
2233 chat_history: OneOrMany::one(CompletionMessage::User {
2234 content: OneOrMany::one(UserContent::Text(Text::new("Hello!".to_string()))),
2235 }),
2236 documents: vec![],
2237 tools: vec![],
2238 temperature: None,
2239 max_tokens: None,
2240 tool_choice: None,
2241 additional_params: None,
2242 output_schema: None,
2243 record_telemetry_content: false,
2244 };
2245
2246 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2247 .expect("Failed to create Ollama request");
2248 let serialized =
2249 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2250
2251 assert_eq!(serialized["options"], json!({}));
2252 }
2253
2254 #[test]
2255 fn test_completion_request_with_output_schema() {
2256 use crate::OneOrMany;
2257 use crate::completion::Message as CompletionMessage;
2258 use crate::message::{Text, UserContent};
2259
2260 let schema: schemars::Schema = serde_json::from_value(json!({
2261 "type": "object",
2262 "properties": {
2263 "age": { "type": "integer" },
2264 "available": { "type": "boolean" }
2265 },
2266 "required": ["age", "available"]
2267 }))
2268 .expect("Failed to parse schema");
2269
2270 let completion_request = CompletionRequest {
2271 model: Some("llama3.1".to_string()),
2272 preamble: None,
2273 chat_history: OneOrMany::one(CompletionMessage::User {
2274 content: OneOrMany::one(UserContent::Text(Text::new(
2275 "How old is Ollama?".to_string(),
2276 ))),
2277 }),
2278 documents: vec![],
2279 tools: vec![],
2280 temperature: None,
2281 max_tokens: None,
2282 tool_choice: None,
2283 additional_params: None,
2284 output_schema: Some(schema),
2285 record_telemetry_content: false,
2286 };
2287
2288 let ollama_request = OllamaCompletionRequest::try_from(("llama3.1", completion_request))
2289 .expect("Failed to create Ollama request");
2290
2291 let serialized =
2292 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2293
2294 let format = serialized
2295 .get("format")
2296 .expect("format field should be present");
2297 assert_eq!(
2298 *format,
2299 json!({
2300 "type": "object",
2301 "properties": {
2302 "age": { "type": "integer" },
2303 "available": { "type": "boolean" }
2304 },
2305 "required": ["age", "available"]
2306 })
2307 );
2308 }
2309
2310 #[test]
2311 fn test_completion_request_without_output_schema() {
2312 use crate::OneOrMany;
2313 use crate::completion::Message as CompletionMessage;
2314 use crate::message::{Text, UserContent};
2315
2316 let completion_request = CompletionRequest {
2317 model: Some("llama3.1".to_string()),
2318 preamble: None,
2319 chat_history: OneOrMany::one(CompletionMessage::User {
2320 content: OneOrMany::one(UserContent::Text(Text::new("Hello!".to_string()))),
2321 }),
2322 documents: vec![],
2323 tools: vec![],
2324 temperature: None,
2325 max_tokens: None,
2326 tool_choice: None,
2327 additional_params: None,
2328 output_schema: None,
2329 record_telemetry_content: false,
2330 };
2331
2332 let ollama_request = OllamaCompletionRequest::try_from(("llama3.1", completion_request))
2333 .expect("Failed to create Ollama request");
2334
2335 let serialized =
2336 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2337
2338 assert!(
2339 serialized.get("format").is_none(),
2340 "format field should be absent when output_schema is None"
2341 );
2342 }
2343
2344 #[test]
2345 fn test_client_initialization() {
2346 let _client = crate::providers::ollama::Client::new(Nothing).expect("Client::new() failed");
2347 let _client_from_builder = crate::providers::ollama::Client::builder()
2348 .api_key(Nothing)
2349 .build()
2350 .expect("Client::builder() failed");
2351 }
2352
2353 #[test]
2354 fn ndjson_buffer_returns_complete_lines_in_single_chunk() {
2355 let mut buf = NdjsonBuffer::new();
2356 let lines = buf.decode(b"{\"a\":1}\n{\"b\":2}\n");
2357 assert_eq!(lines, vec![b"{\"a\":1}".to_vec(), b"{\"b\":2}".to_vec()]);
2358 }
2359
2360 #[test]
2361 fn ndjson_buffer_reassembles_line_split_across_chunks() {
2362 let mut buf = NdjsonBuffer::new();
2363
2364 assert!(buf.decode(b"{\"model\":\"llama\",\"mes").is_empty());
2365
2366 let lines = buf.decode(b"sage\":\"hi\"}\n{\"done\"");
2367 assert_eq!(
2368 lines,
2369 vec![b"{\"model\":\"llama\",\"message\":\"hi\"}".to_vec()]
2370 );
2371
2372 let lines = buf.decode(b":true}\n");
2373 assert_eq!(lines, vec![b"{\"done\":true}".to_vec()]);
2374 }
2375
2376 #[test]
2377 fn ndjson_buffer_skips_blank_lines() {
2378 let mut buf = NdjsonBuffer::new();
2379 let lines = buf.decode(b"\n{\"a\":1}\n\n");
2380 assert_eq!(lines, vec![b"{\"a\":1}".to_vec()]);
2381 }
2382
2383 #[test]
2384 fn ndjson_buffer_retains_unterminated_trailing_data() {
2385 let mut buf = NdjsonBuffer::new();
2386 let lines = buf.decode(b"{\"a\":1}\n{\"b\":2");
2387 assert_eq!(lines, vec![b"{\"a\":1}".to_vec()]);
2388 let lines = buf.decode(b"}\n");
2389 assert_eq!(lines, vec![b"{\"b\":2}".to_vec()]);
2390 }
2391
2392 #[test]
2393 fn ndjson_buffer_handles_empty_chunk() {
2394 let mut buf = NdjsonBuffer::new();
2395 assert!(buf.decode(b"").is_empty());
2396
2397 buf.decode(b"{\"a\":1");
2398 assert!(buf.decode(b"").is_empty());
2399
2400 let lines = buf.decode(b"}\n");
2401 assert_eq!(lines, vec![b"{\"a\":1}".to_vec()]);
2402 }
2403
2404 #[test]
2405 fn ndjson_buffer_handles_multi_byte_utf8_split_across_chunks() {
2406 let mut buf = NdjsonBuffer::new();
2410 assert!(buf.decode(&[0xd0]).is_empty());
2411 assert!(buf.decode(&[0xb8, 0xd0, 0xb7, 0xd0]).is_empty());
2412 assert!(
2413 buf.decode(&[
2414 0xb2, 0xd0, 0xb5, 0xd1, 0x81, 0xd1, 0x82, 0xd0, 0xbd, 0xd0, 0xb8
2415 ])
2416 .is_empty()
2417 );
2418
2419 let lines = buf.decode(b"\n");
2420 assert_eq!(lines.len(), 1);
2421 assert_eq!(std::str::from_utf8(&lines[0]).unwrap(), "известни");
2422 }
2423
2424 #[test]
2425 fn ndjson_buffer_yields_parseable_chunks_when_split_arbitrarily() {
2426 let original = concat!(
2427 "{\"model\":\"llama3.2\",\"message\":{\"role\":\"assistant\",\"content\":\"hi\"},\"done\":false}\n",
2428 "{\"model\":\"llama3.2\",\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"done\":true}\n",
2429 );
2430
2431 let mut buf = NdjsonBuffer::new();
2432 let mut received = Vec::new();
2433 for byte in original.as_bytes() {
2434 for line in buf.decode(std::slice::from_ref(byte)) {
2435 let parsed: serde_json::Value =
2436 serde_json::from_slice(&line).expect("each drained line must be valid JSON");
2437 received.push(parsed);
2438 }
2439 }
2440
2441 assert_eq!(received.len(), 2);
2442 assert_eq!(received[0]["message"]["content"], "hi");
2443 assert_eq!(received[1]["done"], true);
2444 }
2445
2446 #[tokio::test]
2450 async fn completion_non_success_preserves_status_and_body() {
2451 use crate::client::CompletionClient;
2452 use crate::completion::CompletionModel;
2453 use crate::test_utils::RecordingHttpClient;
2454
2455 let body = r#"{"error":"model not found"}"#;
2456 let http_client =
2457 RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
2458 let client = Client::builder()
2459 .api_key("test-key")
2460 .http_client(http_client)
2461 .build()
2462 .expect("build client");
2463 let model = client.completion_model(LLAMA3_2);
2464 let request = model.completion_request("hello").build();
2465
2466 let error = model
2467 .completion(request)
2468 .await
2469 .expect_err("should fail with non-success status");
2470
2471 assert!(matches!(error, CompletionError::HttpError(_)));
2472 assert_eq!(
2473 error.provider_response_status(),
2474 Some(http::StatusCode::SERVICE_UNAVAILABLE)
2475 );
2476 assert_eq!(error.provider_response_body(), Some(body));
2477 }
2478
2479 #[tokio::test]
2483 async fn embeddings_non_success_preserves_status_and_body() {
2484 use crate::client::EmbeddingsClient;
2485 use crate::embeddings::EmbeddingModel;
2486 use crate::test_utils::RecordingHttpClient;
2487
2488 let body = r#"{"error":"model not found"}"#;
2489 let http_client =
2490 RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
2491 let client = Client::builder()
2492 .api_key("test-key")
2493 .http_client(http_client)
2494 .build()
2495 .expect("build client");
2496 let model = client.embedding_model(ALL_MINILM);
2497
2498 let error = model
2499 .embed_texts(vec!["hello".to_string()])
2500 .await
2501 .expect_err("should fail with non-success status");
2502
2503 assert!(matches!(error, EmbeddingError::HttpError(_)));
2504 assert_eq!(
2505 error.provider_response_status(),
2506 Some(http::StatusCode::SERVICE_UNAVAILABLE)
2507 );
2508 assert_eq!(error.provider_response_body(), Some(body));
2509 }
2510}