quvyta-cli 0.1.0-alpha.5

//! A conversation with the model, kept in the Messages shape it is sent in.
//!
//! The messages are kept exactly as the API reads them, so what is saved is what was said. What a
//! request sends is a *view* of them: a long conversation goes out trimmed to a size, with a note
//! standing in where something was left out. Nothing kept is translated on the way.

use std::collections::HashMap;
use std::ops::Range;

use serde_json::{Value, json};

/// A tool call the model made.
#[derive(Debug, Clone, PartialEq)]
pub struct Call {
    /// The id the result must name.
    pub id: String,
    /// The tool.
    pub name: String,
    /// What it was given.
    pub input: Value,
}

/// The messages of one conversation.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Conversation {
    messages: Vec<Value>,
}

impl Conversation {
    /// The messages, oldest first.
    #[must_use]
    pub fn messages(&self) -> &[Value] {
        &self.messages
    }

    /// The messages as they go out in a request that is at most `budget` bytes: the conversation
    /// as it is when it fits, and a trimmed view of it when it does not. What is kept is kept
    /// whole, and nothing is asked of the model to make the view smaller.
    ///
    /// Three parts are never shortened: the first message, which usually holds the task; the
    /// newest message, which holds the tool results the model is waiting for an answer to; and
    /// the shape of a call and its result, so that every call still has its answer. What is
    /// shortened first is the tool results of the turns in between, oldest first, since they are
    /// the biggest part of an old turn and the least needed of it; the text of those turns
    /// follows in the same order, only if the results alone are not enough. Wherever something
    /// is left out, an English note stands in its place saying what was left out, and the
    /// conversation kept here is not touched.
    #[must_use]
    pub fn view(&self, budget: usize) -> Vec<Value> {
        let mut view = self.messages.clone();
        if bytes(&view) <= budget {
            return view;
        }
        let names = tool_names(&self.messages);
        let old = 1..view.len().saturating_sub(1);
        if shorten(&mut view, &old, budget, "tool_result", &names) {
            return view;
        }
        shorten(&mut view, &old, budget, "text", &names);
        view
    }

    /// Whether nothing has been said yet.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    /// Adds what the person wrote. Right after another user message (a reply that was stopped
    /// before it said anything) the text joins that message, since the API wants the roles to
    /// take turns.
    pub fn say(&mut self, text: &str) {
        let block = json!({"type": "text", "text": text});
        self.push_user(vec![block]);
    }

    /// Adds the model's reply. Nothing is added for a reply with no content.
    pub fn reply(&mut self, content: Vec<Value>) {
        if !content.is_empty() {
            self.messages.push(json!({"role": "assistant", "content": Value::Array(content)}));
        }
    }

    /// Adds the results of tool calls, each a `tool_result` block.
    pub fn results(&mut self, results: Vec<Value>) {
        if !results.is_empty() {
            self.push_user(results);
        }
    }

    fn push_user(&mut self, blocks: Vec<Value>) {
        if let Some(last) = self.messages.last_mut()
            && last["role"] == "user"
            && let Some(content) = last.get_mut("content").and_then(Value::as_array_mut)
        {
            content.extend(blocks);
            return;
        }
        self.messages.push(json!({"role": "user", "content": blocks}));
    }

    /// The tool calls of the last message that have no result yet, in the order they were made.
    #[must_use]
    pub fn open_calls(&self) -> Vec<Call> {
        let Some(last) = self.messages.last().filter(|message| message["role"] == "assistant") else {
            return Vec::new();
        };
        blocks(last)
            .filter(|block| block["type"] == "tool_use")
            .map(|block| Call {
                id: block["id"].as_str().unwrap_or_default().to_owned(),
                name: block["name"].as_str().unwrap_or_default().to_owned(),
                input: block.get("input").cloned().unwrap_or(Value::Null),
            })
            .collect()
    }

    /// The conversation as it is saved.
    #[must_use]
    pub fn to_json(&self) -> Value {
        json!({"version": 1, "messages": self.messages})
    }

    /// Whether `value` is a conversation as this program keeps one. A file that is not one holds
    /// somebody's thread in a shape this program cannot show, which is not the same as a file
    /// that holds nothing.
    #[must_use]
    pub fn holds(value: &Value) -> bool {
        value["messages"].is_array()
    }

