src/analyzer/llm_client.py

# src/analyzer/llm_client.py
from __future__ import annotations
import logging
import os
import time
from typing import Any, Dict, List, Optional, Callable, Literal

from .utils.errors import LLMError, ConfigError

ModelType = Literal["router", "translator", "analyzer"]
VerbosityLevel = Literal["low", "medium", "high"]
ReasoningEffort = Literal["minimal", "medium", "high"]

try:
    from openai import OpenAI
    import httpx
except ImportError:
    OpenAI = None
    httpx = None

logger = logging.getLogger(__name__)


class LLMClient:
    """
    Lightweight wrapper around OpenAI Chat Completions with:
      - per-call max_tokens
      - retry with exponential backoff
      - proper error handling and timeouts
    """

    def __init__(self, cfg: Any = None):
        """
        Initialize LLM client with config.

        Args:
            cfg: Configuration object/dict (deprecated - uses get_settings() if None)

        Raises:
            ConfigError: If provider is not OpenAI or client initialization fails
        """
        # Use new settings if no config provided
        if cfg is None:
            from .config import get_settings
            settings = get_settings()
            self.provider = settings.LLM_PROVIDER
            self.model = settings.LLM_MODEL_QA
            self.disable_llm = settings.DISABLE_LLM
            self.model_router = settings.LLM_MODEL_ROUTER
            self.model_translator = settings.LLM_MODEL_TRANSLATE
            self.model_analyzer = settings.LLM_MODEL_QA  # Use QA model for analysis
            api_key = settings.OPENAI_API_KEY
            base_url = os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1")
            timeout_s = settings.TIMEOUT_S
        else:
            # Legacy config object/dict support
            self.provider = self._get_cfg(cfg, "provider", "openai")
            self.model = self._get_cfg(cfg, "model", "gpt-5-mini")
            self.disable_llm = self._get_cfg(cfg, "disable_llm", False)
            self.model_router = self._get_cfg(cfg, "model_router", "gpt-5-nano")
            self.model_translator = self._get_cfg(cfg, "model_translator", "gpt-5-mini")
            self.model_analyzer = self._get_cfg(cfg, "model_analyzer", "gpt-5")
            api_key = self._get_cfg(cfg, "api_key") or os.getenv("OPENAI_API_KEY")
            base_url = self._get_cfg(cfg, "base_url") or os.getenv(
                "OPENAI_API_BASE",
                "https://api.openai.com/v1"
            )
            timeout_s = self._get_cfg(cfg, "timeout_s", 30.0)

        self.max_retries = int(os.getenv("LLM_MAX_RETRIES", "3"))
        self.retry_delay = float(os.getenv("LLM_RETRY_DELAY", "1.0"))
        self.client = None

        # CRITICAL: Only OpenAI is supported - fail fast
        if self.provider != "openai":
            raise ConfigError(
                f"Unsupported LLM_PROVIDER: {self.provider}. "
                f"Only 'openai' is currently supported."
            )

        # Initialize client (if not disabled)
        if self.disable_llm:
            logger.info("LLM disabled via config")
            return

        if OpenAI is None:
            raise ConfigError(
                "openai package not installed. Install with: pip install openai"
            )

        if not api_key:
            raise ConfigError(
                "OPENAI_API_KEY is required when LLM_PROVIDER=openai. "
                "Set the OPENAI_API_KEY environment variable."
            )

        try:
            # Create client with timeout
            timeout = httpx.Timeout(timeout_s, connect=10.0) if httpx else None
            self.client = OpenAI(
                api_key=api_key,
                base_url=base_url,
                timeout=timeout
            )
            logger.info("LLMClient initialized: %s/%s", self.provider, self.model)
        except Exception as e:
            raise ConfigError(f"Failed to initialize OpenAI client: {e}") from e

    def _get_cfg(self, cfg: Any, key: str, default: Any = None) -> Any:
        """Extract config value (works with both dict and object)."""
        if hasattr(cfg, key):
            return getattr(cfg, key) or default
        if isinstance(cfg, dict):
            return cfg.get(key, default)
        return default

    def is_ready(self) -> bool:
        """Check if LLM client is ready to use."""
        return self.client is not None and not self.disable_llm

    def _get_model_for_type(self, model_type: Optional[ModelType] = None) -> str:
        """Get the appropriate model for the given type."""
        if model_type == "router":
            return self.model_router
        elif model_type == "translator":
            return self.model_translator
        elif model_type == "analyzer":
            return self.model_analyzer
        else:
            return self.model

    @staticmethod
    def get_recommended_params(task_type: str) -> Dict[str, Any]:
        """
        Get recommended verbosity and reasoning_effort for common tasks.

        Args:
            task_type: One of "translation", "analysis", "routing", "summary", "comparison"

        Returns:
            Dict with verbosity and reasoning_effort settings
        """
        presets = {
            "translation": {"verbosity": "medium", "reasoning_effort": "minimal"},
            "analysis": {"verbosity": "high", "reasoning_effort": "high"},
            "routing": {"verbosity": "low", "reasoning_effort": "minimal"},
            "summary": {"verbosity": "medium", "reasoning_effort": "medium"},
            "comparison": {"verbosity": "high", "reasoning_effort": "high"}
        }
        return presets.get(task_type, {"verbosity": "medium", "reasoning_effort": "medium"})

    def _retry_with_backoff(
        self,
        fn: Callable[[], Any],
        *,
        max_attempts: int = 3,
        initial_delay: float = 1.0
    ) -> Any:
        """
        Retry function with exponential backoff.

        Retries on transient errors only (timeouts, rate limits).
        Re-raises immediately on permanent errors (auth, invalid model).
        """
        last_error = None
        delay = initial_delay

        for attempt in range(1, max_attempts + 1):
            try:
                return fn()

            except Exception as e:
                last_error = e
                error_str = str(e).lower()

                # Don't retry on permanent errors
                if any(x in error_str for x in [
                    'api key', 'auth', 'forbidden', 'unauthorized',
                    'invalid model', 'model not found'
                ]):
                    raise LLMError(f"Permanent LLM error: {e}") from e

                # Don't retry on programming errors
                if isinstance(e, (ValueError, TypeError, KeyError)):
                    raise LLMError(f"LLM call failed (bad input): {e}") from e

                # Retry on transient errors
                if attempt < max_attempts:
                    logger.warning(
                        "LLM call failed (attempt %d/%d): %s. Retrying in %.1fs...",
                        attempt, max_attempts, e, delay
                    )
                    time.sleep(delay)
                    delay *= 2  # Exponential backoff
                    continue

        # All retries exhausted
        raise LLMError(
            f"LLM call failed after {max_attempts} attempts: {last_error}"
        ) from last_error

    def chat(
        self,
        messages: List[Dict[str, str]],
        *,
        max_tokens: int = 900,
        temperature: float = 0.2,
        top_p: float = 1.0,
        stream: bool = False,
        model_type: Optional[ModelType] = None,
        verbosity: Optional[VerbosityLevel] = None,
        reasoning_effort: Optional[ReasoningEffort] = None,
    ) -> str:
        """
        Single chat completion call with optional streaming and GPT-5 features.

        Args:
            messages: List of message dicts with 'role' and 'content'
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature (0-2)
            top_p: Nucleus sampling parameter
            stream: If True, returns generator yielding tokens (else full response)
            model_type: Type of model to use ('router', 'translator', or 'analyzer')
            verbosity: Response length control (GPT-5 feature)
                - 'low': Brief, concise responses
                - 'medium': Standard length responses
                - 'high': Detailed, comprehensive responses
            reasoning_effort: Thinking time control (GPT-5 feature)
                - 'minimal': Quick, straightforward responses
                - 'medium': Moderate analysis and reasoning
                - 'high': Deep analysis and careful reasoning

        Returns:
            Generated text response (or generator if stream=True)

        Raises:
            LLMError: On API failures or invalid responses
            ConfigError: If client not initialized
        """
        if not self.is_ready():
            if self.disable_llm:
                return self._fallback_response(messages)
            raise LLMError(
                "LLM client not initialized. Check API key and configuration."
            )

        def _make_call():
            # Select model based on model_type
            model = self._get_model_for_type(model_type)

            # Build API parameters
            api_params = {
                "model": model,
                "messages": messages,
                "max_completion_tokens": max_tokens,  # GPT-5 uses max_completion_tokens
                "stream": stream,
            }

            # GPT-5 only supports temperature=1.0 (default), so don't set it
            # For other models, we could add temperature back if needed

            # Add GPT-5 specific parameters if provided
            if verbosity is not None:
                api_params["verbosity"] = verbosity
            if reasoning_effort is not None:
                api_params["reasoning_effort"] = reasoning_effort

            try:
                resp = self.client.chat.completions.create(**api_params)
            except Exception as e:
                # Handle httpx exceptions if available
                if httpx and isinstance(e, httpx.TimeoutException):
                    raise LLMError(f"LLM request timed out") from e
                elif httpx and isinstance(e, httpx.HTTPStatusError):
                    status = e.response.status_code
                    raise LLMError(
                        f"LLM API error (HTTP {status}): {e.response.text[:200]}"
                    ) from e
                else:
                    raise LLMError(f"LLM API call failed: {e}") from e

            if stream:
                # Return generator for streaming
                return self._stream_response(resp)
            else:
                # Validate response
                if not resp.choices:
                    raise LLMError("LLM returned no choices")

                content = resp.choices[0].message.content
                if not content:
                    raise LLMError("LLM returned empty response")

                return content.strip()

        return self._retry_with_backoff(_make_call, max_attempts=self.max_retries)

    def _stream_response(self, stream_resp):
        """
        Handle streaming response from OpenAI API.
        Yields chunks of text as they arrive.
        """
        try:
            full_text = ""
            for chunk in stream_resp:
                if chunk.choices and chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_text += token
                    yield token
        except Exception as e:
            raise LLMError(f"Stream error: {e}") from e

    def summarize(
        self,
        user_text: str,
        *,
        system_text: str = "You are an expert UK grant analyst. Be precise, grounded, concise.",
        max_tokens: int = 900,
        temperature: float = 0.2,
    ) -> str:
        """
        Convenience: single-turn chat.

        Args:
            user_text: User message content
            system_text: System prompt
            max_tokens: Maximum tokens
            temperature: Sampling temperature

        Returns:
            Generated summary

        Raises:
            LLMError: On API failures
        """
        messages = [
            {"role": "system", "content": system_text},
            {"role": "user", "content": user_text},
        ]
        return self.chat(messages, max_tokens=max_tokens, temperature=temperature)

    def summarize_long(
        self,
        user_text: str,
        *,
        system_text: str = "You are an expert UK grant analyst. Be precise, grounded, concise.",
        max_tokens: int = 1300,
        temperature: float = 0.2,
        hard_clip_chars: int = 60_000,
    ) -> str:
        """
        Larger token budget with optional hard clip on input.

        Prevents 400 errors from overly long prompts.
        """
        if len(user_text) > hard_clip_chars:
            logger.warning(
                "Input text truncated from %d to %d chars",
                len(user_text), hard_clip_chars
            )
            user_text = user_text[:hard_clip_chars] + "\n\n[...truncated...]"

        return self.summarize(
            user_text,
            system_text=system_text,
            max_tokens=max_tokens,
            temperature=temperature
        )

    def _fallback_response(self, messages: List[Dict]) -> str:
        """Return safe fallback when LLM is disabled."""
        user_msg = next(
            (m['content'] for m in reversed(messages) if m['role'] == 'user'),
            ""
        )
        return f"[LLM disabled] Your query: {user_msg[:100]}..."