Commit Diff


commit - 39bde005520a131c0400a5d689e05bb05e4cbf29
commit + 58a690fcd50ec0816b9b0f9d23b57441197e9acb
blob - a0b812570e109c7eedba2c307af82fa9d813c0f5
blob + 34137e38327a1922d4715acfe772acec97a76ebf
--- src/ollama.rs
+++ src/ollama.rs
@@ -2,6 +2,7 @@
 //!
 //! Reference: <https://docs.ollama.com/api/chat>
 
+use std::fmt;
 use std::time::Duration;
 
 use crate::Error;
@@ -10,7 +11,8 @@ use chrono::{DateTime, Utc};
 use derive_builder::Builder;
 use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
-use serde_json::{Value, json};
+use serde_json::{Map, Value, json};
+use thiserror::Error;
 
 /// Convenience Client for the Ollama Chat API
 pub struct Client {
@@ -54,6 +56,16 @@ pub struct APIError {
     pub message: String,
 }
 
+/// Errors that can occur when converting an Ollama API response into a
+/// normalized [`Response`].
+#[derive(Debug, Clone, PartialEq, Eq, Error)]
+pub enum ResponseError {
+    #[error("response did not contain a message")]
+    NoMessage,
+    #[error("response message missing text content")]
+    MissingContent,
+}
+
 /// Messages in Chat API request and response
 #[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
 pub struct Message {
@@ -85,12 +97,12 @@ pub enum ReasoningEffort {
     High,
 }
 
-impl ReasoningEffort {
-    fn as_str(self) -> &'static str {
+impl fmt::Display for ReasoningEffort {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            ReasoningEffort::Low => "low",
-            ReasoningEffort::Medium => "medium",
-            ReasoningEffort::High => "high",
+            ReasoningEffort::Low => write!(f, "low"),
+            ReasoningEffort::Medium => write!(f, "medium"),
+            ReasoningEffort::High => write!(f, "high"),
         }
     }
 }
@@ -141,7 +153,8 @@ impl Client {
 
         if resp.status() == StatusCode::OK {
             let chat: ChatAPIResponse = resp.json().await?;
-            Response::try_from(chat).map_err(Error::UnexpectedResponse)
+            Response::try_from(chat)
+                .map_err(|e| Error::UnexpectedResponse(e.to_string()))
         } else {
             let error = resp.json::<APIError>().await?;
             Err(Error::OllamaError { error })
@@ -149,7 +162,8 @@ impl Client {
     }
 
     fn chat_endpoint(&self) -> String {
-        format!("{}{}", self.base_url, "/api/chat")
+        let base = self.base_url.trim_end_matches('/');
+        format!("{}/api/chat", base)
     }
 }
 
@@ -190,25 +204,23 @@ impl ChatRequest {
             .map(|m| serde_json::to_value(m).expect("Message always serializes"))
             .collect::<Vec<_>>();
 
-        let mut payload = json!({
-            "model": self.model,
-            "messages": messages,
-            "stream": false,
-        });
+        let mut payload = Map::new();
+        payload.insert("model".to_string(), json!(self.model));
+        payload.insert("messages".to_string(), json!(messages));
+        payload.insert("stream".to_string(), json!(false));
 
-        if let Some(temperature) = self.temperature
-            && let Some(obj) = payload.as_object_mut()
-        {
-            obj.insert("options".to_string(), json!({ "temperature": temperature }));
+        if let Some(temperature) = self.temperature {
+            payload.insert(
+                "options".to_string(),
+                json!({ "temperature": temperature }),
+            );
         }
 
-        if let Some(think) = &self.think
-            && let Some(obj) = payload.as_object_mut()
-        {
-            obj.insert("think".to_string(), json!(think.as_str()));
+        if let Some(think) = self.think {
+            payload.insert("think".to_string(), json!(think.to_string()));
         }
 
-        payload
+        Value::Object(payload)
     }
 }
 
@@ -245,15 +257,13 @@ fn default_role() -> Role {
 }
 
 impl TryFrom<ChatAPIResponse> for Response {
-    type Error = String;
+    type Error = ResponseError;
 
     fn try_from(value: ChatAPIResponse) -> Result<Self, Self::Error> {
-        let message = value
-            .message
-            .ok_or_else(|| "response did not contain a message".to_string())?;
+        let message = value.message.ok_or(ResponseError::NoMessage)?;
 
         if message.content.is_empty() {
-            return Err("response message missing text content".to_string());
+            return Err(ResponseError::MissingContent);
         }
 
         let finish_reason = match value.done_reason.as_deref() {
@@ -384,8 +394,8 @@ mod test {
         let err = Response::try_from(resp).expect_err("should error");
 
         assert!(
-            err.contains("did not contain a message"),
-            "unexpected error text: {err}"
+            matches!(err, ResponseError::NoMessage),
+            "unexpected error: {err:?}"
         );
 
         Ok(())