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