1use 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
21const OLLAMA_CODEC_ID: CodecId = CodecId(0);
26
27#[derive(Clone, Debug, PartialEq, Eq)]
30pub struct OllamaRequestOptions {
31 pub model: String,
33 pub stream: bool,
35 pub tools: bool,
37}
38
39pub 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
84pub struct OllamaCodecLib {
86 symbol: Symbol,
87 codec_id: CodecId,
88}
89
90impl OllamaCodecLib {
91 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 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
132pub 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
194pub 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
209pub 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
242pub 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
258pub 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 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}