honey-biscuit-workshop/workshop-llm-detector/src/llm/openai.rs

use async_trait::async_trait;
use anyhow::Result;
use openai_api_rust::*;
use openai_api_rust::chat::*;

use super::LLMClient;

pub struct OpenAIClient {
    client: OpenAI,
    model: String,
}

impl OpenAIClient {
    pub fn new(api_key: String, base_url: Option<String>, model: String) -> Self {
        let auth = Auth::new(api_key.as_str());
        let client = if let Some(url) = base_url {
            OpenAI::new(auth, url.as_str())
        } else {
            OpenAI::new(auth, "https://api.openai.com/v1/")
        };
        Self { client, model }
    }
}

#[async_trait]
impl LLMClient for OpenAIClient {
    async fn generate(&self, prompt: &str) -> Result<String> {
        let message = Message {
            role: Role::User,
            content: prompt.to_string(),
        };

        let body = ChatBody {
            model: self.model.clone(),
            messages: vec![message],
            max_tokens: Some(2048),
            temperature: Some(0.1),
            top_p: Some(0.1),
            n: Some(1),
            stream: Some(false),
            stop: None,
            presence_penalty: None,
            frequency_penalty: None,
            logit_bias: None,
            user: None,
        };

        let response = self.client.chat_completion_create(&body).unwrap();

        if let Some(choice) = response.choices.first() {
            Ok(choice.message.as_ref().unwrap().content.clone())
        } else {
            anyhow::bail!("No response from OpenAI")
        }
    }
}