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

use std::ops::Deref;

use async_openai::{
    Client as OpenAIClient,
    config::OpenAIConfig,
    types::{
        audio::{AudioInput, CreateSpeechRequestArgs, CreateTranscriptionRequestArgs},
        chat::{ChatCompletionRequestMessage, CreateChatCompletionRequestArgs},
        images::{
            CreateImageEditRequestArgs, CreateImageRequestArgs,
            Image, ImageInput, ImageModel, ImageResponseFormat,
        },
    },
};

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

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,
            content: LLMMessageContent::Text("Hello!".to_string()),
            timestamp: chrono::Utc::now(),
        }];

        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,
                content: LLMMessageContent::Text(prompt_text),
                timestamp: chrono::Utc::now(),
            })
        };

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

        if let Some(max_completion_tokens) = text_generation_config.max_completion_tokens {
            request_builder.max_completion_tokens(max_completion_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(AudioInput::from_vec_u8(filename, 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().transcription().create(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 {
                ImageModel::DallE2 => ImageModel::DallE2,
                ImageModel::DallE3 => ImageModel::DallE2,
                ImageModel::Other(_) => {
                    ImageModel::DallE2
                }
                _ => original_model.clone(),
            }
        } else {
            original_model
        };

        let quality = if params.cheaper_quality_switching_allowed {
            // Switch to a cheaper quality
            match &image_generation_config.quality {
                Some(quality) => match quality {
                    async_openai::types::images::ImageQuality::Standard => {
                        Some(async_openai::types::images::ImageQuality::Standard)
                    }
                    async_openai::types::images::ImageQuality::HD => {
                        Some(async_openai::types::images::ImageQuality::Standard)
                    }
                    // New quality levels - keep as-is or downgrade to Standard
                    async_openai::types::images::ImageQuality::High => {
                        Some(async_openai::types::images::ImageQuality::Standard)
                    }
                    async_openai::types::images::ImageQuality::Medium => {
                        Some(async_openai::types::images::ImageQuality::Medium)
                    }
                    async_openai::types::images::ImageQuality::Low => {
                        Some(async_openai::types::images::ImageQuality::Low)
                    }
                    async_openai::types::images::ImageQuality::Auto => {
                        Some(async_openai::types::images::ImageQuality::Auto)
                    }
                },
                None => None,
            }
        } else {
            image_generation_config.quality.clone()
        };

        let size = params
            .size_override
            .map(|s| convert_string_to_enum::<async_openai::types::images::ImageSize>(&s).unwrap())
            .or(image_generation_config.size);

        let response_format = match model.clone() {
            ImageModel::DallE2 => Some(ImageResponseFormat::B64Json),
            ImageModel::DallE3 => Some(ImageResponseFormat::B64Json),
            // gpt-image-1 only outputs base64 and we don't need to specify the response format.
            // In fact, specifying the response format results in an error.
            ImageModel::GptImage1 => None,
            ImageModel::GptImage1Mini => None,
            ImageModel::GptImage1dot5 => None,
            ImageModel::Other(_) => Some(ImageResponseFormat::B64Json),
        };

        let mut request_builder = CreateImageRequestArgs::default();

        request_builder.model(model).prompt(prompt.to_owned());

        if let Some(response_format) = response_format {
            request_builder.response_format(response_format);
        }

        if let Some(style) = &image_generation_config.style {
            request_builder.style(style.clone());
        }

        if let Some(quality) = quality {
            request_builder.quality(quality.clone());
        }

        if let Some(size) = size {
            request_builder.size(size);
        }

        let request = request_builder.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().generate(request).await?;

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

                    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 create_image_edit(
        &self,
        prompt: &str,
        images: Vec<ImageSource>,
        _params: ImageEditParams,
    ) -> anyhow::Result<ImageEditResult> {
        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
                ),
            ));
        };

        if images.is_empty() {
            return Err(anyhow::anyhow!("No image sources provided"));
        }

        let mut image_inputs: Vec<ImageInput> = Vec::new();
        for image in images {
            image_inputs.push(image.into());
        }

        let dalle2_size = match image_generation_config.size {
            Some(async_openai::types::images::ImageSize::S256x256) => Some(async_openai::types::images::ImageSize::S256x256),
            Some(async_openai::types::images::ImageSize::S512x512) => Some(async_openai::types::images::ImageSize::S512x512),
            Some(async_openai::types::images::ImageSize::S1024x1024) => Some(async_openai::types::images::ImageSize::S1024x1024),
            _ => None,
        };

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

        let response_format = match model.clone() {
            ImageModel::DallE2 => {
                Some(ImageResponseFormat::B64Json)
            }
            ImageModel::DallE3 => {
                Some(ImageResponseFormat::B64Json)
            }
            // gpt-image-1 only outputs base64 and we don't need to specify the response format.
            // In fact, specifying the response format results in an error.
            ImageModel::GptImage1 => None,
            ImageModel::GptImage1Mini => None,
            ImageModel::GptImage1dot5 => None,
            ImageModel::Other(_) => Some(ImageResponseFormat::B64Json),
        };

        let mut request_builder = CreateImageEditRequestArgs::default();

        request_builder
            .image(image_inputs)
            .prompt(prompt.to_owned())
            .model(model);

        if let Some(size) = dalle2_size {
            request_builder.size(size);
        }

        if let Some(response_format) = response_format {
            request_builder.response_format(response_format);
        }

        let request = request_builder
            .build()
            .map_err(|e| anyhow::anyhow!("Failed to build CreateImageEditRequest: {}", e))?;

        tracing::trace!(
            model = format!("{:?}", request.model),
            size = format!("{:?}", request.size),
            response_format = format!("{:?}", request.response_format),
            "Sending OpenAI image edit API request"
        );

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

        if let Some(image_data) = response.data.into_iter().next() {
            match image_data.deref() {
                Image::B64Json { b64_json, .. } => {
                    let bytes = base64_decode(b64_json.as_ref())?;
                    return Ok(ImageEditResult {
                        bytes,
                        mime_type: mxlink::mime::IMAGE_PNG,
                    });
                }
                Image::Url { url, .. } => {
                    tracing::warn!(?url, "Received URL instead of B64Json for image edit");
                    return Err(anyhow::anyhow!(
                        "Unexpected image type (URL) when B64Json was requested"
                    ));
                }
            }
        }

        Err(anyhow::anyhow!(
            "The OpenAI image edit 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::audio::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().create(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::audio::SpeechResponseFormat,
) -> Option<mxlink::mime::Mime> {
    let content_type = match response_format {
        async_openai::types::audio::SpeechResponseFormat::Mp3 => "audio/mp3".to_owned(),
        async_openai::types::audio::SpeechResponseFormat::Wav => "audio/wav".to_owned(),
        async_openai::types::audio::SpeechResponseFormat::Opus => "audio/ogg".to_owned(),
        async_openai::types::audio::SpeechResponseFormat::Aac => "audio/aac".to_owned(),
        async_openai::types::audio::SpeechResponseFormat::Flac => "audio/flac".to_owned(),
        async_openai::types::audio::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))
}