monothink 0.3.1

/*
* 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/>.

*/

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

use std::{
    collections::HashMap,
    env,
    path::Path,
    pin::Pin,
    sync::{Arc, Mutex},
};

use anyhow::Result;
use tokio::sync::mpsc::{Receiver, Sender};

use crate::{
    config::{self},
    core::conversation::{ConversationItem, ConversationItemKind},
    llm::{
        compaction::CompactionHook, compactors::SharedObservations, estimators::SimpleEstimator,
        tools::RetrieveObservation,
    },
    session::SessionMessage,
};
use futures::{Stream, StreamExt};
use rig::{
    Agent,
    agent::{MultiTurnStreamItem, StreamingError},
    client::AgentClientExt,
    completion::Message,
    message::{ToolCall, ToolFunction, ToolResultContent},
    providers::openai::{self},
    streaming::{StreamedAssistantContent, StreamingChat},
};

#[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, StreamingError>> + 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>,
    model: Option<String>,
    observations: SharedObservations,
}

impl Llm {
    pub fn new(
        session_tx: Sender<SessionMessage>,
        rx: Receiver<LlmMessage>,
        config: config::General,
    ) -> Self {
        Self {
            session_tx,
            rx,
            config,
            context_window: config::default_context_window(),
            endpoint: None,
            model: None,
            observations: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub async fn run(&mut self) {
        // @todo issue an error system message and do not publish any tools on run
        // @todo get this from session together with rest of 'General' config
        let base_dir = env::current_dir().expect("Cannot get current directory");

        let mut agent = None;
        let mut state = LlmState::Idle;

        loop {
            state = match state {
                LlmState::Unconnected => self.handle_unconnected(&mut agent, &base_dir).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,
            };
        }
    }

    // note: We need the responses completion model as this is the new stateless completion api by openai
    // for some reason, it does not implement CompletionModel, so M: CompletionModel does not work
    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,
        }
    }

    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
                return LlmState::Streaming {
                    stream: a.stream_chat(user_prompt, history).await,
                };
            }
            None => LlmState::Unconnected,
        }
    }

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

    async fn handle_llm_item(&mut self, chunk: MultiTurnStreamItem) {
        match chunk {
            MultiTurnStreamItem::StreamAssistantItem(streamed_assistant_content) => {
                match streamed_assistant_content {
                    StreamedAssistantContent::Text(text) => {
                        // @todo handle error by logging it (send does not work) and quitting the actor
                        self.session_tx
                            .send(SessionMessage::AgentResponse(text.to_string()))
                            .await
                            .unwrap();
                    }
                    StreamedAssistantContent::ToolCall {
                        tool_call,
                        internal_call_id: _,
                    } => {
                        self.session_tx
                            .send(SessionMessage::ToolCall {
                                name: tool_call.function.name,
                                args: tool_call.function.arguments,
                                call_id: tool_call.id,
                                provider_id: tool_call.provider,
                            })
                            .await
                            .unwrap();
                    }
                    StreamedAssistantContent::ToolCallDelta {
                        internal_call_id: _,
                        content: _,
                    } => {
                        // ignored because we are only interested in complete tool calls
                        // Rationale: Tool calls are short and streaming them does not take long
                        // There is no reason not to wait until a complete call is finished.
                    }
                    StreamedAssistantContent::Reasoning {
                        reasoning: _,
                        id: _,
                    } => {
                        // ignored (if provided at all) as we process reasoning as a stream already
                    }
                    StreamedAssistantContent::ReasoningDelta {
                        id: _,
                        provider_id: _,
                        reasoning,
                    } => {
                        // @todo handle error by logging it (send does not work) and quitting the actor
                        self.session_tx
                            .send(SessionMessage::AgentReasoning(reasoning))
                            .await
                            .unwrap();
                    }
                    StreamedAssistantContent::Final(_message) => {
                        // ignored as we process reasoning as we handle usage elsewhere
                    }
                    StreamedAssistantContent::Unknown(_value) => {
                        // @todo report this as system so that the user knows what happened on the provider side
                        // at least log this to the application log
                    }
                }
            }
            MultiTurnStreamItem::ToolExecutionCommitted {
                tool_call: _,
                internal_call_id: _,
            } => {
                // tools are observed via StreamAssistantItem to provide real-time insights
            }
            MultiTurnStreamItem::StreamUserItem(
                rig::streaming::StreamedUserContent::ToolResult {
                    tool_result,
                    internal_call_id: _,
                },
            ) => {
                self.session_tx
                    .send(SessionMessage::ToolResult {
                        name: tool_result.name,
                        output: format!("{:?}", tool_result.content),
                        // output: tool_result
                        //     .content
                        //     .first()
                        //     .as_text()
                        //     .unwrap_or("Tool result is not text")
                        //     .to_string(), // @todo handle multiple results
                        success: true,
                        call_id: tool_result.call,
                        provider_id: tool_result.provider,
                        // @todo is there no success for streamed user content
                    })
                    .await
                    .unwrap();
            }
            MultiTurnStreamItem::CompletionCall(completion_call)
                if completion_call.usage.has_values() =>
            {
                self.session_tx
                    .send(SessionMessage::AgentCompletion(completion_call.usage))
                    .await
                    .unwrap();
            }
            MultiTurnStreamItem::ModelTurnRetried { turn: _ } => {
                // we assume that this does not happen
                // report it as a system message to be safe (or simply log it)
            }
            MultiTurnStreamItem::FinalResponse(prompt_response) => {
                self.session_tx
                    .send(SessionMessage::FinalAgentResponse(prompt_response.output))
                    .await
                    .unwrap();
            }
            _ => {
                // @todo report an error for developers (log?)
            }
        }
    }

    async fn handle_unconnected(&mut self, agent: &mut Option<Agent>, base_dir: &Path) -> LlmState {
        if let Some(endpoint) = &self.endpoint
            && let Some(model) = &self.model
        {
            // @todo report and log error instead of unwraps
            let api_token = endpoint
                .secrets
                .clone()
                .unwrap_or(config::EndpointSecrets {
                    api_token: Some("".to_string()),
                })
                .api_token
                .unwrap_or("".to_string());

            let new_agent = match endpoint.api_type {
                // @todo find a way to deduplicate this code
                config::EndpointApi::OpenAICompletions => {
                    let client = openai::Client::builder()
                        .base_url(&endpoint.base_url)
                        .api_key(api_token)
                        .build()
                        .unwrap()
                        .completions_api();

                    client
                        .agent(model)
                        .preamble("You are a experienced software developer advising a colleague.")
                        .tool(tools::ListDir::new(
                            base_dir,
                            self.config.max_list_dir_results,
                        ))
                        .tool(tools::ReadFile::new(base_dir))
                        .tool(tools::SearchText::new(
                            base_dir,
                            self.config.max_search_text_results,
                        ))
                        .tool(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(),
                        ))
                        .default_max_turns(self.config.max_iterations_per_request)
                        .max_tokens(self.config.max_tokens)
                        .build()
                }
                config::EndpointApi::OpenAIResponses => {
                    let client = openai::Client::builder()
                        .base_url(&endpoint.base_url)
                        .api_key(api_token)
                        .build()
                        .unwrap();

                    client
                        .agent(model)
                        .preamble("You are a experienced software developer advising a colleague.")
                        .tool(tools::ListDir::new(
                            base_dir,
                            self.config.max_list_dir_results,
                        ))
                        .tool(tools::ReadFile::new(base_dir))
                        .tool(tools::SearchText::new(
                            base_dir,
                            self.config.max_search_text_results,
                        ))
                        .tool(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(),
                        ))
                        .default_max_turns(self.config.max_iterations_per_request)
                        .max_tokens(self.config.max_tokens)
                        .build()
                }
            };
            *agent = Some(new_agent);

            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,
        }
    }
}

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) => {
                Some(rig::completion::Message::Assistant {
                    id: None,
                    // do it manuall as there is no `tool_call_for()` function
                    content: vec![rig::message::AssistantContent::ToolCall(ToolCall {
                        id: tool_invocation_content.call_id.clone(),
                        provider: tool_invocation_content.provider_id.clone(),
                        function: ToolFunction {
                            name: tool_invocation_content.name.clone(),
                            arguments: tool_invocation_content.arguments.clone(),
                        },
                        signature: None,
                        additional_params: None,
                    })],
                })
            }
            ConversationItemKind::ToolOutput(tool_output_content) => {
                let content = rig::message::UserContent::tool_result_for(
                    tool_output_content.call_id.clone(),
                    tool_output_content.provider_id.clone(),
                    tool_output_content.name.clone(),
                    vec![ToolResultContent::Text(
                        tool_output_content.output.clone().into(),
                    )],
                );
                Some(rig::completion::Message::User {
                    content: vec![content],
                })
            }
            ConversationItemKind::Notification(_notification_type, _) => None,
        })
        .collect()
}

// from https://users.rust-lang.org/t/optional-future-for-optional-connections/77005
// Thanks, alice.
// async fn next_if_some<S>(s: &mut Option<S>) -> Option<S::Item>
// where
//     S: Stream + Unpin,
// {
//     match s.as_mut() {
//         Some(stream) => stream.next().await,
//         None => std::future::pending().await,
//     }
// }