llm_client.py

"""
Multi-provider LLM client with ultra-low latency and streaming support.
Supports Groq and Cerebras for fast inference with A/B comparison capability.

Security: API keys are accessed lazily through SecureConfig,
only when an actual LLM call is made.
"""
import time
from dataclasses import dataclass
from typing import Optional, Generator

from config import SECURE_CONFIG


@dataclass
class LLMResponse:
    """Structured response from LLM call."""
    text: str
    provider: str
    model: str
    latency_ms: float
    had_context: bool
    tokens_used: Optional[int] = None


class LLMClient:
    """
    Multi-provider LLM client with ultra-low latency and streaming.
    
    Optimizations:
    - Streaming responses for real-time feel
    - Faster models (8b for speed, 70b for complex tasks)
    - Reduced max_tokens for faster first-token latency
    - Lazy initialization of API clients
    """
    
    # Fast models for low latency
    FAST_MODELS = {
        "groq": "llama-3.1-8b-instant",      # ~100ms latency
        "cerebras": "llama3.1-8b"             # ~50ms latency
    }
    
    # Full models for complex tasks
    FULL_MODELS = {
        "groq": "llama-3.3-70b-versatile",   # ~500ms latency
        "cerebras": "llama3.1-70b"
    }
    
    def __init__(self, use_fast_model: bool = True):
        self._groq = None
        self._cerebras = None
        self._initialized = False
        self._use_fast_model = use_fast_model
    
    def _ensure_initialized(self):
        """Lazy-initialize API clients only when first LLM call is made."""
        if self._initialized:
            return
        
        is_valid, message = SECURE_CONFIG.validate_keys()
        if not is_valid:
            raise ValueError(message)
        
        if SECURE_CONFIG.has_groq_key():
            from groq import Groq
            self._groq = Groq(api_key=SECURE_CONFIG.groq_api_key)
        
        if SECURE_CONFIG.has_cerebras_key():
            from cerebras.cloud.sdk import Cerebras
            self._cerebras = Cerebras(api_key=SECURE_CONFIG.cerebras_api_key)
        
        self._initialized = True
    
    @property
    def primary_provider(self) -> str:
        return SECURE_CONFIG.primary_provider
    
    @property
    def available_providers(self) -> list[str]:
        providers = []
        if SECURE_CONFIG.has_groq_key():
            providers.append("groq")
        if SECURE_CONFIG.has_cerebras_key():
            providers.append("cerebras")
        return providers
    
    def _get_model(self, provider: str) -> str:
        """Get appropriate model based on speed preference."""
        models = self.FAST_MODELS if self._use_fast_model else self.FULL_MODELS
        return models.get(provider, self.FAST_MODELS.get(provider))
    
    def generate_stream(
        self,
        prompt: str,
        context: str = "",
        provider: Optional[str] = None
    ) -> Generator[str, None, LLMResponse]:
        """
        Generate response with streaming for real-time display.
        
        Yields chunks of text as they arrive.
        Returns final LLMResponse when complete.
        """
        self._ensure_initialized()
        
        provider = provider or self.primary_provider
        full_prompt = f"{context}\n\n{prompt}" if context else prompt
        model = self._get_model(provider)
        
        start_time = time.perf_counter()
        full_text = ""
        tokens = None
        
        if provider == "groq" and self._groq:
            stream = self._groq.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": full_prompt}],
                temperature=0.7,
                max_tokens=512,  # Reduced for faster responses
                stream=True
            )
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    text_chunk = chunk.choices[0].delta.content
                    full_text += text_chunk
                    yield text_chunk
        
        elif provider == "cerebras" and self._cerebras:
            stream = self._cerebras.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": full_prompt}],
                max_tokens=512,
                stream=True
            )
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    text_chunk = chunk.choices[0].delta.content
                    full_text += text_chunk
                    yield text_chunk
        else:
            raise ValueError(f"Provider '{provider}' not available.")
        
        latency_ms = (time.perf_counter() - start_time) * 1000
        
        return LLMResponse(
            text=full_text,
            provider=provider,
            model=model,
            latency_ms=latency_ms,
            had_context=bool(context),
            tokens_used=tokens
        )
    
    def generate(
        self,
        prompt: str,
        context: str = "",
        provider: Optional[str] = None
    ) -> LLMResponse:
        """
        Generate response (non-streaming, optimized for speed).
        """
        self._ensure_initialized()
        
        provider = provider or self.primary_provider
        full_prompt = f"{context}\n\n{prompt}" if context else prompt
        model = self._get_model(provider)
        
        start_time = time.perf_counter()
        
        if provider == "groq" and self._groq:
            response = self._groq.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": full_prompt}],
                temperature=0.7,
                max_tokens=512  # Reduced for faster responses
            )
            text = response.choices[0].message.content
            tokens = getattr(response.usage, 'total_tokens', None)
        
        elif provider == "cerebras" and self._cerebras:
            response = self._cerebras.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": full_prompt}],
                max_tokens=512
            )
            text = response.choices[0].message.content
            tokens = None
        
        else:
            raise ValueError(
                f"Provider '{provider}' not available. "
                f"Available: {self.available_providers}."
            )
        
        latency_ms = (time.perf_counter() - start_time) * 1000
        
        return LLMResponse(
            text=text,
            provider=provider,
            model=model,
            latency_ms=latency_ms,
            had_context=bool(context),
            tokens_used=tokens
        )
    
    def generate_raw(self, prompt: str) -> LLMResponse:
        """Generate without any protection context."""
        return self.generate(prompt, context="")
    
    def generate_protected(self, prompt: str, context: str) -> LLMResponse:
        """Generate with protection context."""
        return self.generate(prompt, context=context)