monothink 0.2.1

/*
* 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, path::Path, pin::Pin};

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

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

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

// abstract different OpenAI APIs
pub enum LlmProviderType {
    Completion(Agent<openai::GenericCompletionModel>),
    Responses(Agent<GenericResponsesCompletionModel>),
}

// @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
pub enum LlmStream {
    Completion(
        Pin<
            Box<
                dyn Stream<
                        Item = Result<
                            MultiTurnStreamItem<openai::streaming::StreamingCompletionResponse>,
                            StreamingError,
                        >,
                    > + Send
                    + 'static,
            >,
        >,
    ),
    Responses(
        Pin<
            Box<
                dyn Stream<
                        Item = Result<
                            MultiTurnStreamItem<StreamingCompletionResponse>,
                            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,
    endpoint: Option<config::Endpoint>,
    model: Option<String>,
}

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

    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::SelectModel(model)) => {
                self.model = Some(model);
                LlmState::Unconnected
            }
            Some(LlmMessage::SetEndpoint(endpoint)) => {
                self.endpoint = Some(endpoint);
                LlmState::Unconnected
            }
            Some(LlmMessage::SubmitPrompt { prompt, history }) => {
                let history = rig_history_from_internal(&history);
                LlmState::Requesting {
                    user_prompt: prompt,
                    history,
                }
            }
            Some(LlmMessage::Quit) => LlmState::Quitting,
            None => LlmState::Quitting,
        }
    }

    async fn handle_requesting(
        &mut self,
        agent: &Option<LlmProviderType>,
        user_prompt: String,
        history: Vec<Message>,
    ) -> LlmState {
        let stream = match agent {
            Some(LlmProviderType::Responses(a)) => {
                // @todo handle new request while request is running
                let s = a.stream_chat(user_prompt, history).await;
                LlmStream::Responses(s)
            }
            Some(LlmProviderType::Completion(a)) => {
                // @todo handle new request while request is running
                let s = a.stream_chat(user_prompt, history).await;
                LlmStream::Completion(s)
            }
            None => return LlmState::Unconnected,
        };
        LlmState::Streaming { stream }
    }

    async fn handle_streaming(&mut self, mut stream: LlmStream) -> LlmState {
        match stream {
            LlmStream::Completion(s) => self.handle_stream_inner(s).await,
            LlmStream::Responses(s) => self.handle_stream_inner(s).await,
        }
    }

    async fn handle_stream_inner<S, T>(&mut self, mut s: S) -> LlmState
    where
        S: Stream<Item = Result<MultiTurnStreamItem<T>, StreamingError>> + Unpin,
    {
        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<R>(&mut self, chunk: MultiTurnStreamItem<R>) {
        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,
                                id: tool_call.id,
                                call_id: tool_call.call_id,
                            })
                            .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,
                        id: tool_result.id,
                        call_id: tool_result.call_id,
                        // @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?)
            }
        }
    }

    async fn handle_unconnected(
        &mut self,
        agent: &mut Option<LlmProviderType>,
        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 {
                    config::EndpointApi::OpenAICompletions => {
                        let client = openai::Client::builder()
                            .base_url(&endpoint.base_url)
                            .api_key(api_token)
                            .build()
                            .unwrap()
                            .completions_api();

                        LlmProviderType::Completion(
                        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,
                            ))
                            .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();

                        LlmProviderType::Responses(
                    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,
                        ))
                        .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::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> {
    let rig_history = 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) => None,
            // ConversationItemKind::ToolInvocation(tool_invocation_content) => {
            //     Some(rig::completion::Message::Assistant {
            //         id: None,
            //         content: rig::OneOrMany::one(match &tool_invocation_content.call_id {
            //             Some(call_id) => rig::message::AssistantContent::tool_call_with_call_id(
            //                 tool_invocation_content.id.clone(),
            //                 call_id.to_string(),
            //                 tool_invocation_content.name.clone(),
            //                 tool_invocation_content.arguments.clone(),
            //             ),
            //             None => rig::message::AssistantContent::tool_call(
            //                 tool_invocation_content.id.clone(),
            //                 tool_invocation_content.name.clone(),
            //                 tool_invocation_content.arguments.clone(),
            //             ),
            //         }),
            //     })
            // }
            ConversationItemKind::ToolOutput(tool_output_content) => None,
            // ConversationItemKind::ToolOutput(tool_output_content) => {
            //     Some(rig::completion::Message::tool_result_with_call_id(
            //         tool_output_content.id.clone(),
            //         tool_output_content.call_id.clone(),
            //         tool_output_content.output.clone(),
            //     ))
            // }
            ConversationItemKind::Notification(_notification_type, _) => None,
        })
        .collect();

    rig_history
}

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