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

*/

mod tools;

use std::{env, pin::Pin};

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

use crate::{
    config::{Config, Endpoint, EndpointSecrets},
    core::conversation::ConversationItemKind,
    session::SessionMessage,
};
use futures::{Stream, StreamExt};
use rig::{
    Agent,
    agent::{MultiTurnStreamItem, StreamingError},
    client::{AgentClientExt, Client},
    providers::openai::{
        self, OpenAIResponsesExt,
        responses_api::{GenericResponsesCompletionModel, streaming::StreamingCompletionResponse},
    },
    streaming::{StreamedAssistantContent, StreamingChat},
};

#[derive(Clone, Debug, PartialEq)]
pub enum LlmMessage {
    SubmitPrompt {
        prompt: String,
        history: Vec<ConversationItemKind>,
    },
    //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 ResponseStream = Pin<
    Box<
        dyn Stream<Item = Result<MultiTurnStreamItem<StreamingCompletionResponse>, StreamingError>>
            + Send
            + 'static,
    >,
>;

pub enum LlmState {
    // @todo Unconnected,
    Idle,
    Requesting {
        user_prompt: String,
        history: Vec<rig::completion::Message>,
    },
    Streaming {
        stream: ResponseStream,
    },
    Quitting,
}

#[derive(Debug)]
pub struct Llm {
    session_tx: Sender<SessionMessage>,
    rx: Receiver<LlmMessage>,
    config: Config,
}

impl Llm {
    pub fn new(
        session_tx: Sender<SessionMessage>,
        rx: Receiver<LlmMessage>,
        config: Config,
    ) -> Self {
        Self {
            session_tx,
            rx,
            config,
        }
    }

    pub async fn run(&mut self) {
        // @todo later don't pass the config to Llm but have a Connect command called from Session
        // also have multiple models per endpoint
        // @todo report and log error
        let endpoint = self.get_active_endpoint().expect("Endpoint missing");
        let client = self.connect(&endpoint.base_url, endpoint.secrets).unwrap();

        // @todo issue an error system message and do not publish any tools on run
        let base_dir = env::current_dir().expect("Cannot get current directory");

        // @todo build the agent on prompt request (extra state: Initializing where we wait for model and system prompt)
        let agent = client
            .agent(endpoint.default_model)
            .preamble("You are a experienced software developer advising a colleague.")
            .tool(tools::ListDir::new(
                &base_dir,
                self.config.general.max_list_dir_results,
            ))
            .tool(tools::ReadFile::new(&base_dir))
            .tool(tools::SearchText::new(
                &base_dir,
                self.config.general.max_search_text_results,
            ))
            .default_max_turns(self.config.general.max_iterations_per_request)
            .max_tokens(self.config.general.max_tokens)
            .build();

        let mut state = LlmState::Idle;

        loop {
            state = match state {
                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,
            };
        }
    }

    fn connect(
        &self,
        base_url: &str,
        secrets: Option<EndpointSecrets>,
    ) -> Result<Client<OpenAIResponsesExt>> {
        // at least in rig 0.41 api_key is mandatory. Therefore provide one in any case.
        let api_token = secrets
            .unwrap_or(EndpointSecrets {
                api_token: Some("".to_string()),
            })
            .api_token
            .unwrap_or("".to_string());

        let client = openai::Client::builder()
            .base_url(base_url)
            .api_key(api_token)
            .build()?;

        Ok(client)
    }

    // @todo this should happen in session and send via Connect(endpoint, system prompt) request
    fn get_active_endpoint(&self) -> Result<Endpoint> {
        self.config
            .endpoint
            .get(&self.config.general.primary_endpoint)
            .map(|val| val.to_owned())
            .ok_or(Error::msg(format!(
                "Endpoint {} not found in endpoints.",
                self.config.general.primary_endpoint
            )))
    }

    // 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::SubmitPrompt { prompt, history }) => {
                //let history = rig_history_from_internal(history);
                let history: Vec<rig::completion::Message> = Vec::new();
                LlmState::Requesting {
                    user_prompt: prompt,
                    history,
                }
            }
            Some(LlmMessage::Quit) => LlmState::Quitting,
            None => LlmState::Quitting,
        }
    }

    async fn handle_requesting(
        &mut self,
        agent: &Agent<GenericResponsesCompletionModel>,
        user_prompt: String,
        history: Vec<rig::completion::Message>,
    ) -> LlmState {
        let request = agent.stream_chat(user_prompt, history);
        tokio::select! {
            stream = request => { LlmState::Streaming { stream } }
            // @todo handle new request while request is running
            Some(LlmMessage::Quit) = self.rx.recv() => { LlmState::Idle }
        }
    }

    async fn handle_streaming(&mut self, mut stream: ResponseStream) -> LlmState {
        loop {
            tokio::select! {
                maybe_chunk = stream.next() => {
                    if let Some(result) = maybe_chunk{
                        if let Ok(chunk) = result {
                            self.handle_llm_item(chunk).await;

                        } else {
                            // @todo why does this error occur if the endpoint is not a valid address?
                            // it should fail earlier, imho
                            // println!("Buuhuuhuu! Error: {}", result.unwrap_err());
                            self.session_tx.send(SessionMessage::Error(format!("{:?}", result))).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<StreamingCompletionResponse>) {
        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,
                            })
                            .await
                            .unwrap();
                    }
                    StreamedAssistantContent::ToolCallDelta {
                        id: _,
                        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) => {
                        // ignored (if provided at all) as we process reasoning as a stream already
                    }
                    StreamedAssistantContent::ReasoningDelta { 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 {
                        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, // @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) => {
                // we assemble the prompt and usage from other items, ignore the blocking version
                // @todo use it to indicate completion of a request
            }
            _ => {
                // @todo report an error for developers (log?)
            }
        }
    }
}

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