baibot-withmcp/src/agent/provider/openai/controller.rs

use std::ops::Deref;

use async_openai::{
    config::OpenAIConfig,
    types::{
        ChatCompletionRequestMessage, CreateChatCompletionRequestArgs, CreateImageRequestArgs,
        CreateSpeechRequestArgs, CreateTranscriptionRequestArgs,
    },
    Client as OpenAIClient,
};

use super::super::ControllerTrait;
use crate::{
    agent::{
        provider::{
            entity::{ImageGenerationResult, PingResult, TextToSpeechParams, TextToSpeechResult},
            openai::utils::convert_string_to_enum,
        },
        AgentPurpose,
    },
    strings,
};
use crate::{
    agent::{
        provider::{
            entity::{TextGenerationParams, TextGenerationResult},
            ImageGenerationParams, SpeechToTextParams, SpeechToTextResult,
        },
        utils::base64_decode,
    },
    conversation::llm::{
        shorten_messages_list_to_context_size, Author as LLMAuthor,
        Conversation as LLMConversation, Message as LLMMessage,
    },
};

use super::config::Config;

#[derive(Debug, Clone)]
pub struct Controller {
    config: Config,
    client: OpenAIClient<OpenAIConfig>,
}

impl Controller {
    pub fn new(config: Config) -> Self {
        let openai_config = OpenAIConfig::new()
            .with_api_base(config.base_url.clone())
            .with_api_key(config.api_key.clone());

        let client = OpenAIClient::with_config(openai_config);

        Self { config, client }
    }
}

impl ControllerTrait for Controller {
    async fn ping(&self) -> anyhow::Result<PingResult> {
        if !self.supports_purpose(AgentPurpose::TextGeneration) {
            return Ok(PingResult::Inconclusive);
        }

        let messages = vec![LLMMessage {
            author: LLMAuthor::User,
            message_text: "Hello!".to_string(),
        }];

        let conversation = LLMConversation { messages };

        self.generate_text(conversation, TextGenerationParams::default())
            .await?;

        Ok(PingResult::Successful)
    }

    async fn generate_text(
        &self,
        conversation: LLMConversation,
        params: TextGenerationParams,
    ) -> anyhow::Result<TextGenerationResult> {
        let Some(text_generation_config) = &self.config.text_generation else {
            return Err(anyhow::anyhow!(
                strings::agent::no_configuration_for_purpose_so_cannot_be_used(
                    &AgentPurpose::TextGeneration
                ),
            ));
        };

        let prompt_text = params.prompt_variables.format(
            params
                .prompt_override
                .unwrap_or(self.text_generation_prompt().unwrap_or("".to_owned()))
                .trim(),
        );

        let prompt_message = if prompt_text.is_empty() {
            None
        } else {
            Some(LLMMessage {
                author: LLMAuthor::Prompt,
                message_text: prompt_text,
            })
        };

        let mut conversation_messages = conversation.messages;

        if params.context_management_enabled {
            tracing::trace!("Shortening messages list to context size");

            conversation_messages = shorten_messages_list_to_context_size(
                &text_generation_config.model_id,
                &prompt_message,
                conversation_messages,
                text_generation_config.max_response_tokens,
                text_generation_config.max_context_tokens,
            );

            tracing::trace!("Finished shortening messages list to context size");
        };

        if let Some(prompt_message) = prompt_message {
            conversation_messages.insert(0, prompt_message);
        }

        let openai_conversation_messages: Vec<ChatCompletionRequestMessage> =
            super::utils::convert_llm_messages_to_openai_messages(conversation_messages);

        let messages_count = openai_conversation_messages.len();

        let temperature = params
            .temperature_override
            .unwrap_or(text_generation_config.temperature);

        let mut request_builder = CreateChatCompletionRequestArgs::default();

        request_builder
            .model(&text_generation_config.model_id)
            .temperature(temperature)
            .messages(openai_conversation_messages);

        if let Some(max_response_tokens) = text_generation_config.max_response_tokens {
            request_builder.max_tokens(max_response_tokens);
        }

        let request = request_builder.build()?;

        if let Ok(request_as_json) = serde_json::to_string(&request) {
            tracing::trace!(
                model = format!("{:?}", request.model),
                ?messages_count,
                request = request_as_json,
                "Sending OpenAI chat completion API request"
            );
        }

        let response = self.client.chat().create(request).await?;

        tracing::trace!(
            ?response,
            "Got response from the OpenAI chat completion API"
        );

        // We only request 1 result, so there should only be 1 choice.
        if let Some(choice) = response.choices.into_iter().next() {
            match choice.message.content {
                Some(text) => {
                    return Ok(TextGenerationResult { text });
                }
                None => {
                    return Err(anyhow::anyhow!(
                        "No content was found in the response choice from the OpenAI chat completion API"
                    ));
                }
            }
        }

        Err(anyhow::anyhow!(
            "No response messages choices were returned from the OpenAI chat completion API"
        ))
    }

