simple-groq-rs 0.1.3

use reqwest::{Client, header};
use serde::{Deserialize, Serialize};
use std::env;
use thiserror::Error;

/// Custom error type for the crate
#[derive(Error, Debug)]
pub enum GroqError {
    #[error("Reqwest error: {0}")]
    Reqwest(#[from] reqwest::Error),

    #[error("API error: {code} - {message}")]
    Api { code: String, message: String },

    #[error("Environment error: {0}")]
    Env(String),

    #[error("No response choices returned")]
    NoChoices,
}

/// A single message in the conversation
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct Message {
    pub role: String,
    pub content: String,
}

impl Message {
    pub fn new(role: &str, content: &str) -> Self {
        Self {
            role: role.to_string(),
            content: content.to_string(),
        }
    }

    pub fn system(content: &str) -> Self {
        Self::new("system", content)
    }

    pub fn user(content: &str) -> Self {
        Self::new("user", content)
    }

    pub fn assistant(content: &str) -> Self {
        Self::new("assistant", content)
    }
}

/// Request body for chat completions
#[derive(Serialize)]
struct ChatRequest {
    model: String,
    messages: Vec<Message>,
    temperature: Option<f32>,
    max_tokens: Option<u32>,
    stream: Option<bool>,
}

/// Part of the response
#[derive(Deserialize)]
pub struct Choice {
    pub message: Message,
    pub finish_reason: Option<String>,
}

/// Main response from the API
#[derive(Deserialize)]
pub struct ChatResponse {
    pub choices: Vec<Choice>,
    // You can add usage, id, created, etc. later
}

/// The main client
#[derive(Clone)]
pub struct GroqClient {
    client: Client,
    api_key: String,
    base_url: String,
}

impl GroqClient {
    /// Create a new client with an API key
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            client: Client::new(),
            api_key: api_key.into(),
            base_url: "https://api.groq.com/openai/v1".to_string(),
        }
    }

    /// Create from Groq_API_KEY environment variable
    pub fn from_env() -> Result<Self, GroqError> {
        let api_key = env::var("GROQ_API_KEY")
            .map_err(|_| GroqError::Env("GROQ_API_KEY not set".to_string()))?;
        Ok(Self::new(api_key))
    }

    /// Send a chat completion request
    pub async fn chat_completion(
        &self,
        model: &str,
        messages: Vec<Message>,
        temperature: Option<f32>,
        max_tokens: Option<u32>,
    ) -> Result<String, GroqError> {
        let request = ChatRequest {
            model: model.to_string(),
            messages,
            temperature,
            max_tokens,
            stream: Some(false),
        };

        let response = self
            .client
            .post(format!("{}/chat/completions", self.base_url))
            .header(header::AUTHORIZATION, format!("Bearer {}", self.api_key))
            .header(header::CONTENT_TYPE, "application/json")
            .json(&request)
            .send()
            .await?;

        // Better error handling for API errors
        if !response.status().is_success() {
            let status = response.status();
            let text = response.text().await.unwrap_or_default();
            return Err(GroqError::Api {
                code: status.to_string(),
                message: text,
            });
        }

        let chat_response: ChatResponse = response.json().await?;

        let content = chat_response
            .choices
            .into_iter()
            .next()
            .ok_or(GroqError::NoChoices)?
            .message
            .content;

        Ok(content)
    }
}