PTSRaiAgent/app/llm/giga.py

"""GigaChat provider (152-FZ: the only allowed LLM). Transport only — prompts live in `prompts.py`.

Logging policy: only attempt number, latency, token counts and error class. Never prompt or answer text.
"""

import asyncio
import json
import logging
import re
import time
from typing import Any, Protocol, TypeVar

from pydantic import BaseModel, ValidationError

from app.config import Settings
from app.core.types import ScrubbedText

log = logging.getLogger(__name__)

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

TEMPERATURE = 0.1
TIMEOUT_S = 8.0
MAX_RETRIES = 2  # SPEC §7: up to 2 retries → at most 3 attempts.
MAX_TOKENS = 512

_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*|\s*```\s*$", re.IGNORECASE)


class ChatClient(Protocol):
    """The subset of `gigachat.GigaChat` we use; lets tests inject a fake without network."""

    async def achat(self, payload: Any) -> Any: ...


def parse_json_object(raw: str) -> dict[str, Any] | None:
    """Parse a JSON object from model output, tolerating ```json fences and surrounding chatter."""
    text = _FENCE_RE.sub("", raw.strip())
    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        start, end = text.find("{"), text.rfind("}")
        if start == -1 or end <= start:
            return None
        try:
            data = json.loads(text[start : end + 1])
        except json.JSONDecodeError:
            return None
    return data if isinstance(data, dict) else None


class GigaChatProvider:
    def __init__(self, settings: Settings, client: ChatClient | None = None) -> None:
        self._settings = settings
        self._client = client

    def _get_client(self) -> ChatClient:
        if self._client is None:
            from gigachat import GigaChat

            self._client = GigaChat(
                credentials=self._settings.gigachat_credentials,
                scope=self._settings.gigachat_scope,
                model=self._settings.gigachat_model,
                verify_ssl_certs=self._settings.gigachat_verify_ssl,
                timeout=TIMEOUT_S,
            )
        return self._client

    async def _call(self, system: str, user: str, attempt: int) -> str | None:
        from gigachat.models import Chat, Messages, MessagesRole

        payload = Chat(
            messages=[
                Messages(role=MessagesRole.SYSTEM, content=system),
                Messages(role=MessagesRole.USER, content=user),
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
        )
        started = time.monotonic()
        try:
            resp = await asyncio.wait_for(self._get_client().achat(payload), timeout=TIMEOUT_S)
            content = resp.choices[0].message.content
        except Exception as exc:  # network, auth, timeout, malformed response — all mean "no answer"
            log.warning(
                "gigachat call failed attempt=%d latency_ms=%d error=%s",
                attempt,
                (time.monotonic() - started) * 1000,
                type(exc).__name__,
            )
            return None
        usage = getattr(resp, "usage", None)
        log.info(
            "gigachat call ok attempt=%d latency_ms=%d prompt_tokens=%s completion_tokens=%s",
            attempt,
            (time.monotonic() - started) * 1000,
            getattr(usage, "prompt_tokens", "?"),
            getattr(usage, "completion_tokens", "?"),
        )
        return content if isinstance(content, str) else None

    async def complete_json(
        self,
        system: str,
        user: ScrubbedText,
        schema: type[M],
        validation_context: dict[str, Any] | None = None,
    ) -> M | None:
        for attempt in range(1, MAX_RETRIES + 2):
            raw = await self._call(system, user, attempt)
            if raw is None:
                continue
            data = parse_json_object(raw)
            if data is None:
                log.warning("gigachat json parse failed attempt=%d schema=%s", attempt, schema.__name__)
                continue
            try:
                return schema.model_validate(data, context=validation_context)
            except ValidationError as exc:
                log.warning(
                    "gigachat json invalid attempt=%d schema=%s errors=%d",
                    attempt,
                    schema.__name__,
                    exc.error_count(),
                )
        return None

    async def complete_text(self, system: str, user: ScrubbedText) -> str | None:
        for attempt in range(1, MAX_RETRIES + 2):
            raw = await self._call(system, user, attempt)
            if raw and raw.strip():
                return raw.strip()
        return None