pryscraper/llm_providers/registry.py

# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Rug Munch Media LLC
#
# Part of Pry — https://git.rugmunch.io/RugMunchMedia/pryscraper
# Licensed under MIT. See LICENSE.
"""LLM provider registry with fallback chain and cost tracking."""

import json
import logging
import time
from collections import defaultdict
from datetime import UTC, datetime
from typing import Any

from llm_providers.base import LLMProvider, LLMResponse, ReferralConfig
from paths import PRY_DATA_DIR

logger = logging.getLogger(__name__)

USAGE_DIR = PRY_DATA_DIR / "llm_usage"
USAGE_DIR.mkdir(parents=True, exist_ok=True)


class LLMRegistry:
    """Registry of LLM providers with fallback chain, cost tracking, and referral config."""

    def __init__(self, referral: ReferralConfig | None = None):
        self.providers: dict[str, LLMProvider] = {}
        self.referral = referral or ReferralConfig()
        self.fallback_chain: list[str] = []
        # Usage tracking
        self.usage: dict[str, dict[str, Any]] = defaultdict(
            lambda: {
                "calls": 0,
                "input_tokens": 0,
                "output_tokens": 0,
                "cost_usd": 0.0,
                "last_used": None,
            }
        )
        self._load_usage()

    def register(self, provider: LLMProvider) -> None:
        self.providers[provider.name] = provider
        if provider.name not in self.fallback_chain:
            self.fallback_chain.append(provider.name)
        logger.info("provider_registered", extra={"provider_name": provider.name})

    def set_fallback_chain(self, chain: list[str]) -> None:
        self.fallback_chain = chain

    async def complete(
        self,
        prompt: str,
        system: str = "",
        provider_name: str = "",
        max_tokens: int = 1024,
        temperature: float = 0.7,
        model: str = "",
        fallback: bool = True,
    ) -> LLMResponse:
        """Complete via specified provider or fallback chain."""
        names = [provider_name] if provider_name else list(self.fallback_chain)
        if not fallback:
            names = (
                [provider_name]
                if provider_name
                else [self.fallback_chain[0]]
                if self.fallback_chain
                else []
            )

        last_error = ""
        for name in names:
            provider = self.providers.get(name)
            if not provider:
                continue
            try:
                start = time.time()
                response = await provider.complete(prompt, system, max_tokens, temperature, model)
                response.latency_ms = int((time.time() - start) * 1000)
                response.cost_usd = provider.estimate_cost(
                    response.input_tokens, response.output_tokens
                )
                response.referral_id = self.referral.program_id
                self._track(response)
                return response
            except Exception as e:  # noqa: BLE001
                last_error = str(e)
                logger.warning(
                    "llm_provider_failed", extra={"provider": name, "error": str(e)[:100]}
                )
                if not fallback:
                    break
        raise Exception(f"All LLM providers failed. Last: {last_error}")

    async def complete_or_empty(
        self,
        prompt: str,
        system: str = "",
        provider_name: str = "",
        max_tokens: int = 1024,
        temperature: float = 0.7,
        model: str = "",
        fallback: bool = True,
    ) -> LLMResponse:
        """Complete via LLM, returning an empty response on total failure.

        Use this at API boundaries where a 500 is worse than degraded output.
        """
        try:
            return await self.complete(
                prompt, system, provider_name, max_tokens, temperature, model, fallback
            )
        except Exception as e:  # noqa: BLE001
            logger.warning("llm_complete_failed_all_providers", extra={"error": str(e)[:120]})
            return LLMResponse(
                text="",
                provider="none",
                model="",
                input_tokens=0,
                output_tokens=0,
                cost_usd=0.0,
            )

    async def embed(self, text: str, provider_name: str = "", model: str = "") -> list[float]:
        names = [provider_name] if provider_name else list(self.fallback_chain)
        for name in names:
            provider = self.providers.get(name)
            if not provider:
                continue
            try:
                return await provider.embed(text, model)
            except Exception as e:  # noqa: BLE001
                logger.warning(
                    "embed_provider_failed", extra={"provider": name, "error": str(e)[:80]}
                )
        raise Exception("All embedding providers failed")

    def _track(self, response: LLMResponse) -> None:
        u = self.usage[response.provider]
        u["calls"] += 1
        u["input_tokens"] += response.input_tokens
        u["output_tokens"] += response.output_tokens
        u["cost_usd"] += response.cost_usd
        u["last_used"] = datetime.now(UTC).isoformat()
        # NEW: track in-app LLM usage as referral revenue
        try:
            from pry_referrals import ReferralTracker

            if not hasattr(self, "_referral_tracker"):
                self._referral_tracker = ReferralTracker()
            # Estimate revenue at 5% of user cost (typical affiliate share)
            self._referral_tracker.track_in_app_usage(response.provider, response.cost_usd * 0.05)
        except ImportError:
            pass
        # Persist daily
        today = datetime.now(UTC).strftime("%Y-%m-%d")
        path = USAGE_DIR / f"usage_{today}.jsonl"
        try:
            with open(path, "a") as f:
                f.write(
                    json.dumps(
                        {
                            "ts": datetime.now(UTC).isoformat(),
                            "provider": response.provider,
                            "model": response.model,
                            "input_tokens": response.input_tokens,
                            "output_tokens": response.output_tokens,
                            "cost_usd": response.cost_usd,
                            "referral_id": response.referral_id,
                        }
                    )
                    + "\n"
                )
        except OSError:
            pass

    def _load_usage(self) -> None:
        for path in USAGE_DIR.glob("usage_*.jsonl"):
            try:
                for line in path.read_text().splitlines():
                    if line.strip():
                        r = json.loads(line)
                        prov = r.get("provider", "unknown")
                        u = self.usage[prov]
                        u["calls"] += 1
                        u["input_tokens"] += r.get("input_tokens", 0)
                        u["output_tokens"] += r.get("output_tokens", 0)
                        u["cost_usd"] += r.get("cost_usd", 0.0)
            except (json.JSONDecodeError, OSError):
                pass

    def get_stats(self) -> dict[str, Any]:
        return {
            "providers": {name: self.usage[p.name] for name, p in self.providers.items()},
            "referral": {
                "enabled": self.referral.enabled,
                "program_id": self.referral.program_id,
                "links": self.referral.referral_links,
            },
            "fallback_chain": self.fallback_chain,
        }


# Global registry
_registry: LLMRegistry | None = None


def get_registry() -> LLMRegistry:
    """Get or create the global LLM registry with default providers."""
    global _registry
    if _registry is None:
        _registry = LLMRegistry()
        # Register providers lazily (only if API keys are set)
        from llm_providers.providers import register_default_providers

        register_default_providers(_registry)
    return _registry