    /// A saved conversation. Anything that is not a message is left out, and calls that were
    /// waiting when the program closed are answered as not done, so the conversation can go on.
    #[must_use]
    pub fn from_json(value: &Value) -> Self {
        let messages = value["messages"]
            .as_array()
            .map(|messages| {
                messages
                    .iter()
                    .filter(|message| {
                        matches!(message["role"].as_str(), Some("user" | "assistant")) && message["content"].is_array()
                    })
                    .cloned()
                    .collect()
            })
            .unwrap_or_default();
        let mut conversation = Self { messages: mend(messages) };
        let open = conversation.open_calls();
        conversation.results(open.iter().map(|call| result(&call.id, CLOSED, true)).collect());
        conversation
    }
}

/// `messages` as the API accepts them, however an older program or a crash left them: every
/// call's input is an object, and every call in the middle of the thread has its result first
/// thing in the message after it, where a result that was never written is added as not done.
/// Nothing that was said is taken out. The calls of the last message are left to the caller,
/// which answers them with its own sentence.
fn mend(messages: Vec<Value>) -> Vec<Value> {
    let mut mended: Vec<Value> = Vec::with_capacity(messages.len());
    let mut waiting: Vec<String> = Vec::new();
    for mut message in messages {
        if !waiting.is_empty() {
            let mut calls = std::mem::take(&mut waiting);
            if message["role"] == "user" {
                // Each result already there answers one call of its id.
                for block in blocks(&message).filter(|block| block["type"] == "tool_result") {
                    let id = block["tool_use_id"].as_str().unwrap_or_default();
                    if let Some(at) = calls.iter().position(|call| call == id) {
                        calls.remove(at);
                    }
                }
            }
            let answers: Vec<Value> = calls.iter().map(|id| result(id, CLOSED, true)).collect();
            if answers.is_empty() {
            } else if message["role"] == "user"
                && let Some(content) = message.get_mut("content").and_then(Value::as_array_mut)
            {
                content.splice(0..0, answers);
            } else {
                mended.push(json!({"role": "user", "content": answers}));
            }
        }
        if message["role"] == "assistant" {
            waiting = calls_of(&mut message);
        }
        mended.push(message);
    }
    mended
}

/// The ids of the calls in `message`, with each call's input made an object on the way.
fn calls_of(message: &mut Value) -> Vec<String> {
    let Some(content) = message.get_mut("content").and_then(Value::as_array_mut) else { return Vec::new() };
    let mut ids = Vec::new();
    for block in content.iter_mut().filter(|block| block["type"] == "tool_use") {
        if !block["input"].is_object() {
            let kept = match &block["input"] {
                Value::String(text) => text.clone(),
                Value::Null => String::new(),
                other => other.to_string(),
            };
            block["input"] = if kept.is_empty() { json!({}) } else { json!({ crate::provider::UNPARSED: kept }) };
        }
        ids.push(block["id"].as_str().unwrap_or_default().to_owned());
    }
    ids
}

/// What a call that was waiting when the program closed is answered with.
const CLOSED: &str = "Not done: the program was closed before this ran.";

/// The blocks of `message`.
pub fn blocks(message: &Value) -> impl Iterator<Item = &Value> {
    message["content"].as_array().into_iter().flatten()
}

/// A `tool_result` block for the call `id`.
#[must_use]
pub fn result(id: &str, text: &str, is_error: bool) -> Value {
    json!({"type": "tool_result", "tool_use_id": id, "content": text, "is_error": is_error})
}

/// How many bytes `messages` takes in a request.
fn bytes(messages: &[Value]) -> usize {
    // Counted as written rather than written into a string: a long conversation is megabytes of
    // JSON, and the view measures it more than once per request.
    struct Count(usize);
    impl std::io::Write for Count {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            self.0 += buf.len();
            Ok(buf.len())
        }
        fn flush(&mut self) -> std::io::Result<()> {
            Ok(())
        }
    }
    let mut count = Count(0);
    serde_json::to_writer(&mut count, messages).map_or(0, |()| count.0)
}

/// What each call was called, by its id: a note can then name the tool whose output it left out,
/// which is read rather than guessed.
fn tool_names(messages: &[Value]) -> HashMap<&str, &str> {
    messages
        .iter()
        .flat_map(blocks)
        .filter(|block| block["type"] == "tool_use")
        .filter_map(|block| Some((block["id"].as_str()?, block["name"].as_str()?)))
        .collect()
}

