1use 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#[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#[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#[derive(Debug, Clone, Serialize, Deserialize)]
74#[serde(untagged)]
75pub enum Format {
76 Json,
77 Schema(serde_json::Value),
78}
79
80#[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#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct Tool {
133 #[serde(rename = "type")]
134 pub tool_type: String,
135 pub function: ToolFunction,
136}
137
138#[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#[derive(Debug, Clone, Serialize, Deserialize)]
148pub struct ToolCall {
149 pub function: ToolCallFunction,
150}
151
152#[derive(Debug, Clone, Serialize, Deserialize)]
154pub struct ToolCallFunction {
155 pub name: String,
156 pub arguments: HashMap<String, serde_json::Value>,
157}
158
159#[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 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 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}