1use crate::client::{self, ApiKey, DebugExt, ModelLister, Nothing, Provider, ProviderClient};
42use crate::completion::Usage;
43use crate::http_client::{self, HttpClientExt};
44use crate::message::DocumentSourceKind;
45use crate::model::{Model, ModelList, ModelListingError};
46use crate::providers::internal;
47use crate::streaming::{RawStreamingChoice, RawStreamingResult, StreamFinal};
48use crate::telemetry::{CompletionOperation, CompletionSpanBuilder, SpanCombinator};
49use crate::{
50 completion::{self, CompletionError, CompletionRequest},
51 embeddings::{self, EmbeddingError},
52 json_utils, message,
53 message::Text,
54 streaming,
55 wasm_compat::{WasmCompatSend, WasmCompatSync},
56};
57use async_stream::stream;
58use futures::StreamExt;
59use serde::{Deserialize, Serialize};
60use serde_json::{Value, json};
61use std::convert::TryFrom;
62use tracing_futures::Instrument;
63const OLLAMA_API_BASE_URL: &str = "http://localhost:11434";
66
67const PROVIDER_NAME: &str = "ollama";
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
121client::impl_capabilities!(
122 OllamaExt,
123 completion = CompletionModel<H>,
124 embeddings = EmbeddingModel<H>,
125 model_listing = OllamaModelLister<H>,
126);
127
128impl DebugExt for OllamaExt {}
129
130client::impl_default_provider_builder!(
131 OllamaBuilder => OllamaExt,
132 api_key = OllamaApiKey,
133 base_url = OLLAMA_API_BASE_URL,
134);
135
136pub type Client<H = reqwest::Client> = client::Client<OllamaExt, H>;
137pub type ClientBuilder<H = crate::markers::Missing> =
138 client::ClientBuilder<OllamaBuilder, OllamaApiKey, H>;
139
140impl ProviderClient for Client {
141 type Input = OllamaApiKey;
142 type Error = crate::client::ProviderClientError;
143
144 fn from_env() -> Result<Self, Self::Error> {
145 let api_base = crate::client::optional_env_var("OLLAMA_API_BASE_URL")?
146 .unwrap_or_else(|| OLLAMA_API_BASE_URL.to_string());
147
148 let api_key = crate::client::optional_env_var("OLLAMA_API_KEY")?
149 .map(OllamaApiKey::from)
150 .unwrap_or_default();
151
152 Self::builder()
153 .api_key(api_key)
154 .base_url(&api_base)
155 .build()
156 .map_err(Into::into)
157 }
158
159 fn from_val(api_key: Self::Input) -> Result<Self, Self::Error> {
160 Self::builder().api_key(api_key).build().map_err(Into::into)
161 }
162}
163
164pub const ALL_MINILM: &str = "all-minilm";
167pub const NOMIC_EMBED_TEXT: &str = "nomic-embed-text";
168
169fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
170 match identifier {
171 ALL_MINILM => Some(384),
172 NOMIC_EMBED_TEXT => Some(768),
173 _ => None,
174 }
175}
176
177#[derive(Debug, Serialize, Deserialize)]
178pub struct EmbeddingResponse {
179 pub model: String,
180 pub embeddings: Vec<Vec<f64>>,
181 #[serde(default)]
182 pub total_duration: Option<u64>,
183 #[serde(default)]
184 pub load_duration: Option<u64>,
185 #[serde(default)]
186 pub prompt_eval_count: Option<u64>,
187}
188
189#[derive(Clone)]
192pub struct EmbeddingModel<T = reqwest::Client> {
193 client: Client<T>,
194 pub model: String,
195 ndims: usize,
196}
197
198impl<T> EmbeddingModel<T> {
199 pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
200 Self {
201 client,
202 model: model.into(),
203 ndims,
204 }
205 }
206
207 pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
208 Self {
209 client,
210 model: model.into(),
211 ndims,
212 }
213 }
214}
215
216impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
217where
218 T: HttpClientExt + Clone + 'static,
219{
220 type Client = Client<T>;
221
222 fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
223 let model = model.into();
224 let dims = dims
225 .or(model_dimensions_from_identifier(&model))
226 .unwrap_or_default();
227 Self::new(client.clone(), model, dims)
228 }
229
230 const MAX_DOCUMENTS: usize = 1024;
231 fn ndims(&self) -> usize {
232 self.ndims
233 }
234
235 async fn embed_texts(
236 &self,
237 documents: impl IntoIterator<Item = String>,
238 ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
239 let docs: Vec<String> = documents.into_iter().collect();
240
241 let body = serde_json::to_vec(&json!({
242 "model": self.model,
243 "input": docs
244 }))?;
245
246 let req = self
247 .client
248 .post("api/embed")?
249 .body(body)
250 .map_err(|e| EmbeddingError::HttpError(e.into()))?;
251
252 let response = self.client.send::<_, Vec<u8>>(req).await?;
253
254 let status = response.status();
255 if !status.is_success() {
256 let text = http_client::text(response).await?;
257 return Err(EmbeddingError::from_http_response(status, text));
258 }
259
260 let bytes: Vec<u8> = response.into_body().await?;
261
262 let api_resp: EmbeddingResponse = serde_json::from_slice(&bytes)?;
263
264 if api_resp.embeddings.len() != docs.len() {
265 return Err(EmbeddingError::ResponseError(
266 "Number of returned embeddings does not match input".into(),
267 ));
268 }
269 Ok(api_resp
270 .embeddings
271 .into_iter()
272 .zip(docs.into_iter())
273 .map(|(vec, document)| embeddings::Embedding { document, vec })
274 .collect())
275 }
276}
277
278pub const LLAMA3_2: &str = "llama3.2";
281pub const LLAVA: &str = "llava";
282pub const MISTRAL: &str = "mistral";
283
284#[derive(Debug, Serialize, Deserialize)]
285pub struct CompletionResponse {
286 pub model: String,
287 pub created_at: String,
288 pub message: Message,
289 pub done: bool,
290 #[serde(default)]
291 pub done_reason: Option<String>,
292 #[serde(default)]
293 pub total_duration: Option<u64>,
294 #[serde(default)]
295 pub load_duration: Option<u64>,
296 #[serde(default)]
297 pub prompt_eval_count: Option<u64>,
298 #[serde(default)]
299 pub prompt_eval_duration: Option<u64>,
300 #[serde(default)]
301 pub eval_count: Option<u64>,
302 #[serde(default)]
303 pub eval_duration: Option<u64>,
304}
305pub(crate) fn map_done_reason(reason: &str) -> completion::FinishReason {
311 match reason {
312 "stop" => completion::FinishReason::Stop,
313 "length" => completion::FinishReason::Length,
314 other => completion::FinishReason::Other(other.to_owned()),
315 }
316}
317
318impl From<&CompletionResponse> for Usage {
319 fn from(response: &CompletionResponse) -> Usage {
320 let input_tokens = response.prompt_eval_count.unwrap_or(0);
321 let output_tokens = response.eval_count.unwrap_or(0);
322 crate::providers::internal::completion_usage(
323 input_tokens,
324 output_tokens,
325 input_tokens + output_tokens,
326 0,
327 )
328 }
329}
330
331impl crate::telemetry::ProviderResponseExt for CompletionResponse {
332 type Usage = Usage;
333
334 fn get_response_id(&self) -> Option<String> {
336 None
337 }
338
339 fn get_response_model_name(&self) -> Option<String> {
340 Some(self.model.clone())
341 }
342
343 fn get_text_response(&self) -> Option<String> {
344 match &self.message {
345 Message::Assistant { content, .. } if !content.is_empty() => Some(content.clone()),
346 _ => None,
347 }
348 }
349
350 fn get_usage(&self) -> Option<Self::Usage> {
351 Some(Usage::from(self))
352 }
353}
354
355impl TryFrom<CompletionResponse> for completion::CompletionResponse {
356 type Error = CompletionError;
357 fn try_from(resp: CompletionResponse) -> Result<Self, Self::Error> {
358 let usage = Usage::from(&resp);
359 let finish_reason = resp.done_reason.as_deref().map(map_done_reason);
360 let model = resp.model.clone();
361 let permits_omitted_think_start = resp.model.to_ascii_lowercase().contains("qwen3");
362
363 let Message::Assistant {
365 content,
366 thinking,
367 tool_calls,
368 ..
369 } = resp.message
370 else {
371 return Err(CompletionError::ResponseError(
372 "Chat response does not include an assistant message".into(),
373 ));
374 };
375
376 let mut assistant_contents = Vec::new();
377 let (legacy_thinking, visible_content) = if matches!(thinking.as_deref(), None | Some("")) {
378 split_legacy_thinking(&content, permits_omitted_think_start)
379 } else {
380 (None, content.as_str())
381 };
382 if let Some(thinking) = thinking.as_deref().filter(|t| !t.is_empty()) {
389 assistant_contents.push(completion::AssistantContent::reasoning(thinking));
390 }
391 if let Some(legacy_thinking) = legacy_thinking {
392 assistant_contents.push(completion::AssistantContent::reasoning(legacy_thinking));
393 }
394 if !visible_content.is_empty() {
396 assistant_contents.push(completion::AssistantContent::text(visible_content));
397 }
398 for tc in tool_calls.iter() {
406 assistant_contents.push(completion::AssistantContent::tool_call(
407 tc.id.as_deref().unwrap_or(""),
408 tc.function.name.clone(),
409 tc.function.arguments.clone(),
410 ));
411 }
412 let choice = crate::message::require_non_empty_response(assistant_contents)?;
413
414 Ok(
415 completion::CompletionResponse::new(choice, usage, PROVIDER_NAME)
416 .with_model(model)
417 .with_optional_finish_reason(finish_reason),
418 )
419 }
420}
421
422fn split_legacy_thinking(content: &str, permits_omitted_start: bool) -> (Option<&str>, &str) {
427 let trimmed = content.trim_start();
428 let split = if let Some(reasoning_start) = trimmed.strip_prefix("<think>") {
429 reasoning_start.split_once("</think>")
430 } else if permits_omitted_start {
431 trimmed.split_once("\n</think>\n\n")
435 } else {
436 None
437 };
438 let Some((reasoning, visible)) = split else {
439 return (None, content);
440 };
441
442 let reasoning = reasoning.trim();
443 if reasoning.is_empty() {
444 return (None, visible.trim_start());
445 }
446
447 (Some(reasoning), visible.trim_start())
448}
449
450#[derive(Debug, Serialize, Deserialize)]
451pub(super) struct OllamaCompletionRequest {
452 model: String,
453 pub messages: Vec<Message>,
454 #[serde(skip_serializing_if = "Vec::is_empty")]
455 tools: Vec<ToolDefinition>,
456 pub stream: bool,
457 #[serde(skip_serializing_if = "Option::is_none")]
458 think: Option<Think>,
459 #[serde(skip_serializing_if = "Option::is_none")]
460 keep_alive: Option<String>,
461 #[serde(skip_serializing_if = "Option::is_none")]
462 format: Option<schemars::Schema>,
463 options: serde_json::Value,
464}
465
466impl TryFrom<(&str, CompletionRequest)> for OllamaCompletionRequest {
467 type Error = CompletionError;
468
469 fn try_from((model, req): (&str, CompletionRequest)) -> Result<Self, Self::Error> {
470 let chat_history = req.chat_history_with_documents();
471 let model = req.model.clone().unwrap_or_else(|| model.to_string());
472 if req.tool_choice.is_some() {
473 tracing::warn!("WARNING: `tool_choice` not supported for Ollama");
474 }
475 let mut partial_history = vec![];
477 partial_history.extend(chat_history);
478 crate::providers::internal::resolve_empty_tool_result_names(&mut partial_history);
481
482 let mut full_history: Vec<Message> = match &req.preamble {
484 Some(preamble) => vec![Message::system(preamble)],
485 None => vec![],
486 };
487
488 full_history.extend(
490 partial_history
491 .into_iter()
492 .map(message::Message::try_into)
493 .collect::<Result<Vec<Vec<Message>>, _>>()?
494 .into_iter()
495 .flatten()
496 .collect::<Vec<_>>(),
497 );
498
499 let mut think: Option<Think> = None;
500 let mut keep_alive: Option<String> = None;
501
502 let mut base_options = serde_json::Map::new();
506 if let Some(temperature) = req.temperature {
507 base_options.insert("temperature".to_string(), json!(temperature));
508 }
509 if let Some(max_tokens) = req.max_tokens {
510 base_options.insert("num_predict".to_string(), json!(max_tokens));
511 }
512 let base_options = Value::Object(base_options);
513
514 let options = if let Some(mut extra) = req.additional_params {
515 if let Some(obj) = extra.as_object_mut() {
517 if let Some(think_val) = obj.remove("think") {
519 think = Some(match think_val {
520 Value::Bool(think) => Think::Bool(think),
521 Value::String(think) => Think::Level(match think.to_lowercase().as_str() {
522 "low" => Level::Low,
523 "medium" => Level::Medium,
524 "high" => Level::High,
525 "max" => Level::Max,
526 _ => {
527 return Err(CompletionError::RequestError(
528 "`think` must be a 'low', 'medium', 'high', 'max' or bool"
529 .into(),
530 ));
531 }
532 }),
533 _ => {
534 return Err(CompletionError::RequestError(
535 "`think` must be a 'low', 'medium', 'high', 'max' or bool".into(),
536 ));
537 }
538 });
539 }
540
541 if let Some(keep_alive_val) = obj.remove("keep_alive") {
543 keep_alive = Some(
544 keep_alive_val
545 .as_str()
546 .ok_or_else(|| {
547 CompletionError::RequestError(
548 "`keep_alive` must be a string".into(),
549 )
550 })?
551 .to_string(),
552 );
553 }
554 }
555
556 json_utils::merge(base_options, extra)
557 } else {
558 base_options
559 };
560
561 Ok(Self {
562 model: model.to_string(),
563 messages: full_history,
564 stream: false,
565 think,
566 keep_alive,
567 format: req.output_schema,
568 tools: req
569 .tools
570 .clone()
571 .into_iter()
572 .map(ToolDefinition::from)
573 .collect::<Vec<_>>(),
574 options,
575 })
576 }
577}
578
579#[derive(Clone)]
580pub struct CompletionModel<T = reqwest::Client> {
581 client: Client<T>,
582 pub model: String,
583}
584
585impl<T> CompletionModel<T> {
586 pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
587 Self {
588 client,
589 model: model.into(),
590 }
591 }
592}
593
594impl<T> crate::client::ConstructCompletionModel<Client<T>> for CompletionModel<T>
595where
596 Client<T>: Clone,
597{
598 fn construct(client: &Client<T>, model: String) -> Self {
599 Self::new(client.clone(), model)
600 }
601}
602
603#[derive(Debug, Clone, Serialize, Deserialize)]
604#[serde(untagged)]
605enum Think {
606 Bool(bool),
607 Level(Level),
608}
609
610#[derive(Debug, Clone, Serialize, Deserialize)]
611#[serde(rename_all = "lowercase")]
612enum Level {
613 Low,
614 Medium,
615 High,
616 Max,
617}
618
619#[derive(Clone, Serialize, Deserialize, Debug)]
624pub struct StreamingCompletionResponse {
625 pub model: String,
627 pub done_reason: Option<String>,
628 pub total_duration: Option<u64>,
629 pub load_duration: Option<u64>,
630 pub prompt_eval_count: Option<u64>,
631 pub prompt_eval_duration: Option<u64>,
632 pub eval_count: Option<u64>,
633 pub eval_duration: Option<u64>,
634}
635
636impl From<&StreamingCompletionResponse> for Usage {
637 fn from(response: &StreamingCompletionResponse) -> Usage {
638 let input_tokens = response.prompt_eval_count.unwrap_or_default();
639 let output_tokens = response.eval_count.unwrap_or_default();
640 crate::providers::internal::completion_usage(
641 input_tokens,
642 output_tokens,
643 input_tokens + output_tokens,
644 0,
645 )
646 }
647}
648
649impl From<StreamingCompletionResponse> for StreamFinal {
650 fn from(response: StreamingCompletionResponse) -> StreamFinal {
651 StreamFinal::new(PROVIDER_NAME, Usage::from(&response))
654 .with_optional_finish_reason(response.done_reason.as_deref().map(map_done_reason))
655 .with_model(response.model)
656 }
657}
658
659#[derive(Default)]
665struct NdjsonBuffer {
666 buf: Vec<u8>,
667}
668
669impl NdjsonBuffer {
670 fn new() -> Self {
671 Self::default()
672 }
673
674 fn decode(&mut self, chunk: &[u8]) -> Vec<Vec<u8>> {
677 self.buf.extend_from_slice(chunk);
678
679 let mut lines = Vec::new();
680 while let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
681 let mut line: Vec<u8> = self.buf.drain(..=pos).collect();
682 line.pop();
683 if !line.is_empty() {
684 lines.push(line);
685 }
686 }
687 lines
688 }
689}
690
691impl<T> CompletionModel<T>
692where
693 T: HttpClientExt + Clone + Default + std::fmt::Debug + Send + 'static,
694{
695 pub async fn raw_completion(
704 &self,
705 completion_request: CompletionRequest,
706 ) -> Result<CompletionResponse, CompletionError> {
707 let system_instructions = completion_request.preamble.clone();
708 let record_telemetry_content = completion_request.record_telemetry_content;
709 let request = OllamaCompletionRequest::try_from((self.model.as_ref(), completion_request))?;
710 let span =
711 CompletionSpanBuilder::new(PROVIDER_NAME, &request.model, CompletionOperation::Chat)
712 .system_instructions(system_instructions.as_deref(), record_telemetry_content)
713 .build();
714
715 internal::trace_json(
716 crate::providers::internal::LogTarget::Completions,
717 "Ollama completion request",
718 &request,
719 );
720
721 let body = serde_json::to_vec(&request)?;
722
723 let req = self
724 .client
725 .post("api/chat")?
726 .body(body)
727 .map_err(http_client::Error::from)?;
728
729 let async_block = internal::completion_send::send_completion::<
730 _,
731 internal::envelope::DirectPayload<CompletionResponse>,
732 _,
733 >(
734 &self.client,
735 req,
736 "Ollama completion",
737 None,
739 |response| {
740 let span = tracing::Span::current();
741 span.record_response_metadata(response);
742 span.record_token_usage(&Usage::from(response));
743 },
744 );
745
746 tracing::Instrument::instrument(async_block, span)
747 .await
748 .map(|(payload, _)| payload)
749 }
750
751 pub async fn raw_stream(
759 &self,
760 request: CompletionRequest,
761 ) -> Result<RawStreamingResult<StreamingCompletionResponse>, CompletionError> {
762 let system_instructions = request.preamble.clone();
763 let record_telemetry_content = request.record_telemetry_content;
764 let mut request = OllamaCompletionRequest::try_from((self.model.as_ref(), request))?;
765 let span = CompletionSpanBuilder::new(
766 PROVIDER_NAME,
767 &request.model,
768 CompletionOperation::ChatStreaming,
769 )
770 .system_instructions(system_instructions.as_deref(), record_telemetry_content)
771 .build();
772 request.stream = true;
773
774 internal::trace_json(
775 crate::providers::internal::LogTarget::Completions,
776 "Ollama streaming completion request",
777 &request,
778 );
779
780 let body = serde_json::to_vec(&request)?;
781
782 let req = self
783 .client
784 .post("api/chat")?
785 .body(body)
786 .map_err(http_client::Error::from)?;
787
788 let response = self
789 .client
790 .send_streaming(req)
791 .instrument(span.clone())
792 .await?;
793 let status = response.status();
794 let mut byte_stream = response.into_body();
795
796 if !status.is_success() {
797 let mut body = Vec::new();
798 while let Some(chunk) = byte_stream.next().await {
799 match chunk {
800 Ok(bytes) => body.extend_from_slice(&bytes),
801 Err(e) => {
802 tracing::warn!(error = %e, "failed reading Ollama error-response body; preserving partial body");
803 break;
804 }
805 }
806 }
807 return Err(CompletionError::from_http_response(
808 status,
809 String::from_utf8_lossy(&body),
810 ));
811 }
812
813 let transport = stream! {
817 let mut line_buf = NdjsonBuffer::new();
818 while let Some(chunk) = byte_stream.next().await {
819 let bytes = match chunk {
820 Ok(bytes) => bytes,
821 Err(e) => {
822 yield Err(CompletionError::from(http_client::Error::Instance(e.into())));
823 break;
824 }
825 };
826
827 for line in line_buf.decode(&bytes) {
828 tracing::debug!(target: "rig", "Received NDJSON line from Ollama: {}", String::from_utf8_lossy(&line));
829 yield Ok(internal::adapter::WireFrame::Bytes(line));
830 }
831 }
832 };
833
834 let stream: RawStreamingResult<StreamingCompletionResponse> = Box::pin(
835 internal::adapter::run_wire_stream(transport, OllamaAdapter::default())
836 .instrument(span),
837 );
838
839 Ok(stream)
840 }
841}
842
843struct OllamaAdapter {
852 reasoning: internal::chunk_lifecycle::MintedReasoningLifecycle,
856 tool_ids: crate::streaming::SyntheticIds,
861}
862
863impl Default for OllamaAdapter {
864 fn default() -> Self {
865 Self {
866 reasoning: internal::chunk_lifecycle::MintedReasoningLifecycle::new(
867 crate::streaming::StreamPartId::minted(crate::streaming::MintKind::Reasoning, 0),
868 ),
869 tool_ids: crate::streaming::SyntheticIds::tool(),
870 }
871 }
872}
873
874impl internal::adapter::WireAdapter for OllamaAdapter {
875 type Frame = internal::adapter::WireFrame;
876 type Event = CompletionResponse;
877 type Response = StreamingCompletionResponse;
878
879 fn classify(&self, frame: Self::Frame) -> internal::wire::WireEvent<CompletionResponse> {
880 match frame {
881 internal::adapter::WireFrame::Bytes(line) => {
882 internal::wire::classify_untyped_line(&line)
883 }
884 internal::adapter::WireFrame::Text(line) => {
885 internal::wire::classify_untyped_line(line.as_bytes())
886 }
887 }
888 }
889
890 fn interpret(
891 &mut self,
892 response: CompletionResponse,
893 out: &mut internal::adapter::AdapterOutput<Self::Response>,
894 ) {
895 let span = tracing::Span::current();
896 if response.done {
897 span.record("gen_ai.response.model", &response.model);
898 }
899
900 if let Message::Assistant {
901 content,
902 thinking,
903 tool_calls,
904 ..
905 } = response.message
906 {
907 let mut tool_events = Vec::with_capacity(tool_calls.len());
913 for tool_call in tool_calls {
914 let key = match tool_call
915 .id
916 .as_deref()
917 .and_then(crate::streaming::WireId::new)
918 {
919 Some(wire_id) => crate::streaming::StreamPartId::wire(wire_id.as_str()),
920 None => self.tool_ids.mint(),
921 };
922 tool_events.push(RawStreamingChoice::ToolCall(
923 crate::streaming::RawStreamingToolCall::new(
924 key,
925 tool_call.function.name,
926 tool_call.function.arguments,
927 ),
928 ));
929 }
930
931 self.reasoning.emit_chunk(
934 internal::chunk_lifecycle::ChunkParts {
935 reasoning: thinking,
936 reasoning_signature: None,
937 text: Some(content),
938 tool_events,
939 },
940 out,
941 );
942 }
943
944 if response.done {
947 span.record("gen_ai.usage.input_tokens", response.prompt_eval_count);
948 span.record("gen_ai.usage.output_tokens", response.eval_count);
949 out.push(Ok(RawStreamingChoice::FinalResponse(
950 StreamingCompletionResponse {
951 model: response.model,
952 total_duration: response.total_duration,
953 load_duration: response.load_duration,
954 prompt_eval_count: response.prompt_eval_count,
955 prompt_eval_duration: response.prompt_eval_duration,
956 eval_count: response.eval_count,
957 eval_duration: response.eval_duration,
958 done_reason: response.done_reason,
959 },
960 )));
961 }
962 }
963
964 fn finish(&mut self, _out: &mut internal::adapter::AdapterOutput<Self::Response>) {
965 }
968}
969
970impl<T> completion::CompletionModel for CompletionModel<T>
971where
972 T: HttpClientExt + Clone + Default + std::fmt::Debug + Send + 'static,
973{
974 async fn completion(
975 &self,
976 completion_request: CompletionRequest,
977 ) -> Result<completion::CompletionResponse, CompletionError> {
978 let raw = self.raw_completion(completion_request).await?;
980 let captured = serde_json::to_value(&raw)?;
981 let response: completion::CompletionResponse = raw.try_into()?;
982 Ok(response.with_raw(captured))
983 }
984
985 async fn stream(
986 &self,
987 request: CompletionRequest,
988 ) -> Result<streaming::StreamingCompletionResponse, CompletionError> {
989 let stream = self.raw_stream(request).await?;
990 let normalized =
991 streaming::normalize_stream(stream, |response: StreamingCompletionResponse| {
992 Ok(response.into())
993 });
994
995 Ok(streaming::StreamingCompletionResponse::stream(
996 PROVIDER_NAME,
997 normalized,
998 ))
999 }
1000}
1001
1002#[derive(Debug, Deserialize)]
1005struct ListModelsResponse {
1006 models: Vec<ListModelEntry>,
1007}
1008
1009#[derive(Debug, Deserialize)]
1010struct ListModelEntry {
1011 name: String,
1012 model: String,
1013}
1014
1015impl From<ListModelEntry> for Model {
1016 fn from(value: ListModelEntry) -> Self {
1017 Model::new(value.model, value.name)
1018 }
1019}
1020
1021#[derive(Clone)]
1023pub struct OllamaModelLister<H = reqwest::Client> {
1024 client: Client<H>,
1025}
1026
1027impl<H> ModelLister<H> for OllamaModelLister<H>
1028where
1029 H: HttpClientExt + WasmCompatSend + WasmCompatSync + 'static,
1030{
1031 type Client = Client<H>;
1032
1033 fn new(client: Self::Client) -> Self {
1034 Self { client }
1035 }
1036
1037 async fn list_all(&self) -> Result<ModelList, ModelListingError> {
1038 let api_resp: ListModelsResponse = crate::providers::internal::model_listing::get_json(
1039 &self.client,
1040 "Ollama",
1041 "/api/tags",
1042 )
1043 .await?;
1044 let models = api_resp.models.into_iter().map(Model::from).collect();
1045
1046 Ok(ModelList::new(models))
1047 }
1048}
1049
1050#[derive(Clone, Debug, Deserialize, Serialize)]
1054pub struct ToolDefinition {
1055 #[serde(rename = "type")]
1056 pub type_field: String, pub function: completion::ToolDefinition,
1058}
1059
1060impl From<crate::completion::ToolDefinition> for ToolDefinition {
1062 fn from(tool: crate::completion::ToolDefinition) -> Self {
1063 ToolDefinition {
1064 type_field: "function".to_owned(),
1065 function: completion::ToolDefinition {
1066 name: tool.name,
1067 description: tool.description,
1068 parameters: tool.parameters,
1069 },
1070 }
1071 }
1072}
1073
1074#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
1075pub struct ToolCall {
1076 #[serde(default, skip_serializing)]
1082 pub id: Option<String>,
1083 #[serde(default, rename = "type")]
1084 pub r#type: ToolType,
1085 pub function: Function,
1086}
1087#[derive(Default, Debug, Serialize, Deserialize, PartialEq, Clone)]
1088#[serde(rename_all = "lowercase")]
1089pub enum ToolType {
1090 #[default]
1091 Function,
1092}
1093#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
1094pub struct Function {
1095 pub name: String,
1096 pub arguments: Value,
1097}
1098
1099#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
1102#[serde(tag = "role", rename_all = "lowercase")]
1103pub enum Message {
1104 User {
1105 content: String,
1106 #[serde(skip_serializing_if = "Option::is_none")]
1107 images: Option<Vec<String>>,
1108 #[serde(skip_serializing_if = "Option::is_none")]
1109 name: Option<String>,
1110 },
1111 Assistant {
1112 #[serde(default)]
1113 content: String,
1114 #[serde(skip_serializing_if = "Option::is_none")]
1115 thinking: Option<String>,
1116 #[serde(skip_serializing_if = "Option::is_none")]
1117 images: Option<Vec<String>>,
1118 #[serde(skip_serializing_if = "Option::is_none")]
1119 name: Option<String>,
1120 #[serde(default, deserialize_with = "json_utils::null_or_default")]
1121 tool_calls: Vec<ToolCall>,
1122 },
1123 System {
1124 content: String,
1125 #[serde(skip_serializing_if = "Option::is_none")]
1126 images: Option<Vec<String>>,
1127 #[serde(skip_serializing_if = "Option::is_none")]
1128 name: Option<String>,
1129 },
1130 #[serde(rename = "tool")]
1131 ToolResult {
1132 #[serde(rename = "tool_name")]
1133 name: String,
1134 content: String,
1135 },
1136}
1137
1138fn user_message_from_content(
1142 content: Vec<crate::message::UserContent>,
1143) -> Result<Message, crate::message::MessageError> {
1144 let mut texts = Vec::new();
1145 let mut images = Vec::new();
1146
1147 for content in content {
1148 match content {
1149 crate::message::UserContent::Text(crate::message::Text { text, .. }) => {
1150 texts.push(text);
1151 }
1152 crate::message::UserContent::Image(crate::message::Image {
1153 data: DocumentSourceKind::Base64(data),
1154 ..
1155 }) => images.push(data),
1156 crate::message::UserContent::Image(_) => {
1157 return Err(crate::message::MessageError::ConversionError(
1158 "Ollama images must be base64 encoded data".into(),
1159 ));
1160 }
1161 crate::message::UserContent::Document(crate::message::Document {
1162 data: DocumentSourceKind::Base64(data) | DocumentSourceKind::String(data),
1163 ..
1164 }) => texts.push(data),
1165 crate::message::UserContent::Document(_) => {
1166 return Err(crate::message::MessageError::ConversionError(
1167 "Ollama documents must be string or base64 encoded data".into(),
1168 ));
1169 }
1170 crate::message::UserContent::Audio(_) => {
1171 return Err(crate::message::MessageError::ConversionError(
1172 "Ollama does not support audio user content".into(),
1173 ));
1174 }
1175 crate::message::UserContent::Video(_) => {
1176 return Err(crate::message::MessageError::ConversionError(
1177 "Ollama does not support video user content".into(),
1178 ));
1179 }
1180 crate::message::UserContent::ToolResult(_) => {
1181 return Err(crate::message::MessageError::ConversionError(
1182 "tool results must be converted to a separate Ollama message".into(),
1183 ));
1184 }
1185 }
1186 }
1187
1188 Ok(Message::User {
1189 content: texts.join(" "),
1190 images: (!images.is_empty()).then_some(images),
1191 name: None,
1192 })
1193}
1194
1195impl TryFrom<crate::message::Message> for Vec<Message> {
1198 type Error = crate::message::MessageError;
1199 fn try_from(internal_msg: crate::message::Message) -> Result<Self, Self::Error> {
1200 use crate::message::Message as InternalMessage;
1201 match internal_msg {
1202 InternalMessage::System { content } => Ok(vec![Message::System {
1203 content,
1204 images: None,
1205 name: None,
1206 }]),
1207 InternalMessage::User { content, .. } => {
1208 let mut messages = Vec::new();
1209 let mut pending_user_content = Vec::new();
1210
1211 for content in content {
1212 match content {
1213 crate::message::UserContent::ToolResult(crate::message::ToolResult {
1214 name,
1215 content,
1216 ..
1217 }) => {
1218 let function_name = name;
1220 if !pending_user_content.is_empty() {
1221 messages.push(user_message_from_content(std::mem::take(
1222 &mut pending_user_content,
1223 ))?);
1224 }
1225
1226 let content = content
1227 .into_iter()
1228 .map(|content| match content {
1229 crate::message::ToolResultContent::Text(text) => Ok(text.text),
1230 crate::message::ToolResultContent::Json { value } => {
1231 Ok(value.to_string())
1232 }
1233 crate::message::ToolResultContent::Image(_) => {
1234 Err(crate::message::MessageError::ConversionError(
1235 "Ollama does not support images in tool results".into(),
1236 ))
1237 }
1238 })
1239 .collect::<Result<Vec<_>, _>>()?
1240 .join("\n");
1241 messages.push(Message::ToolResult {
1242 name: function_name,
1243 content,
1244 });
1245 }
1246 content => pending_user_content.push(content),
1247 }
1248 }
1249
1250 if !pending_user_content.is_empty() {
1251 messages.push(user_message_from_content(pending_user_content)?);
1252 }
1253
1254 Ok(messages)
1255 }
1256 InternalMessage::Assistant { content, .. } => {
1257 let mut thinking: Option<String> = None;
1258 let mut text_content = Vec::new();
1259 let mut tool_calls = Vec::new();
1260
1261 for content in content.into_iter() {
1262 match content {
1263 crate::message::AssistantContent::Text(text) => {
1264 text_content.push(text.text)
1265 }
1266 crate::message::AssistantContent::ToolCall(tool_call) => {
1267 tool_calls.push(tool_call)
1268 }
1269 crate::message::AssistantContent::Reasoning(reasoning) => {
1270 let display = reasoning.display_text();
1271 if !display.is_empty() {
1272 thinking = Some(display);
1273 }
1274 }
1275 crate::message::AssistantContent::Image(_) => {
1276 return Err(crate::message::MessageError::ConversionError(
1277 "Ollama currently doesn't support images.".into(),
1278 ));
1279 }
1280 }
1281 }
1282
1283 Ok(vec![Message::Assistant {
1289 content: text_content.join(" "),
1290 thinking,
1291 images: None,
1292 name: None,
1293 tool_calls: tool_calls
1294 .into_iter()
1295 .map(|tool_call| tool_call.into())
1296 .collect::<Vec<_>>(),
1297 }])
1298 }
1299 }
1300 }
1301}
1302
1303impl From<Message> for crate::completion::Message {
1314 fn from(msg: Message) -> Self {
1315 match msg {
1316 Message::User { content, .. } => crate::completion::Message::User {
1317 content: vec![crate::completion::message::UserContent::Text(Text::new(
1318 content,
1319 ))],
1320 },
1321 Message::Assistant {
1322 content,
1323 thinking,
1324 tool_calls,
1325 ..
1326 } => {
1327 let mut assistant_contents = Vec::new();
1328 if let Some(thinking) = thinking.filter(|t| !t.is_empty()) {
1330 assistant_contents.push(
1331 crate::completion::message::AssistantContent::reasoning(thinking),
1332 );
1333 }
1334 if !content.is_empty() {
1341 assistant_contents.push(crate::completion::message::AssistantContent::Text(
1342 Text::new(content),
1343 ));
1344 }
1345 for tc in tool_calls {
1348 assistant_contents.push(
1349 crate::completion::message::AssistantContent::tool_call(
1350 tc.id.as_deref().unwrap_or(""),
1351 tc.function.name,
1352 tc.function.arguments,
1353 ),
1354 );
1355 }
1356 crate::completion::Message::Assistant {
1357 id: None,
1358 content: assistant_contents,
1359 }
1360 }
1361 Message::System { content, .. } => crate::completion::Message::User {
1363 content: vec![crate::completion::message::UserContent::Text(Text::new(
1364 content,
1365 ))],
1366 },
1367 Message::ToolResult { name, content } => crate::completion::Message::User {
1368 content: vec![message::UserContent::tool_result_from_wire(
1371 "",
1372 name,
1373 vec![message::ToolResultContent::text(content)],
1374 )],
1375 },
1376 }
1377 }
1378}
1379
1380impl Message {
1381 pub fn system(content: &str) -> Self {
1383 Message::System {
1384 content: content.to_owned(),
1385 images: None,
1386 name: None,
1387 }
1388 }
1389}
1390
1391impl From<crate::message::ToolCall> for ToolCall {
1394 fn from(tool_call: crate::message::ToolCall) -> Self {
1395 Self {
1396 id: None,
1399 r#type: ToolType::Function,
1400 function: Function {
1401 name: tool_call.function.name,
1402 arguments: tool_call.function.arguments,
1403 },
1404 }
1405 }
1406}
1407
1408#[cfg(test)]
1413mod tests {
1414 use super::*;
1415 use serde_json::json;
1416
1417 #[test]
1420 fn classify_ndjson_line_is_known_or_corrupt() {
1421 let line = json!({
1422 "model": "llama3.2",
1423 "created_at": "2024-01-01T00:00:00Z",
1424 "message": {"role": "assistant", "content": "hi"},
1425 "done": false,
1426 })
1427 .to_string();
1428 assert!(matches!(
1429 internal::wire::classify_untyped_line::<CompletionResponse>(line.as_bytes()),
1430 internal::wire::WireEvent::Known(_)
1431 ));
1432 assert!(matches!(
1433 internal::wire::classify_untyped_line::<CompletionResponse>(b"{not json"),
1434 internal::wire::WireEvent::Corrupt(_)
1435 ));
1436 assert!(matches!(
1437 internal::wire::classify_untyped_line::<CompletionResponse>(br#"{"done": 42}"#),
1438 internal::wire::WireEvent::Corrupt(_)
1439 ));
1440 }
1441
1442 #[test]
1443 fn splits_legacy_reasoning_with_or_without_opening_marker() {
1444 assert_eq!(
1445 split_legacy_thinking("<think>private reasoning</think>\n\nvisible answer", false),
1446 (Some("private reasoning"), "visible answer")
1447 );
1448 assert_eq!(
1449 split_legacy_thinking("private reasoning\n</think>\n\nvisible answer", true),
1450 (Some("private reasoning"), "visible answer")
1451 );
1452 }
1453
1454 #[test]
1455 fn leaves_unterminated_or_inline_reasoning_markers_visible() {
1456 assert_eq!(
1457 split_legacy_thinking("<think>unterminated", true),
1458 (None, "<think>unterminated")
1459 );
1460 assert_eq!(
1461 split_legacy_thinking("The literal marker is <think>.", true),
1462 (None, "The literal marker is <think>.")
1463 );
1464 assert_eq!(
1465 split_legacy_thinking(" visible indentation", true),
1466 (None, " visible indentation")
1467 );
1468 assert_eq!(
1469 split_legacy_thinking("The closing token </think> is XML-like.", true),
1470 (None, "The closing token </think> is XML-like.")
1471 );
1472 assert_eq!(
1473 split_legacy_thinking("Example:\n</think>\nis a closing tag.", true),
1474 (None, "Example:\n</think>\nis a closing tag.")
1475 );
1476 }
1477
1478 #[tokio::test]
1480 async fn test_chat_completion() {
1481 let sample_chat_response = json!({
1483 "model": "llama3.2",
1484 "created_at": "2023-08-04T19:22:45.499127Z",
1485 "message": {
1486 "role": "assistant",
1487 "content": "The sky is blue because of Rayleigh scattering.",
1488 "images": null,
1489 "tool_calls": [
1490 {
1491 "type": "function",
1492 "function": {
1493 "name": "get_current_weather",
1494 "arguments": {
1495 "location": "San Francisco, CA",
1496 "format": "celsius"
1497 }
1498 }
1499 }
1500 ]
1501 },
1502 "done": true,
1503 "total_duration": 8000000000u64,
1504 "load_duration": 6000000u64,
1505 "prompt_eval_count": 61u64,
1506 "prompt_eval_duration": 400000000u64,
1507 "eval_count": 468u64,
1508 "eval_duration": 7700000000u64
1509 });
1510 let sample_text = sample_chat_response.to_string();
1511
1512 let chat_resp: CompletionResponse =
1513 serde_json::from_str(&sample_text).expect("Invalid JSON structure");
1514 let conv: completion::CompletionResponse = chat_resp.try_into().unwrap();
1515 assert!(
1516 !conv.choice.is_empty(),
1517 "Expected non-empty choice in chat response"
1518 );
1519 }
1520
1521 #[test]
1522 fn done_reason_maps_documented_values_and_preserves_the_rest() {
1523 assert_eq!(map_done_reason("stop"), completion::FinishReason::Stop);
1524 assert_eq!(map_done_reason("length"), completion::FinishReason::Length);
1525 assert_eq!(
1528 map_done_reason("load"),
1529 completion::FinishReason::Other("load".to_owned())
1530 );
1531 assert_eq!(
1532 map_done_reason("unload"),
1533 completion::FinishReason::Other("unload".to_owned())
1534 );
1535 }
1536
1537 #[test]
1538 fn response_metadata_is_normalized() {
1539 let response: CompletionResponse = serde_json::from_value(json!({
1540 "model": "llama3.2",
1541 "created_at": "2023-08-04T19:22:45.499127Z",
1542 "message": {"role": "assistant", "content": "Hi!", "tool_calls": []},
1543 "done": true,
1544 "done_reason": "length",
1545 "prompt_eval_count": 12u64,
1546 "eval_count": 3u64
1547 }))
1548 .expect("fixture should deserialize");
1549
1550 let normalized: completion::CompletionResponse =
1551 response.try_into().expect("normalization should succeed");
1552
1553 assert_eq!(normalized.provider, PROVIDER_NAME);
1554 assert_eq!(normalized.model.as_deref(), Some("llama3.2"));
1555 assert_eq!(
1556 normalized.finish_reason(),
1557 Some(completion::FinishReason::Length)
1558 );
1559 assert_eq!(normalized.message_id, None);
1561 assert_eq!(normalized.usage.input_tokens, 12);
1562 assert_eq!(normalized.usage.output_tokens, 3);
1563 assert_eq!(normalized.usage.total_tokens, 15);
1564 }
1565
1566 #[test]
1569 fn tool_call_turn_upgrades_a_plain_stop_to_tool_calls() {
1570 let response: CompletionResponse = serde_json::from_value(json!({
1571 "model": "qwen3:4b",
1572 "created_at": "2023-08-04T19:22:45.499127Z",
1573 "message": {
1574 "role": "assistant",
1575 "content": "",
1576 "tool_calls": [
1577 {"type": "function", "function": {"name": "get_weather", "arguments": {"location": "Berlin"}}}
1578 ]
1579 },
1580 "done": true,
1581 "done_reason": "stop"
1582 }))
1583 .expect("fixture should deserialize");
1584
1585 let normalized: completion::CompletionResponse =
1586 response.try_into().expect("normalization should succeed");
1587
1588 assert_eq!(
1589 normalized.finish_reason(),
1590 Some(completion::FinishReason::ToolCalls)
1591 );
1592 }
1593
1594 #[test]
1595 fn streaming_terminal_record_is_normalized() {
1596 let terminal = StreamingCompletionResponse {
1597 model: "llama3.2".to_string(),
1598 done_reason: Some("dragons".to_string()),
1599 total_duration: None,
1600 load_duration: None,
1601 prompt_eval_count: Some(7),
1602 prompt_eval_duration: None,
1603 eval_count: Some(5),
1604 eval_duration: None,
1605 };
1606
1607 let final_record = StreamFinal::from(terminal);
1608 assert_eq!(final_record.provider, PROVIDER_NAME);
1609 assert_eq!(final_record.model.as_deref(), Some("llama3.2"));
1610 assert_eq!(
1611 final_record.finish_reason,
1612 Some(completion::FinishReason::Other("dragons".to_owned()))
1613 );
1614 assert_eq!(final_record.usage.total_tokens, 12);
1615 }
1616
1617 #[test]
1619 fn test_message_conversion() {
1620 let provider_msg = Message::User {
1622 content: "Test message".to_owned(),
1623 images: None,
1624 name: None,
1625 };
1626 let comp_msg: crate::completion::Message = provider_msg.into();
1628 match comp_msg {
1629 crate::completion::Message::User { content } => {
1630 let first_content = content.first();
1631 match first_content {
1633 Some(crate::completion::message::UserContent::Text(text_struct)) => {
1634 assert_eq!(text_struct.text, "Test message");
1635 }
1636 _ => panic!("Expected text content in conversion"),
1637 }
1638 }
1639 _ => panic!("Conversion from provider Message to completion Message failed"),
1640 }
1641 }
1642
1643 #[test]
1644 fn empty_assistant_history_converts_to_empty_content_not_a_sentinel() {
1645 let provider_msg = Message::Assistant {
1651 content: String::new(),
1652 thinking: None,
1653 images: None,
1654 name: None,
1655 tool_calls: Vec::new(),
1656 };
1657 let comp_msg: crate::completion::Message = provider_msg.into();
1658 match comp_msg {
1659 crate::completion::Message::Assistant { content, .. } => {
1660 assert!(content.is_empty(), "expected empty content: {content:?}");
1661 }
1662 other => panic!("expected an assistant message, got {other:?}"),
1663 }
1664
1665 let provider_msg = Message::Assistant {
1667 content: "hello".to_owned(),
1668 thinking: None,
1669 images: None,
1670 name: None,
1671 tool_calls: Vec::new(),
1672 };
1673 let comp_msg: crate::completion::Message = provider_msg.into();
1674 match comp_msg {
1675 crate::completion::Message::Assistant { content, .. } => {
1676 assert!(
1677 matches!(
1678 content.as_slice(),
1679 [crate::completion::message::AssistantContent::Text(text)]
1680 if text.text == "hello"
1681 ),
1682 "unexpected content: {content:?}"
1683 );
1684 }
1685 other => panic!("expected an assistant message, got {other:?}"),
1686 }
1687 }
1688
1689 #[test]
1690 fn mixed_user_content_preserves_message_order() {
1691 use crate::message::{Message as RigMessage, ToolResultContent, UserContent};
1692
1693 let message = RigMessage::User {
1694 content: vec![
1695 UserContent::text("before"),
1696 UserContent::tool_result(
1697 "",
1698 "lookup",
1699 vec![ToolResultContent::json(json!({ "ok": true }))],
1700 ),
1701 UserContent::text("after"),
1702 ],
1703 };
1704
1705 let messages = Vec::<Message>::try_from(message).expect("mixed content should convert");
1706 assert_eq!(messages.len(), 3);
1707 assert!(matches!(
1708 &messages[0],
1709 Message::User { content, .. } if content == "before"
1710 ));
1711 assert!(matches!(
1712 &messages[1],
1713 Message::ToolResult { name, content }
1714 if name == "lookup" && content == r#"{"ok":true}"#
1715 ));
1716 assert!(matches!(
1717 &messages[2],
1718 Message::User { content, .. } if content == "after"
1719 ));
1720 }
1721
1722 #[test]
1723 fn unsupported_user_content_returns_a_conversion_error() {
1724 use crate::message::{ImageMediaType, Message as RigMessage, UserContent};
1725
1726 let message = RigMessage::User {
1727 content: vec![UserContent::image_url(
1728 "https://example.com/image.png",
1729 Some(ImageMediaType::PNG),
1730 None,
1731 )],
1732 };
1733
1734 let error = Vec::<Message>::try_from(message).expect_err("URL image should be rejected");
1735 assert!(error.to_string().contains("base64"));
1736 }
1737
1738 #[test]
1740 fn test_tool_definition_conversion() {
1741 let internal_tool = crate::completion::ToolDefinition {
1743 name: "get_current_weather".to_owned(),
1744 description: "Get the current weather for a location".to_owned(),
1745 parameters: json!({
1746 "type": "object",
1747 "properties": {
1748 "location": {
1749 "type": "string",
1750 "description": "The location to get the weather for, e.g. San Francisco, CA"
1751 },
1752 "format": {
1753 "type": "string",
1754 "description": "The format to return the weather in, e.g. 'celsius' or 'fahrenheit'",
1755 "enum": ["celsius", "fahrenheit"]
1756 }
1757 },
1758 "required": ["location", "format"]
1759 }),
1760 };
1761 let ollama_tool: ToolDefinition = internal_tool.into();
1763 assert_eq!(ollama_tool.type_field, "function");
1764 assert_eq!(ollama_tool.function.name, "get_current_weather");
1765 assert_eq!(
1766 ollama_tool.function.description,
1767 "Get the current weather for a location"
1768 );
1769 let params = &ollama_tool.function.parameters;
1771 assert_eq!(params["properties"]["location"]["type"], "string");
1772 }
1773
1774 #[tokio::test]
1776 async fn test_chat_completion_with_thinking() {
1777 let sample_response = json!({
1778 "model": "qwen-thinking",
1779 "created_at": "2023-08-04T19:22:45.499127Z",
1780 "message": {
1781 "role": "assistant",
1782 "content": "The answer is 42.",
1783 "thinking": "Let me think about this carefully. The question asks for the meaning of life...",
1784 "images": null,
1785 "tool_calls": []
1786 },
1787 "done": true,
1788 "total_duration": 8000000000u64,
1789 "load_duration": 6000000u64,
1790 "prompt_eval_count": 61u64,
1791 "prompt_eval_duration": 400000000u64,
1792 "eval_count": 468u64,
1793 "eval_duration": 7700000000u64
1794 });
1795
1796 let chat_resp: CompletionResponse =
1797 serde_json::from_value(sample_response).expect("Failed to deserialize");
1798
1799 if let Message::Assistant {
1801 thinking, content, ..
1802 } = &chat_resp.message
1803 {
1804 assert_eq!(
1805 thinking.as_ref().unwrap(),
1806 "Let me think about this carefully. The question asks for the meaning of life..."
1807 );
1808 assert_eq!(content, "The answer is 42.");
1809 } else {
1810 panic!("Expected Assistant message");
1811 }
1812 }
1813
1814 #[tokio::test]
1816 async fn test_chat_completion_without_thinking() {
1817 let sample_response = json!({
1818 "model": "llama3.2",
1819 "created_at": "2023-08-04T19:22:45.499127Z",
1820 "message": {
1821 "role": "assistant",
1822 "content": "Hello!",
1823 "images": null,
1824 "tool_calls": []
1825 },
1826 "done": true,
1827 "total_duration": 8000000000u64,
1828 "load_duration": 6000000u64,
1829 "prompt_eval_count": 10u64,
1830 "prompt_eval_duration": 400000000u64,
1831 "eval_count": 5u64,
1832 "eval_duration": 7700000000u64
1833 });
1834
1835 let chat_resp: CompletionResponse =
1836 serde_json::from_value(sample_response).expect("Failed to deserialize");
1837
1838 if let Message::Assistant {
1840 thinking, content, ..
1841 } = &chat_resp.message
1842 {
1843 assert!(thinking.is_none());
1844 assert_eq!(content, "Hello!");
1845 } else {
1846 panic!("Expected Assistant message");
1847 }
1848 }
1849
1850 #[test]
1852 fn test_streaming_response_with_thinking() {
1853 let sample_chunk = json!({
1854 "model": "qwen-thinking",
1855 "created_at": "2023-08-04T19:22:45.499127Z",
1856 "message": {
1857 "role": "assistant",
1858 "content": "",
1859 "thinking": "Analyzing the problem...",
1860 "images": null,
1861 "tool_calls": []
1862 },
1863 "done": false
1864 });
1865
1866 let chunk: CompletionResponse =
1867 serde_json::from_value(sample_chunk).expect("Failed to deserialize");
1868
1869 if let Message::Assistant {
1870 thinking, content, ..
1871 } = &chunk.message
1872 {
1873 assert_eq!(thinking.as_ref().unwrap(), "Analyzing the problem...");
1874 assert_eq!(content, "");
1875 } else {
1876 panic!("Expected Assistant message");
1877 }
1878 }
1879
1880 #[test]
1882 fn test_message_conversion_with_thinking() {
1883 let reasoning_content = crate::message::Reasoning::new("Step 1: Consider the problem");
1885
1886 let internal_msg = crate::message::Message::Assistant {
1887 id: None,
1888 content: vec![
1889 crate::message::AssistantContent::Reasoning(reasoning_content),
1890 crate::message::AssistantContent::Text(crate::message::Text::new(
1891 "The answer is X".to_string(),
1892 )),
1893 ],
1894 };
1895
1896 let provider_msgs: Vec<Message> = internal_msg.try_into().unwrap();
1898 assert_eq!(provider_msgs.len(), 1);
1899
1900 if let Message::Assistant {
1901 thinking, content, ..
1902 } = &provider_msgs[0]
1903 {
1904 assert_eq!(thinking.as_ref().unwrap(), "Step 1: Consider the problem");
1905 assert_eq!(content, "The answer is X");
1906 } else {
1907 panic!("Expected Assistant message with thinking");
1908 }
1909 }
1910
1911 #[test]
1915 fn wire_message_conversion_preserves_the_daemon_tool_call_id() {
1916 let wire = Message::Assistant {
1917 content: String::new(),
1918 thinking: None,
1919 images: None,
1920 name: None,
1921 tool_calls: vec![ToolCall {
1922 id: Some("call_abc".to_owned()),
1923 r#type: ToolType::default(),
1924 function: Function {
1925 name: "get_weather".to_owned(),
1926 arguments: json!({}),
1927 },
1928 }],
1929 };
1930
1931 let converted: crate::completion::Message = wire.into();
1932 let crate::completion::Message::Assistant { content, .. } = converted else {
1933 panic!("Expected Assistant message");
1934 };
1935 let ids: Vec<String> = content
1936 .iter()
1937 .filter_map(|item| match item {
1938 crate::message::AssistantContent::ToolCall(call) => {
1939 Some(call.id.as_str().to_owned())
1940 }
1941 _ => None,
1942 })
1943 .collect();
1944 assert_eq!(ids, vec!["call_abc".to_owned()]);
1945 }
1946
1947 #[tokio::test]
1954 async fn nonstreaming_response_preserves_thinking_as_reasoning() {
1955 let sample_response = json!({
1956 "model": "qwen3:4b",
1957 "created_at": "2023-08-04T19:22:45.499127Z",
1958 "message": {
1959 "role": "assistant",
1960 "content": "",
1961 "thinking": "The user asked for the weather in Berlin. I should call get_weather with location=Berlin.",
1962 "images": null,
1963 "tool_calls": [
1964 { "type": "function", "function": { "name": "get_weather", "arguments": { "location": "Berlin" } } }
1965 ]
1966 },
1967 "done": true,
1968 "done_reason": "stop",
1969 "total_duration": 8000000000u64,
1970 "load_duration": 6000000u64,
1971 "prompt_eval_count": 61u64,
1972 "prompt_eval_duration": 400000000u64,
1973 "eval_count": 468u64,
1974 "eval_duration": 7700000000u64
1975 });
1976
1977 let raw: CompletionResponse =
1978 serde_json::from_value(sample_response).expect("deserialize ollama response");
1979 let completed: completion::CompletionResponse =
1980 raw.try_into().expect("convert to completion response");
1981
1982 let reasoning = completed.choice.iter().find_map(|c| match c {
1983 completion::AssistantContent::Reasoning(r) => Some(r.clone()),
1984 _ => None,
1985 });
1986 let has_tool_call = completed
1987 .choice
1988 .iter()
1989 .any(|c| matches!(c, completion::AssistantContent::ToolCall(_)));
1990
1991 assert!(has_tool_call, "tool call should survive the conversion");
1992 let reasoning = reasoning.expect(
1993 "non-streaming response must surface `thinking` as AssistantContent::Reasoning (issue #1926)",
1994 );
1995 assert_eq!(
1996 reasoning.display_text(),
1997 "The user asked for the weather in Berlin. I should call get_weather with location=Berlin.",
1998 );
1999 }
2000
2001 #[test]
2003 fn test_empty_thinking_content() {
2004 let sample_response = json!({
2005 "model": "llama3.2",
2006 "created_at": "2023-08-04T19:22:45.499127Z",
2007 "message": {
2008 "role": "assistant",
2009 "content": "Response",
2010 "thinking": "",
2011 "images": null,
2012 "tool_calls": []
2013 },
2014 "done": true,
2015 "total_duration": 8000000000u64,
2016 "load_duration": 6000000u64,
2017 "prompt_eval_count": 10u64,
2018 "prompt_eval_duration": 400000000u64,
2019 "eval_count": 5u64,
2020 "eval_duration": 7700000000u64
2021 });
2022
2023 let chat_resp: CompletionResponse =
2024 serde_json::from_value(sample_response).expect("Failed to deserialize");
2025
2026 if let Message::Assistant {
2027 thinking, content, ..
2028 } = &chat_resp.message
2029 {
2030 assert_eq!(thinking.as_ref().unwrap(), "");
2032 assert_eq!(content, "Response");
2033 } else {
2034 panic!("Expected Assistant message");
2035 }
2036 }
2037
2038 #[test]
2040 fn test_thinking_with_tool_calls() {
2041 let sample_response = json!({
2042 "model": "qwen-thinking",
2043 "created_at": "2023-08-04T19:22:45.499127Z",
2044 "message": {
2045 "role": "assistant",
2046 "content": "Let me check the weather.",
2047 "thinking": "User wants weather info, I should use the weather tool",
2048 "images": null,
2049 "tool_calls": [
2050 {
2051 "type": "function",
2052 "function": {
2053 "name": "get_weather",
2054 "arguments": {
2055 "location": "San Francisco"
2056 }
2057 }
2058 }
2059 ]
2060 },
2061 "done": true,
2062 "total_duration": 8000000000u64,
2063 "load_duration": 6000000u64,
2064 "prompt_eval_count": 30u64,
2065 "prompt_eval_duration": 400000000u64,
2066 "eval_count": 50u64,
2067 "eval_duration": 7700000000u64
2068 });
2069
2070 let chat_resp: CompletionResponse =
2071 serde_json::from_value(sample_response).expect("Failed to deserialize");
2072
2073 if let Message::Assistant {
2074 thinking,
2075 content,
2076 tool_calls,
2077 ..
2078 } = &chat_resp.message
2079 {
2080 assert_eq!(
2081 thinking.as_ref().unwrap(),
2082 "User wants weather info, I should use the weather tool"
2083 );
2084 assert_eq!(content, "Let me check the weather.");
2085 assert_eq!(tool_calls.len(), 1);
2086 assert_eq!(tool_calls[0].function.name, "get_weather");
2087 } else {
2088 panic!("Expected Assistant message with thinking and tool calls");
2089 }
2090 }
2091
2092 #[test]
2094 fn test_completion_request_with_think_param() {
2095 use crate::completion::Message as CompletionMessage;
2096 use crate::message::{Text, UserContent};
2097
2098 let completion_request = CompletionRequest {
2100 model: None,
2101 preamble: Some("You are a helpful assistant.".to_string()),
2102 chat_history: vec![CompletionMessage::User {
2103 content: vec![UserContent::Text(Text::new("What is 2 + 2?".to_string()))],
2104 }],
2105 documents: vec![],
2106 tools: vec![],
2107 temperature: Some(0.7),
2108 max_tokens: Some(1024),
2109 tool_choice: None,
2110 additional_params: Some(json!({
2111 "think": true,
2112 "keep_alive": "-1m",
2113 "num_ctx": 4096
2114 })),
2115 output_schema: None,
2116 record_telemetry_content: false,
2117 };
2118
2119 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
2121 .expect("Failed to create Ollama request");
2122
2123 let serialized =
2125 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2126
2127 let expected = json!({
2133 "model": "qwen3:8b",
2134 "messages": [
2135 {
2136 "role": "system",
2137 "content": "You are a helpful assistant."
2138 },
2139 {
2140 "role": "user",
2141 "content": "What is 2 + 2?"
2142 }
2143 ],
2144 "stream": false,
2145 "think": true,
2146 "keep_alive": "-1m",
2147 "options": {
2148 "temperature": 0.7,
2149 "num_predict": 1024,
2150 "num_ctx": 4096
2151 }
2152 });
2153
2154 assert_eq!(serialized, expected);
2155 }
2156
2157 #[test]
2159 fn test_completion_request_with_level_low_think_param() {
2160 use crate::completion::Message as CompletionMessage;
2161 use crate::message::{Text, UserContent};
2162
2163 let completion_request = CompletionRequest {
2165 model: None,
2166 preamble: Some("You are a helpful assistant.".to_string()),
2167 chat_history: vec![CompletionMessage::User {
2168 content: vec![UserContent::Text(Text::new("What is 2 + 2?".to_string()))],
2169 }],
2170 documents: vec![],
2171 tools: vec![],
2172 temperature: Some(0.7),
2173 max_tokens: Some(1024),
2174 tool_choice: None,
2175 additional_params: Some(json!({
2176 "think": "low",
2177 "keep_alive": "-1m",
2178 "num_ctx": 4096
2179 })),
2180 output_schema: None,
2181 record_telemetry_content: false,
2182 };
2183
2184 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
2186 .expect("Failed to create Ollama request");
2187
2188 let serialized =
2190 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2191
2192 let expected = json!({
2198 "model": "qwen3:8b",
2199 "messages": [
2200 {
2201 "role": "system",
2202 "content": "You are a helpful assistant."
2203 },
2204 {
2205 "role": "user",
2206 "content": "What is 2 + 2?"
2207 }
2208 ],
2209 "stream": false,
2210 "think": "low",
2211 "keep_alive": "-1m",
2212 "options": {
2213 "temperature": 0.7,
2214 "num_predict": 1024,
2215 "num_ctx": 4096
2216 }
2217 });
2218
2219 assert_eq!(serialized, expected);
2220 }
2221
2222 #[test]
2224 fn test_completion_request_with_level_medium_think_param() {
2225 use crate::completion::Message as CompletionMessage;
2226 use crate::message::{Text, UserContent};
2227
2228 let completion_request = CompletionRequest {
2230 model: None,
2231 preamble: Some("You are a helpful assistant.".to_string()),
2232 chat_history: vec![CompletionMessage::User {
2233 content: vec![UserContent::Text(Text::new("What is 2 + 2?".to_string()))],
2234 }],
2235 documents: vec![],
2236 tools: vec![],
2237 temperature: Some(0.7),
2238 max_tokens: Some(1024),
2239 tool_choice: None,
2240 additional_params: Some(json!({
2241 "think": "medium",
2242 "keep_alive": "-1m",
2243 "num_ctx": 4096
2244 })),
2245 output_schema: None,
2246 record_telemetry_content: false,
2247 };
2248
2249 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
2251 .expect("Failed to create Ollama request");
2252
2253 let serialized =
2255 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2256
2257 let expected = json!({
2263 "model": "qwen3:8b",
2264 "messages": [
2265 {
2266 "role": "system",
2267 "content": "You are a helpful assistant."
2268 },
2269 {
2270 "role": "user",
2271 "content": "What is 2 + 2?"
2272 }
2273 ],
2274 "stream": false,
2275 "think": "medium",
2276 "keep_alive": "-1m",
2277 "options": {
2278 "temperature": 0.7,
2279 "num_predict": 1024,
2280 "num_ctx": 4096
2281 }
2282 });
2283
2284 assert_eq!(serialized, expected);
2285 }
2286
2287 #[test]
2289 fn test_completion_request_with_level_high_think_param() {
2290 use crate::completion::Message as CompletionMessage;
2291 use crate::message::{Text, UserContent};
2292
2293 let completion_request = CompletionRequest {
2295 model: None,
2296 preamble: Some("You are a helpful assistant.".to_string()),
2297 chat_history: vec![CompletionMessage::User {
2298 content: vec![UserContent::Text(Text::new("What is 2 + 2?".to_string()))],
2299 }],
2300 documents: vec![],
2301 tools: vec![],
2302 temperature: Some(0.7),
2303 max_tokens: Some(1024),
2304 tool_choice: None,
2305 additional_params: Some(json!({
2306 "think": "high",
2307 "keep_alive": "-1m",
2308 "num_ctx": 4096
2309 })),
2310 output_schema: None,
2311 record_telemetry_content: false,
2312 };
2313
2314 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request))
2316 .expect("Failed to create Ollama request");
2317
2318 let serialized =
2320 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2321
2322 let expected = json!({
2328 "model": "qwen3:8b",
2329 "messages": [
2330 {
2331 "role": "system",
2332 "content": "You are a helpful assistant."
2333 },
2334 {
2335 "role": "user",
2336 "content": "What is 2 + 2?"
2337 }
2338 ],
2339 "stream": false,
2340 "think": "high",
2341 "keep_alive": "-1m",
2342 "options": {
2343 "temperature": 0.7,
2344 "num_predict": 1024,
2345 "num_ctx": 4096
2346 }
2347 });
2348
2349 assert_eq!(serialized, expected);
2350 }
2351
2352 #[test]
2354 fn test_completion_request_with_level_invalid_think_param() {
2355 use crate::completion::Message as CompletionMessage;
2356 use crate::message::{Text, UserContent};
2357
2358 let completion_request = CompletionRequest {
2360 model: None,
2361 preamble: Some("You are a helpful assistant.".to_string()),
2362 chat_history: vec![CompletionMessage::User {
2363 content: vec![UserContent::Text(Text::new("What is 2 + 2?".to_string()))],
2364 }],
2365 documents: vec![],
2366 tools: vec![],
2367 temperature: Some(0.7),
2368 max_tokens: Some(1024),
2369 tool_choice: None,
2370 additional_params: Some(json!({
2371 "think": "invalid",
2372 "keep_alive": "-1m",
2373 "num_ctx": 4096
2374 })),
2375 output_schema: None,
2376 record_telemetry_content: false,
2377 };
2378
2379 let ollama_request = OllamaCompletionRequest::try_from(("qwen3:8b", completion_request));
2381
2382 assert!(ollama_request.is_err())
2383 }
2384
2385 #[test]
2388 fn test_completion_request_with_think_omitted_by_default() {
2389 use crate::completion::Message as CompletionMessage;
2390 use crate::message::{Text, UserContent};
2391
2392 let completion_request = CompletionRequest {
2394 model: None,
2395 preamble: Some("You are a helpful assistant.".to_string()),
2396 chat_history: vec![CompletionMessage::User {
2397 content: vec![UserContent::Text(Text::new("Hello!".to_string()))],
2398 }],
2399 documents: vec![],
2400 tools: vec![],
2401 temperature: Some(0.5),
2402 max_tokens: None,
2403 tool_choice: None,
2404 additional_params: None,
2405 output_schema: None,
2406 record_telemetry_content: false,
2407 };
2408
2409 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2411 .expect("Failed to create Ollama request");
2412
2413 let serialized =
2415 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2416
2417 let expected = json!({
2420 "model": "llama3.2",
2421 "messages": [
2422 {
2423 "role": "system",
2424 "content": "You are a helpful assistant."
2425 },
2426 {
2427 "role": "user",
2428 "content": "Hello!"
2429 }
2430 ],
2431 "stream": false,
2432 "options": {
2433 "temperature": 0.5
2434 }
2435 });
2436
2437 assert_eq!(serialized, expected);
2438 }
2439
2440 #[test]
2444 fn test_completion_request_num_predict_from_additional_params_wins() {
2445 use crate::completion::Message as CompletionMessage;
2446 use crate::message::{Text, UserContent};
2447
2448 let completion_request = CompletionRequest {
2449 model: None,
2450 preamble: None,
2451 chat_history: vec![CompletionMessage::User {
2452 content: vec![UserContent::Text(Text::new("Hello!".to_string()))],
2453 }],
2454 documents: vec![],
2455 tools: vec![],
2456 temperature: None,
2457 max_tokens: Some(1024),
2458 tool_choice: None,
2459 additional_params: Some(json!({ "num_predict": 42 })),
2460 output_schema: None,
2461 record_telemetry_content: false,
2462 };
2463
2464 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2465 .expect("Failed to create Ollama request");
2466 let serialized =
2467 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2468
2469 assert_eq!(serialized["options"], json!({ "num_predict": 42 }));
2470 assert_eq!(serialized.get("max_tokens"), None);
2471 }
2472
2473 #[test]
2478 fn test_completion_request_num_predict_without_additional_params() {
2479 use crate::completion::Message as CompletionMessage;
2480 use crate::message::{Text, UserContent};
2481
2482 let completion_request = CompletionRequest {
2483 model: None,
2484 preamble: None,
2485 chat_history: vec![CompletionMessage::User {
2486 content: vec![UserContent::Text(Text::new("Hello!".to_string()))],
2487 }],
2488 documents: vec![],
2489 tools: vec![],
2490 temperature: Some(0.7),
2491 max_tokens: Some(1024),
2492 tool_choice: None,
2493 additional_params: None,
2494 output_schema: None,
2495 record_telemetry_content: false,
2496 };
2497
2498 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2499 .expect("Failed to create Ollama request");
2500 let serialized =
2501 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2502
2503 assert_eq!(
2504 serialized["options"],
2505 json!({ "temperature": 0.7, "num_predict": 1024 })
2506 );
2507 assert_eq!(serialized.get("max_tokens"), None);
2509 assert_eq!(serialized.get("temperature"), None);
2510 }
2511
2512 #[test]
2516 fn test_completion_request_options_omit_unset_parameters() {
2517 use crate::completion::Message as CompletionMessage;
2518 use crate::message::{Text, UserContent};
2519
2520 let completion_request = CompletionRequest {
2521 model: None,
2522 preamble: None,
2523 chat_history: vec![CompletionMessage::User {
2524 content: vec![UserContent::Text(Text::new("Hello!".to_string()))],
2525 }],
2526 documents: vec![],
2527 tools: vec![],
2528 temperature: None,
2529 max_tokens: None,
2530 tool_choice: None,
2531 additional_params: None,
2532 output_schema: None,
2533 record_telemetry_content: false,
2534 };
2535
2536 let ollama_request = OllamaCompletionRequest::try_from(("llama3.2", completion_request))
2537 .expect("Failed to create Ollama request");
2538 let serialized =
2539 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2540
2541 assert_eq!(serialized["options"], json!({}));
2542 }
2543
2544 #[test]
2545 fn test_completion_request_with_output_schema() {
2546 use crate::completion::Message as CompletionMessage;
2547 use crate::message::{Text, UserContent};
2548
2549 let schema: schemars::Schema = serde_json::from_value(json!({
2550 "type": "object",
2551 "properties": {
2552 "age": { "type": "integer" },
2553 "available": { "type": "boolean" }
2554 },
2555 "required": ["age", "available"]
2556 }))
2557 .expect("Failed to parse schema");
2558
2559 let completion_request = CompletionRequest {
2560 model: Some("llama3.1".to_string()),
2561 preamble: None,
2562 chat_history: vec![CompletionMessage::User {
2563 content: vec![UserContent::Text(Text::new(
2564 "How old is Ollama?".to_string(),
2565 ))],
2566 }],
2567 documents: vec![],
2568 tools: vec![],
2569 temperature: None,
2570 max_tokens: None,
2571 tool_choice: None,
2572 additional_params: None,
2573 output_schema: Some(schema),
2574 record_telemetry_content: false,
2575 };
2576
2577 let ollama_request = OllamaCompletionRequest::try_from(("llama3.1", completion_request))
2578 .expect("Failed to create Ollama request");
2579
2580 let serialized =
2581 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2582
2583 let format = serialized
2584 .get("format")
2585 .expect("format field should be present");
2586 assert_eq!(
2587 *format,
2588 json!({
2589 "type": "object",
2590 "properties": {
2591 "age": { "type": "integer" },
2592 "available": { "type": "boolean" }
2593 },
2594 "required": ["age", "available"]
2595 })
2596 );
2597 }
2598
2599 #[test]
2600 fn test_completion_request_without_output_schema() {
2601 use crate::completion::Message as CompletionMessage;
2602 use crate::message::{Text, UserContent};
2603
2604 let completion_request = CompletionRequest {
2605 model: Some("llama3.1".to_string()),
2606 preamble: None,
2607 chat_history: vec![CompletionMessage::User {
2608 content: vec![UserContent::Text(Text::new("Hello!".to_string()))],
2609 }],
2610 documents: vec![],
2611 tools: vec![],
2612 temperature: None,
2613 max_tokens: None,
2614 tool_choice: None,
2615 additional_params: None,
2616 output_schema: None,
2617 record_telemetry_content: false,
2618 };
2619
2620 let ollama_request = OllamaCompletionRequest::try_from(("llama3.1", completion_request))
2621 .expect("Failed to create Ollama request");
2622
2623 let serialized =
2624 serde_json::to_value(&ollama_request).expect("Failed to serialize request");
2625
2626 assert!(
2627 serialized.get("format").is_none(),
2628 "format field should be absent when output_schema is None"
2629 );
2630 }
2631
2632 #[test]
2633 fn test_client_initialization() {
2634 let _client = crate::providers::ollama::Client::new(Nothing).expect("Client::new() failed");
2635 let _client_from_builder = crate::providers::ollama::Client::builder()
2636 .api_key(Nothing)
2637 .build()
2638 .expect("Client::builder() failed");
2639 }
2640
2641 #[test]
2642 fn ndjson_buffer_returns_complete_lines_in_single_chunk() {
2643 let mut buf = NdjsonBuffer::new();
2644 let lines = buf.decode(b"{\"a\":1}\n{\"b\":2}\n");
2645 assert_eq!(lines, vec![b"{\"a\":1}".to_vec(), b"{\"b\":2}".to_vec()]);
2646 }
2647
2648 #[test]
2649 fn ndjson_buffer_reassembles_line_split_across_chunks() {
2650 let mut buf = NdjsonBuffer::new();
2651
2652 assert!(buf.decode(b"{\"model\":\"llama\",\"mes").is_empty());
2653
2654 let lines = buf.decode(b"sage\":\"hi\"}\n{\"done\"");
2655 assert_eq!(
2656 lines,
2657 vec![b"{\"model\":\"llama\",\"message\":\"hi\"}".to_vec()]
2658 );
2659
2660 let lines = buf.decode(b":true}\n");
2661 assert_eq!(lines, vec![b"{\"done\":true}".to_vec()]);
2662 }
2663
2664 #[test]
2665 fn ndjson_buffer_skips_blank_lines() {
2666 let mut buf = NdjsonBuffer::new();
2667 let lines = buf.decode(b"\n{\"a\":1}\n\n");
2668 assert_eq!(lines, vec![b"{\"a\":1}".to_vec()]);
2669 }
2670
2671 #[test]
2672 fn ndjson_buffer_retains_unterminated_trailing_data() {
2673 let mut buf = NdjsonBuffer::new();
2674 let lines = buf.decode(b"{\"a\":1}\n{\"b\":2");
2675 assert_eq!(lines, vec![b"{\"a\":1}".to_vec()]);
2676 let lines = buf.decode(b"}\n");
2677 assert_eq!(lines, vec![b"{\"b\":2}".to_vec()]);
2678 }
2679
2680 #[test]
2681 fn ndjson_buffer_handles_empty_chunk() {
2682 let mut buf = NdjsonBuffer::new();
2683 assert!(buf.decode(b"").is_empty());
2684
2685 buf.decode(b"{\"a\":1");
2686 assert!(buf.decode(b"").is_empty());
2687
2688 let lines = buf.decode(b"}\n");
2689 assert_eq!(lines, vec![b"{\"a\":1}".to_vec()]);
2690 }
2691
2692 #[test]
2693 fn ndjson_buffer_handles_multi_byte_utf8_split_across_chunks() {
2694 let mut buf = NdjsonBuffer::new();
2698 assert!(buf.decode(&[0xd0]).is_empty());
2699 assert!(buf.decode(&[0xb8, 0xd0, 0xb7, 0xd0]).is_empty());
2700 assert!(
2701 buf.decode(&[
2702 0xb2, 0xd0, 0xb5, 0xd1, 0x81, 0xd1, 0x82, 0xd0, 0xbd, 0xd0, 0xb8
2703 ])
2704 .is_empty()
2705 );
2706
2707 let lines = buf.decode(b"\n");
2708 assert_eq!(lines.len(), 1);
2709 assert_eq!(std::str::from_utf8(&lines[0]).unwrap(), "известни");
2710 }
2711
2712 #[test]
2713 fn ndjson_buffer_yields_parseable_chunks_when_split_arbitrarily() {
2714 let original = concat!(
2715 "{\"model\":\"llama3.2\",\"message\":{\"role\":\"assistant\",\"content\":\"hi\"},\"done\":false}\n",
2716 "{\"model\":\"llama3.2\",\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"done\":true}\n",
2717 );
2718
2719 let mut buf = NdjsonBuffer::new();
2720 let mut received = Vec::new();
2721 for byte in original.as_bytes() {
2722 for line in buf.decode(std::slice::from_ref(byte)) {
2723 let parsed: serde_json::Value =
2724 serde_json::from_slice(&line).expect("each drained line must be valid JSON");
2725 received.push(parsed);
2726 }
2727 }
2728
2729 assert_eq!(received.len(), 2);
2730 assert_eq!(received[0]["message"]["content"], "hi");
2731 assert_eq!(received[1]["done"], true);
2732 }
2733
2734 #[tokio::test]
2738 async fn truncated_stream_does_not_synthesize_a_terminal_record() {
2739 use crate::client::CompletionClient;
2740 use crate::completion::CompletionModel;
2741 use crate::streaming::StreamedAssistantContent;
2742 use crate::test_utils::MockStreamingClient;
2743 use futures::StreamExt;
2744
2745 let ndjson = concat!(
2746 r#"{"model":"llama3.2","created_at":"2023-08-04T19:22:45.499127Z","message":{"role":"assistant","content":"hi"},"done":false}"#,
2747 "\n",
2748 );
2749 let client = Client::builder()
2750 .api_key("test-key")
2751 .http_client(MockStreamingClient {
2752 sse_bytes: bytes::Bytes::from(ndjson),
2753 })
2754 .build()
2755 .expect("build client");
2756 let model = client.completion_model(LLAMA3_2);
2757 let request = model.completion_request("hello").build();
2758
2759 let mut stream = model.stream(request).await.expect("stream should open");
2760
2761 let mut texts = Vec::new();
2762 let mut saw_terminal = false;
2763 while let Some(item) = stream.next().await {
2764 match item.expect("stream item should be Ok") {
2765 StreamedAssistantContent::Text(text) => texts.push(text.text),
2766 StreamedAssistantContent::Final(_) => saw_terminal = true,
2767 _ => {}
2768 }
2769 }
2770
2771 assert_eq!(texts, ["hi"]);
2772 assert!(
2773 !saw_terminal,
2774 "EOF without a done record must not synthesize a terminal record"
2775 );
2776 assert!(stream.response.is_none());
2777 }
2778
2779 #[tokio::test]
2783 async fn malformed_line_is_surfaced_and_the_terminal_still_arrives() {
2784 use crate::client::CompletionClient;
2785 use crate::completion::CompletionModel;
2786 use crate::streaming::StreamedAssistantContent;
2787 use crate::test_utils::MockStreamingClient;
2788 use futures::StreamExt;
2789
2790 let ndjson = concat!(
2791 r#"{"model":"llama3.2","created_at":"2023-08-04T19:22:45.499127Z","message":{"role":"assistant","content":"hi"},"done":false}"#,
2792 "\n",
2793 "{not json\n",
2794 r#"{"model":"llama3.2","created_at":"2023-08-04T19:22:46.499127Z","message":{"role":"assistant","content":" there"},"done":false}"#,
2795 "\n",
2796 r#"{"model":"llama3.2","created_at":"2023-08-04T19:22:47.499127Z","message":{"role":"assistant","content":""},"done":true,"done_reason":"stop","prompt_eval_count":10,"eval_count":4}"#,
2797 "\n",
2798 );
2799 let client = Client::builder()
2800 .api_key("test-key")
2801 .http_client(MockStreamingClient {
2802 sse_bytes: bytes::Bytes::from(ndjson),
2803 })
2804 .build()
2805 .expect("build client");
2806 let model = client.completion_model(LLAMA3_2);
2807 let request = model.completion_request("hello").build();
2808
2809 let mut stream = model.stream(request).await.expect("stream should open");
2810
2811 let mut texts = Vec::new();
2812 let mut saw_error = false;
2813 let mut terminal = None;
2814 while let Some(item) = stream.next().await {
2815 match item {
2816 Ok(StreamedAssistantContent::Text(text)) => texts.push(text.text),
2817 Ok(StreamedAssistantContent::Final(final_response)) => {
2818 terminal = Some(final_response)
2819 }
2820 Ok(_) => {}
2821 Err(_) => saw_error = true,
2822 }
2823 }
2824
2825 assert_eq!(texts, ["hi", " there"]);
2826 assert!(saw_error, "the malformed line must reach the consumer");
2827 let terminal = terminal.expect("the genuine done record must still arrive");
2828 assert_eq!(terminal.usage.input_tokens, 10);
2829 assert_eq!(terminal.usage.output_tokens, 4);
2830 }
2831
2832 #[tokio::test]
2836 async fn content_after_the_done_record_is_not_yielded() {
2837 use crate::client::CompletionClient;
2838 use crate::completion::CompletionModel;
2839 use crate::streaming::StreamedAssistantContent;
2840 use crate::test_utils::MockStreamingClient;
2841 use futures::StreamExt;
2842
2843 let ndjson = concat!(
2844 r#"{"model":"llama3.2","created_at":"2023-08-04T19:22:45.499127Z","message":{"role":"assistant","content":"hi"},"done":false}"#,
2845 "\n",
2846 r#"{"model":"llama3.2","created_at":"2023-08-04T19:22:46.499127Z","message":{"role":"assistant","content":""},"done":true,"done_reason":"stop","prompt_eval_count":10,"eval_count":4}"#,
2847 "\n",
2848 r#"{"model":"llama3.2","created_at":"2023-08-04T19:22:47.499127Z","message":{"role":"assistant","content":"stray"},"done":false}"#,
2849 "\n",
2850 );
2851 let client = Client::builder()
2852 .api_key("test-key")
2853 .http_client(MockStreamingClient {
2854 sse_bytes: bytes::Bytes::from(ndjson),
2855 })
2856 .build()
2857 .expect("build client");
2858 let model = client.completion_model(LLAMA3_2);
2859 let request = model.completion_request("hello").build();
2860
2861 let mut stream = model.stream(request).await.expect("stream should open");
2862
2863 let mut texts = Vec::new();
2864 let mut terminal = None;
2865 while let Some(item) = stream.next().await {
2866 match item.expect("stream item should be Ok") {
2867 StreamedAssistantContent::Text(text) => texts.push(text.text),
2868 StreamedAssistantContent::Final(final_response) => {
2869 assert!(
2870 terminal.is_none(),
2871 "the terminal record must be yielded exactly once"
2872 );
2873 terminal = Some(final_response);
2874 }
2875 other => panic!("unexpected stream item: {other:?}"),
2876 }
2877 }
2878
2879 assert_eq!(
2880 texts,
2881 ["hi"],
2882 "content after the done record must not be yielded"
2883 );
2884 let terminal = terminal.expect("the done record must yield the terminal record");
2885 assert_eq!(terminal.usage.input_tokens, 10);
2886 assert_eq!(terminal.usage.output_tokens, 4);
2887 }
2888
2889 #[tokio::test]
2893 async fn completion_non_success_preserves_status_and_body() {
2894 use crate::client::CompletionClient;
2895 use crate::completion::CompletionModel;
2896 use crate::test_utils::RecordingHttpClient;
2897
2898 let body = r#"{"error":"model not found"}"#;
2899 let http_client =
2900 RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
2901 let client = Client::builder()
2902 .api_key("test-key")
2903 .http_client(http_client)
2904 .build()
2905 .expect("build client");
2906 let model = client.completion_model(LLAMA3_2);
2907 let request = model.completion_request("hello").build();
2908
2909 let error = model
2910 .completion(request)
2911 .await
2912 .expect_err("should fail with non-success status");
2913
2914 assert!(matches!(error, CompletionError::HttpError(_)));
2915 assert_eq!(
2916 error.provider_response_status(),
2917 Some(http::StatusCode::SERVICE_UNAVAILABLE)
2918 );
2919 assert_eq!(error.provider_response_body(), Some(body));
2920 }
2921
2922 #[tokio::test]
2926 async fn embeddings_non_success_preserves_status_and_body() {
2927 use crate::client::EmbeddingsClient;
2928 use crate::embeddings::EmbeddingModel;
2929 use crate::test_utils::RecordingHttpClient;
2930
2931 let body = r#"{"error":"model not found"}"#;
2932 let http_client =
2933 RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
2934 let client = Client::builder()
2935 .api_key("test-key")
2936 .http_client(http_client)
2937 .build()
2938 .expect("build client");
2939 let model = client.embedding_model(ALL_MINILM);
2940
2941 let error = model
2942 .embed_texts(vec!["hello".to_string()])
2943 .await
2944 .expect_err("should fail with non-success status");
2945
2946 assert!(matches!(error, EmbeddingError::HttpError(_)));
2947 assert_eq!(
2948 error.provider_response_status(),
2949 Some(http::StatusCode::SERVICE_UNAVAILABLE)
2950 );
2951 assert_eq!(error.provider_response_body(), Some(body));
2952 }
2953
2954 mod raw_capture {
2962 use super::*;
2963 use crate::client::CompletionClient;
2964 use crate::completion::CompletionModel as _;
2965 use crate::test_utils::RecordingHttpClient;
2966
2967 const BODY: &str = r#"{
2968 "model": "llama3.2",
2969 "created_at": "2023-08-04T19:22:45.499127Z",
2970 "message": {"role": "assistant", "content": "hello"},
2971 "done": true,
2972 "done_reason": "stop",
2973 "total_duration": 5043500667,
2974 "load_duration": 5025959,
2975 "prompt_eval_count": 26,
2976 "prompt_eval_duration": 325953000,
2977 "eval_count": 5,
2978 "eval_duration": 4709213000
2979 }"#;
2980
2981 fn model() -> CompletionModel<RecordingHttpClient> {
2982 let client = Client::builder()
2983 .api_key("test-key")
2984 .http_client(RecordingHttpClient::new(BODY))
2985 .build()
2986 .expect("build client");
2987 client.completion_model(LLAMA3_2)
2988 }
2989
2990 #[tokio::test]
2998 async fn completion_captures_raw_that_round_trips_into_the_wire_type() {
2999 let model = model();
3000
3001 let response = model
3002 .completion(model.completion_request("hello").build())
3003 .await
3004 .expect("completion");
3005
3006 let raw = &response.raw;
3007 let typed: CompletionResponse =
3008 serde_json::from_value(raw.clone()).expect("raw must deserialize");
3009 assert_eq!(
3010 serde_json::to_value(&typed).expect("re-serialize"),
3011 *raw,
3012 "the capture must be exactly what the wire type serializes to"
3013 );
3014 assert_eq!(typed.total_duration, Some(5_043_500_667));
3015 assert_eq!(typed.eval_duration, Some(4_709_213_000));
3016 assert_eq!(raw["total_duration"], 5_043_500_667_u64);
3017 assert_eq!(typed.done_reason.as_deref(), Some("stop"));
3018
3019 let renormalized: completion::CompletionResponse =
3020 typed.try_into().expect("re-normalize the capture");
3021 assert_eq!(response.identity(), renormalized.identity());
3022 assert_eq!(response.finish_reason(), renormalized.finish_reason());
3023 assert_eq!(response.model, renormalized.model);
3024 assert_eq!(response.usage, renormalized.usage);
3025 assert_eq!(response.choice, renormalized.choice);
3026 assert_eq!(
3027 response.finish_reason(),
3028 Some(completion::FinishReason::Stop)
3029 );
3030 assert_eq!(response.model.as_deref(), Some("llama3.2"));
3031 assert_eq!(response.usage.total_tokens, 31);
3032 }
3033 }
3034}