app/llm/client.py

from groq import Groq
from typing import List, Dict, Optional
import json

from app.config.settings import settings


class LLMClient:
    """Client for interacting with Groq LLM API."""
    
    def __init__(self):
        """Initialize Groq client."""
        self.client = Groq(api_key=settings.GROQ_API_KEY)
        self.model = settings.GROQ_MODEL
    
    def get_completion(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 1024,
        json_mode: bool = False
    ) -> str:
        """
        Get completion from Groq LLM.
        
        Args:
            messages: List of message dictionaries with 'role' and 'content'
            temperature: Sampling temperature (0-2)
            max_tokens: Maximum tokens in response
            json_mode: Whether to request JSON output
            
        Returns:
            Response content string
        """
        try:
            response_format = {"type": "json_object"} if json_mode else None
            
            chat_completion = self.client.chat.completions.create(
                messages=messages,
                model=self.model,
                temperature=temperature,
                max_tokens=max_tokens,
                response_format=response_format
            )
            
            return chat_completion.choices[0].message.content
            
        except Exception as e:
            print(f"Error getting LLM completion: {e}")
            raise
    
    def get_completion_with_retry(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 1024,
        json_mode: bool = False,
        max_retries: int = 3
    ) -> str:
        """
        Get completion with retry logic.
        
        Args:
            messages: List of message dictionaries
            temperature: Sampling temperature
            max_tokens: Maximum tokens
            json_mode: Whether to request JSON output
            max_retries: Maximum number of retries
            
        Returns:
            Response content string
        """
        for attempt in range(max_retries):
            try:
                return self.get_completion(messages, temperature, max_tokens, json_mode)
            except Exception as e:
                if attempt == max_retries - 1:
                    raise
                print(f"Retry {attempt + 1}/{max_retries} after error: {e}")
                continue


# Global LLM client instance
llm_client = LLMClient()