backend/app/providers/base.py

"""Provider abstraction: one interface, many model backends.

The whole point: the pipeline depends only on `LLMProvider.complete()`. Swapping
mock ↔ Gemma ↔ Claude ↔ Gemini changes nothing upstream. Each response carries a
full token/cost/cache usage record so every call is measurable.
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Optional

from .pricing import compute_cost


@dataclass
class CacheBlock:
    """A prompt segment marked cacheable (placed at the head of the payload)."""

    text: str
    cacheable: bool = True


@dataclass
class LLMRequest:
    """A model-agnostic completion request.

    `system_blocks` are ordered; cacheable blocks form the stable, cached prefix
    (system prompt + tool defs + schema). `user_content` is the dynamic suffix
    (the document) and is never cached. Keeping these separate is what makes
    prompt caching actually work.
    """

    system_blocks: list[CacheBlock] = field(default_factory=list)
    user_content: str = ""
    max_tokens: int = 1024
    temperature: float = 0.0
    task: str = "generic"          # classify|extract|normalize|validate|agent|...
    json_schema: Optional[dict] = None  # if set, request structured output
    context: dict = field(default_factory=dict)  # provider hints (doc_type, etc.)

    def full_prompt(self) -> str:
        """Concatenation of all blocks + user content (for token estimation/hash)."""
        return "\n".join(b.text for b in self.system_blocks) + "\n" + self.user_content

    def cacheable_prefix(self) -> str:
        return "\n".join(b.text for b in self.system_blocks if b.cacheable)


@dataclass
class Usage:
    input_tokens: int = 0
    output_tokens: int = 0
    cache_read_tokens: int = 0
    cache_write_tokens: int = 0

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


@dataclass
class LLMResponse:
    text: str
    model: str
    provider: str
    usage: Usage
    latency_ms: float
    cache_hit: bool = False
    cost_usd: float = 0.0
    routing_reason: str = ""
    error: Optional[str] = None

    @classmethod
    def build(
        cls,
        text: str,
        model: str,
        provider: str,
        usage: Usage,
        latency_ms: float,
        cache_hit: bool = False,
        routing_reason: str = "",
        error: Optional[str] = None,
    ) -> "LLMResponse":
        cost = compute_cost(
            model,
            usage.input_tokens,
            usage.output_tokens,
            usage.cache_read_tokens,
            usage.cache_write_tokens,
        )
        return cls(
            text=text,
            model=model,
            provider=provider,
            usage=usage,
            latency_ms=round(latency_ms, 1),
            cache_hit=cache_hit,
            cost_usd=cost,
            routing_reason=routing_reason,
            error=error,
        )


def estimate_tokens(text: str) -> int:
    """Cheap heuristic: ~4 chars/token. Used by mock/local providers and as a
    fallback when a real API doesn't return usage. Good enough for cost demos."""
    if not text:
        return 0
    return max(1, len(text) // 4)


class LLMProvider:
    """Base class. Subclasses implement `complete`."""

    name: str = "base"
    tier: str = "base"  # offline|local|hosted

    def available(self) -> bool:  # pragma: no cover - trivial
        return True

    def complete(self, req: LLMRequest, model: str) -> LLMResponse:  # pragma: no cover
        raise NotImplementedError

    # Helper so subclasses can time themselves consistently.
    @staticmethod
    def _now() -> float:
        return time.perf_counter()