monothink 0.5.2

/*
* 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 monothink.
* If not, see <https://www.gnu.org/licenses/>.

*/

mod compaction;
mod compactors;
mod estimators;
mod tools;

use crate::{
    common::{
        conversation::{ConversationItem, ConversationItemKind},
        package_resolver::PackageResolver,
        path_traversal::PathTraversal,
    },
    config::{self},
    llm::{
        compaction::CompactionHook, compactors::SharedObservations, estimators::SimpleEstimator,
        tools::find_package::FindPackage,
    },
    session::SessionMessage,
};
use anyhow::Result;
use futures::{Stream, StreamExt};
use rig::{
    Agent as RigAgent, AgentBuilder, DynModel,
    agent::{AgentHook, MultiTurnStreamItem},
    completion::Message,
    effect::Outcome,
    message::{
        AssistantMessage, LocalCallId, ToolCall, ToolFunction, ToolName, ToolResultContent,
        UserContent,
    },
    providers::openai::{self},
    run::PromptError,
    tool::rmcp::tools_from_server,
};
use std::{collections::HashMap, path::PathBuf, pin::Pin, sync::Arc};
use tokio::sync::{
    Mutex,
    mpsc::{Receiver, Sender},
};

struct Agent {
    agent: RigAgent,
}

#[derive(Clone, Debug, PartialEq)]
pub enum LlmMessage {
    SubmitPrompt {
        prompt: String,
        history: Vec<ConversationItem>,
    },
    SelectModel(String),
    SetContextWindow(u64),
    SetEndpoint(config::Endpoint),
    //CancelGeneration,
    Quit,
}

// @todo does it have to be this detailed? Can we abbreviate or at least make it look cleaner?
// This is the Response type of the streaming completion API of rig
type LlmStream =
    Pin<Box<dyn Stream<Item = Result<MultiTurnStreamItem, PromptError>> + Send + 'static>>;

pub enum LlmState {
    Unconnected,
    Idle,
    Requesting {
        user_prompt: String,
        history: Vec<Message>,
    },
    Streaming {
        stream: LlmStream,
    },
    Quitting,
}

#[derive(Debug)]
pub struct Llm {
    session_tx: Sender<SessionMessage>,
    rx: Receiver<LlmMessage>,
    config: config::General,
    context_window: u64,
    endpoint: Option<config::Endpoint>,
    fatal_error: bool,
    tools: config::Tools,
    model: Option<String>,
    observations: SharedObservations,
    package_resolver: Arc<Mutex<PackageResolver>>,
    path_traversal: Arc<Mutex<PathTraversal>>,
}

impl Llm {
    pub async fn new(
        session_tx: Sender<SessionMessage>,
        rx: Receiver<LlmMessage>,
        config: config::General,
        tools: config::Tools,
        workspace_root: PathBuf,
    ) -> Self {
        let observe_hidden = config.observe_hidden;
        let observe_ignore = config.observe_ignore;
        Self {
            session_tx,
            rx,
            config,
            context_window: config::default_context_window(),
            endpoint: None,
            fatal_error: false,
            tools,
            model: None,
            observations: Arc::new(Mutex::new(HashMap::new())),
            package_resolver: Arc::new(Mutex::new(
                PackageResolver::new(workspace_root.clone()).await,
            )),
            path_traversal: Arc::new(Mutex::new(PathTraversal::new(
                &workspace_root,
                observe_hidden,
                observe_ignore,
            ))),
        }
    }

    #[tracing::instrument(skip(self))]
    pub async fn run(&mut self) {
        tracing::info!("Starting llm actor");
        let mut agent = None;
        let mut state = LlmState::Idle;
        self.session_tx
            .send(SessionMessage::Info(format!(
                "Package managers found: {:?}",
                self.package_resolver.lock().await.list_pacmans()
            )))
            .await
            .unwrap();

        loop {
            if self.fatal_error {
                break;
            }
            state = match state {
                LlmState::Unconnected => self.handle_unconnected(&mut agent).await,
                LlmState::Idle => self.handle_idle().await,
                LlmState::Requesting {
                    user_prompt,
                    history,
                } => self.handle_requesting(&agent, user_prompt, history).await,
                LlmState::Streaming { stream } => self.handle_streaming(stream).await,
                LlmState::Quitting => break,
            };
        }
    }

