reasonkit/llm/ollama/
client.rs1use 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 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 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}