/// Stands `block` in for itself, with a note in place of its text, when it is of type `wanted`
/// and too big to go out whole. `None` when there is nothing to gain, so a block that is already
/// small is left alone rather than made bigger.
fn note_for(block: &Value, wanted: &str, names: &HashMap<&str, &str>) -> Option<Value> {
    if block["type"].as_str()? != wanted {
        return None;
    }
    let (field, what) = match wanted {
        "tool_result" => {
            let id = block["tool_use_id"].as_str().unwrap_or_default();
            ("content", format!("output from {}", names.get(id).copied().unwrap_or("an earlier call")))
        }
        "text" => ("text", "earlier text".to_owned()),
        _ => return None,
    };
    let mut note = block.clone();
    note[field] = Value::String(format!(
        "(Left out: {} bytes of {what}. The whole of it is in the conversation saved for this folder.)",
        block[field].as_str()?.len()
    ));
    (bytes(std::slice::from_ref(&note)) < bytes(std::slice::from_ref(block))).then_some(note)
}

/// Replaces the blocks of type `wanted` in the messages `old` holds, oldest first, until `view`
/// fits `budget` or there is nothing left that can be shortened. `true` when it fits.
fn shorten(view: &mut [Value], old: &Range<usize>, budget: usize, wanted: &str, names: &HashMap<&str, &str>) -> bool {
    let mut over = bytes(view).saturating_sub(budget);
    for message in &mut view[old.clone()] {
        if over == 0 {
            return true;
        }
        let Some(blocks) = message.get_mut("content").and_then(Value::as_array_mut) else { continue };
        for block in blocks.iter_mut() {
            if over == 0 {
                return true;
            }
            let Some(note) = note_for(block, wanted, names) else { continue };
            over = over.saturating_sub(bytes(std::slice::from_ref(block)) - bytes(std::slice::from_ref(&note)));
            *block = note;
        }
    }
    over == 0
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn roles_take_turns_even_after_a_reply_that_said_nothing() {
        let mut conversation = Conversation::default();
        conversation.say("first");
        conversation.reply(Vec::new());
        conversation.say("second");
        assert_eq!(conversation.messages().len(), 1, "the two texts share one user message");
        assert_eq!(conversation.messages()[0]["content"][1]["text"], "second");
        conversation.reply(vec![json!({"type": "text", "text": "answer"})]);
        assert_eq!(conversation.messages()[1]["role"], "assistant");
    }

    #[test]
    fn open_calls_are_the_last_replys_tool_uses() {
        let mut conversation = Conversation::default();
        conversation.say("look");
        conversation.reply(vec![
            json!({"type": "text", "text": "Looking."}),
            json!({"type": "tool_use", "id": "a", "name": "list_dir", "input": {"path": "."}}),
            json!({"type": "tool_use", "id": "b", "name": "read_file", "input": {"path": "x"}}),
        ]);
        let open = conversation.open_calls();
        assert_eq!(open.iter().map(|call| call.id.as_str()).collect::<Vec<_>>(), ["a", "b"]);
        assert_eq!(open[1].input, json!({"path": "x"}));
        conversation.results(vec![result("a", "x/", false), result("b", "text", false)]);
        assert!(conversation.open_calls().is_empty());
    }

    #[test]
    fn a_saved_conversation_comes_back_and_waiting_calls_are_answered() {
        let mut conversation = Conversation::default();
        conversation.say("run it");
        conversation
            .reply(vec![json!({"type": "tool_use", "id": "r", "name": "run_command", "input": {"command": "ls"}})]);
        let saved = conversation.to_json();
        let back = Conversation::from_json(&saved);
        assert!(back.open_calls().is_empty(), "the waiting call is answered");
        let last = back.messages().last().expect("a message");
        assert_eq!(last["content"][0]["tool_use_id"], "r");
        assert_eq!(last["content"][0]["is_error"], true);
        assert_eq!(Conversation::from_json(&json!({"messages": [{"role": "system"}, 3]})), Conversation::default());
    }

    #[test]
    fn a_message_whose_content_is_not_a_list_is_left_out_and_the_rest_stays() {
        let back = Conversation::from_json(&json!({"messages": [
            {"role": "user", "content": [{"type": "text", "text": "kept"}]},
            {"role": "assistant", "content": 7},
            {"role": "user", "content": "a plain string"},
            {"role": "user", "content": [{"type": "text", "text": "kept too"}]},
        ]}));
        let texts: Vec<_> = back.messages().iter().map(|message| message["content"][0]["text"].clone()).collect();
        assert_eq!(texts, [json!("kept"), json!("kept too")], "the two good messages, in order");
    }

    #[test]
    fn a_call_with_no_id_gets_one_and_never_panics_on_a_missing_field() {
        let mut conversation = Conversation::default();
        conversation.reply(vec![
            json!({"type": "tool_use", "name": "list_dir", "input": {"path": "."}}),
            json!({"type": "tool_use", "id": 5, "name": null}),
        ]);
        let open = conversation.open_calls();
        assert_eq!(open.len(), 2, "both blocks are calls to be taken");
        assert_eq!(open[0].id, "", "a call that names no id gets none");
        assert_eq!(open[1].name, "", "nor a call that names no tool");
        assert_eq!(open[1].input, Value::Null, "and no input either");

        // The same blocks coming back from a file are answered, empty id and all.
        let back = Conversation::from_json(&conversation.to_json());
        let answered: Vec<_> = blocks(back.messages().last().expect("the results"))
            .map(|block| block["tool_use_id"].as_str().unwrap_or_default().to_owned())
            .collect();
        assert_eq!(answered, ["", ""], "both are answered, so the next request is one the API accepts");
    }

    #[test]
    fn two_calls_with_one_id_are_answered_one_for_one() {
        let back = Conversation::from_json(&json!({"messages": [
            {"role": "assistant", "content": [
                {"type": "tool_use", "id": "same", "name": "list_dir", "input": {}},
                {"type": "tool_use", "id": "same", "name": "read_file", "input": {}},
            ]},
        ]}));
        assert!(back.open_calls().is_empty());
        let answered: Vec<_> = blocks(back.messages().last().expect("the results"))
            .map(|block| block["tool_use_id"].as_str().unwrap_or_default().to_owned())
            .collect();
        assert_eq!(answered, ["same", "same"], "both calls are answered, each in its own place");
    }

    #[test]
    fn a_result_of_a_call_that_is_not_there_is_kept_and_the_call_it_left_waiting_is_answered() {
        let back = Conversation::from_json(&json!({"messages": [
            {"role": "assistant", "content": [{"type": "tool_use", "id": "here", "name": "list_dir", "input": {}}]},
            {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "elsewhere", "content": "stray"}]},
        ]}));
        // The call of the file had no result, which the API refuses on every request; it is
        // answered first thing in the message after it, and what the file said stays behind it.
        assert_eq!(back.messages().len(), 2);
        let last = back.messages().last().expect("a message");
        let said: Vec<_> = blocks(last)
            .map(|block| {
                (block["tool_use_id"].as_str().unwrap_or_default(), block["content"].as_str().unwrap_or_default())
            })
            .collect();
        assert_eq!(said, [("here", CLOSED), ("elsewhere", "stray")]);
    }

    #[test]
    fn a_call_answered_in_the_file_is_not_answered_twice() {
        let saved = json!({"messages": [
            {"role": "user", "content": [{"type": "text", "text": "look"}]},
            {"role": "assistant", "content": [{"type": "tool_use", "id": "a", "name": "list_dir", "input": {}}]},
            {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "a", "content": "x/"}]},
        ]});
        let back = Conversation::from_json(&saved);
        assert_eq!(
            back.messages(),
            saved["messages"].as_array().expect("messages").as_slice(),
            "a whole thread is kept as it is"
        );
    }
}