    async fn build_agent<M>(&mut self, model: M) -> Agent
    where
        M: Into<DynModel<rig::core::operation::Completion>>,
    {
        let mut builder = AgentBuilder::new(model)
                        .preamble(format!(
                            "You are a experienced software developer advising a colleague.\nFor your information: The current date is {}.\n\
                            The workspace root directory is {}. Tools that take a workspace path operate relative to this directory.",
                            chrono::Local::now().format("%A, %Y-%m-%d %H:%M %z"),
                            self.package_resolver.lock().await.workspace_dir.to_string_lossy()
                        ))
                        .tool(tools::list_dir::ListDir::new(
                            self.package_resolver.clone(),
                            self.path_traversal.clone(),
                            Some(self.config.max_list_dir_results),
                        ).await)
                        .tool(tools::read_file::ReadFile::new(
                            self.package_resolver.clone(),
                            self.path_traversal.clone(),
                        ).await)
                        .tool(tools::search_text::SearchText::new(
                            self.package_resolver.clone(),
                            self.path_traversal.clone(),
                            self.config.max_search_text_results,
                            400
                        ).await)
                        .tool(tools::retrieve_observation::RetrieveObservation::new(self.observations.clone()))
                        .add_hook(CompactionHook::new(
                            self.session_tx.clone(),
                            self.context_window,
                            SimpleEstimator::new(4.0),
                            self.config.max_tokens,
                            self.config.max_tool_observations,
                            self.config.min_reasoning_observations,
                            self.observations.clone(),
                        ))
                        .add_hook(ToolResultHook { session_tx: self.session_tx.clone() })
                        .default_max_turns(self.config.max_iterations_per_request)
                        .max_tokens(self.config.max_tokens);

        if !self.package_resolver.lock().await.active_pacmans.is_empty() {
            builder = builder.tool(FindPackage::new(self.package_resolver.clone()));
        }
        if let Some(ws_config) = &self.tools.web_search
            && let Some(web_search_tool) = tools::web_search::WebSearch::new(ws_config).await
        {
            builder = builder.tool(web_search_tool);
        }
        let mcp_tools = tools::get_mcp_tools(&self.tools.mcp).await;
        for (client, tools) in mcp_tools {
            let rmcp_tools = tools_from_server(tools, &client);
            for tool in rmcp_tools {
                if let Ok(dyn_tool) = rig::tool::DynamicTool::try_from(tool) {
                    builder = builder.dynamic_tool(dyn_tool);
                }
            }
        }
        Agent {
            agent: builder.build(),
        }
    }

    #[tracing::instrument(skip(self))]
    async fn handle_idle(&mut self) -> LlmState {
        match self.rx.recv().await {
            Some(LlmMessage::SetContextWindow(context_window)) => {
                self.context_window = context_window;
                LlmState::Unconnected
            }
            Some(LlmMessage::SelectModel(model)) => {
                self.model = Some(model);
                LlmState::Unconnected
            }
            Some(LlmMessage::SetEndpoint(endpoint)) => {
                self.endpoint = Some(endpoint);
                LlmState::Unconnected
            }
            Some(LlmMessage::SubmitPrompt { prompt, history }) => {
                let rig_history = rig_history_from_internal(&history);
                LlmState::Requesting {
                    user_prompt: prompt,
                    history: rig_history,
                }
            }
            Some(LlmMessage::Quit) => LlmState::Quitting,
            None => LlmState::Quitting,
        }
    }

    #[tracing::instrument(skip(self, agent))]
    async fn handle_requesting(
        &mut self,
        agent: &Option<Agent>,
        user_prompt: String,
        history: Vec<Message>,
    ) -> LlmState {
        match agent {
            Some(a) => {
                // @todo handle new request while request is running
                LlmState::Streaming {
                    stream: a.agent.prompt(user_prompt).history(&history).stream(),
                }
            }
            None => LlmState::Unconnected,
        }
    }

