monothink 0.4.7

/*
 * 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 std::{
    env, fs,
    io::Write,
    path::{Path, PathBuf},
};

use anyhow::{Error, Result};
use directories::ProjectDirs;
use rig::{
    completion::Usage,
    message::{ProviderCallId, ToolCallId},
};
use tokio::sync::mpsc::{Receiver, Sender};

use crate::{
    config::{Config, Endpoint, ViewMode, default_context_window},
    core::conversation::{Conversation, ConversationItemKind, MessageState, 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.
    ContextPressure {
        initial: u64,
        updated: u64,
        window: u64,
    },
    Debug(String),
    FinalAgentResponse(String),
    Timeout,
    ToolCall {
        name: String,
        args: serde_json::Value,
        call_id: ToolCallId,
        provider_id: Option<ProviderCallId>,
    },
    ToolResult {
        name: String,
        output: String,
        success: bool,
        call_id: ToolCallId,
        provider_id: Option<ProviderCallId>,
    },
    Error(String),
    Info(String),
    Warning(String),
}

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

impl Session {
    pub fn new(
        ui_tx: Sender<UiMessage>,
        llm_tx: Sender<LlmMessage>,
        rx: Receiver<SessionMessage>,
        config: Config,
        workspace_root: PathBuf,
    ) -> Self {
        let endpoint = Self::get_active_endpoint(&config);
        let active_model = match endpoint {
            Ok(endpoint) => Some(endpoint.default_model),
            Err(_) => None,
        };

        let context_window = match &active_model {
            Some(model_name) => config
                .model
                .get(model_name)
                .map(|model| model.context_window)
                .unwrap_or_else(default_context_window),
            None => default_context_window(),
        };

        Self {
            config,
            active_model: active_model.clone(),
            ui_tx,
            llm_tx,
            rx,
            context_window,
            conversation: Conversation::new(active_model, workspace_root),
            token_usage: TokenUsage::default(),
            quit: false,
        }
    }

    pub async fn run(&mut self) {
        self.ui_tx
            .send(UiMessage::ActiveModel(self.active_model.clone()))
            .await
            .unwrap();

        // @todo proper error handling: Tell user to configure an endpoint
        let endpoint = Self::get_active_endpoint(&self.config).unwrap();
        self.llm_tx
            .send(LlmMessage::SetEndpoint(endpoint))
            .await
            .unwrap();
        self.llm_tx
            .send(LlmMessage::SetContextWindow(self.context_window))
            .await
            .unwrap();
        // @todo proper error handling: Tell user to configure a model
        self.llm_tx
            .send(LlmMessage::SelectModel(self.active_model.clone().unwrap()))
            .await
            .unwrap();
        self.ui_tx
            .send(UiMessage::ConversationViewMode(
                self.config.general.view_mode,
            ))
            .await
            .unwrap();
        self.ui_tx
            .send(UiMessage::ContextState {
                fill_state: 0,
                total: self.context_window,
            })
            .await
            .unwrap();

        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).clone();
                        self.llm_tx
                            .send(LlmMessage::SubmitPrompt {
                                prompt,
                                history: self.conversation.messages().clone(),
                            })
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.clone()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(true))
                            .await
                            .unwrap();
                    }
                    SessionMessage::AgentResponse(chunk) => {
                        let item = self.conversation.add_assistant_response_chunk(&chunk);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::AgentReasoning(chunk) => {
                        let item = self.conversation.add_assistant_reasoning_chunk(&chunk);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(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 = {},\n \
                            cached input tokens = {}, cached creation input tokens = {},\n \
                            reasoning tokens = {}, tool use prompt tokens = {},\n 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::Debug, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ContextPressure {
                        initial,
                        updated,
                        window,
                    } => {
                        let msg = format!(
                            "Context window: {window}, initial pressure: {initial}, updated pressure: {updated}.",
                        ).to_string();
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Debug, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::ContextState {
                                fill_state: updated,
                                total: window,
                            })
                            .await
                            .unwrap();
                    }
                    SessionMessage::FinalAgentResponse(msg) => {
                        let item = self.conversation.set_final_assistant_response(&msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(false))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Error(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Error, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(false))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Warning(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Warn, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Timeout => {
                        let item = self.conversation.add_notification(
                            NotificationType::Error,
                            "The connection to the LLM provider timed out.".to_string(),
                        );
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(false))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ToolCall {
                        name,
                        args,
                        call_id,
                        provider_id,
                    } => {
                        let item =
                            self.conversation
                                .add_tool_invocation(name, args, call_id, provider_id);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ToolResult {
                        name,
                        output,
                        success,
                        call_id,
                        provider_id,
                    } => {
                        let item = self.conversation.add_tool_result(
                            name,
                            output,
                            success,
                            call_id,
                            provider_id,
                        );
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Debug(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Debug, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Info(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Info, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(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);
                    }
                }
                "v" | "view_mode" => {
                    self.handle_view_mode_command(params).await;
                }
                _ => {
                    // @todo command not understood
                }
            }
        } else {
            // @todo report error: received a command that does not start with ':". Should never happen.
        }
    }

    // @todo this should be part of config
    fn get_active_endpoint(config: &Config) -> Result<Endpoint> {
        config
            .endpoint
            .get(&config.general.primary_endpoint)
            .map(|val| val.to_owned())
            .ok_or(Error::msg(format!(
                "Endpoint {} not found in endpoints.",
                config.general.primary_endpoint
            )))
    }

    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, message_kind) => {
                        if !message_content.raw().trim().is_empty()
                            && *message_kind == MessageState::Final
                        {
                            // 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_pretty(&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;
    }

    async fn handle_view_mode_command(&mut self, params: &str) {
        let view_mode = match params {
            "compact" => ViewMode::Compact,
            "debug" => ViewMode::Debug,
            "detailed" => ViewMode::Detailed,
            "full" => ViewMode::Full,
            "minimal" => ViewMode::Minimal,
            _ => return, // @todo report error
        };
        self.ui_tx
            .send(UiMessage::ConversationViewMode(view_mode))
            .await
            .unwrap();
    }

    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);
        }
    }
}