#[cfg(test)]
mod view_tests {
    use super::*;

    /// A conversation of `turns` turns, each an answer of `answer` bytes and a result of
    /// `result` bytes, after the task that started it.
    fn long(answer: usize, result: usize, turns: usize) -> Conversation {
        let mut conversation = Conversation::default();
        conversation.say("the task");
        for n in 0..turns {
            let id = format!("c{n}");
            conversation.reply(vec![
                json!({"type": "text", "text": "a".repeat(answer)}),
                json!({"type": "tool_use", "id": id, "name": "read_file", "input": {"path": "a.rs"}}),
            ]);
            conversation.results(vec![result_of(&id, &"x".repeat(result))]);
        }
        conversation
    }

    /// A `tool_result` block for the call `id`, holding `text`.
    fn result_of(id: &str, text: &str) -> Value {
        super::result(id, text, false)
    }

    /// What each result of a view holds, in the order they go out.
    fn results(view: &[Value]) -> Vec<String> {
        view.iter()
            .flat_map(blocks)
            .filter(|block| block["type"] == "tool_result")
            .map(|block| block["content"].as_str().unwrap_or_default().to_owned())
            .collect()
    }

    /// Whether `text` is a note standing in for what was left out.
    fn left_out(text: &str) -> bool {
        text.starts_with("(Left out:")
    }