    #[tracing::instrument(skip(self, s))]
    async fn handle_streaming(&mut self, mut s: LlmStream) -> LlmState {
        loop {
            tokio::select! {
                maybe_chunk = s.next() => {
                    if let Some(result) = maybe_chunk {
                        match result {
                            Ok(chunk) => self.handle_llm_item(chunk).await,
                            Err(e) => {
                                    // @todo why does this error occur if the endpoint is not a valid address?
                                    // it should fail earlier, imho
                                self.session_tx
                                    .send(SessionMessage::Error(e.to_string()))
                                    .await
                                    .unwrap();
                                    // @todo handle error
                            }
                        }
                    } else {
                        // it is cleaner to additionally allow handle_llm_item()
                        // to report a completion and change the state.
                        // @todo implement that
                        return LlmState::Idle;
                    }
                }
                // @todo handle new request while stream is running
                // for this, drop stream and go back to Requesting with new prompt
                // alternative block input in ui and only allow command :cancel (hint for user)
                // @todo in this case do not go to idle but quit
                // have an extra message to cancle a request and go to idle
                Some(LlmMessage::Quit) = self.rx.recv() => {
                    return LlmState::Idle;
                }
            }
        }
    }

    #[tracing::instrument(skip(self))]
    async fn handle_llm_item(&mut self, chunk: MultiTurnStreamItem) {
        match chunk {
            MultiTurnStreamItem::StreamAssistantItem(rig::streaming::Item::Event(
                assistant_content,
            )) => {
                match assistant_content {
                    rig::streaming::StreamEvent::Start {
                        part: _,
                        kind: _,
                        name: _,
                    } => {}
                    rig::streaming::StreamEvent::Text { part: _, text } => {
                        self.send_to_session(SessionMessage::AgentResponse(text))
                            .await;
                    }
                    rig::streaming::StreamEvent::Reasoning { part: _, text } => {
                        self.send_to_session(SessionMessage::AgentReasoning(text))
                            .await;
                    }
                    rig::streaming::StreamEvent::Arguments { part: _, json: _ } => {
                        // only necessary if we want to inspect the data sent by the provider
                    }
                    rig::streaming::StreamEvent::End {
                        part: _,
                        content: _,
                    } => {}
                }
            }
            MultiTurnStreamItem::StreamAssistantItem(rig::streaming::Item::Unknown(
                _unknown_payload,
            )) => {}
            MultiTurnStreamItem::ToolCall { tool_call } => {
                tracing::info!(
                    tool = %tool_call.function.name,
                    call_id = %tool_call.id,
                    "tool call"
                );
                self.send_to_session(SessionMessage::ToolCall {
                    name: tool_call.function.name.to_string(),
                    args: tool_call.function.arguments,
                    call_id: tool_call.id,
                })
                .await;
            }
            MultiTurnStreamItem::ToolExecutionCommitted { tool_call: _ } => {
                // tools are observed via StreamAssistantItem to provide real-time insights
            }
            MultiTurnStreamItem::ToolResult { tool_result: _ } => {
                // tool results are observed via hook to get more insights
            }
            MultiTurnStreamItem::CompletionCall(completion_call) => {
                self.send_to_session(SessionMessage::AgentCompletion(completion_call.usage))
                    .await;
            }
            MultiTurnStreamItem::ModelTurnRetried { turn } => {
                tracing::warn!(turn = turn, "Model turn has been retried");
            }
            MultiTurnStreamItem::FinalResponse(prompt_response) => {
                self.send_to_session(SessionMessage::FinalAgentResponse(prompt_response.output()))
                    .await;
            }
            _ => {}
        }
    }

    #[tracing::instrument(skip(self, agent))]
    async fn handle_unconnected(&mut self, agent: &mut Option<Agent>) -> LlmState {
        if let Some(endpoint) = &self.endpoint
            && let Some(model) = &self.model
        {
            let api_key = endpoint
                .secrets
                .clone()
                .unwrap_or_else(|| {
                    tracing::debug!("Endpoint secrets not set, falling back to default");
                    config::EndpointSecrets {
                        api_key: Some("".to_string()),
                    }
                })
                .api_key
                .unwrap_or_else(|| {
                    tracing::debug!("API key not set, falling back to default");
                    "".to_string()
                });

            let new_agent = match endpoint.api_type {
                config::EndpointApi::OpenAICompletions => {
                    let model = openai::OpenAIConfig::new(api_key)
                        .with_base_url(&endpoint.base_url)
                        .client()
                        .completion(model);
                    self.build_agent(model).await
                }
                config::EndpointApi::OpenAIResponses => {
                    let model = openai::OpenAIConfig::new(api_key)
                        .with_base_url(&endpoint.base_url)
                        .client()
                        .responses(model);
                    self.build_agent(model).await
                }
            };
            *agent = Some(new_agent);
            self.send_to_session(SessionMessage::Debug(
                "Connected to LLM provider".to_string(),
            ))
            .await;

            return LlmState::Idle;
        }

        match self.rx.recv().await {
            Some(LlmMessage::SelectModel(model)) => {
                self.model = Some(model);
                LlmState::Unconnected
            }
            Some(LlmMessage::SetContextWindow(context_window)) => {
                self.context_window = context_window;
                LlmState::Unconnected
            }
            Some(LlmMessage::SetEndpoint(endpoint)) => {
                self.endpoint = Some(endpoint);
                LlmState::Unconnected
            }
            Some(LlmMessage::SubmitPrompt {
                prompt: _,
                history: _,
            }) => {
                // @todo report error
                LlmState::Unconnected
            }
            Some(LlmMessage::Quit) => LlmState::Quitting,
            None => LlmState::Quitting,
        }
    }

