src/gcmd_classifier/llm/base.py

"""Provider-neutral model interface and retry helpers."""

from __future__ import annotations

import time
from collections.abc import Mapping
from dataclasses import dataclass, field, replace
from enum import Enum
from typing import Any, Generic, Protocol, TypeVar

from pydantic import BaseModel

from gcmd_classifier.config import ModelSettings
from gcmd_classifier.errors import ModelRetriesExhaustedError, RetryableModelError

StructuredResponseT = TypeVar("StructuredResponseT", bound=BaseModel)


class ModelStage(str, Enum):  # noqa: UP042
    """Supported model-call stages for the MVP classifier."""

    TOPIC = "topic"
    TERM = "term"
    VARIABLE = "variable"


@dataclass(frozen=True)
class TokenUsage:
    """Token usage metadata returned by a provider when available."""

    input_tokens: int | None = None
    output_tokens: int | None = None
    total_tokens: int | None = None


@dataclass(frozen=True)
class ModelRequest(Generic[StructuredResponseT]):
    """Provider-neutral structured model request."""

    stage: ModelStage
    prompt: str
    response_schema: type[StructuredResponseT]
    provider: str
    model_name: str
    prompt_version: str
    temperature: float
    timeout_seconds: float
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @classmethod
    def from_settings(
        cls,
        *,
        stage: ModelStage,
        prompt: str,
        response_schema: type[StructuredResponseT],
        settings: ModelSettings,
        metadata: Mapping[str, Any] | None = None,
    ) -> ModelRequest[StructuredResponseT]:
        """Create a request from model settings for a specific stage."""
        return cls(
            stage=stage,
            prompt=prompt,
            response_schema=response_schema,
            provider=settings.provider,
            model_name=settings.model_name,
            prompt_version=settings.prompt_version_for_stage(stage.value),
            temperature=settings.temperature,
            timeout_seconds=settings.timeout_seconds,
            metadata={} if metadata is None else dict(metadata),
        )


@dataclass(frozen=True)
class ModelResponse(Generic[StructuredResponseT]):
    """Provider-neutral structured model response."""

    parsed: StructuredResponseT
    provider: str
    model_name: str
    prompt_version: str
    retry_count: int = 0
    duration_seconds: float | None = None
    token_usage: TokenUsage | None = None
    estimated_cost: float | None = None
    raw_response: str | None = None
    warnings: tuple[str, ...] = ()
    diagnostics: Mapping[str, Any] = field(default_factory=dict)


class ModelClient(Protocol):
    """Provider-neutral interface used by future classification components."""

    def generate_structured(
        self,
        request: ModelRequest[StructuredResponseT],
    ) -> ModelResponse[StructuredResponseT]:
        """Return a typed structured response for a model request."""


@dataclass(frozen=True)
class RetryPolicy:
    """Retry settings for temporary model failures."""

    max_retries: int = 2

    @classmethod
    def from_settings(cls, settings: ModelSettings) -> RetryPolicy:
        """Create a retry policy from model settings."""
        return cls(max_retries=settings.max_retries)


def generate_with_retries(
    client: ModelClient,
    request: ModelRequest[StructuredResponseT],
    policy: RetryPolicy,
) -> ModelResponse[StructuredResponseT]:
    """Call a model client, retrying only explicitly retryable failures."""
    retry_count = 0
    started = time.perf_counter()
    while True:
        try:
            response = client.generate_structured(request)
            duration = response.duration_seconds
            if duration is None:
                duration = time.perf_counter() - started
            return replace(response, retry_count=retry_count, duration_seconds=duration)
        except RetryableModelError as exc:
            if retry_count >= policy.max_retries:
                raise ModelRetriesExhaustedError(
                    "Retryable model failure exceeded configured max retries.",
                    retry_count=retry_count,
                ) from exc
            retry_count += 1