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