Skip to main content

sim_codec_chat/
ollama.rs

1//! Ollama provider bridge: encode a chat model-request transcript into an
2//! Ollama JSON request, and decode Ollama JSON responses and streamed chunks
3//! back into canonical chat transcript `Expr` values.
4
5use serde_json::{Map, Value, json};
6use std::sync::Arc;
7
8use sim_codec::{
9    DecodeBudget, DecodeLimits, Decoder, DomainCodecLib, Encoder, Input, Output, ReadCx,
10    domain_input_text,
11};
12use sim_codec_json::{JsonProjectionMode, json_number_to_u64, project_json_to_expr_budgeted};
13use sim_kernel::{CodecId, Error, Expr, Lib, LibManifest, Linker, LoadCx, Result, Symbol, WriteCx};
14use sim_value::access;
15
16use crate::{
17    is_model_request_expr, model_response_expr, text_part, usage_record, validate_chat_transcript,
18};
19
20/// Codec id used to tag decode-budget errors raised while projecting an Ollama
21/// provider response. The Ollama bridge is a set of free functions with no
22/// registered codec id of its own; the value only appears in budget-exceeded
23/// error messages on hostile input.
24const OLLAMA_CODEC_ID: CodecId = CodecId(0);
25
26/// Options controlling how a chat model-request transcript is projected into an
27/// Ollama JSON request body.
28#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct OllamaRequestOptions {
30    /// The Ollama model name to send in the request.
31    pub model: String,
32    /// Whether to request a streamed response.
33    pub stream: bool,
34    /// Whether to include a (currently empty) `tools` array in the request.
35    pub tools: bool,
36}
37
38/// Runtime codec for Ollama chat JSON and NDJSON response bodies.
39pub struct OllamaCodec;
40
41impl Decoder for OllamaCodec {
42    fn decode(&self, cx: &mut ReadCx<'_>, input: Input) -> Result<Expr> {
43        let source = domain_input_text(cx.codec, input)?;
44        let model = first_model_name(&source).unwrap_or_else(|| "ollama".to_owned());
45        let body = source.as_bytes();
46        if source
47            .lines()
48            .filter(|line| !line.trim().is_empty())
49            .count()
50            > 1
51        {
52            decode_ollama_stream(Symbol::qualified("runner", "ollama"), &model, body, false)
53        } else {
54            decode_ollama_response(Symbol::qualified("runner", "ollama"), &model, body, false)
55        }
56    }
57}
58
59impl Encoder for OllamaCodec {
60    fn encode(&self, _cx: &mut WriteCx<'_>, expr: &Expr) -> Result<Output> {
61        let options = OllamaRequestOptions::new(request_model(expr), false, false);
62        let bytes = encode_ollama_request(expr, &options)?;
63        String::from_utf8(bytes)
64            .map(Output::Text)
65            .map_err(|err| Error::Eval(format!("ollama codec encoded invalid utf-8: {err}")))
66    }
67}
68
69/// Host-registered lib for `codec:ollama`.
70pub struct OllamaCodecLib {
71    symbol: Symbol,
72    codec_id: CodecId,
73}
74
75impl OllamaCodecLib {
76    /// Creates the lib bound to the given runtime-assigned codec id.
77    pub fn new(id: CodecId) -> Self {
78        Self {
79            symbol: Symbol::qualified("codec", "ollama"),
80            codec_id: id,
81        }
82    }
83
84    fn domain_lib(&self) -> DomainCodecLib {
85        DomainCodecLib::new(
86            self.symbol.clone(),
87            self.codec_id,
88            Arc::new(OllamaCodec),
89            Arc::new(OllamaCodec),
90            Symbol::qualified("codec", "OllamaTranscript"),
91        )
92    }
93}
94
95impl Lib for OllamaCodecLib {
96    fn manifest(&self) -> LibManifest {
97        self.domain_lib().manifest()
98    }
99
100    fn load(&self, cx: &mut LoadCx, linker: &mut Linker<'_>) -> Result<()> {
101        self.domain_lib().load(cx, linker)
102    }
103}
104
105impl OllamaRequestOptions {
106    /// Creates request options from a `model` name and the `stream`/`tools`
107    /// flags.
108    pub fn new(model: impl Into<String>, stream: bool, tools: bool) -> Self {
109        Self {
110            model: model.into(),
111            stream,
112            tools,
113        }
114    }
115}
116
117/// Encodes a chat model-request transcript into an Ollama JSON request body.
118///
119/// Fails closed unless `expr` is a valid `model-request` transcript: prior
120/// messages plus the request `task` are flattened into Ollama `messages`, and
121/// the `model`/`stream`/`tools` options from `options` are applied.
122///
123/// # Examples
124///
125/// ```
126/// use sim_codec_chat::{encode_ollama_request, model_card_expr, OllamaRequestOptions};
127/// use sim_kernel::Symbol;
128///
129/// // A model-card is not a model-request, so the codec fails closed.
130/// let card = model_card_expr(
131///     Symbol::new("local-reasoner"),
132///     "qwen2.5-coder:14b",
133///     Symbol::new("ollama"),
134///     Symbol::new("local"),
135/// );
136/// let options = OllamaRequestOptions::new("qwen2.5-coder:14b", false, false);
137/// assert!(encode_ollama_request(&card, &options).is_err());
138/// ```
139pub fn encode_ollama_request(expr: &Expr, options: &OllamaRequestOptions) -> Result<Vec<u8>> {
140    if !is_model_request_expr(expr) {
141        return Err(Error::Eval(
142            "ollama codec expects a model-request transcript".to_owned(),
143        ));
144    }
145    validate_chat_transcript(expr)?;
146    let mut payload = Map::new();
147    payload.insert("model".to_owned(), Value::String(options.model.clone()));
148    payload.insert("stream".to_owned(), Value::Bool(options.stream));
149    payload.insert(
150        "messages".to_owned(),
151        Value::Array(transcript_messages(expr)?),
152    );
153    if options.tools {
154        payload.insert("tools".to_owned(), Value::Array(Vec::new()));
155    }
156    serde_json::to_vec(&Value::Object(payload))
157        .map_err(|err| Error::Eval(format!("ollama codec failed to encode request: {err}")))
158}
159
160/// Decodes a non-streamed Ollama JSON response `body` into a canonical
161/// model-response transcript attributed to `runner` and `model`.
162///
163/// Extracts the response text and any token-usage counts; when `include_raw`
164/// is set, the original JSON is also attached as a `raw-provider-response`
165/// field.
166pub fn decode_ollama_response(
167    runner: Symbol,
168    model: &str,
169    body: &[u8],
170    include_raw: bool,
171) -> Result<Expr> {
172    let mut budget = DecodeBudget::new(DecodeLimits::default());
173    budget.check_input_bytes(OLLAMA_CODEC_ID, body.len())?;
174    let value: Value = serde_json::from_slice(body)
175        .map_err(|err| Error::Eval(format!("ollama codec returned invalid json: {err}")))?;
176    response_expr_from_json(runner, model, &value, include_raw, &mut budget)
177}
178
179/// Decodes a newline-delimited Ollama streaming response `body` into a single
180/// canonical model-response transcript attributed to `runner` and `model`.
181///
182/// Concatenates the text of every chunk, derives the stop reason from the final
183/// chunk, and folds in token-usage counts; when `include_raw` is set, the raw
184/// chunks are attached as a `raw-provider-response` list. Errors if no chunks
185/// are present.
186pub fn decode_ollama_stream(
187    runner: Symbol,
188    model: &str,
189    body: &[u8],
190    include_raw: bool,
191) -> Result<Expr> {
192    let mut budget = DecodeBudget::new(DecodeLimits::default());
193    budget.check_input_bytes(OLLAMA_CODEC_ID, body.len())?;
194    let text = std::str::from_utf8(body)
195        .map_err(|err| Error::Eval(format!("ollama stream is not valid utf-8: {err}")))?;
196    let mut chunks = Vec::new();
197    let mut combined = String::new();
198    let mut usage_source = None;
199    let mut stop_reason = Symbol::new("stop");
200    for line in text.lines() {
201        let trimmed = line.trim();
202        if trimmed.is_empty() {
203            continue;
204        }
205        let value: Value = serde_json::from_str(trimmed)
206            .map_err(|err| Error::Eval(format!("ollama stream returned invalid json: {err}")))?;
207        combined.push_str(&response_chunk_text(&value)?);
208        if usage_expr_from_value(&value)?.is_some() {
209            usage_source = Some(value.clone());
210        }
211        if let Some(reason) = value.get("done_reason").and_then(Value::as_str) {
212            stop_reason = Symbol::new(reason);
213        } else if value.get("done").and_then(Value::as_bool).unwrap_or(false) {
214            stop_reason = Symbol::new("stop");
215        }
216        chunks.push(value);
217    }
218    if chunks.is_empty() {
219        return Err(Error::Eval(
220            "ollama stream did not contain any response chunks".to_owned(),
221        ));
222    }
223    let mut entries =
224        match model_response_expr(runner, model, vec![text_part(&combined)], stop_reason) {
225            Expr::Map(entries) => entries,
226            _ => unreachable!("model_response_expr always returns a map"),
227        };
228    if let Some(source) = usage_source.as_ref()
229        && let Some(usage) = usage_expr_from_value(source)?
230    {
231        entries.push((Expr::Symbol(Symbol::new("usage")), usage));
232    }
233    if include_raw {
234        let raw = chunks
235            .iter()
236            .map(|chunk| {
237                project_json_to_expr_budgeted(
238                    chunk,
239                    JsonProjectionMode::UntaggedInterop,
240                    OLLAMA_CODEC_ID,
241                    &mut budget,
242                    0,
243                )
244            })
245            .collect::<Result<Vec<_>>>()?;
246        entries.push((
247            Expr::Symbol(Symbol::new("raw-provider-response")),
248            Expr::List(raw),
249        ));
250    }
251    Ok(Expr::Map(entries))
252}
253
254fn response_expr_from_json(
255    runner: Symbol,
256    model: &str,
257    value: &Value,
258    include_raw: bool,
259    budget: &mut DecodeBudget,
260) -> Result<Expr> {
261    let response = value
262        .as_object()
263        .ok_or_else(|| Error::Eval("ollama response must be a json object".to_owned()))?;
264    let content = response_content(response)?;
265    let stop_reason = response
266        .get("done_reason")
267        .and_then(Value::as_str)
268        .unwrap_or("stop");
269    let mut entries = match model_response_expr(
270        runner,
271        model,
272        vec![text_part(&content)],
273        Symbol::new(stop_reason),
274    ) {
275        Expr::Map(entries) => entries,
276        _ => unreachable!("model_response_expr always returns a map"),
277    };
278    if let Some(usage) = usage_expr_from_value(value)? {
279        entries.push((Expr::Symbol(Symbol::new("usage")), usage));
280    }
281    if include_raw {
282        entries.push((
283            Expr::Symbol(Symbol::new("raw-provider-response")),
284            project_json_to_expr_budgeted(
285                value,
286                JsonProjectionMode::UntaggedInterop,
287                OLLAMA_CODEC_ID,
288                budget,
289                0,
290            )?,
291        ));
292    }
293    Ok(Expr::Map(entries))
294}
295
296fn transcript_messages(expr: &Expr) -> Result<Vec<Value>> {
297    let Expr::Map(entries) = expr else {
298        return Err(Error::Eval(
299            "ollama codec expects request transcript as a map".to_owned(),
300        ));
301    };
302    let mut messages = access::entry_required_list_any(entries, "messages", "ollama messages")?
303        .iter()
304        .map(message_to_json)
305        .collect::<Result<Vec<_>>>()?;
306    messages.push(json!({
307        "role": "user",
308        "content": flatten_expr(map_field(entries, "task")?),
309    }));
310    Ok(messages)
311}
312
313fn message_to_json(expr: &Expr) -> Result<Value> {
314    let Expr::Map(entries) = expr else {
315        return Err(Error::Eval("ollama codec message must be a map".to_owned()));
316    };
317    let role = access::entry_required_sym_any(entries, "role", "ollama message role")?;
318    let content = access::entry_required_list_any(entries, "content", "ollama message content")?
319        .iter()
320        .map(content_part_to_text)
321        .collect::<Result<Vec<_>>>()?
322        .join(" ");
323    Ok(json!({
324        "role": role.name.as_ref(),
325        "content": content,
326    }))
327}
328
329fn content_part_to_text(expr: &Expr) -> Result<String> {
330    let Expr::Map(entries) = expr else {
331        return Err(Error::Eval(
332            "ollama codec content part must be a map".to_owned(),
333        ));
334    };
335    match access::entry_required_sym_any(entries, "type", "ollama content part type")?
336        .name
337        .as_ref()
338    {
339        "text" => Ok(
340            access::entry_required_str_any(entries, "text", "ollama content part text")?.to_owned(),
341        ),
342        other => Err(Error::Eval(format!(
343            "ollama codec does not support content part type {other}"
344        ))),
345    }
346}
347
348fn response_content(response: &Map<String, Value>) -> Result<String> {
349    if let Some(content) = response
350        .get("message")
351        .and_then(Value::as_object)
352        .and_then(|message| message.get("content"))
353        .and_then(Value::as_str)
354    {
355        return Ok(content.to_owned());
356    }
357    if let Some(content) = response.get("response").and_then(Value::as_str) {
358        return Ok(content.to_owned());
359    }
360    Err(Error::Eval(
361        "ollama response missing message.content or response".to_owned(),
362    ))
363}
364
365fn response_chunk_text(value: &Value) -> Result<String> {
366    let object = value
367        .as_object()
368        .ok_or_else(|| Error::Eval("ollama stream chunk must be a json object".to_owned()))?;
369    if let Some(content) = object
370        .get("message")
371        .and_then(Value::as_object)
372        .and_then(|message| message.get("content"))
373        .and_then(Value::as_str)
374    {
375        return Ok(content.to_owned());
376    }
377    if let Some(content) = object.get("response").and_then(Value::as_str) {
378        return Ok(content.to_owned());
379    }
380    Ok(String::new())
381}
382
383fn usage_expr_from_value(value: &Value) -> Result<Option<Expr>> {
384    let response = value
385        .as_object()
386        .ok_or_else(|| Error::Eval("ollama usage source must be a json object".to_owned()))?;
387    let input = response
388        .get("prompt_eval_count")
389        .and_then(json_number_to_u64);
390    let output = response.get("eval_count").and_then(json_number_to_u64);
391    // Ollama reports no total; it is derived only when both counts are present.
392    // Saturate rather than wrap so a hostile body cannot overflow the u64 add.
393    let total = input
394        .zip(output)
395        .map(|(input, output)| input.saturating_add(output));
396    let fields = usage_record(input, output, total);
397    Ok((!fields.is_empty()).then_some(Expr::Map(fields)))
398}
399
400fn map_field<'a>(entries: &'a [(Expr, Expr)], key: &str) -> Result<&'a Expr> {
401    entries
402        .iter()
403        .find_map(|(field, value)| match field {
404            Expr::Symbol(symbol) if symbol.name.as_ref() == key => Some(value),
405            _ => None,
406        })
407        .ok_or_else(|| Error::Eval(format!("ollama codec missing {key} field")))
408}
409
410fn request_model(expr: &Expr) -> String {
411    let Expr::Map(entries) = expr else {
412        return "ollama".to_owned();
413    };
414    entries
415        .iter()
416        .find_map(|(field, value)| match (field, value) {
417            (Expr::Symbol(symbol), Expr::String(model)) if symbol.name.as_ref() == "model" => {
418                Some(model.clone())
419            }
420            _ => None,
421        })
422        .unwrap_or_else(|| "ollama".to_owned())
423}
424
425fn first_model_name(source: &str) -> Option<String> {
426    source
427        .lines()
428        .map(str::trim)
429        .find(|line| !line.is_empty())
430        .and_then(|line| serde_json::from_str::<Value>(line).ok())
431        .and_then(|value| {
432            value
433                .get("model")
434                .and_then(Value::as_str)
435                .map(str::to_owned)
436        })
437}
438
439fn flatten_expr(expr: &Expr) -> String {
440    match expr {
441        Expr::Nil => "nil".to_owned(),
442        Expr::Bool(flag) => flag.to_string(),
443        Expr::Number(number) => number.canonical.clone(),
444        Expr::Symbol(symbol) | Expr::Local(symbol) => symbol.to_string(),
445        Expr::String(text) => text.clone(),
446        Expr::Bytes(bytes) => format!("{bytes:?}"),
447        Expr::List(items) | Expr::Vector(items) | Expr::Set(items) | Expr::Block(items) => {
448            items.iter().map(flatten_expr).collect::<Vec<_>>().join(" ")
449        }
450        Expr::Map(entries) => entries
451            .iter()
452            .map(|(key, value)| format!("{} {}", flatten_expr(key), flatten_expr(value)))
453            .collect::<Vec<_>>()
454            .join(" "),
455        Expr::Call { operator, args } => std::iter::once(flatten_expr(operator))
456            .chain(args.iter().map(flatten_expr))
457            .collect::<Vec<_>>()
458            .join(" "),
459        Expr::Infix {
460            operator,
461            left,
462            right,
463        } => format!(
464            "{} {} {}",
465            flatten_expr(left),
466            operator,
467            flatten_expr(right)
468        ),
469        Expr::Prefix { operator, arg } => format!("{operator} {}", flatten_expr(arg)),
470        Expr::Postfix { operator, arg } => format!("{} {operator}", flatten_expr(arg)),
471        Expr::Quote { expr, .. } | Expr::Annotated { expr, .. } => flatten_expr(expr),
472        Expr::Extension { tag, payload } => format!("{tag} {}", flatten_expr(payload)),
473    }
474}