llm.py

from __future__ import annotations

import json
import re
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Any, Optional

from loguru import logger
from openai import BadRequestError, OpenAI


@dataclass(frozen=True)
class AgentLLM:
    """Thin OpenRouter chat wrapper that returns content and usage for token accounting."""

    model_id: str
    base_url: str
    api_key_path: str
    temperature: float
    max_tokens_per_call: int
    max_retries: int = 3
    extra_body: Optional[Dict[str, Any]] = None

    def __post_init__(self) -> None:  # type: ignore[override]
        assert isinstance(self.model_id, str) and len(self.model_id.strip()) > 0
        assert isinstance(self.base_url, str) and self.base_url.startswith("http")
        assert isinstance(self.api_key_path, str) and len(self.api_key_path.strip()) > 0
        assert isinstance(self.temperature, float)
        assert (
            isinstance(self.max_tokens_per_call, int) and self.max_tokens_per_call > 0
        )
        key_path = Path(self.api_key_path)
        assert key_path.exists() and key_path.is_file()
        api_key = key_path.read_text(encoding="utf-8").strip()
        assert len(api_key) > 0
        object.__setattr__(
            self, "_client", OpenAI(base_url=self.base_url, api_key=api_key)
        )

    def _call_api(self, call_params: dict) -> Any:
        prev_input_tokens: Optional[int] = None
        while True:
            try:
                return self._client.chat.completions.create(**call_params)
            except BadRequestError as e:
                if e.status_code == 400:
                    match = re.search(
                        r"passed (\d+) input tokens.*context length is only (\d+)",
                        str(e),
                        re.DOTALL,
                    )
                    if match:
                        input_tokens = int(match.group(1))
                        context_length = int(match.group(2))
                        new_max = context_length - input_tokens - 1
                        if (
                            new_max > 0
                            and new_max < call_params["max_tokens"]
                            and (prev_input_tokens is None or input_tokens <= prev_input_tokens)
                        ):
                            logger.warning(
                                f"Input too long ({input_tokens} tokens), reducing max_tokens "
                                f"from {call_params['max_tokens']} to {new_max}"
                            )
                            call_params["max_tokens"] = new_max
                            prev_input_tokens = input_tokens
                            continue
                        if prev_input_tokens is not None and input_tokens > prev_input_tokens:
                            raise RuntimeError(
                                f"Prompt too long and growing between retries "
                                f"({prev_input_tokens}→{input_tokens} tokens)."
                            ) from e
                raise

    def chat(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
        assert isinstance(messages, list) and len(messages) >= 1

        for attempt in range(self.max_retries):
            try:
                call_params = dict(
                    model=self.model_id,
                    messages=messages,
                    temperature=self.temperature,
                    max_tokens=self.max_tokens_per_call,
                )
                if self.extra_body:
                    call_params["extra_body"] = self.extra_body
                completion = self._call_api(call_params)
                content = completion.choices[0].message.content or ""

                usage = getattr(completion, "usage", None)
                usage_dict: Dict[str, int] = {
                    "prompt_tokens": (
                        int(getattr(usage, "prompt_tokens", 0))
                        if usage is not None
                        else 0
                    ),
                    "completion_tokens": (
                        int(getattr(usage, "completion_tokens", 0))
                        if usage is not None
                        else 0
                    ),
                    "total_tokens": (
                        int(getattr(usage, "total_tokens", 0))
                        if usage is not None
                        else 0
                    ),
                }
                return {"content": content, "usage": usage_dict}

            except json.JSONDecodeError as e:
                logger.warning(
                    f"JSONDecodeError on attempt {attempt + 1}/{self.max_retries}: {e}"
                )
                if attempt == self.max_retries - 1:
                    raise e
                time.sleep(2**attempt)  # Exponential backoff


__all__ = ["AgentLLM"]