    #[test]
    fn a_conversation_under_the_budget_is_sent_whole() {
        let conversation = long(7, 100, 3);
        let budget = bytes(conversation.messages());
        assert_eq!(conversation.view(budget), conversation.messages());
        assert_eq!(conversation.view(budget + 1), conversation.messages());
    }

    #[test]
    fn the_oldest_results_are_left_out_first_and_the_newest_turn_is_whole() {
        let conversation = long(7, 4_000, 4);
        let view = conversation.view(5_500);
        assert!(bytes(&view) <= 5_500, "{} bytes", bytes(&view));

        let sent = results(&view);
        assert_eq!(sent.len(), 4, "no result is dropped, only shortened");
        assert!(sent[..3].iter().all(|text| left_out(text)), "the oldest three go first: {sent:?}");
        assert!(sent[3..].iter().all(|text| text.len() == 4_000), "and what still fits stays whole: {sent:?}");

        let note = &sent[0];
        assert!(note.starts_with("(Left out: 4000 bytes of output from read_file."), "{note}");
        assert!(note.ends_with("The whole of it is in the conversation saved for this folder.)"), "{note}");
        assert_eq!(view.first(), conversation.messages().first(), "the first message is never touched");
    }

    #[test]
    fn a_result_the_model_waits_for_keeps_its_answer_and_its_call() {
        let view = long(7, 4_000, 4).view(1_000);
        for wanted in ["tool_use", "tool_result"] {
            let found = view.iter().flat_map(blocks).filter(|block| block["type"] == wanted).count();
            assert_eq!(found, 4, "every {wanted} is still there");
        }
        for call in view.iter().flat_map(blocks).filter(|block| block["type"] == "tool_use") {
            let id = call["id"].as_str().expect("an id");
            let answered =
                view.iter().flat_map(blocks).any(|block| block["type"] == "tool_result" && block["tool_use_id"] == id);
            assert!(answered, "{id} is still answered");
        }
    }

    #[test]
    fn a_result_too_small_to_gain_from_a_note_is_left_whole() {
        let conversation = long(7, 20, 3);
        let sent = results(&conversation.view(0));
        assert!(sent.iter().all(|text| !left_out(text)), "{sent:?}: a note would be bigger than the result");
    }

    #[test]
    fn old_text_is_shortened_when_there_is_nothing_to_gain_in_the_results() {
        let conversation = long(4_000, 20, 4);
        let view = conversation.view(5_500);
        assert!(bytes(&view) <= 5_500, "{} bytes", bytes(&view));

        let answers: Vec<&str> = view
            .iter()
            .filter(|message| message["role"] == "assistant")
            .flat_map(blocks)
            .filter(|block| block["type"] == "text")
            .map(|block| block["text"].as_str().unwrap_or_default())
            .collect();
        assert_eq!(answers.len(), 4, "no text is dropped, only shortened");
        assert!(answers[..3].iter().all(|text| left_out(text)), "the oldest three go first: {answers:?}");
        assert_eq!(answers[3].len(), 4_000, "and the newest answer is whole");
        for note in &answers[..3] {
            assert!(note.contains("4000 bytes of earlier text"), "{note}");
        }
    }

    #[test]
    fn a_thinking_block_is_never_touched() {
        let mut conversation = long(4_000, 20, 3);
        conversation.reply(vec![json!({"type": "thinking", "thinking": "z".repeat(4_000), "signature": "sig"})]);
        let view = conversation.view(1_000);
        assert!(bytes(&view) > 1_000, "what cannot be shortened goes out as it is: {} bytes", bytes(&view));
        let thinking = view.iter().flat_map(blocks).find(|block| block["type"] == "thinking").expect("the thinking");
        assert_eq!(thinking["thinking"], json!("z".repeat(4_000)), "it is whole");
        assert_eq!(thinking["signature"], json!("sig"), "and its signature is not broken");
    }

    #[test]
    fn what_cannot_be_shortened_goes_out_whole_even_past_the_budget() {
        let conversation = long(7, 4_000, 3);
        let view = conversation.view(500);
        assert!(bytes(&view) > 500, "the newest result alone is over the budget: {} bytes", bytes(&view));
        assert_eq!(view.last(), conversation.messages().last());
    }

    #[test]
    fn the_stored_conversation_is_never_changed_by_a_view() {
        let conversation = long(7, 4_000, 4);
        let kept = conversation.clone();
        let _ = conversation.view(500);
        assert_eq!(conversation, kept);
    }
}