gyandoot-v2/crates/gyandoot-providers/src/groq.rs

//! Groq API provider implementation.
//!
//! OpenAI-compatible chat completions endpoint (`/v1/chat/completions`) with
//! Bearer token auth. Supports streaming and function calling identically
//! to OpenAI. Base URL: `https://api.groq.com/openai/v1`.

use async_trait::async_trait;
use gyandoot_core::{ConversationContext, ProviderResult};
use std::time::Duration;
use tracing::debug;

use crate::openai::OpenAIProvider;
use crate::{Provider, ProviderResponse};

const GROQ_BASE_URL: &str = "https://api.groq.com/openai/v1";

/// Provider for the Groq API.
///
/// Groq uses an OpenAI-compatible chat completions endpoint.
/// This implementation wraps [`OpenAIProvider`] with the correct base URL.
pub struct GroqProvider {
    inner: OpenAIProvider,
}

impl GroqProvider {
    /// Create a new Groq provider.
    ///
    /// - `api_key`: Groq API key (sent as Bearer token).
    /// - `model`: Model identifier (e.g. `"llama-3.3-70b-versatile"`).
    /// - `base_url`: Optional override for the API base URL.
    pub fn new(api_key: String, model: String, base_url: Option<String>) -> Self {
        let base = base_url.unwrap_or_else(|| GROQ_BASE_URL.to_string());
        Self {
            inner: OpenAIProvider::new(api_key, model, Some(base)),
        }
    }

    /// Set a custom request timeout.
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.inner = self.inner.with_timeout(timeout);
        self
    }
}

#[async_trait]
impl Provider for GroqProvider {
    async fn chat(&self, context: ConversationContext) -> ProviderResult<ProviderResponse> {
        debug!(
            "Groq chat request (model: {}, messages: {})",
            self.inner.model,
            context.messages.len()
        );
        self.inner.chat(context).await
    }

    async fn chat_stream(
        &self,
        context: ConversationContext,
    ) -> ProviderResult<Box<dyn tokio_stream::Stream<Item = ProviderResult<String>> + Send + Unpin>>
    {
        debug!(
            "Groq streaming chat (model: {}, messages: {})",
            self.inner.model,
            context.messages.len()
        );
        self.inner.chat_stream(context).await
    }

    async fn models(&self) -> ProviderResult<Vec<String>> {
        Ok(vec![
            "llama-3.3-70b-versatile".to_string(),
            "llama-3.1-8b-instant".to_string(),
            "mixtral-8x7b-32768".to_string(),
            "gemma2-9b-it".to_string(),
            "deepseek-r1-distill-llama-70b".to_string(),
        ])
    }

    fn supports_tools(&self) -> bool {
        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_provider() -> GroqProvider {
        GroqProvider::new(
            "gsk-test".to_string(),
            "llama-3.3-70b-versatile".to_string(),
            None,
        )
    }

    #[test]
    fn test_creation() {
        let provider = test_provider();
        assert_eq!(provider.inner.model, "llama-3.3-70b-versatile");
    }

    #[test]
    fn test_supports_tools() {
        let provider = test_provider();
        assert!(provider.supports_tools());
    }
}