    async fn send_to_session(&mut self, message: SessionMessage) {
        match self.session_tx.send(message).await {
            Ok(_) => {}
            Err(e) => {
                tracing::error!(
                    error = e.to_string(),
                    "Cannot send message to session. Terminating Llm actor."
                );
                self.fatal_error = true;
            }
        }
    }
}

fn rig_history_from_internal(history: &[ConversationItem]) -> Vec<Message> {
    history
        .iter()
        .filter_map(|item| match &item.content {
            ConversationItemKind::UserMessage(message_content) => {
                Some(Message::user(message_content.raw().to_string()))
            }
            ConversationItemKind::AssistantMessage(message_content, _) => {
                if message_content.raw().trim().is_empty() {
                    None
                } else {
                    Some(Message::assistant(message_content.raw().to_string()))
                }
            }
            ConversationItemKind::AssistantReasoning(_message_content) => None,
            ConversationItemKind::ToolInvocation(tool_invocation_content) => {
                let tool_call = ToolCall::new(
                    tool_invocation_content.call_id.clone(),
                    ToolFunction::new(
                        ToolName::new(tool_invocation_content.name.clone()).unwrap(),
                        serde_json::Value::Null,
                    ),
                );
                Some(rig::completion::Message::Assistant(AssistantMessage::new(
                    vec![rig::message::AssistantContent::ToolCall(tool_call)],
                )))
            }
            ConversationItemKind::ToolOutput(tool_output_content) => {
                let tool_call = ToolCall::new(
                    tool_output_content.call_id.clone(),
                    ToolFunction::new(
                        ToolName::new(tool_output_content.name.clone()).unwrap(),
                        serde_json::Value::Null,
                    ),
                );
                let tool_result = tool_call.result(vec![ToolResultContent::Text(
                    tool_output_content.output.clone().into(),
                )]);
                Some(rig::completion::Message::User {
                    content: vec![UserContent::ToolResult(tool_result)],
                })
            }
            ConversationItemKind::Notification(_notification_type, _) => None,
        })
        .collect()
}

pub struct ToolResultHook {
    session_tx: Sender<SessionMessage>,
}

impl AgentHook for ToolResultHook {
    async fn on_outcome(
        &self,
        _ctx: &rig::agent::HookContext,
        event: rig::agent::OutcomeEvent<'_>,
    ) -> rig::agent::OutcomeAction {
        if let Ok(Outcome::ToolResult { result: outcome }) = event.outcome {
            let content = if outcome.is_success() {
                outcome
                    .output()
                    .as_content()
                    .iter()
                    .map(|content| match content {
                        ToolResultContent::Text(text) => text.to_string(),
                        ToolResultContent::Image(_) => "".to_string(),
                        ToolResultContent::Json { value } => {
                            serde_json::to_string_pretty(&value).unwrap_or_default()
                        }
                    })
                    .fold(String::new(), |result, output| {
                        let separator = if !result.is_empty() { "\n" } else { "" };
                        result + separator + &output
                    })
            } else if outcome.is_error() {
                outcome
                    .error()
                    .expect("There should be an error in the error case")
                    .to_string()
            } else {
                String::new()
            };
            self.session_tx
                .send(SessionMessage::ToolResult {
                    name: event.tool_name().unwrap_or_default().to_string(),
                    output: content,
                    success: outcome.is_success(),
                    call_id: event
                        .call_id
                        .unwrap_or(&rig::message::CallId::Local(LocalCallId::new()))
                        .clone(),
                })
                .await
                .unwrap();
        }
        rig::agent::OutcomeAction::Proceed
    }
}