    async fn speech_to_text(
        &self,
        mime_type: &mxlink::mime::Mime,
        media: Vec<u8>,
        params: SpeechToTextParams,
    ) -> anyhow::Result<SpeechToTextResult> {
        let Some(speech_to_text_config) = &self.config.speech_to_text else {
            return Err(anyhow::anyhow!(
                strings::agent::no_configuration_for_purpose_so_cannot_be_used(
                    &AgentPurpose::SpeechToText
                ),
            ));
        };

        let filename = audio_mime_type_to_file_name(mime_type).unwrap_or("audio.ogg".to_string());

        let language = params.language_override.unwrap_or("".to_string());

        let request = CreateTranscriptionRequestArgs::default()
            .model(&speech_to_text_config.model_id)
            .file(async_openai::types::AudioInput {
                source: async_openai::types::InputSource::VecU8 {
                    filename,
                    vec: media,
                },
            })
            .language(language.clone())
            .build()?;

        tracing::trace!(
            model_id = speech_to_text_config.model_id,
            ?language,
            "Sending OpenAI speech-to-text API request"
        );

        let response = self.client.audio().transcribe(request).await?;

        tracing::trace!(
            ?response,
            "Got response from the OpenAI audio transcription API"
        );

        Ok(SpeechToTextResult {
            text: response.text,
        })
    }

    async fn generate_image(
        &self,
        prompt: &str,
        params: ImageGenerationParams,
    ) -> anyhow::Result<ImageGenerationResult> {
        let Some(image_generation_config) = &self.config.image_generation else {
            return Err(anyhow::anyhow!(
                strings::agent::no_configuration_for_purpose_so_cannot_be_used(
                    &AgentPurpose::ImageGeneration
                ),
            ));
        };

        let original_model = image_generation_config
            .model_id_as_openai_image_model()
            .map_err(|err| anyhow::anyhow!(err))?;

        let model = if params.cheaper_model_switching_allowed {
            // Switch to a cheaper model
            match original_model {
                async_openai::types::ImageModel::DallE2 => async_openai::types::ImageModel::DallE2,
                async_openai::types::ImageModel::DallE3 => async_openai::types::ImageModel::DallE2,
                async_openai::types::ImageModel::Other(_) => {
                    async_openai::types::ImageModel::DallE2
                }
            }
        } else {
            original_model
        };

        let quality = if params.cheaper_quality_switching_allowed {
            // Switch to a cheaper quality
            match &image_generation_config.quality {
                async_openai::types::ImageQuality::Standard => {
                    async_openai::types::ImageQuality::Standard
                }
                async_openai::types::ImageQuality::HD => {
                    async_openai::types::ImageQuality::Standard
                }
            }
        } else {
            image_generation_config.quality.clone()
        };

        let size = params
            .size_override
            .map(|s| {
                convert_string_to_enum::<async_openai::types::ImageSize>(&s)
                    .unwrap_or(image_generation_config.size)
            })
            .unwrap_or(image_generation_config.size);

        let request = CreateImageRequestArgs::default()
            .model(model)
            .prompt(prompt.to_owned())
            .response_format(async_openai::types::ImageResponseFormat::B64Json)
            .size(size)
            .style(image_generation_config.style.clone())
            .quality(quality)
            .build()?;

        tracing::trace!(
            ?prompt,
            model = format!("{:?}", request.model),
            size = format!("{:?}", request.size),
            style = format!("{:?}", request.style),
            quality = format!("{:?}", request.quality),
            "Sending OpenAI image generation API request"
        );

        let response = self.client.images().create(request).await?;

        if let Some(image) = response.data.into_iter().next() {
            match image.deref() {
                async_openai::types::Image::B64Json {
                    b64_json,
                    revised_prompt,
                } => {
                    let bytes = base64_decode(b64_json)?;

                    return Ok(ImageGenerationResult {
                        bytes,
                        mime_type: mxlink::mime::IMAGE_PNG,
                        revised_prompt: revised_prompt.clone(),
                    });
                }
                _ => {
                    return Err(anyhow::anyhow!("Unexpected image type"));
                }
            }
        }

        Err(anyhow::anyhow!(
            "The OpenAI image generation API returned no images"
        ))
    }

