rig_core/providers/ollama/
streaming.rs

1//! The decoder of native Ollama chat replies: a whole `/api/chat` body, or a
2//! stream of NDJSON records of the same shape, the last one `done`.
3//!
4//! A record's `thinking` is reasoning, its `content` text and its
5//! `tool_calls` whole calls. A reply with no `thinking` may carry its
6//! reasoning as a leading `<think>` block in `content`; that block is split
7//! out as reasoning (see [`ChatDecoder`]).
8//!
9//! ```
10//! use rig_core::providers::ollama::streaming::ChatDecoder;
11//!
12//! let decoder = ChatDecoder::default();
13//! # let _ = decoder;
14//! ```
15
16use serde::Deserialize;
17use serde_json::{Map, Value};
18
19use crate::completion::{FinishReason, Usage};
20use crate::error::ProviderError;
21use crate::json_utils::Lenient;
22use crate::message::{CallId, ToolName};
23use crate::observe::ObservedError;
24use crate::operation::{Block, Completion, Finish};
25use crate::providers::internal::wire;
26use crate::wire::{AdapterVerdict, Decoder, Flow, ObservationSink, Out, WireEvent, WireFrame};
27
28/// The keys that make a JSON line a chat record: every record carries
29/// `message` or `done`, and an in-band failure carries `error`.
30const RECORD_KEYS: &[&str] = &["message", "done", "error"];
31
32/// The tag that opens an inline reasoning block.
33const THINK_OPEN: &str = "<think>";
34/// The tag that closes it.
35const THINK_CLOSE: &str = "</think>";
36
37/// One `/api/chat` record, as the daemon sent it.
38#[derive(Debug, Default, Deserialize)]
39#[serde(transparent)]
40pub struct ChatRecord(pub Map<String, Value>);
41
42/// Decode native Ollama chat replies.
43///
44/// `thinking` and `content` each continue the block of the record before
45/// while the kind stays the same, and each call is a whole block. The
46/// reply ends at the record that says `done`; its `done_reason` is the
47/// finish reason, `tool_calls` when a `stop` reply called a tool, and its
48/// `prompt_eval_count`, `prompt_eval_cached_count` (the part of the prompt
49/// read from the daemon's cache) and `eval_count` are the usage.
50///
51/// A reply with no `thinking` whose `content`, after leading whitespace,
52/// starts with `<think>` has that block split out as reasoning, by its tags
53/// alone. Content is held only while its start could still become
54/// `<think>`; content that cannot is text at once. Inside the block,
55/// reasoning streams as it arrives, holding back only trailing whitespace
56/// and a partial `</think>`. After `</think>` the rest is text. A block
57/// that never closes is reasoning to the end of the reply.
58#[derive(Debug, Default)]
59pub struct ChatDecoder {
60    /// Where the inline split stands.
61    split: Split,
62    /// Whether the reply called a tool.
63    called: bool,
64    /// The model the records name.
65    model: Option<String>,
66}
67
68/// Where a reply's inline reasoning split stands.
69#[derive(Debug)]
70enum Split {
71    /// Content held while its start, after whitespace, could still become
72    /// `<think>`.
73    Opening(String),
74    /// Inside a leading `<think>` block. `held` is trailing whitespace and a
75    /// partial `</think>` not yet written; `started` is whether any
76    /// reasoning was, so leading whitespace is dropped until then.
77    Inside { held: String, started: bool },
78    /// Content is text. `trim` drops leading whitespace until visible text
79    /// arrives after a split block.
80    Text { trim: bool },
81}
82
83impl Default for Split {
84    fn default() -> Self {
85        Self::Opening(String::new())
86    }
87}
88
89/// The length of the longest proper prefix of `tag` that `text` ends with.
90/// The tags are ASCII, so the cut is a character boundary.
91fn partial_suffix(text: &str, tag: &str) -> usize {
92    (1..tag.len())
93        .rev()
94        .find(|&len| text.ends_with(&tag[..len]))
95        .unwrap_or(0)
96}
97
98impl<'id> Decoder<'id, Completion> for ChatDecoder {
99    type Event = ChatRecord;
100
101    fn classify(&self, frame: WireFrame) -> WireEvent<ChatRecord> {
102        wire::classify_marker_keyed_frame(&frame.as_str(), RECORD_KEYS)
103    }
104
105    /// Fields are read leniently: one the record leaves out or sends with
106    /// another type is absent, so only `error` and `done` steer the reply.
107    fn decode(
108        &mut self,
109        ChatRecord(fields): ChatRecord,
110        mut out: Out<'id, Completion>,
111    ) -> Result<Flow, ProviderError> {
112        let record = Value::Object(fields);
113        if let Some(error) = record.get("error").filter(|error| !error.is_null()) {
114            let body = serde_json::json!({ "error": error }).to_string();
115            return Err(ProviderError::from_provider_body(body));
116        }
117        if let Some(model) = record.str("model").filter(|model| !model.is_empty())
118            && self.model.as_deref() != Some(model)
119        {
120            tracing::Span::current().record("gen_ai.response.model", model);
121            self.model = Some(model.to_owned());
122        }
123        let message = record.get("message").unwrap_or(&Value::Null);
124        if let Some(thinking) = message
125            .str("thinking")
126            .filter(|thinking| !thinking.is_empty())
127        {
128            // A reply with native thinking splits nothing out of its content.
129            self.release(&mut out)?;
130            reason(thinking, &mut out)?;
131        }
132        if let Some(content) = message.str("content").filter(|content| !content.is_empty()) {
133            self.content(content, &mut out)?;
134        }
135        for call in message.arr("tool_calls") {
136            self.call(call, &mut out)?;
137        }
138        if record.bool("done") != Some(true) {
139            return Ok(Flow::More);
140        }
141        self.release(&mut out)?;
142        out.end_run()?;
143        let reason = record.str("done_reason").map(|reason| match reason {
144            "stop" if self.called => FinishReason::ToolCalls,
145            "stop" => FinishReason::Stop,
146            "length" => FinishReason::Length,
147            other => FinishReason::Other(other.to_owned()),
148        });
149        let (input, output) = (record.u64("prompt_eval_count"), record.u64("eval_count"));
150        let usage = Usage {
151            input_tokens: input,
152            output_tokens: output,
153            total_tokens: input.zip(output).map(|(input, output)| input + output),
154            cached_input_tokens: record.u64("prompt_eval_cached_count"),
155            ..Usage::default()
156        };
157        Ok(out.end(Finish {
158            usage,
159            reason,
160            model: self.model.take(),
161            ..Finish::default()
162        }))
163    }
164}
165
166impl ChatDecoder {
167    /// Write a fragment of `content`, through the inline split.
168    fn content(
169        &mut self,
170        content: &str,
171        out: &mut Out<'_, Completion>,
172    ) -> Result<(), ProviderError> {
173        match std::mem::replace(&mut self.split, Split::Text { trim: false }) {
174            Split::Text { trim } => {
175                let text = if trim { content.trim_start() } else { content };
176                self.split = Split::Text {
177                    trim: trim && text.is_empty(),
178                };
179                if !text.is_empty() {
180                    out.run(Block::Text, text)?;
181                }
182                Ok(())
183            }
184            Split::Opening(mut held) => {
185                held.push_str(content);
186                let trimmed = held.trim_start();
187                if let Some(rest) = trimmed.strip_prefix(THINK_OPEN) {
188                    let rest = rest.to_owned();
189                    self.split = Split::Inside {
190                        held: String::new(),
191                        started: false,
192                    };
193                    self.content(&rest, out)
194                } else if THINK_OPEN.starts_with(trimmed) {
195                    self.split = Split::Opening(held);
196                    Ok(())
197                } else {
198                    out.run(Block::Text, &held)?;
199                    Ok(())
200                }
201            }
202            Split::Inside {
203                mut held,
204                mut started,
205            } => {
206                held.push_str(content);
207                if let Some((reasoning, rest)) = held.split_once(THINK_CLOSE) {
208                    write_reasoning(reasoning.trim_end(), started, out)?;
209                    self.split = Split::Text { trim: true };
210                    return self.content(rest, out);
211                }
212                let cut = held.len() - partial_suffix(&held, THINK_CLOSE);
213                let cut = held[..cut].trim_end().len();
214                started |= write_reasoning(&held[..cut], started, out)?;
215                self.split = Split::Inside {
216                    held: held.split_off(cut),
217                    started,
218                };
219                Ok(())
220            }
221        }
222    }
223
224    /// End the split: content held at the start is text as it arrived, and
225    /// what an open block held is reasoning.
226    fn release(&mut self, out: &mut Out<'_, Completion>) -> Result<(), ProviderError> {
227        match std::mem::replace(&mut self.split, Split::Text { trim: false }) {
228            Split::Opening(held) if !held.is_empty() => {
229                out.run(Block::Text, &held)?;
230            }
231            Split::Inside { held, started } => {
232                write_reasoning(held.trim_end(), started, out)?;
233            }
234            Split::Opening(_) | Split::Text { .. } => {}
235        }
236        Ok(())
237    }
238
239    /// Write one whole call. A call without a name is dropped: nothing can
240    /// answer it.
241    fn call(&mut self, call: &Value, out: &mut Out<'_, Completion>) -> Result<(), ProviderError> {
242        let Ok(name) = ToolName::new(
243            call.at("/function/name")
244                .and_then(Value::as_str)
245                .unwrap_or_default(),
246        ) else {
247            tracing::warn!("Ollama sent a tool call without a name; nothing can answer it");
248            return Ok(());
249        };
250        let arguments = match call.at("/function/arguments") {
251            None | Some(Value::Null) => "{}".to_owned(),
252            Some(Value::String(arguments)) => arguments.clone(),
253            Some(arguments) => arguments.to_string(),
254        };
255        self.release(out)?;
256        out.end_run()?;
257        self.called = true;
258        let id = CallId::from_wire(call.str("id").unwrap_or_default());
259        let index = out.fresh_index();
260        out.open(index, Block::Call { id, name }, call.clone())?;
261        out.push(index, &arguments)?;
262        out.finish(index)
263    }
264
265    /// Project the verdict and an error envelope off one raw record
266    /// before normalization discards them. A payload that is not JSON
267    /// projects nothing.
268    pub(crate) fn project(payload: &[u8], sink: &mut ObservationSink<'_>) {
269        let Ok(record) = serde_json::from_slice::<Value>(payload) else {
270            return;
271        };
272        let verdict = match record.str("done_reason") {
273            Some(reason) => AdapterVerdict {
274                finish_reason: Some(sink.scrub(reason)),
275                block_reason: None,
276                detail: None,
277                model: record.str("model").map(|model| sink.scrub(model)),
278            },
279            None => AdapterVerdict::default(),
280        };
281        sink.provider(verdict, None);
282        if let Some(error) = ObservedError::of(&record) {
283            error.emit(sink);
284        }
285    }
286}
287
288/// Write a fragment of an inline block's reasoning, its leading whitespace
289/// dropped until the block has `started`. Returns whether it wrote any.
290fn write_reasoning(
291    text: &str,
292    started: bool,
293    out: &mut Out<'_, Completion>,
294) -> Result<bool, ProviderError> {
295    let text = if started { text } else { text.trim_start() };
296    if text.is_empty() {
297        return Ok(false);
298    }
299    reason(text, out)?;
300    Ok(true)
301}
302
303/// Write a fragment of reasoning. Its block's item is the `thinking` it
304/// joins, so a same-model replay sends it back.
305fn reason(text: &str, out: &mut Out<'_, Completion>) -> Result<(), ProviderError> {
306    let index = out.run(Block::Reasoning { redacted: false }, text)?;
307    out.edit(index, |item| match item {
308        Value::Object(fields) => {
309            if let Some(Value::String(thinking)) = fields.get_mut("thinking") {
310                thinking.push_str(text);
311            }
312        }
313        _ => *item = serde_json::json!({ "thinking": text }),
314    })
315}
316
317pub(crate) mod document;
318
319#[cfg(test)]
320mod tests;
rig_core/providers/ollama/streaming.rs

rig_core/providers/ollama/
streaming.rs
