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::{
17 is_model_request_expr, model_response_expr, text_part, usage_record, validate_chat_transcript,
18};
19
20const OLLAMA_CODEC_ID: CodecId = CodecId(0);
25
26#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct OllamaRequestOptions {
30 pub model: String,
32 pub stream: bool,
34 pub tools: bool,
36}
37
38pub 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
69pub struct OllamaCodecLib {
71 symbol: Symbol,
72 codec_id: CodecId,
73}
74
75impl OllamaCodecLib {
76 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 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
117pub 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
160pub 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
179pub 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 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}