monothink 0.4.0

/*
 * This file is part of monothink.
 *
 * monothink is free software: you can redistribute it and/or modify it under the terms
 * of the GNU General Public License as published by the Free Software Foundation,
 * either version 3 of the License, or (at your option) any later version.
 *
 * monothink is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with Wavething.
 * If not, see <https://www.gnu.org/licenses/>.
 */

use anyhow::Result;
use markdown::mdast::Node;
use rig::message::{ProviderCallId, ToolCallId};
use serde::{Deserialize, Serialize};
use std::{fs::File, io::Write, path::Path};
use uuid::Uuid;

#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)]
pub struct Document {
    pub blocks: Vec<Block>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ListItemMarker {
    Bullet,
    Checkbox(bool),
    Number(u32),
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ListItem {
    pub marker: ListItemMarker,
    pub content: Vec<Block>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct TableCell {
    pub content: Vec<Inline>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct TableRow {
    pub content: Vec<TableCell>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Block {
    Code {
        lang: String,
        content: String,
        copy_id: u32,
    },
    Heading {
        depth: u8,
        content: Vec<Inline>,
    },
    List {
        content: Vec<ListItem>,
    },
    Paragraph(Vec<Inline>),
    Quote(Vec<Block>),
    Table {
        content: Vec<TableRow>,
    },
    ThematicBreak,
    Unsupported(String),
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Inline {
    Code(String),
    Emphasis(Vec<Inline>),
    StrongEmphasis(Vec<Inline>),
    Link {
        url: String,
        title: Option<String>,
        label: Vec<Inline>,
    },
    Image(String), // title, url, description
    LineBreak,
    Text(String),
    Unsupported(String),
}

#[derive(Debug, Clone, PartialEq)]
pub struct MessageContent {
    unprocessed: String,
    structured: Document,
}

impl Serialize for MessageContent {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        self.unprocessed.serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for MessageContent {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let unprocessed = String::deserialize(deserializer)?;

        Ok(Self {
            unprocessed,
            structured: Document::default(),
        })
    }
}

impl MessageContent {
    pub fn new(data: &str) -> Self {
        Self {
            unprocessed: data.to_string(),
            structured: parse_from_markdown(data).unwrap(),
        }
    }

    pub fn append(&mut self, data: &str) {
        self.unprocessed.push_str(data);
        self.structured = parse_from_markdown(&self.unprocessed).unwrap();
    }

    pub fn content(&self) -> &Document {
        &self.structured
    }

    pub fn raw(&self) -> &String {
        &self.unprocessed
    }

    fn write(&mut self, text: &str) {
        self.unprocessed = text.to_string();
        self.structured = parse_from_markdown(&self.unprocessed).unwrap();
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolInvocationContent {
    pub name: String,
    pub arguments: serde_json::Value,
    // this is a tool invocation id assigned by rig
    pub call_id: ToolCallId,
    // this is an optional tool call id assigned by the provider
    // we keep it as the provider needs it for matching calls and responses
    pub provider_id: Option<ProviderCallId>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolOutputContent {
    // name of the tool that produced the output
    pub name: String,
    pub output: String,
    pub success: bool,
    // this is a tool invocation id assigned by rig
    pub call_id: ToolCallId,
    // this is an optional tool call id assigned by the provider
    // we keep it as the provider needs it for matching calls and responses
    pub provider_id: Option<ProviderCallId>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum MessageState {
    // an intermediate message while working on an answer
    Intermediate,

    // the final response to a message
    Final,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum NotificationType {
    // information that is relevant for someone who is interested in technical details
    Debug,

    // information that is relevant for the user to know
    Info,

    // warn the user that an action was completed but not yield the expected result
    Warn,

    // inform the user that an action failed to complete successfully
    Error,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ConversationItemKind {
    UserMessage(MessageContent),
    AssistantMessage(MessageContent, MessageState),
    AssistantReasoning(MessageContent),
    ToolInvocation(ToolInvocationContent),
    ToolOutput(ToolOutputContent),
    Notification(NotificationType, String),
}

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct ConversationItem {
    #[serde(flatten)]
    pub content: ConversationItemKind,
    pub id: Uuid,
    // started_at: DateTime<Utc>,
    // completed_at: Option<DateTime<Utc>>,
}

impl ConversationItem {
    pub fn new(content: ConversationItemKind) -> Self {
        Self {
            content,
            id: Uuid::new_v4(),
        }
    }

    pub fn is_user_message(&self) -> bool {
        matches!(self.content, ConversationItemKind::UserMessage(_))
    }

    pub fn is_assistant_message(&self, state: Option<MessageState>) -> bool {
        if let ConversationItemKind::AssistantMessage(_, ref msg_state) = self.content {
            if let Some(compare_state) = state {
                compare_state == *msg_state
            } else {
                true
            }
        } else {
            false
        }
    }

    pub fn is_notification(&self, notification_type: Option<NotificationType>) -> bool {
        if let ConversationItemKind::Notification(ref msg_nt, _) = self.content {
            if let Some(compare_nt) = notification_type {
                compare_nt == *msg_nt
            } else {
                true
            }
        } else {
            false
        }
    }

    pub fn is_reasoning(&self) -> bool {
        matches!(self.content, ConversationItemKind::AssistantReasoning(_))
    }

    pub fn is_tool_invocation(&self) -> bool {
        matches!(self.content, ConversationItemKind::ToolInvocation(_))
    }

    pub fn is_tool_output(&self) -> bool {
        matches!(self.content, ConversationItemKind::ToolOutput(_))
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Conversation {
    id: Uuid,
    title: Option<String>,
    model: Option<String>,
    messages: Vec<ConversationItem>,
}

impl Conversation {
    pub fn new(model: Option<String>) -> Self {
        Conversation {
            id: Uuid::new_v4(),
            title: None,
            model,
            messages: Vec::new(),
        }
    }

    pub fn add_assistant_reasoning_chunk(&mut self, chunk: &str) -> &ConversationItem {
        if let Some(ConversationItem {
            content: ConversationItemKind::AssistantReasoning(msg),
            id: _,
        }) = self.messages.last_mut()
        {
            msg.append(chunk);
        } else {
            let item = ConversationItem::new(ConversationItemKind::AssistantReasoning(
                MessageContent::new(chunk),
            ));
            self.messages.push(item);
        }
        self.assign_copy_ids();
        self.messages.last().expect(
            "Getting last message after pushing or modifying existing agent reasoning failed",
        )
    }

    pub fn add_assistant_response_chunk(&mut self, chunk: &str) -> &ConversationItem {
        if let Some(ConversationItem {
            content: ConversationItemKind::AssistantMessage(msg, MessageState::Intermediate),
            id: _,
        }) = self.messages.last_mut()
        {
            msg.append(chunk);
        } else {
            let item = ConversationItem::new(ConversationItemKind::AssistantMessage(
                MessageContent::new(chunk),
                MessageState::Intermediate,
            ));
            self.messages.push(item);
        }
        self.assign_copy_ids();
        self.messages.last().expect(
            "Getting last message after pushing or modifying existing agent response failed",
        )
    }

    pub fn add_user_prompt(&mut self, prompt: &str) -> &ConversationItem {
        let item = ConversationItem::new(ConversationItemKind::UserMessage(MessageContent::new(
            prompt,
        )));
        self.messages.push(item);
        self.assign_copy_ids();

        self.messages
            .last()
            .expect("Getting last message after pushing a user prompt failed")
    }

    pub fn add_notification(
        &mut self,
        notification_type: NotificationType,
        msg: String,
    ) -> &ConversationItem {
        let item =
            ConversationItem::new(ConversationItemKind::Notification(notification_type, msg));
        self.messages.push(item);

        self.messages
            .last()
            .expect("Getting last message after pushing a user prompt failed")
    }

    pub fn add_tool_invocation(
        &mut self,
        name: String,
        args: serde_json::Value,
        call_id: ToolCallId,
        provider_id: Option<ProviderCallId>,
    ) -> &ConversationItem {
        let item = ConversationItem::new(ConversationItemKind::ToolInvocation(
            ToolInvocationContent {
                name,
                arguments: args,
                call_id,
                provider_id,
            },
        ));
        self.messages.push(item);

        self.messages
            .last()
            .expect("Getting last message after pushing a user prompt failed")
    }

    pub fn add_tool_result(
        &mut self,
        name: String,
        output: String,
        success: bool,
        call_id: ToolCallId,
        provider_id: Option<ProviderCallId>,
    ) -> &ConversationItem {
        let item = ConversationItem::new(ConversationItemKind::ToolOutput(ToolOutputContent {
            output,
            success,
            call_id,
            provider_id,
            name,
        }));
        self.messages.push(item);

        self.messages
            .last()
            .expect("Getting last message after pushing a user prompt failed")
    }

    pub fn id(&self) -> Uuid {
        self.id
    }

    pub fn messages(&self) -> &Vec<ConversationItem> {
        &self.messages
    }

    pub fn set_final_assistant_response(&mut self, text: &str) -> &ConversationItem {
        // assumption: the last streamed assistant message equals the final response
        // according to documentation for rig 0.42 this assumption holds
        let outdated_item = self
            .messages
            .iter_mut()
            .rev()
            .find(|it| matches!(it.content, ConversationItemKind::AssistantMessage(_, _)));
        if let Some(ConversationItem {
            content: ConversationItemKind::AssistantMessage(msg, kind),
            id: _,
        }) = outdated_item
            && *kind == MessageState::Intermediate
        {
            msg.write(text);
            *kind = MessageState::Final;
        } else {
            // note: this should never happen. If I ever observe it, I should adjust my assumptions
            // it might happen, if a system message comes before the final response
            let item = ConversationItem::new(ConversationItemKind::AssistantMessage(
                MessageContent::new(text),
                MessageState::Final,
            ));
            self.messages.push(item);
        }
        self.assign_copy_ids();

        self.messages
            .iter()
            .rev()
            .find(|it| {
                matches!(
                    it.content,
                    ConversationItemKind::AssistantMessage(_, MessageState::Final)
                )
            })
            .expect("After setting a final message, there should be one in 'messages'")
    }

    pub fn title(&self) -> Option<String> {
        self.title.clone()
    }

    pub fn set_title(&mut self, title: &str) {
        self.title = Some(title.to_string());
    }

    fn assign_copy_ids(&mut self) {
        let mut current_copy_id = 1;
        for message in &mut self.messages {
            if let ConversationItemKind::AssistantMessage(message_content, _) = &mut message.content
            {
                for block in &mut message_content.structured.blocks {
                    Self::assign_copy_ids_to_block(block, &mut current_copy_id);
                }
            }
        }
    }

    fn assign_copy_ids_to_block(block: &mut Block, current_copy_id: &mut u32) {
        match block {
            Block::Code {
                lang: _,
                content: _,
                copy_id,
            } => {
                *copy_id = *current_copy_id;
                *current_copy_id += 1;
            }
            Block::List { content } => {
                for item in content {
                    for block in &mut item.content {
                        Self::assign_copy_ids_to_block(block, current_copy_id);
                    }
                }
            }
            Block::Quote(blocks) => {
                for block in blocks {
                    Self::assign_copy_ids_to_block(block, current_copy_id);
                }
            }
            _ => {}
        }
    }
}

fn parse_from_markdown(text: &str) -> Result<Document> {
    let ast =
        markdown::to_mdast(text, &markdown::ParseOptions::gfm()).map_err(anyhow::Error::msg)?;

    Ok(transform_document(ast))
}

// always returns a document if it starts at root
// forgiving parsing
// @todo provide an error function that collects all Unsupported elements
fn transform_document(nd: Node) -> Document {
    if let Node::Root(rt) = nd {
        Document {
            blocks: rt.children.iter().map(transform_block).collect(),
        }
    } else {
        Document {
            blocks: vec![Block::Unsupported(
                "Documents must start with a root node.".to_string(),
            )],
        }
    }
}

// always returns a block, error are wrappe in unsupported
// reason: the result is a partially usable tree
fn transform_block(nd: &Node) -> Block {
    match nd {
        Node::Blockquote(blockquote) => {
            Block::Quote(blockquote.children.iter().map(transform_block).collect())
        }
        Node::List(list) => tansform_list(list),
        Node::Code(code) => Block::Code {
            lang: code.lang.clone().unwrap_or("".to_string()).to_string(),
            content: code.value.clone(),
            copy_id: 0,
        },
        Node::Heading(heading) => Block::Heading {
            depth: heading.depth,
            content: heading.children.iter().map(transform_phrasing).collect(),
        },
        Node::Table(table) => transform_table(table),
        Node::ThematicBreak(_thematic_break) => Block::ThematicBreak,
        Node::Paragraph(paragraph) => {
            Block::Paragraph(paragraph.children.iter().map(transform_phrasing).collect())
        }
        x => Block::Unsupported(format!(
            "The following block is not supported when parsing blocks: {}",
            x.to_string()
        )),
    }
}

fn transform_phrasing(nd: &Node) -> Inline {
    match nd {
        Node::Break(_) => Inline::LineBreak,
        Node::InlineCode(inline_code) => Inline::Code(inline_code.value.clone()),
        //Node::InlineMath(inline_math) => todo!(),
        //Node::Delete(delete) => todo!(),
        Node::Emphasis(emphasis) => {
            Inline::Emphasis(emphasis.children.iter().map(transform_phrasing).collect())
        }
        //Node::Image(image) => todo!(),
        //Node::ImageReference(image_reference) => todo!(),
        Node::Link(link) => Inline::Link {
            url: link.url.clone(),
            title: link.title.clone(),
            label: link.children.iter().map(transform_phrasing).collect(),
        },
        Node::Strong(strong) => {
            Inline::StrongEmphasis(strong.children.iter().map(transform_phrasing).collect())
        }
        Node::Text(text) => Inline::Text(text.value.clone()),
        //Node::Math(math) => todo!(),
        x => Inline::Unsupported(format!(
            "The following block is not supported when parsing blocks: {}",
            x.to_string()
        )),
        //Node::TableRow(table_row) => todo!(),
        //Node::TableCell(table_cell) => todo!(),
    }
}

// this function supports lists that mix ordered and unordered elements
// while this should never happen in practice (it is not allowed by CommonMark)
// the data model of mdast allows it.
// Therefore, it is assumed that markdown-rs does not produce such a list.
// If it does, this function is forgiving and allows us to render it.
fn tansform_list(list: &markdown::mdast::List) -> Block {
    let mut number = list.start.unwrap_or(1);
    let mut items = Vec::new();
    for item in list.children.iter() {
        if let Node::ListItem(li) = item {
            if list.ordered {
                items.push(ListItem {
                    marker: ListItemMarker::Number(number),
                    content: li.children.iter().map(transform_block).collect(),
                });
                number += 1;
            } else {
                let marker = if let Some(checked) = li.checked {
                    ListItemMarker::Checkbox(checked)
                } else {
                    ListItemMarker::Bullet
                };
                items.push(ListItem {
                    marker,
                    content: li.children.iter().map(transform_block).collect(),
                });
            }
        } else {
            // @todo add error element or log error
            // However, this branch should never be taken
        }
    }
    Block::List { content: items }
}

fn transform_table(table: &markdown::mdast::Table) -> Block {
    fn to_cell(nd: &Node) -> Option<TableCell> {
        match nd {
            Node::TableCell(table_cell) => {
                let cell_content = table_cell.children.iter().map(transform_phrasing).collect();
                Some(TableCell {
                    content: cell_content,
                })
            }
            _ => None, // silently drop non-matching node (should not be there anyhow)
        }
    }

    fn to_row(nd: &Node) -> Option<TableRow> {
        match nd {
            Node::TableRow(row) => {
                let row_content = row.children.iter().filter_map(to_cell).collect();
                Some(TableRow {
                    content: row_content,
                })
            }
            _ => None, // silently drop non-matching node (should not be there anyhow)
        }
    }

    let rows: Vec<TableRow> = table.children.iter().filter_map(to_row).collect();

    Block::Table { content: rows }
}

pub fn generate_dummy_conversation() {
    let conv = Conversation {
        id: Uuid::new_v4(),
        title: Some("Test".to_string()),
        model: None,
        messages: vec![
            ConversationItem::new(ConversationItemKind::Notification(
                NotificationType::Info,
                "Starting...".to_string(),
            )),
            ConversationItem::new(ConversationItemKind::UserMessage(MessageContent {
                unprocessed: "".to_string(),
                structured: Document::default(),
            })),
            ConversationItem::new(ConversationItemKind::AssistantReasoning(MessageContent {
                unprocessed: "Hmm...".to_string(),
                structured: Document::default(),
            })),
            ConversationItem::new(ConversationItemKind::ToolInvocation(
                ToolInvocationContent {
                    name: "help".to_string(),
                    arguments: serde_json::Value::String("What shall I do".to_string()),
                    call_id: ToolCallId::mint(),
                    provider_id: None,
                },
            )),
            ConversationItem::new(ConversationItemKind::ToolOutput(ToolOutputContent {
                name: "help".to_string(),
                output: "Figure it out".to_string(),
                success: true,
                call_id: ToolCallId::mint(),
                provider_id: None,
            })),
            ConversationItem::new(ConversationItemKind::AssistantMessage(
                MessageContent {
                    unprocessed: "I don't know.".to_string(),
                    structured: Document::default(),
                },
                MessageState::Final,
            )),
        ],
    };

    let mut file = File::create(Path::new("test.json")).expect("Should be able to open testfile");
    file.write_fmt(format_args!(
        "{}",
        serde_json::to_string(&conv).expect("Test structure should be serializable.")
    ))
    .expect("Should be able to write to test file.");
}