ollama_sdk/types/
chat.rs

1//! Contains all data structures that are particularly used for Ollama Chat API
2
3use std::pin::Pin;
4
5use crate::parser::{GenericStreamParser, StreamEventExt};
6use crate::types::Thinking;
7use crate::Result;
8use bytes::Bytes;
9use futures::Stream;
10use ollama_sdk_macros::FromBytes;
11use serde::{Deserialize, Serialize};
12
13use super::{Role, ThinkingLevel};
14
15/// Represents a chat request to the Ollama API.
16///
17/// This struct is used to send a series of messages to a specified model
18/// and receive a chat completion. It supports both streaming and non-streaming
19/// responses, as well as tool integration.
20#[derive(Serialize, Default, Debug, Clone)]
21pub struct ChatRequest {
22    /// The name of the model to use for the chat completion (e.g., "llama2").
23    pub model: String,
24    /// A list of messages exchanged in the chat.
25    pub messages: Vec<ChatRequestMessage>,
26    /// If `true`, the response will be streamed back as a series of [`ChatStreamEvent`]s.
27    /// If `false` or `None`, a single [`ChatResponse`] will be returned.
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub stream: Option<bool>,
30    /// An optional list of tools that the model can use.
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub tools: Option<Vec<ToolSpec>>,
33    /// Configuration for the model's "thinking" process.
34    #[serde(default)]
35    pub think: Thinking,
36}
37
38/// Represents a single message in a chat request.
39#[derive(Deserialize, Serialize, Debug, Clone)]
40#[serde(untagged)]
41pub enum ChatRequestMessage {
42    /// A standard chat message.
43    Message(RegularChatRequestMessage),
44    /// A tool call result
45    ToolCallResult(ToolCallResultMessage),
46}
47
48/// Represents a standard chat message in a chat request.
49#[derive(Deserialize, Serialize, Debug, Clone)]
50pub struct RegularChatRequestMessage {
51    /// The role of the sender (e.g., `User`, `Assistant`, `System`).
52    pub role: Role,
53    /// The content of the message.
54    pub content: String,
55    /// An optional list of tool calls made by the assistant.
56    #[serde(default)]
57    pub tool_calls: Vec<FunctionalTool>,
58}
59
60impl RegularChatRequestMessage {
61    /// Creates a new [`RegularChatRequestMessage`].
62    pub fn new(role: Role, content: String) -> Self {
63        Self {
64            role,
65            content,
66            tool_calls: Vec::new(),
67        }
68    }
69
70    /// Adds a tool call to the message.
71    pub fn add_tool_call(mut self, tool: FunctionalTool) -> Self {
72        self.tool_calls.push(tool);
73        self
74    }
75
76    /// Converts the [`RegularChatRequestMessage`] into a [`ChatRequestMessage`].
77    pub fn to_chat_request_message(self) -> ChatRequestMessage {
78        ChatRequestMessage::Message(self)
79    }
80}
81
82/// Represents a tool call result message in a chat request.
83#[derive(Deserialize, Serialize, Debug, Clone)]
84pub struct ToolCallResultMessage {
85    /// The role of the sender (e.g., `User`, `Assistant`, `System`).
86    pub role: Role,
87    /// The name of the tool.
88    pub name: String,
89    /// The content of the tool call result.
90    pub content: String,
91    /// The unique identifier for the tool call.
92    pub tool_call_id: String,
93}
94
95impl ToolCallResultMessage {
96    /// Creates a new [`ToolCallResultMessage`].
97    pub fn new(name: String, content: String, tool_call_id: String) -> Self {
98        Self {
99            role: Role::Tool,
100            name,
101            content,
102            tool_call_id,
103        }
104    }
105
106    /// Converts the [`ToolCallResultMessage`] into a [`ChatRequestMessage`].
107    pub fn to_chat_request_message(self) -> ChatRequestMessage {
108        ChatRequestMessage::ToolCallResult(self)
109    }
110}
111
112/// Specifies a tool that the model can use.
113#[derive(Serialize, Debug, Clone)]
114#[serde(tag = "type")]
115pub enum ToolSpec {
116    /// A functional tool definition.
117    Function { function: FunctionalTool },
118}
119
120/// Represents a functional tool that the model can call.
121#[derive(Deserialize, Serialize, Debug, Clone)]
122pub struct FunctionalTool {
123    /// The name of the tool.
124    pub name: String,
125    /// An optional description of the tool's purpose.
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub description: Option<String>,
128    /// The JSON schema for the tool's parameters.
129    pub parameters: serde_json::Value,
130}
131
132/// Represents a chat response from the Ollama API.
133///
134/// This struct is used for non-streaming chat completions.
135#[derive(Deserialize, Serialize, Default, FromBytes, Debug, Clone)]
136pub struct ChatResponse {
137    /// The name of the model that generated the response.
138    pub model: String,
139    /// The timestamp when the response was created.
140    #[serde(default)]
141    pub created_at: String,
142    /// The message content from the model.
143    pub message: ChatResponseMessage,
144    /// Indicates if the chat completion is complete.
145    pub done: bool,
146}
147
148/// Represents a single message in a chat response.
149#[derive(Deserialize, Serialize, Default, Debug, Clone)]
150pub struct ChatResponseMessage {
151    /// The role of the sender (e.g., `User`, `Assistant`).
152    pub role: Role,
153    /// The content of the message.
154    pub content: String,
155    /// The model's internal "thinking" process, if enabled.
156    #[serde(default)]
157    pub thinking: String,
158    // A list of tool calls made by the assistant.
159    #[serde(default)]
160    pub tool_calls: Vec<ToolCall>,
161}
162
163/// Represents a tool call made by the model in a chat response.
164#[derive(Deserialize, Serialize, Debug, Clone)]
165pub struct ToolCall {
166    /// Unique identifier for this tool call.
167    pub id: String,
168    /// The function invocation details.
169    pub function: FunctionInvocation,
170}
171
172/// Represents the details of a function invocation in a tool call.
173#[derive(Deserialize, Serialize, Debug, Clone)]
174pub struct FunctionInvocation {
175    /// An optional index indicating the position of the function in a list.
176    pub index: Option<usize>,
177    /// The name of the function being invoked.
178    pub name: String,
179    /// The arguments passed to the function as a JSON value.
180    pub arguments: serde_json::Value,
181}
182
183/// A simplified chat request for non-streaming responses.
184///
185/// This struct is a convenience wrapper for creating a [`ChatRequest`]
186/// that explicitly requests a non-streaming response.
187#[derive(Serialize, Default, Debug, Clone)]
188pub struct SimpleChatRequest {
189    /// The name of the model to use.
190    pub model: String,
191    /// The messages in the chat history.
192    pub messages: Vec<ChatRequestMessage>,
193    /// Configuration for the model's "thinking" process.
194    pub think: Thinking,
195}
196
197impl SimpleChatRequest {
198    /// Creates a new [`SimpleChatRequest`].
199    pub fn new(model: String) -> Self {
200        Self {
201            model,
202            messages: Vec::new(),
203            think: Thinking::default(),
204        }
205    }
206
207    /// Adds a message to the chat request.
208    pub fn add_message(mut self, message: RegularChatRequestMessage) -> Self {
209        self.messages.push(ChatRequestMessage::Message(message));
210        self
211    }
212
213    /// Adds a tool call result message to the chat request.
214    pub fn add_tool_call_result(mut self, message: ToolCallResultMessage) -> Self {
215        self.messages
216            .push(ChatRequestMessage::ToolCallResult(message));
217        self
218    }
219
220    /// Sets `think` param in the API call to `true`.
221    pub fn enable_thinking(mut self) -> Self {
222        self.think = Thinking::Boolean(true);
223        self
224    }
225
226    /// Sets `think` param in the API call to `false`.
227    pub fn disable_thinking(mut self) -> Self {
228        self.think = Thinking::Boolean(false);
229        self
230    }
231
232    /// Sets `think` param in the API call to specified level (`high`, `medium`, `low`).
233    pub fn set_thinking_level(mut self, level: ThinkingLevel) -> Self {
234        self.think = Thinking::Level(level);
235        self
236    }
237}
238
239/// A simplified chat request for streaming responses.
240///
241/// This struct is a convenience wrapper for creating a [`ChatRequest`]
242/// that explicitly requests a streaming response.
243#[derive(Serialize, Default, Debug, Clone)]
244pub struct StreamingChatRequest {
245    /// The name of the model to use.
246    pub model: String,
247    /// The messages in the chat history.
248    pub messages: Vec<ChatRequestMessage>,
249    /// An optional list of tools that the model can use.
250    #[serde(skip_serializing_if = "Option::is_none")]
251    pub tools: Option<Vec<ToolSpec>>,
252    /// Configuration for the model's "thinking" process.
253    pub think: Thinking,
254}
255
256impl StreamingChatRequest {
257    /// Creates a new [`StreamingChatRequest`].
258    pub fn new(model: String) -> Self {
259        Self {
260            model,
261            messages: Vec::new(),
262            tools: None,
263            think: Thinking::default(),
264        }
265    }
266
267    /// Adds a message to the chat request.
268    pub fn add_message(mut self, message: ChatRequestMessage) -> Self {
269        self.messages.push(message);
270        self
271    }
272
273    /// Adds a [`RegularChatRequestMessage`] to the chat request.
274    pub fn add_regular_message(mut self, message: RegularChatRequestMessage) -> Self {
275        self.messages.push(ChatRequestMessage::Message(message));
276        self
277    }
278
279    /// Adds a tool call result message to the chat request.
280    pub fn add_tool_call_result(mut self, message: ToolCallResultMessage) -> Self {
281        self.messages
282            .push(ChatRequestMessage::ToolCallResult(message));
283        self
284    }
285
286    /// Sets `think` param in the API call to `true`.
287    pub fn enable_thinking(mut self) -> Self {
288        self.think = Thinking::Boolean(true);
289        self
290    }
291
292    /// Sets `think` param in the API call to `false`.
293    pub fn disable_thinking(mut self) -> Self {
294        self.think = Thinking::Boolean(false);
295        self
296    }
297
298    /// Sets `think` param in the API call to specified level (`high`, `medium`, `low`).
299    pub fn set_thinking_level(mut self, level: ThinkingLevel) -> Self {
300        self.think = Thinking::Level(level);
301        self
302    }
303
304    /// Sets the tools for the streaming chat request.
305    pub fn tools(mut self, tools: Vec<ToolSpec>) -> Self {
306        self.tools = Some(tools);
307        self
308    }
309}
310
311impl From<SimpleChatRequest> for ChatRequest {
312    /// Converts a [`SimpleChatRequest`] into a [`ChatRequest`] with `stream` set to `false`.
313    fn from(value: SimpleChatRequest) -> Self {
314        ChatRequest {
315            model: value.model,
316            messages: value.messages,
317            stream: Some(false),
318            think: value.think,
319            tools: None,
320        }
321    }
322}
323
324impl From<StreamingChatRequest> for ChatRequest {
325    /// Converts a [`StreamingChatRequest`] into a [`ChatRequest`] with `stream` set to `true`.
326    fn from(value: StreamingChatRequest) -> Self {
327        ChatRequest {
328            model: value.model,
329            messages: value.messages,
330            stream: Some(true),
331            think: value.think,
332            tools: value.tools,
333        }
334    }
335}
336
337/// Represents an event received from a streaming chat response.
338#[derive(Deserialize, Serialize, Debug)]
339pub enum ChatStreamEvent {
340    /// A complete chat response message.
341    Message(ChatResponse),
342    /// An error occurred during the streaming process.
343    Error(String),
344    /// A partial response, returned when the content was un-parseable
345    Partial {
346        /// The un-parseable content.
347        partial: String,
348        /// An optional error message associated with the partial response.
349        error: Option<String>,
350    },
351}
352
353/// A stream of [`ChatStreamEvent`]s for streaming chat completions.
354pub struct ChatStream {
355    pub inner: Pin<Box<dyn Stream<Item = Result<ChatStreamEvent>> + Send>>,
356}
357
358impl Stream for ChatStream {
359    type Item = Result<ChatStreamEvent>;
360
361    fn poll_next(
362        mut self: std::pin::Pin<&mut Self>,
363        cx: &mut std::task::Context<'_>,
364    ) -> std::task::Poll<Option<Self::Item>> {
365        self.inner.as_mut().poll_next(cx)
366    }
367}
368
369impl ChatStream {
370    pub fn from_bytes_stream<S>(stream: S) -> Self
371    where
372        S: Stream<Item = Result<Bytes>> + Send + Unpin + 'static,
373    {
374        let parser = GenericStreamParser::<S, ChatResponse, ChatStreamEvent>::new(stream);
375        ChatStream {
376            inner: Box::pin(parser),
377        }
378    }
379}
380
381impl StreamEventExt<ChatResponse> for ChatStreamEvent {
382    fn from_message(msg: ChatResponse) -> Self {
383        ChatStreamEvent::Message(msg)
384    }
385
386    fn from_error(err: String) -> Self {
387        ChatStreamEvent::Error(err)
388    }
389
390    fn partial(partial: String, error: Option<String>) -> Self {
391        ChatStreamEvent::Partial { partial, error }
392    }
393}
ollama_sdk/types/chat.rs

ollama_sdk/types/
chat.rs
