llm_gateway.py

"""Provider-neutral, testable LLM gateway primitives.

The module intentionally has no third-party dependencies and makes no network
calls. Real provider adapters can implement ``LLMProvider.complete``.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
import uuid
from collections import defaultdict, deque
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Mapping, Protocol


logger = logging.getLogger("llm_gateway")


class GatewayError(RuntimeError):
    """Base class for safe, application-facing gateway errors."""


class BudgetExceeded(GatewayError):
    """The request would exceed its conversation token budget."""


class InvalidProviderResponse(GatewayError):
    """The model response did not satisfy the application contract."""


class TransientProviderError(GatewayError):
    """A retryable provider or transport failure."""


class ToolRejected(GatewayError):
    """A tool call was unknown or had invalid arguments."""


@dataclass(frozen=True)
class ProviderResponse:
    content: str
    input_tokens: int
    output_tokens: int
    tool_calls: tuple[Mapping[str, Any], ...] = ()


class LLMProvider(Protocol):
    async def complete(
        self,
        *,
        messages: tuple[Mapping[str, str], ...],
        max_tokens: int,
        request_id: str,
    ) -> ProviderResponse:
        """Return one provider response without leaking credentials to callers."""


@dataclass(frozen=True)
class TriageResult:
    category: str
    priority: str
    summary: str
    needs_human: bool


def validate_triage_json(raw: str) -> TriageResult:
    """Fail closed when a model response differs from the mobile API schema."""

    try:
        value = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise InvalidProviderResponse("provider returned invalid JSON") from exc

    if not isinstance(value, dict):
        raise InvalidProviderResponse("provider response must be a JSON object")

    expected = {"category", "priority", "summary", "needs_human"}
    if set(value) != expected:
        raise InvalidProviderResponse("provider response has unexpected fields")

    if value["category"] not in {"billing", "technical", "account", "other"}:
        raise InvalidProviderResponse("unsupported category")
    if value["priority"] not in {"low", "normal", "high"}:
        raise InvalidProviderResponse("unsupported priority")
    if not isinstance(value["summary"], str) or not value["summary"].strip():
        raise InvalidProviderResponse("summary must be a non-empty string")
    if len(value["summary"]) > 280:
        raise InvalidProviderResponse("summary exceeds 280 characters")
    if not isinstance(value["needs_human"], bool):
        raise InvalidProviderResponse("needs_human must be boolean")

    return TriageResult(
        category=value["category"],
        priority=value["priority"],
        summary=value["summary"].strip(),
        needs_human=value["needs_human"],
    )


class ConversationMemory:
    """Small bounded memory; a Redis adapter can preserve the same interface."""

    def __init__(self, max_messages: int = 8) -> None:
        if max_messages < 1:
            raise ValueError("max_messages must be positive")
        self._messages: dict[str, deque[Mapping[str, str]]] = defaultdict(
            lambda: deque(maxlen=max_messages)
        )

    def read(self, conversation_id: str) -> tuple[Mapping[str, str], ...]:
        return tuple(self._messages[conversation_id])

    def append(self, conversation_id: str, role: str, content: str) -> None:
        if role not in {"user", "assistant"}:
            raise ValueError("unsupported role")
        self._messages[conversation_id].append({"role": role, "content": content})


class UsageLedger:
    """Atomic per-conversation token accounting for concurrent requests."""

    def __init__(self, conversation_limit: int = 10_000) -> None:
        self._limit = conversation_limit
        self._used: dict[str, int] = defaultdict(int)
        self._lock = asyncio.Lock()

    async def ensure_available(self, conversation_id: str, requested: int) -> None:
        if requested < 1:
            raise ValueError("requested tokens must be positive")
        async with self._lock:
            if self._used[conversation_id] + requested > self._limit:
                raise BudgetExceeded("conversation token budget exhausted")

    async def commit(self, conversation_id: str, actual: int) -> None:
        if actual < 0:
            raise ValueError("actual tokens cannot be negative")
        async with self._lock:
            if self._used[conversation_id] + actual > self._limit:
                raise BudgetExceeded("provider usage exceeded conversation budget")
            self._used[conversation_id] += actual

    def used(self, conversation_id: str) -> int:
        return self._used[conversation_id]


ToolHandler = Callable[[Mapping[str, Any]], Awaitable[Mapping[str, Any]]]


@dataclass(frozen=True)
class RegisteredTool:
    required: frozenset[str]
    allowed: frozenset[str]
    handler: ToolHandler


class ToolRegistry:
    """Allow-list model tool calls and reject extra or missing arguments."""

    def __init__(self) -> None:
        self._tools: dict[str, RegisteredTool] = {}

    def register(
        self,
        name: str,
        *,
        required: set[str],
        allowed: set[str],
        handler: ToolHandler,
    ) -> None:
        if not re.fullmatch(r"[a-z][a-z0-9_]{1,63}", name):
            raise ValueError("invalid tool name")
        if not required.issubset(allowed):
            raise ValueError("required arguments must be allowed")
        self._tools[name] = RegisteredTool(
            required=frozenset(required),
            allowed=frozenset(allowed),
            handler=handler,
        )

    async def execute(self, call: Mapping[str, Any]) -> Mapping[str, Any]:
        if set(call) != {"name", "arguments"}:
            raise ToolRejected("malformed tool call")
        name = call["name"]
        arguments = call["arguments"]
        if not isinstance(name, str) or name not in self._tools:
            raise ToolRejected("tool is not allowed")
        if not isinstance(arguments, dict):
            raise ToolRejected("tool arguments must be an object")

        tool = self._tools[name]
        keys = set(arguments)
        if not tool.required.issubset(keys) or not keys.issubset(tool.allowed):
            raise ToolRejected("tool arguments do not match the schema")
        return await tool.handler(arguments)


class LLMGateway:
    def __init__(
        self,
        provider: LLMProvider,
        *,
        memory: ConversationMemory | None = None,
        ledger: UsageLedger | None = None,
        tools: ToolRegistry | None = None,
        timeout_seconds: float = 8.0,
        max_attempts: int = 3,
    ) -> None:
        if max_attempts < 1:
            raise ValueError("max_attempts must be positive")
        self._provider = provider
        self._memory = memory or ConversationMemory()
        self._ledger = ledger or UsageLedger()
        self._tools = tools or ToolRegistry()
        self._timeout = timeout_seconds
        self._max_attempts = max_attempts

    @property
    def ledger(self) -> UsageLedger:
        return self._ledger

    async def triage(
        self,
        *,
        conversation_id: str,
        message: str,
        max_tokens: int = 300,
    ) -> tuple[TriageResult, str]:
        if not re.fullmatch(r"[A-Za-z0-9._:-]{1,128}", conversation_id):
            raise ValueError("invalid conversation_id")
        if not message.strip() or len(message) > 4_000:
            raise ValueError("message must contain 1..4000 characters")
        if not 32 <= max_tokens <= 1_000:
            raise ValueError("max_tokens must be between 32 and 1000")

        await self._ledger.ensure_available(conversation_id, max_tokens)
        request_id = str(uuid.uuid4())
        messages = self._memory.read(conversation_id) + (
            {"role": "user", "content": message.strip()},
        )

        response = await self._complete_with_retry(
            messages=messages,
            max_tokens=max_tokens,
            request_id=request_id,
        )
        await self._ledger.commit(
            conversation_id, response.input_tokens + response.output_tokens
        )

        for call in response.tool_calls:
            await self._tools.execute(call)

        result = validate_triage_json(response.content)
        self._memory.append(conversation_id, "user", message.strip())
        self._memory.append(conversation_id, "assistant", response.content)
        logger.info(
            "llm_request_succeeded request_id=%s tokens=%s",
            request_id,
            response.input_tokens + response.output_tokens,
        )
        return result, request_id

    async def _complete_with_retry(
        self,
        *,
        messages: tuple[Mapping[str, str], ...],
        max_tokens: int,
        request_id: str,
    ) -> ProviderResponse:
        for attempt in range(1, self._max_attempts + 1):
            try:
                return await asyncio.wait_for(
                    self._provider.complete(
                        messages=messages,
                        max_tokens=max_tokens,
                        request_id=request_id,
                    ),
                    timeout=self._timeout,
                )
            except (TransientProviderError, TimeoutError, asyncio.TimeoutError):
                if attempt == self._max_attempts:
                    raise TransientProviderError("provider temporarily unavailable")
                await asyncio.sleep(0.05 * (2 ** (attempt - 1)))
        raise AssertionError("retry loop exhausted")


class DemoProvider:
    """Deterministic provider used only for local proof and tests."""

    async def complete(
        self,
        *,
        messages: tuple[Mapping[str, str], ...],
        max_tokens: int,
        request_id: str,
    ) -> ProviderResponse:
        del max_tokens, request_id
        last = messages[-1]["content"]
        billing = any(word in last.lower() for word in ("payment", "invoice", "оплат"))
        content = json.dumps(
            {
                "category": "billing" if billing else "technical",
                "priority": "high" if "twice" in last.lower() else "normal",
                "summary": last[:120],
                "needs_human": billing,
            },
            ensure_ascii=False,
        )
        return ProviderResponse(
            content=content,
            input_tokens=max(1, len(last) // 4),
            output_tokens=max(1, len(content) // 4),
        )