rig-core 0.45.0

//! The native `/api/chat` reassembler: a stream of records rebuilds the body
//! a unary call returns.

use serde_json::{Map, Value};

use super::ChatRecord;
use crate::wire::document::Reassemble;

/// The `/api/chat` reply a stream of records adds up to.
///
/// Top-level fields take the last non-null value a record states, so the
/// final `done` record's reason, counts and durations stand. Within
/// `message`, `content` and `thinking` append, `tool_calls` and `images`
/// collect, and other fields take the last non-null value. Top-level
/// `logprobs` collect. A reply that failed in band keeps its `error`
/// beside what arrived before it.
#[derive(Debug, Default)]
pub struct ChatResponse {
    document: Map<String, Value>,
}

/// The `message` fields whose fragments join as text.
const APPENDED: &[&str] = &["content", "thinking"];

/// The fields whose lists join, in `message` and at the top level.
const COLLECTED: &[&str] = &["tool_calls", "images", "logprobs"];

/// Fold `fields` into `into`: text in `appended` joins, lists in
/// [`COLLECTED`] extend, any other non-null value replaces, and `null`
/// only fills an absent key.
fn fold<'a>(
    into: &mut Map<String, Value>,
    fields: impl IntoIterator<Item = (&'a String, &'a Value)>,
    appended: &[&str],
) {
    for (key, value) in fields {
        match (into.get_mut(key), value) {
            (Some(Value::String(text)), Value::String(more))
                if appended.contains(&key.as_str()) =>
            {
                text.push_str(more);
            }
            (Some(Value::Array(items)), Value::Array(more))
                if COLLECTED.contains(&key.as_str()) =>
            {
                items.extend(more.iter().cloned());
            }
            (Some(_), Value::Null) => {}
            (_, value) => {
                into.insert(key.clone(), value.clone());
            }
        }
    }
}

impl crate::wire::document::Serves<crate::operation::Completion> for ChatResponse {}

impl Reassemble<ChatRecord> for ChatResponse {
    fn absorb(&mut self, ChatRecord(fields): &ChatRecord) {
        if let Some(Value::Object(message)) = fields.get("message") {
            match self.document.get_mut("message") {
                Some(Value::Object(into)) => fold(into, message, APPENDED),
                _ => {
                    let message = Value::Object(message.clone());
                    self.document.insert("message".to_owned(), message);
                }
            }
        }
        let rest = fields.iter().filter(|(key, _)| *key != "message");
        fold(&mut self.document, rest, &[]);
    }

    fn finish(self) -> Value {
        if self.document.is_empty() {
            Value::Null
        } else {
            Value::Object(self.document)
        }
    }
}