ollama_sdk/types/
generate.rs

1//! Contains all data structures that are particularly used for Ollama Generate 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::ThinkingLevel;
14
15/// Represents a request to the Ollama API for text generation.
16///
17/// This struct allows specifying the model, prompt, system message,
18/// and various generation options. It supports both streaming and non-streaming
19/// responses.
20#[derive(Serialize, Default, Debug, Clone)]
21pub struct GenerateRequest {
22    /// The name of the model to use for generation (e.g., "llama2").
23    pub model: String,
24    /// The primary prompt for the model.
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub prompt: Option<String>,
27    /// A suffix to be appended to the prompt.
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub suffix: Option<String>,
30    /// Optional base64-encoded images to include in the prompt.
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub images: Option<Vec<String>>,
33    /// A system message to guide the model's behavior.
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub system: Option<String>,
36    /// If `true`, the response will be streamed back as a series of [`GenerateStreamEvent`]s.
37    /// If `false`, a single [`GenerateResponse`] will be returned.
38    pub stream: bool,
39    /// Configuration for the model's "thinking" process.
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub think: Option<Thinking>,
42    /// If `true`, the raw prompt will be used without any templating.
43    #[serde(skip_serializing_if = "Option::is_none")]
44    pub raw: Option<bool>,
45    /// Additional generation options.
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub options: Option<GenerateOptions>,
48}
49
50/// Represents various options that can be configured for text generation.
51#[derive(Serialize, Default, Debug, Clone)]
52pub struct GenerateOptions {
53    /// The random seed for reproducibility.
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub seed: Option<u8>,
56    /// The temperature for sampling, controlling randomness. Higher values mean more random.
57    #[serde(skip_serializing_if = "Option::is_none")]
58    pub temperature: Option<f32>,
59    /// The top-k sampling parameter.
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub top_k: Option<u8>,
62    /// The top-p sampling parameter.
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub top_p: Option<f32>,
65    /// The minimum-p sampling parameter.
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub min_p: Option<f32>,
68    /// A list of strings that, if generated, will stop the generation.
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub stop: Option<Vec<String>>,
71    /// The size of the context window.
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub num_ctx: Option<u16>,
74    /// The maximum number of tokens to predict.
75    #[serde(skip_serializing_if = "Option::is_none")]
76    pub num_predict: Option<u16>,
77}
78
79/// Represents a response from the Ollama API for text generation.
80///
81/// This struct is used for non-streaming generation responses.
82#[derive(Deserialize, Serialize, Default, FromBytes, Debug, Clone)]
83pub struct GenerateResponse {
84    /// The name of the model that generated the response.
85    pub model: String,
86    /// The timestamp when the response was created.
87    pub created_at: String,
88    /// The generated text response.
89    pub response: String,
90    /// The model's internal "thinking" process, if enabled.
91    #[serde(default)]
92    pub thinking: String,
93    /// Indicates if the generation is complete.
94    pub done: bool,
95    /// The reason why the generation finished (e.g., "stop", "length").
96    #[serde(default)]
97    pub done_reason: Option<String>,
98    /// The total duration of the generation process in nanoseconds.
99    #[serde(default)]
100    pub total_duration: u64,
101    /// The duration spent loading the model in nanoseconds.
102    #[serde(default)]
103    pub load_duration: u64,
104    /// The number of tokens in the prompt that were evaluated.
105    #[serde(default)]
106    pub prompt_eval_count: u64,
107    /// The duration spent evaluating the prompt in nanoseconds.
108    #[serde(default)]
109    pub prompt_eval_duration: u64,
110    /// The number of tokens generated.
111    #[serde(default)]
112    pub eval_count: u64,
113    /// The duration spent generating tokens in nanoseconds.
114    #[serde(default)]
115    pub eval_duration: u64,
116}
117
118/// A simplified generation request for non-streaming responses.
119///
120/// This struct is a convenience wrapper for creating a [`GenerateRequest`]
121/// that explicitly requests a non-streaming response.
122#[derive(Serialize, Default, Debug, Clone)]
123pub struct SimpleGenerateRequest {
124    /// The name of the model to use.
125    pub model: String,
126    /// The primary prompt for the model.
127    pub prompt: Option<String>,
128    /// A suffix to be appended to the prompt.
129    pub suffix: Option<String>,
130    /// Optional base64-encoded images to include in the prompt.
131    pub images: Option<Vec<String>>,
132    /// A system message to guide the model's behavior.
133    pub system: Option<String>,
134    /// Configuration for the model's "thinking" process.
135    pub think: Option<Thinking>,
136    /// If `true`, the raw prompt will be used without any templating.
137    pub raw: Option<bool>,
138    /// Additional generation options.
139    pub options: Option<GenerateOptions>,
140}
141
142impl SimpleGenerateRequest {
143    /// Creates a new [`SimpleGenerateRequest`].
144    pub fn new(model: String, prompt: String) -> Self {
145        Self {
146            model,
147            prompt: Some(prompt),
148            ..Default::default()
149        }
150    }
151
152    /// Sets `think` param in the API call to `true`.
153    pub fn enable_thinking(mut self) -> Self {
154        self.think = Thinking::Boolean(true).into();
155        self
156    }
157
158    /// Sets `think` param in the API call to `false`.
159    pub fn disable_thinking(mut self) -> Self {
160        self.think = Thinking::Boolean(false).into();
161        self
162    }
163
164    /// Sets `think` param in the API call to specified level (`high`, `medium`, `low`).
165    pub fn set_thinking_level(mut self, level: ThinkingLevel) -> Self {
166        self.think = Thinking::Level(level).into();
167        self
168    }
169
170    /// Sets the system message for the request.
171    pub fn system(mut self, system: String) -> Self {
172        self.system = Some(system);
173        self
174    }
175
176    /// Sets the images for the request. An image should be a Base64-encoded string
177    pub fn images(mut self, images: Vec<String>) -> Self {
178        self.images = Some(images);
179        self
180    }
181
182    /// Sets the generation options for the request.
183    pub fn options(mut self, options: GenerateOptions) -> Self {
184        self.options = Some(options);
185        self
186    }
187}
188
189impl From<SimpleGenerateRequest> for GenerateRequest {
190    /// Converts a [`SimpleGenerateRequest`] into a [`GenerateRequest`] with `stream` set to `false`.
191    fn from(request: SimpleGenerateRequest) -> GenerateRequest {
192        GenerateRequest {
193            model: request.model,
194            prompt: request.prompt,
195            suffix: request.suffix,
196            images: request.images,
197            system: request.system,
198            think: request.think,
199            raw: request.raw,
200            options: request.options,
201            stream: false,
202        }
203    }
204}
205
206/// A simplified generation request for streaming responses.
207///
208/// This struct is a convenience wrapper for creating a [`GenerateRequest`]
209/// that explicitly requests a streaming response.
210#[derive(Serialize, Default, Debug, Clone)]
211pub struct StreamingGenerateRequest {
212    /// The name of the model to use.
213    pub model: String,
214    /// The primary prompt for the model.
215    pub prompt: Option<String>,
216    /// A suffix to be appended to the prompt.
217    pub suffix: Option<String>,
218    /// Optional base64-encoded images to include in the prompt.
219    pub images: Option<Vec<String>>,
220    /// A system message to guide the model's behavior.
221    pub system: Option<String>,
222    /// Configuration for the model's "thinking" process.
223    pub think: Option<Thinking>,
224    /// If `true`, the raw prompt will be used without any templating.
225    pub raw: Option<bool>,
226    /// Additional generation options.
227    pub options: Option<GenerateOptions>,
228}
229
230impl StreamingGenerateRequest {
231    /// Creates a new [`StreamingGenerateRequest`].
232    pub fn new(model: String, prompt: String) -> Self {
233        Self {
234            model,
235            prompt: Some(prompt),
236            ..Default::default()
237        }
238    }
239
240    /// Sets `think` param in the API call to `true`.
241    pub fn enable_thinking(mut self) -> Self {
242        self.think = Thinking::Boolean(true).into();
243        self
244    }
245
246    /// Sets `think` param in the API call to `false`.
247    pub fn disable_thinking(mut self) -> Self {
248        self.think = Thinking::Boolean(false).into();
249        self
250    }
251
252    /// Sets `think` param in the API call to specified level (`high`, `medium`, `low`).
253    pub fn set_thinking_level(mut self, level: ThinkingLevel) -> Self {
254        self.think = Thinking::Level(level).into();
255        self
256    }
257
258    /// Sets the system message for the request.
259    pub fn system(mut self, system: String) -> Self {
260        self.system = Some(system);
261        self
262    }
263
264    /// Sets the images for the request.
265    pub fn images(mut self, images: Vec<String>) -> Self {
266        self.images = Some(images);
267        self
268    }
269
270    /// Sets the generation options for the request.
271    pub fn options(mut self, options: GenerateOptions) -> Self {
272        self.options = Some(options);
273        self
274    }
275}
276
277impl From<StreamingGenerateRequest> for GenerateRequest {
278    /// Converts a [`StreamingGenerateRequest`] into a [`GenerateRequest`] with `stream` set to `true`.
279    fn from(request: StreamingGenerateRequest) -> GenerateRequest {
280        GenerateRequest {
281            model: request.model,
282            prompt: request.prompt,
283            suffix: request.suffix,
284            images: request.images,
285            system: request.system,
286            think: request.think,
287            raw: request.raw,
288            options: request.options,
289            stream: true,
290        }
291    }
292}
293
294/// Represents an event received from a streaming generation response.
295#[derive(Deserialize, Serialize, Debug)]
296#[serde(untagged)]
297pub enum GenerateStreamEvent {
298    /// A chunk of the generated response.
299    MessageChunk(GenerateResponse),
300    /// An error occurred during the streaming process.
301    Error(String),
302    /// A partial response, returned when the content was un-parseable
303    Partial {
304        /// The un-parseable content.
305        partial: String,
306        /// An optional error message associated with the partial response.
307        error: Option<String>,
308    },
309}
310
311/// A stream of [`GenerateStreamEvent`]s for streaming text generation.
312pub struct GenerateStream {
313    pub inner: Pin<Box<dyn Stream<Item = Result<GenerateStreamEvent>> + Send>>,
314}
315
316impl Stream for GenerateStream {
317    type Item = Result<GenerateStreamEvent>;
318
319    fn poll_next(
320        mut self: std::pin::Pin<&mut Self>,
321        cx: &mut std::task::Context<'_>,
322    ) -> std::task::Poll<Option<Self::Item>> {
323        self.inner.as_mut().poll_next(cx)
324    }
325}
326
327impl GenerateStream {
328    pub fn from_bytes_stream<S>(stream: S) -> Self
329    where
330        S: Stream<Item = Result<Bytes>> + Send + Unpin + 'static,
331    {
332        let parser = GenericStreamParser::<S, GenerateResponse, GenerateStreamEvent>::new(stream);
333        GenerateStream {
334            inner: Box::pin(parser),
335        }
336    }
337}
338
339impl StreamEventExt<GenerateResponse> for GenerateStreamEvent {
340    fn from_message(msg: GenerateResponse) -> Self {
341        GenerateStreamEvent::MessageChunk(msg)
342    }
343
344    fn from_error(err: String) -> Self {
345        GenerateStreamEvent::Error(err)
346    }
347
348    fn partial(partial: String, error: Option<String>) -> Self {
349        GenerateStreamEvent::Partial { partial, error }
350    }
351}
ollama_sdk/types/generate.rs

ollama_sdk/types/
generate.rs
