path: root/crates/asist-llm/src/chat.rs

//! Conversation history in the format the chat API expects.

use serde::{Deserialize, Serialize};
use serde_json::{json, Value};

use asist_core::tools::{ToolCall, ToolOutcome};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
    Tool,
}

impl Role {
    fn as_str(self) -> &'static str {
        match self {
            Role::System => "system",
            Role::User => "user",
            Role::Assistant => "assistant",
            Role::Tool => "tool",
        }
    }
}

#[derive(Debug, Clone)]
pub struct Message {
    pub role: Role,
    pub content: String,
    /// Tools the assistant asked for in this turn.
    pub tool_calls: Vec<ToolCall>,
    /// For `tool` role messages: which call they answer.
    pub tool_call_id: Option<String>,
}

impl Message {
    pub fn system(content: impl Into<String>) -> Self {
        Self::plain(Role::System, content)
    }

    pub fn user(content: impl Into<String>) -> Self {
        Self::plain(Role::User, content)
    }

    pub fn assistant(content: impl Into<String>) -> Self {
        Self::plain(Role::Assistant, content)
    }

    fn plain(role: Role, content: impl Into<String>) -> Self {
        Self {
            role,
            content: content.into(),
            tool_calls: Vec::new(),
            tool_call_id: None,
        }
    }

    /// Assistant turn that asked for tools instead of speaking.
    pub fn tool_request(content: String, tool_calls: Vec<ToolCall>) -> Self {
        Self {
            role: Role::Assistant,
            content,
            tool_calls,
            tool_call_id: None,
        }
    }

    /// Result handed back to the model.
    pub fn tool_result(outcome: &ToolOutcome) -> Self {
        Self {
            role: Role::Tool,
            content: outcome.output.clone(),
            tool_calls: Vec::new(),
            tool_call_id: Some(outcome.id.clone()),
        }
    }

    pub fn to_json(&self) -> Value {
        let mut object = json!({ "role": self.role.as_str(), "content": self.content });
        if !self.tool_calls.is_empty() {
            object["tool_calls"] = Value::Array(
                self.tool_calls
                    .iter()
                    .map(|call| {
                        json!({
                            "id": call.id,
                            "type": "function",
                            "function": { "name": call.name, "arguments": call.arguments }
                        })
                    })
                    .collect(),
            );
        }
        if let Some(id) = &self.tool_call_id {
            object["tool_call_id"] = json!(id);
        }
        object
    }
}

/// History with a fixed system prompt and a sliding window of turns, so
/// the conversation does not grow without end.
#[derive(Debug, Clone)]
pub struct Conversation {
    system: Message,
    turns: Vec<Message>,
    max_turns: usize,
}

impl Conversation {
    pub fn new(system_prompt: impl Into<String>, max_turns: usize) -> Self {
        Self {
            system: Message::system(system_prompt),
            turns: Vec::new(),
            max_turns,
        }
    }

    /// Changes the system prompt without touching the history.
    ///
    /// The turn alternates between the tool guide and the style guide, and the
    /// history must survive the switch: if it were reset, the model would lose
    /// the results it just asked for.
    pub fn set_system(&mut self, prompt: &str) {
        self.system = Message::system(prompt);
    }

    pub fn push(&mut self, message: Message) {
        self.turns.push(message);
        self.trim();
    }

    pub fn extend(&mut self, messages: impl IntoIterator<Item = Message>) {
        self.turns.extend(messages);
        self.trim();
    }

    /// Trims to `max_turns` user utterances, never leaving an orphan `tool`
    /// role message at the start: the API rejects it unless it comes right
    /// after the call that caused it.
    fn trim(&mut self) {
        if self.max_turns == 0 {
            self.turns.clear();
            return;
        }
        let user_positions: Vec<usize> = self
            .turns
            .iter()
            .enumerate()
            .filter(|(_, m)| m.role == Role::User)
            .map(|(i, _)| i)
            .collect();
        if user_positions.len() <= self.max_turns {
            return;
        }
        let cut = user_positions[user_positions.len() - self.max_turns];
        self.turns.drain(..cut);
    }

    pub fn messages(&self) -> Vec<&Message> {
        std::iter::once(&self.system)
            .chain(self.turns.iter())
            .collect()
    }

    pub fn to_json(&self) -> Value {
        Value::Array(self.messages().into_iter().map(Message::to_json).collect())
    }

    pub fn len(&self) -> usize {
        self.turns.len()
    }

    pub fn is_empty(&self) -> bool {
        self.turns.is_empty()
    }

    pub fn clear(&mut self) {
        self.turns.clear();
    }
}

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

    #[test]
    fn system_prompt_always_goes_first() {
        let mut chat = Conversation::new("sé breve", 4);
        chat.push(Message::user("hola"));
        let json = chat.to_json();
        assert_eq!(json[0]["role"], "system");
        assert_eq!(json[0]["content"], "sé breve");
        assert_eq!(json[1]["role"], "user");
    }

    #[test]
    fn changing_the_system_prompt_keeps_the_history() {
        let mut chat = Conversation::new("primera", 4);
        chat.push(Message::user("hola"));
        chat.set_system("segunda");
        let json = chat.to_json();
        assert_eq!(json[0]["content"], "segunda");
        assert_eq!(
            json[1]["content"], "hola",
            "the history must not be lost on the switch"
        );
    }

    #[test]
    fn history_is_trimmed_by_user_turns() {
        let mut chat = Conversation::new("s", 2);
        for i in 0..5 {
            chat.push(Message::user(format!("p{i}")));
            chat.push(Message::assistant(format!("r{i}")));
        }
        let messages = chat.messages();
        let users: Vec<&str> = messages
            .iter()
            .filter(|m| m.role == Role::User)
            .map(|m| m.content.as_str())
            .collect();
        assert_eq!(users, vec!["p3", "p4"]);
    }

    #[test]
    fn trimming_does_not_orphan_a_tool_result() {
        // The API rejects a `tool` message that does not follow the call that
        // asked for it, so the cut must land on a user turn.
        let mut chat = Conversation::new("s", 1);
        chat.push(Message::user("p0"));
        chat.push(Message::tool_request(
            String::new(),
            vec![ToolCall {
                id: "a".into(),
                name: "t".into(),
                arguments: "{}".into(),
            }],
        ));
        chat.push(Message {
            role: Role::Tool,
            content: "ok".into(),
            tool_calls: vec![],
            tool_call_id: Some("a".into()),
        });
        chat.push(Message::user("p1"));

        let roles: Vec<Role> = chat.messages().iter().map(|m| m.role).collect();
        assert_eq!(roles, vec![Role::System, Role::User]);
    }

    #[test]
    fn zero_turns_means_no_memory() {
        let mut chat = Conversation::new("s", 0);
        chat.push(Message::user("hola"));
        assert_eq!(chat.messages().len(), 1);
    }

    #[test]
    fn tool_call_is_serialized_as_the_api_expects() {
        let message = Message::tool_request(
            String::new(),
            vec![ToolCall {
                id: "call_1".into(),
                name: "hora_actual".into(),
                arguments: "{}".into(),
            }],
        );
        let json = message.to_json();
        assert_eq!(json["tool_calls"][0]["type"], "function");
        assert_eq!(json["tool_calls"][0]["function"]["name"], "hora_actual");
    }
}