Skip to main content

oai_sdk/
chat.rs

1// Copyright 2026 Cloudflavor GmbH
2
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6
7// http://www.apache.org/licenses/LICENSE-2.0
8
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::client::ModelClient;
16use crate::client::handle_error_response;
17use crate::client::json_lines_stream;
18use crate::error::{OllamaError, Result};
19use serde::{Deserialize, Serialize};
20use std::collections::HashMap;
21use tokio_stream::Stream;
22
23/// Controls the reasoning level for models that support thinking.
24///
25/// Most models (e.g., qwen3, deepseek-r1) accept a boolean: `true` enables
26/// reasoning traces.  GPT-OSS variants require one of `low`, `medium`, or
27/// `high`.
28#[derive(Debug, Clone, Serialize, Deserialize)]
29#[serde(untagged)]
30pub enum ThinkLevel {
31    Bool(bool),
32    Level(String),
33}
34
35impl Default for ThinkLevel {
36    fn default() -> Self {
37        ThinkLevel::Bool(true)
38    }
39}
40
41impl From<bool> for ThinkLevel {
42    fn from(v: bool) -> Self {
43        ThinkLevel::Bool(v)
44    }
45}
46
47impl From<&str> for ThinkLevel {
48    fn from(v: &str) -> Self {
49        ThinkLevel::Level(v.to_string())
50    }
51}
52
53/// Request for chat completion.
54#[derive(Debug, Clone, Serialize, Deserialize, Default)]
55pub struct ChatRequest {
56    pub model: String,
57    pub messages: Vec<Message>,
58    #[serde(default)]
59    pub stream: bool,
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub format: Option<Format>,
62    #[serde(skip_serializing_if = "Option::is_none")]
63    pub options: Option<HashMap<String, serde_json::Value>>,
64    #[serde(skip_serializing_if = "Option::is_none")]
65    pub keep_alive: Option<String>,
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub tools: Option<Vec<Tool>>,
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub think: Option<ThinkLevel>,
70}
71
72/// Format for the response.
73#[derive(Debug, Clone, Serialize, Deserialize)]
74#[serde(untagged)]
75pub enum Format {
76    Json,
77    Schema(serde_json::Value),
78}
79
80/// A message in a chat.
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct Message {
83    pub role: String,
84    pub content: String,
85    #[serde(skip_serializing_if = "Option::is_none")]
86    pub images: Option<Vec<String>>,
87    #[serde(skip_serializing_if = "Option::is_none")]
88    pub tool_calls: Option<Vec<ToolCall>>,
89    #[serde(skip_serializing_if = "Option::is_none")]
90    pub tool_name: Option<String>,
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub thinking: Option<String>,
93}
94
95impl Message {
96    pub fn user(content: impl Into<String>) -> Self {
97        Self {
98            role: "user".to_string(),
99            content: content.into(),
100            images: None,
101            tool_calls: None,
102            tool_name: None,
103            thinking: None,
104        }
105    }
106
107    pub fn assistant(content: impl Into<String>) -> Self {
108        Self {
109            role: "assistant".to_string(),
110            content: content.into(),
111            images: None,
112            tool_calls: None,
113            tool_name: None,
114            thinking: None,
115        }
116    }
117
118    pub fn system(content: impl Into<String>) -> Self {
119        Self {
120            role: "system".to_string(),
121            content: content.into(),
122            images: None,
123            tool_calls: None,
124            tool_name: None,
125            thinking: None,
126        }
127    }
128}
129
130/// A tool that can be used by the model.
131#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct Tool {
133    #[serde(rename = "type")]
134    pub tool_type: String,
135    pub function: ToolFunction,
136}
137
138/// A tool function.
139#[derive(Debug, Clone, Serialize, Deserialize)]
140pub struct ToolFunction {
141    pub name: String,
142    pub description: String,
143    pub parameters: serde_json::Value,
144}
145
146/// A tool call.
147#[derive(Debug, Clone, Serialize, Deserialize)]
148pub struct ToolCall {
149    pub function: ToolCallFunction,
150}
151
152/// A tool call function.
153#[derive(Debug, Clone, Serialize, Deserialize)]
154pub struct ToolCallFunction {
155    pub name: String,
156    pub arguments: HashMap<String, serde_json::Value>,
157}
158
159/// Response for chat completion.
160#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct ChatResponse {
162    pub model: String,
163    pub created_at: String,
164    pub message: Message,
165    pub done: bool,
166    #[serde(skip_serializing_if = "Option::is_none")]
167    pub done_reason: Option<String>,
168    #[serde(default)]
169    pub total_duration: u64,
170    #[serde(default)]
171    pub load_duration: u64,
172    #[serde(default)]
173    pub prompt_eval_count: u32,
174    #[serde(default)]
175    pub prompt_eval_duration: u64,
176    #[serde(default)]
177    pub eval_count: u32,
178    #[serde(default)]
179    pub eval_duration: u64,
180}
181
182impl ModelClient {
183    /// Generate a chat completion.
184    pub async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
185        let url = self
186            .base_url
187            .join("api/chat")
188            .map_err(OllamaError::UrlError)?;
189        let response = self
190            .client
191            .post(url)
192            .json(&request)
193            .send()
194            .await
195            .map_err(OllamaError::RequestError)?;
196
197        self.handle_response(response, Some(&request.model)).await
198    }
199
200    /// Generate a streaming chat completion.
201    pub async fn chat_stream(
202        &self,
203        request: ChatRequest,
204    ) -> Result<impl Stream<Item = Result<ChatResponse>> + '_> {
205        let url = self
206            .base_url
207            .join("api/chat")
208            .map_err(OllamaError::UrlError)?;
209        let response = self
210            .client
211            .post(url)
212            .json(&request)
213            .send()
214            .await
215            .map_err(OllamaError::RequestError)?;
216
217        if !response.status().is_success() {
218            return Err(handle_error_response(response, Some(&request.model)).await);
219        }
220
221        Ok(json_lines_stream(response))
222    }
223}