src/eval_framework/utils/llm_client.py

"""LLM client utilities for async API calls."""

import asyncio
import logging
import re
from typing import Optional, Any
from abc import ABC, abstractmethod
import json

from openai import AsyncOpenAI, OpenAI
from openai import RateLimitError as OpenAIRateLimitError
from anthropic import Anthropic, AsyncAnthropic
from tenacity import (
    retry,
    retry_if_exception_type,
    wait_exponential,
    stop_after_attempt,
    before_sleep_log,
)

from ..config import get_settings

logger = logging.getLogger(__name__)


def _parse_retry_after(exc: Exception) -> float:
    """Extract 'retry after N seconds' from a Groq 429 error message."""
    msg = str(exc)
    match = re.search(r"try again in ([\d\.]+)s", msg)
    if match:
        return min(float(match.group(1)) + 1.0, 60.0)
    return 5.0


async def _groq_retry_wait(retry_state):
    """Custom wait that honours the 'retry after Xs' hint in the error."""
    exc = retry_state.outcome.exception()
    secs = _parse_retry_after(exc) if exc else 5.0
    logger.warning(f"Rate limited — waiting {secs:.1f}s before retry")
    await asyncio.sleep(secs)


_groq_retry = retry(
    retry=retry_if_exception_type(OpenAIRateLimitError),
    wait=wait_exponential(multiplier=1, min=2, max=60),
    stop=stop_after_attempt(8),
    before_sleep=before_sleep_log(logger, logging.WARNING),
    reraise=False,
)


class LLMClient(ABC):
    """Abstract base class for LLM clients."""
    
    @abstractmethod
    async def generate(
        self,
        system_prompt: str,
        user_message: str,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> str:
        """Generate text from the LLM."""
        pass
    
    @abstractmethod
    async def generate_json(
        self,
        system_prompt: str,
        user_message: str,
        schema: dict,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> dict:
        """Generate structured JSON from the LLM."""
        pass


class OpenAIClient(LLMClient):
    """Async OpenAI client wrapper."""
    
    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
        settings = get_settings()
        self.client = AsyncOpenAI(api_key=api_key or settings.openai_api_key)
        self.model = model or settings.openai_model
    
    async def generate(
        self,
        system_prompt: str,
        user_message: str,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> str:
        """Generate text using OpenAI API."""
        try:
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
            )
            return response.choices[0].message.content or ""
        except Exception as e:
            logger.error(f"OpenAI API error: {e}")
            raise
    
    async def generate_json(
        self,
        system_prompt: str,
        user_message: str,
        schema: dict,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> dict:
        """Generate structured JSON using OpenAI API."""
        try:
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                response_format={
                    "type": "json_schema",
                    "json_schema": {
                        "name": "evaluation_result",
                        "schema": schema,
                        "strict": True,
                    },
                },
            )
            content = response.choices[0].message.content or "{}"
            return json.loads(content)
        except Exception as e:
            logger.error(f"OpenAI JSON generation error: {e}")
            raise


class AnthropicClient(LLMClient):
    """Async Anthropic client wrapper."""
    
    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
        settings = get_settings()
        self.client = AsyncAnthropic(api_key=api_key or settings.anthropic_api_key)
        self.model = model or settings.anthropic_model
    
    async def generate(
        self,
        system_prompt: str,
        user_message: str,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> str:
        """Generate text using Anthropic API."""
        try:
            response = await self.client.messages.create(
                model=self.model,
                max_tokens=max_tokens,
                system=system_prompt,
                messages=[
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                top_p=top_p,
            )
            return response.content[0].text
        except Exception as e:
            logger.error(f"Anthropic API error: {e}")
            raise
    
    async def generate_json(
        self,
        system_prompt: str,
        user_message: str,
        schema: dict,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> dict:
        """Generate structured JSON using Anthropic API."""
        # Anthropic doesn't have native JSON mode, so we format the request carefully
        json_instructions = f"""
Please respond with valid JSON matching this schema:
{json.dumps(schema, indent=2)}

Only return the JSON object, no other text.
"""
        try:
            response = await self.client.messages.create(
                model=self.model,
                max_tokens=max_tokens,
                system=system_prompt + "\n" + json_instructions,
                messages=[
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                top_p=top_p,
            )
            content = response.content[0].text
            # Try to parse JSON from the response
            return json.loads(content)
        except Exception as e:
            logger.error(f"Anthropic JSON generation error: {e}")
            raise


class GroqClient(LLMClient):
    """Async Groq client wrapper.

    Groq exposes an OpenAI-compatible REST API, so we reuse AsyncOpenAI with a
    custom base_url.  Note: Groq does NOT support the ``response_format`` JSON
    schema mode, so ``generate_json`` falls back to prompt-level instructions
    and normal JSON parsing (same approach as AnthropicClient).
    """

    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
        settings = get_settings()
        self.client = AsyncOpenAI(
            api_key=api_key or settings.groq_api_key,
            base_url="https://api.groq.com/openai/v1",
        )
        self.model = model or settings.groq_model

    @_groq_retry
    async def generate(
        self,
        system_prompt: str,
        user_message: str,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> str:
        """Generate text using Groq API (auto-retries on 429)."""
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message},
            ],
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
        )
        return response.choices[0].message.content or ""

    @_groq_retry
    async def generate_json(
        self,
        system_prompt: str,
        user_message: str,
        schema: dict,
        temperature: float = 0.1,
        top_p: float = 0.9,
        max_tokens: int = 2000,
    ) -> dict:
        """Generate structured JSON using Groq API (prompt-based, auto-retries on 429)."""
        json_instructions = (
            f"\nRespond with valid JSON matching this schema:\n"
            f"{json.dumps(schema, indent=2)}\n"
            "Only return the JSON object, no other text."
        )
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system_prompt + json_instructions},
                {"role": "user", "content": user_message},
            ],
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
        )
        content = response.choices[0].message.content or "{}"
        return json.loads(content)


def get_llm_client(provider: str = "groq", model: Optional[str] = None) -> LLMClient:
    """Get an LLM client based on provider name.

    Args:
        provider: One of "openai", "anthropic", "groq".
        model: Optional model override. If None, uses the provider's default
               from settings. Pass a specific model name to use a different
               model than the default (e.g. a cheaper judge model).
    """
    if provider.lower() == "openai":
        return OpenAIClient(model=model)
    elif provider.lower() == "anthropic":
        return AnthropicClient(model=model)
    elif provider.lower() == "groq":
        return GroqClient(model=model)
    else:
        raise ValueError(f"Unknown provider: {provider}. Choose from: openai, anthropic, groq")