monothink 0.2.0

/*
 * This file is part of monothink.
 *
 * Wavething 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.
 *
 * Wavething 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 std::{env, fs, io::Write, path::Path};

use directories::ProjectDirs;
use rig::completion::Usage;
use tokio::sync::mpsc::{Receiver, Sender};

use crate::{
    core::conversation::{Conversation, ConversationItemKind, NotificationType},
    llm::LlmMessage,
    ui::UiMessage,
};

#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct TokenUsage {
    pub input: Option<u64>,
    pub output: Option<u64>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum SessionMessage {
    SubmitCommand(String),
    SubmitPrompt(String),
    AgentResponse(String),
    AgentReasoning(String),
    AgentCompletion(Usage), // Usage is a rig type we we should abstract from. But at this point we need the same information anyhow.
    ToolCall {
        name: String,
        args: serde_json::Value,
    },
    ToolResult {
        output: String,
        success: bool,
    },
    Error(String),
    Warning(String),
}

#[derive(Debug)]
pub struct Session {
    ui_tx: Sender<UiMessage>,
    llm_tx: Sender<LlmMessage>,
    rx: Receiver<SessionMessage>,
    conversation: Conversation,
    token_usage: TokenUsage,
    quit: bool,
}

impl Session {
    pub fn new(
        ui_tx: Sender<UiMessage>,
        llm_tx: Sender<LlmMessage>,
        rx: Receiver<SessionMessage>,
    ) -> Self {
        Self {
            ui_tx,
            llm_tx,
            rx,
            conversation: Conversation::new(),
            token_usage: TokenUsage::default(),
            quit: false,
        }
    }

    pub async fn run(&mut self) {
        while !self.quit {
            let event = self.rx.recv().await;
            if let Some(msg) = event {
                match msg {
                    SessionMessage::SubmitCommand(command) => self.evaluate_command(command).await,
                    SessionMessage::SubmitPrompt(prompt) => {
                        let item = self.conversation.add_user_prompt(&prompt);
                        self.llm_tx
                            .send(LlmMessage::SubmitPrompt {
                                prompt,
                                history: Vec::new(), // @todo send recent history
                            })
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::NewConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::AgentResponse(chunk) => {
                        let (item, new) = self.conversation.add_assistant_response_chunk(&chunk);
                        if new {
                            self.ui_tx
                                .send(UiMessage::NewConversationItem(item.to_owned()))
                                .await
                                .unwrap();
                        } else {
                            self.ui_tx
                                .send(UiMessage::UpdatedConversationItem(item.to_owned()))
                                .await
                                .unwrap();
                        }
                    }
                    SessionMessage::AgentReasoning(chunk) => {
                        let (item, new) = self.conversation.add_assistant_reasoning_chunk(&chunk);
                        if new {
                            self.ui_tx
                                .send(UiMessage::NewConversationItem(item.to_owned()))
                                .await
                                .unwrap();
                        } else {
                            self.ui_tx
                                .send(UiMessage::UpdatedConversationItem(item.to_owned()))
                                .await
                                .unwrap();
                        }
                    }
                    SessionMessage::AgentCompletion(usage) => {
                        self.update_token_usage(usage);
                        self.ui_tx
                            .send(UiMessage::GenerationFinished(TokenUsage {
                                input: self.token_usage.input,
                                output: self.token_usage.output,
                            }))
                            .await
                            .unwrap();

                        // publish report as system message for analyzing provider behavior
                        // bypass conversation as this is only for development
                        let msg = format!(
                            "Usage: input tokens = {}, output tokens = {}, \
                            cached input tokens = {}, cached creation input tokens = {}, \
                            reasoning tokens = {}, tool use prompt tokens = {}, total tokens = {}",
                            usage.input_tokens,
                            usage.output_tokens,
                            usage.cached_input_tokens,
                            usage.cache_creation_input_tokens,
                            usage.reasoning_tokens,
                            usage.tool_use_prompt_tokens,
                            usage.total_tokens
                        );
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Info, msg);
                        self.ui_tx
                            .send(UiMessage::NewConversationItem(item.to_owned()))
                            .await
                            .unwrap();

                        // We could use this event to mark our messages as complete.
                        // But as of now a message is complete when a new message arrives.
                    }
                    SessionMessage::Error(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Error, msg);
                        self.ui_tx
                            .send(UiMessage::NewConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Warning(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Warn, msg);
                        self.ui_tx
                            .send(UiMessage::NewConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ToolCall { name, args } => {
                        let item = self.conversation.add_tool_invocation(name, args);
                        self.ui_tx
                            .send(UiMessage::NewConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ToolResult { output, success } => {
                        let item = self.conversation.add_tool_result(output, success);
                        self.ui_tx
                            .send(UiMessage::NewConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                }
            }
        }
    }

    async fn evaluate_command(&mut self, command: String) {
        if let Some(start) = command.chars().nth(0)
            && start == ':'
        {
            let (name, params) = command[1..]
                .split_once(char::is_whitespace)
                .unwrap_or((&command[1..], ""));
            let params = params.trim_start();
            match name {
                "c" | "cancel" => {
                    // @todo cancel request
                }
                "e" | "export" => {
                    self.handle_export_command(params).await;
                }
                "q" | "quit" => {
                    self.handle_quit_command().await;
                }
                "t" | "title" => {
                    if params.is_empty() {
                        // @todo show title in statusbar (send to ui and there show in info area)
                    } else {
                        self.conversation.set_title(params);
                    }
                }
                _ => {
                    // @todo command not understood
                }
            }
        } else {
            // @todo report error: received a command that does not start with ':". Should never happen.
        }
    }

    async fn handle_export_command(&self, params: &str) {
        let mut path = Path::new(params).to_path_buf();
        // check if it is an absolute or relative path
        if Some(true) != path.parent().and_then(|p| p.try_exists().ok()) {
            let mut relative_path = env::current_dir().expect("Cannot get current directory");
            relative_path.push(params);
            path = relative_path;
        }

        let mut first = true;

        if Some(true) == path.parent().and_then(|p| p.try_exists().ok()) {
            // iterate over messages and write user and assistant messages as md
            let mut output = String::new();
            for msg in self.conversation.messages() {
                match &msg.content {
                    ConversationItemKind::UserMessage(message_content) => {
                        if !first {
                            output.push_str("\n\n");
                        } else {
                            first = false;
                        }
                        output.push_str("# User\n\n");
                        output.push_str(message_content.raw());
                    }
                    ConversationItemKind::AssistantMessage(message_content) => {
                        if !message_content.raw().trim().is_empty() {
                            // note: Assistant can never be the first
                            output.push_str("\n\n");
                            // no newlines after heading needed as the message already contains them
                            output.push_str("# Assistant");
                            output.push_str(message_content.raw());
                        }
                    }
                    ConversationItemKind::AssistantReasoning(_) => {}
                    ConversationItemKind::ToolInvocation(_) => {}
                    ConversationItemKind::ToolOutput(_) => {}
                    ConversationItemKind::Notification(_, _) => {}
                }
            }
            let mut export_file = fs::File::create(path).expect("Cannot open md file for export");
            let _ = export_file
                .write(output.as_bytes())
                .expect("Cannot write md export to file");
        }
    }

    async fn handle_quit_command(&mut self) {
        self.llm_tx.send(LlmMessage::Quit).await.unwrap();
        self.ui_tx.send(UiMessage::Quit).await.unwrap();

        // @todo do not quit if changes are unsafed and command is not force exit
        if !self.conversation.messages().is_empty() {
            // save session
            // @todo make autosafe a config option and ask otherwise
            if let Some(proj_dirs) = ProjectDirs::from("org", "codemonk", "monothink") {
                let mut state_dir = proj_dirs
                    .state_dir()
                    .expect("Running on other OS than Linux is not supported.")
                    .to_path_buf();
                state_dir.push("sessions");
                state_dir.push(format!("{}.json", self.conversation.id()));
                let session_path = state_dir.as_path();
                // @todo stream to a writer
                let output = serde_json::to_string(&self.conversation)
                    .expect("Should be able to serealize a session.");
                let mut session_file =
                    fs::File::create(session_path).expect("Cannot open session file for writing");
                let _ = session_file
                    .write(output.as_bytes())
                    .expect("Cannot write session to file");
            } else {
                // @todo probably should not panic here, as we are quitting anyhow.
                panic!("Could not get state directory for saving session.");
            };
        }

        self.quit = true;
    }

    fn update_token_usage(&mut self, usage: Usage) {
        // assume that OpenAI typical token usage reports,
        // i.e. only input and output are reported
        if usage.input_tokens > 0 {
            self.token_usage.input = Some(self.token_usage.input.unwrap_or(0) + usage.input_tokens);
        }
        if usage.output_tokens > 0 {
            self.token_usage.output =
                Some(self.token_usage.output.unwrap_or(0) + usage.output_tokens);
        }
    }
}