rig_core/providers/ollama/
extension.rs

1//! Ollama's typed request options and reply extras. The shared section
2//! goes to the OpenAI-compatible `/v1` and the native `/api/chat`, and is
3//! refused on the Anthropic-compatible `/v1/messages`, which reads none of
4//! it; the `"ollama.chat"` section, with the model parameters, goes to the
5//! native route only and is skipped on the others.
6//!
7//! ```
8//! use rig_core::completion::CompletionRequest;
9//! use rig_core::providers::ollama::extension::{KeepAlive, OllamaOptions};
10//!
11//! let options = OllamaOptions::default()
12//!     .keep_alive(KeepAlive::duration("5m"))
13//!     .num_ctx(8192);
14//! let request = CompletionRequest::new("hi").provider_option(options);
15//! ```
16
17use serde::{Deserialize, Serialize};
18use serde_json::Value;
19
20use crate::completion::{ExtensionOptions, ProviderExtension, ReplyExtras};
21use crate::message::Api;
22
23/// The `ollama` provider: its key, [`OllamaOptions`] and [`OllamaExtras`].
24#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
25pub struct OllamaExt;
26
27impl ProviderExtension for OllamaExt {
28    const PROVIDER: &'static str = super::PROVIDER_NAME;
29    type Options = OllamaOptions;
30    type Extras = OllamaExtras;
31}
32
33/// Ollama's request options.
34#[non_exhaustive]
35#[derive(Clone, Debug, Default, PartialEq, Serialize)]
36pub struct OllamaOptions {
37    /// The fields the OpenAI-compatible and native routes read.
38    #[serde(rename = "*")]
39    pub shared: OllamaShared,
40    /// The fields only `/api/chat` reads.
41    #[serde(rename = "ollama.chat")]
42    pub chat: OllamaNative,
43}
44
45impl ExtensionOptions for OllamaOptions {
46    type Ext = OllamaExt;
47
48    /// `keep_alive` on the Anthropic-compatible route, whose body has no
49    /// such field.
50    fn unsupported(
51        &self,
52        target: &dyn crate::completion::ReplayTarget,
53        _request: &crate::completion::CompletionRequest,
54    ) -> Vec<(&'static str, String)> {
55        let messages = Api::from_static(crate::providers::anthropic::wire::MESSAGES_API);
56        match self.shared.keep_alive {
57            Some(_) if target.api() == messages => vec![(
58                "keep_alive",
59                "Ollama's Anthropic-compatible API does not read it".to_owned(),
60            )],
61            _ => Vec::new(),
62        }
63    }
64}
65
66impl OllamaOptions {
67    /// How long the daemon keeps the model loaded after the request
68    /// (`keep_alive`).
69    pub fn keep_alive(mut self, keep_alive: KeepAlive) -> Self {
70        self.shared.keep_alive = Some(keep_alive);
71        self
72    }
73
74    /// The context window, in tokens (`options.num_ctx`).
75    pub fn num_ctx(mut self, num_ctx: u32) -> Self {
76        self.chat.options.num_ctx = Some(num_ctx);
77        self
78    }
79
80    /// The prompt tokens kept when the context is shifted
81    /// (`options.num_keep`).
82    pub fn num_keep(mut self, num_keep: u32) -> Self {
83        self.chat.options.num_keep = Some(num_keep);
84        self
85    }
86
87    /// Sample from the `top_k` most likely tokens (`options.top_k`).
88    pub fn top_k(mut self, top_k: u32) -> Self {
89        self.chat.options.top_k = Some(top_k);
90        self
91    }
92
93    /// Drop tokens below `min_p` times the likeliest token's probability
94    /// (`options.min_p`).
95    pub fn min_p(mut self, min_p: f64) -> Self {
96        self.chat.options.min_p = Some(min_p);
97        self
98    }
99
100    /// Penalize repeated tokens (`options.repeat_penalty`).
101    pub fn repeat_penalty(mut self, penalty: f64) -> Self {
102        self.chat.options.repeat_penalty = Some(penalty);
103        self
104    }
105
106    /// How far back the repeat penalty looks, in tokens; `0` disables it
107    /// and `-1` is the context window (`options.repeat_last_n`).
108    pub fn repeat_last_n(mut self, last_n: i32) -> Self {
109        self.chat.options.repeat_last_n = Some(last_n);
110        self
111    }
112
113    /// The model layers offloaded to the GPU (`options.num_gpu`).
114    pub fn num_gpu(mut self, num_gpu: i32) -> Self {
115        self.chat.options.num_gpu = Some(num_gpu);
116        self
117    }
118
119    /// The CPU threads generation uses (`options.num_thread`).
120    pub fn num_thread(mut self, num_thread: u32) -> Self {
121        self.chat.options.num_thread = Some(num_thread);
122        self
123    }
124
125    /// Return each generated token's log probability (`logprobs`), read
126    /// back as [`OllamaExtras::logprobs`].
127    pub fn logprobs(mut self, logprobs: bool) -> Self {
128        self.chat.logprobs = Some(logprobs);
129        self
130    }
131
132    /// The most likely alternatives returned beside each token
133    /// (`top_logprobs`).
134    pub fn top_logprobs(mut self, top_logprobs: u32) -> Self {
135        self.chat.top_logprobs = Some(top_logprobs);
136        self
137    }
138}
139
140/// The fields both Ollama routes read.
141#[non_exhaustive]
142#[derive(Clone, Debug, Default, PartialEq, Serialize)]
143pub struct OllamaShared {
144    /// `keep_alive`.
145    #[serde(skip_serializing_if = "Option::is_none")]
146    pub keep_alive: Option<KeepAlive>,
147}
148
149/// How long the daemon keeps a model loaded after a request.
150#[non_exhaustive]
151#[derive(Clone, Debug, PartialEq, Serialize)]
152#[serde(untagged)]
153pub enum KeepAlive {
154    /// A duration string such as `"5m"`; a negative one keeps the model
155    /// loaded.
156    Duration(String),
157    /// A number of seconds; `0` unloads the model at once and a negative
158    /// number keeps it loaded.
159    Seconds(i64),
160}
161
162impl KeepAlive {
163    /// A duration string such as `"5m"` or `"1h"`.
164    pub fn duration(duration: impl Into<String>) -> Self {
165        Self::Duration(duration.into())
166    }
167
168    /// A number of seconds.
169    pub fn seconds(seconds: i64) -> Self {
170        Self::Seconds(seconds)
171    }
172}
173
174/// The fields only `/api/chat` reads.
175#[non_exhaustive]
176#[derive(Clone, Debug, Default, PartialEq, Serialize)]
177pub struct OllamaNative {
178    /// `options`: model parameters, beside the ones the request and its
179    /// generation options set.
180    #[serde(skip_serializing_if = "ModelOptions::is_empty")]
181    pub options: ModelOptions,
182    /// `logprobs`.
183    #[serde(skip_serializing_if = "Option::is_none")]
184    pub logprobs: Option<bool>,
185    /// `top_logprobs`.
186    #[serde(skip_serializing_if = "Option::is_none")]
187    pub top_logprobs: Option<u32>,
188}
189
190/// The model parameters `/api/chat` reads in `options`. `temperature`,
191/// `num_predict`, `top_p`, `seed` and `stop` are not here: the request and
192/// its generation options set them.
193#[non_exhaustive]
194#[derive(Clone, Debug, Default, PartialEq, Serialize)]
195pub struct ModelOptions {
196    /// `num_ctx`.
197    #[serde(skip_serializing_if = "Option::is_none")]
198    pub num_ctx: Option<u32>,
199    /// `num_keep`.
200    #[serde(skip_serializing_if = "Option::is_none")]
201    pub num_keep: Option<u32>,
202    /// `top_k`.
203    #[serde(skip_serializing_if = "Option::is_none")]
204    pub top_k: Option<u32>,
205    /// `min_p`.
206    #[serde(skip_serializing_if = "Option::is_none")]
207    pub min_p: Option<f64>,
208    /// `repeat_penalty`.
209    #[serde(skip_serializing_if = "Option::is_none")]
210    pub repeat_penalty: Option<f64>,
211    /// `repeat_last_n`.
212    #[serde(skip_serializing_if = "Option::is_none")]
213    pub repeat_last_n: Option<i32>,
214    /// `num_gpu`.
215    #[serde(skip_serializing_if = "Option::is_none")]
216    pub num_gpu: Option<i32>,
217    /// `num_thread`.
218    #[serde(skip_serializing_if = "Option::is_none")]
219    pub num_thread: Option<u32>,
220}
221
222impl ModelOptions {
223    /// Whether no parameter is set.
224    pub fn is_empty(&self) -> bool {
225        self == &Self::default()
226    }
227}
228
229/// The fields of an `/api/chat` reply Rig does not normalize, read from
230/// the whole reply or a stream's final record. A `/v1` reply carries only
231/// [`Self::model`] and leaves the rest `None`.
232#[non_exhaustive]
233#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
234#[serde(default)]
235pub struct OllamaExtras {
236    /// `model`, on both routes.
237    pub model: Option<String>,
238    /// `created_at`.
239    pub created_at: Option<String>,
240    /// `done_reason` as the daemon spells it (`stop`, `length`, ...).
241    pub done_reason: Option<String>,
242    /// `total_duration`, in nanoseconds.
243    pub total_duration: Option<u64>,
244    /// `load_duration`, in nanoseconds.
245    pub load_duration: Option<u64>,
246    /// `prompt_eval_duration`, in nanoseconds.
247    pub prompt_eval_duration: Option<u64>,
248    /// `eval_duration`, in nanoseconds.
249    pub eval_duration: Option<u64>,
250    /// `prompt_eval_count`.
251    pub prompt_eval_count: Option<u64>,
252    /// `prompt_eval_cached_count`: the prompt tokens read from the cache.
253    pub prompt_eval_cached_count: Option<u64>,
254    /// `eval_count`.
255    pub eval_count: Option<u64>,
256    /// `logprobs`, when [`OllamaOptions::logprobs`] asked for them.
257    pub logprobs: Option<Vec<Logprob>>,
258}
259
260/// One generated token's log probability.
261#[non_exhaustive]
262#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
263#[serde(default)]
264pub struct Logprob {
265    /// `token`.
266    pub token: Option<String>,
267    /// `logprob`.
268    pub logprob: Option<f64>,
269    /// `bytes`, the token's UTF-8 bytes.
270    pub bytes: Option<Vec<u8>>,
271    /// `top_logprobs`, the likeliest alternatives, when
272    /// [`OllamaOptions::top_logprobs`] asked for them.
273    pub top_logprobs: Option<Vec<Logprob>>,
274}
275
276impl ReplyExtras for OllamaExtras {
277    fn from_reply(_api: &Api, raw: &Value) -> Result<Self, serde_json::Error> {
278        Self::deserialize(raw)
279    }
280}
281
282#[cfg(test)]
283mod tests;
rig_core/providers/ollama/extension.rs

rig_core/providers/ollama/
extension.rs
