reasonkit/llm/ollama/
client.rs

1use crate::llm::ollama::types::{ChatRequest, ChatResponse, OllamaErrorEnvelope};
2use reqwest::{header, StatusCode};
3use std::time::Duration;
4
5#[derive(Debug, Clone)]
6pub struct OllamaClient {
7    http: reqwest::Client,
8    base_url: String,
9    timeout: Duration,
10}
11
12#[derive(Debug, thiserror::Error)]
13pub enum OllamaClientError {
14    #[error("http error: {0}")]
15    Transport(#[from] reqwest::Error),
16
17    #[error("unexpected status {status}: {body}")]
18    HttpStatus { status: StatusCode, body: String },
19
20    #[error("ollama error: {0}")]
21    Ollama(#[from] OllamaErrorEnvelope),
22
23    #[error("invalid response: {0}")]
24    InvalidResponse(String),
25
26    #[error("streaming not supported by this client; set stream=false")]
27    StreamingNotSupported,
28}
29
30impl OllamaClient {
31    /// `base_url` examples:
32    /// - `http://localhost:11434`
33    /// - `http://127.0.0.1:11434`
34    pub fn new(base_url: impl Into<String>) -> Result<Self, reqwest::Error> {
35        let http = reqwest::Client::builder()
36            .user_agent(concat!(
37                "reasonkit-ollama-client/",
38                env!("CARGO_PKG_VERSION")
39            ))
40            .build()?;
41
42        Ok(Self {
43            http,
44            base_url: base_url.into().trim_end_matches('/').to_string(),
45            timeout: Duration::from_secs(60),
46        })
47    }
48
49    pub fn with_timeout(mut self, timeout: Duration) -> Self {
50        self.timeout = timeout;
51        self
52    }
53
54    fn chat_url(&self) -> String {
55        format!("{}/api/chat", self.base_url)
56    }
57
58    /// Call `/api/chat` with `stream:false`.
59    pub async fn chat(&self, mut req: ChatRequest) -> Result<ChatResponse, OllamaClientError> {
60        if req.stream.is_none() {
61            req.stream = Some(false);
62        }
63        if req.stream != Some(false) {
64            return Err(OllamaClientError::StreamingNotSupported);
65        }
66
67        let resp = self
68            .http
69            .post(self.chat_url())
70            .header(header::ACCEPT, "application/json")
71            .json(&req)
72            .timeout(self.timeout)
73            .send()
74            .await?;
75
76        let status = resp.status();
77        let body = resp.text().await?;
78
79        if !status.is_success() {
80            if let Ok(err_env) = serde_json::from_str::<OllamaErrorEnvelope>(&body) {
81                return Err(OllamaClientError::Ollama(err_env));
82            }
83            return Err(OllamaClientError::HttpStatus { status, body });
84        }
85
86        serde_json::from_str::<ChatResponse>(&body)
87            .map_err(|e| OllamaClientError::InvalidResponse(format!("{e}; body={body}")))
88    }
89}