    async fn text_to_speech(
        &self,
        input: &str,
        params: TextToSpeechParams,
    ) -> anyhow::Result<TextToSpeechResult> {
        let Some(text_to_speech_config) = &self.config.text_to_speech else {
            return Err(anyhow::anyhow!(
                strings::agent::no_configuration_for_purpose_so_cannot_be_used(
                    &AgentPurpose::TextToSpeech
                ),
            ));
        };

        let speed = params.speed_override.unwrap_or(text_to_speech_config.speed);

        let voice = if let Some(voice_string) = params.voice_override {
            // This is a hacky way to construct a Voice enum from the string we have.
            let voice: serde_json::Result<async_openai::types::Voice> =
                serde_json::from_str(&format!("\"{}\"", voice_string));
            match voice {
                Ok(voice) => voice,
                Err(err) => {
                    tracing::debug!(?voice_string, ?err, "Failed to parse voice");

                    return Err(anyhow::anyhow!(
                        "The configured voice ({}) is not supported.",
                        voice_string
                    ));
                }
            }
        } else {
            text_to_speech_config.voice.clone()
        };

        let response_format = text_to_speech_config.response_format;

        let mime_type = response_format_to_mime_type(&response_format).unwrap_or(
            "audio/mp3"
                .parse()
                .expect("Failed parsing default mime type"),
        );

        let request = CreateSpeechRequestArgs::default()
            .model(text_to_speech_config.model_id.clone())
            .voice(voice)
            .speed(speed)
            .response_format(response_format)
            .input(input)
            .build()?;

        tracing::trace!(
            model = format!("{:?}", request.model),
            voice = format!("{:?}", request.voice),
            speed = format!("{:?}", request.speed),
            "Sending OpenAI text-to-speech API request"
        );

        let result = self.client.audio().speech(request).await?;

        Ok(TextToSpeechResult {
            bytes: result.bytes.into(),
            mime_type,
        })
    }

    fn supports_purpose(&self, purpose: AgentPurpose) -> bool {
        match purpose {
            AgentPurpose::TextGeneration => self.config.text_generation.is_some(),
            AgentPurpose::SpeechToText => self.config.speech_to_text.is_some(),
            AgentPurpose::TextToSpeech => self.config.text_to_speech.is_some(),
            AgentPurpose::ImageGeneration => self.config.image_generation.is_some(),
            AgentPurpose::CatchAll => true,
        }
    }

    fn text_generation_model_id(&self) -> Option<String> {
        self.config
            .text_generation
            .as_ref()
            .map(|config| config.model_id.to_owned())
    }

    fn text_generation_prompt(&self) -> Option<String> {
        self.config
            .text_generation
            .as_ref()
            .and_then(|config| config.prompt.clone())
    }

    fn text_generation_temperature(&self) -> Option<f32> {
        self.config
            .text_generation
            .as_ref()
            .map(|config| config.temperature)
    }

    fn text_to_speech_voice(&self) -> Option<String> {
        let Some(text_to_speech_config) = &self.config.text_to_speech else {
            return None;
        };

        // A hacky way to turn this enum to a string
        let voice_as_string = serde_json::to_string(&text_to_speech_config.voice).ok()?;
        Some(voice_as_string.replace("\"", ""))
    }

    fn text_to_speech_speed(&self) -> Option<f32> {
        let Some(text_to_speech_config) = &self.config.text_to_speech else {
            return None;
        };

        Some(text_to_speech_config.speed)
    }
}

fn response_format_to_mime_type(
    response_format: &async_openai::types::SpeechResponseFormat,
) -> Option<mxlink::mime::Mime> {
    let content_type = match response_format {
        async_openai::types::SpeechResponseFormat::Mp3 => "audio/mp3".to_owned(),
        async_openai::types::SpeechResponseFormat::Wav => "audio/wav".to_owned(),
        async_openai::types::SpeechResponseFormat::Opus => "audio/ogg".to_owned(),
        async_openai::types::SpeechResponseFormat::Aac => "audio/aac".to_owned(),
        async_openai::types::SpeechResponseFormat::Flac => "audio/flac".to_owned(),
        async_openai::types::SpeechResponseFormat::Pcm => "audio/L8".to_owned(),
    };

    match content_type.parse() {
        Ok(content_type) => Some(content_type),
        Err(err) => {
            tracing::error!(?err, "Failed to parse content type");
            None
        }
    }
}

fn audio_mime_type_to_file_name(mime_type: &mxlink::mime::Mime) -> Option<String> {
    let mime_type_string = mime_type.to_string();

    let file_extension = match mime_type_string.as_str() {
        "audio/flac" => "flac",
        "audio/x-m4a" | "audio/m4a" => "m4a",
        "audio/mp3" | "audio/mpeg" => "mp3",
        "audio/mp4" => "mp4",
        "application/ogg" | "audio/ogg" => "ogg",
        "audio/wav" | "audio/x-wav" => "wav",
        "audio/webm" => "webm",
        _ => return None,
    };

    Some(format!("audio.{}", file_extension))
}