devplacepy/devplacepy/services/devii/llm.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import logging
import re
from typing import Any, Callable

import httpx

from devplacepy import stealth
from .config import Settings
from .cost import record_cost, record_usage
from .errors import LLMError

logger = logging.getLogger("devii.llm")

# Some upstreams (observed live with openai/gpt-oss-20b via OpenRouter) use
# OpenAI's "Harmony" response format internally and occasionally fail to
# strip its channel-routing special tokens (<|channel|>commentary,
# <|message|>, <|end|>, ...) before filling tool_calls[].function.name -
# the model then tries to call e.g. "plan<|channel|>commentary", gets an
# unrecognized-tool error, and spirals trying to guess a fixed call syntax
# instead of just retrying "plan" cleanly. Sanitized once here, the single
# choke point every completion (main loop, delegate, eval) goes through, so
# neither the dispatcher nor the model's own context ever sees the leak.
_SPECIAL_TOKEN_RE = re.compile(r"<\|[^|>]*\|>")


def _sanitize_tool_name(name: str) -> str:
    match = _SPECIAL_TOKEN_RE.search(name)
    if match is None:
        return name
    cleaned = name[: match.start()].strip()
    return cleaned or name


def _sanitize_tool_calls(message: dict[str, Any]) -> None:
    for call in message.get("tool_calls") or []:
        function = call.get("function")
        if not isinstance(function, dict):
            continue
        raw_name = function.get("name", "")
        if not isinstance(raw_name, str):
            continue
        clean_name = _sanitize_tool_name(raw_name)
        if clean_name != raw_name:
            logger.warning(
                "Upstream tool call name contained a leaked special token, "
                "sanitized %r -> %r",
                raw_name,
                clean_name,
            )
            function["name"] = clean_name


class LLMClient:
    def __init__(
        self,
        settings: Settings,
        key_resolver: Callable[[], str] | None = None,
    ) -> None:
        self._settings = settings
        self._key_resolver = key_resolver
        self._client = stealth.stealth_async_client(
            timeout=settings.timeout_seconds,
            headers={
                "Authorization": f"Bearer {settings.ai_key}",
                "Content-Type": "application/json",
            },
        )

    async def aclose(self) -> None:
        await self._client.aclose()

    async def _post(self, payload: dict[str, Any]) -> httpx.Response:
        response = await self._client.post(self._settings.ai_url, json=payload)
        if response.status_code == 401 and self._refresh_key():
            response = await self._client.post(self._settings.ai_url, json=payload)
        return response

    def _refresh_key(self) -> bool:
        if self._key_resolver is None:
            return False
        try:
            key = self._key_resolver()
        except Exception:  # noqa: BLE001 - refresh is best-effort; the original 401 still surfaces
            logger.exception("Credential refresh failed after 401")
            return False
        if not key or self._client.headers.get("Authorization") == f"Bearer {key}":
            return False
        self._client.headers["Authorization"] = f"Bearer {key}"
        logger.info("Refreshed model endpoint credentials after 401")
        return True

    async def complete(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
    ) -> dict[str, Any]:
        payload = {
            "model": self._settings.ai_model,
            "messages": messages,
            "tools": tools,
            "tool_choice": "auto",
        }
        try:
            logger.debug("LLM request with %d messages", len(messages))
            response = await self._post(payload)
        except httpx.HTTPError as exc:
            raise LLMError(f"Could not reach the model endpoint: {exc}") from exc

        if response.status_code >= 400:
            raise LLMError(
                f"Model endpoint returned {response.status_code}: {self._reason(response)}",
                status=response.status_code,
                body=response.text[:500],
            )

        try:
            data = response.json()
        except ValueError as exc:
            raise LLMError("Model endpoint returned invalid JSON.") from exc

        choices = data.get("choices")
        if not choices:
            raise LLMError("Model response contained no choices.", body=str(data)[:500])

        message = choices[0].get("message")
        if message is None:
            raise LLMError("Model response contained no message.", body=str(data)[:500])

        _sanitize_tool_calls(message)
        record_usage(data.get("usage"))
        self._record_native_cost(response)
        logger.debug(
            "LLM response received (tool_calls=%s)", bool(message.get("tool_calls"))
        )
        return message

    async def complete_text(
        self, messages: list[dict[str, Any]], temperature: float = 0.0
    ) -> str:
        payload = {
            "model": self._settings.ai_model,
            "messages": messages,
            "temperature": temperature,
        }
        try:
            response = await self._post(payload)
        except httpx.HTTPError as exc:
            raise LLMError(f"Could not reach the model endpoint: {exc}") from exc
        if response.status_code >= 400:
            raise LLMError(
                f"Model endpoint returned {response.status_code}: {self._reason(response)}"
            )
        try:
            data = response.json()
            content = data["choices"][0]["message"]["content"] or ""
        except (ValueError, KeyError, IndexError) as exc:
            raise LLMError("Model endpoint returned an unexpected response.") from exc
        record_usage(data.get("usage"))
        self._record_native_cost(response)
        return content

    async def summarize(self, text: str) -> str:
        payload = {
            "model": self._settings.ai_model,
            "messages": [
                {
                    "role": "system",
                    "content": "You are a precise technical summarizer.",
                },
                {"role": "user", "content": text},
            ],
            "temperature": 0.0,
        }
        try:
            response = await self._post(payload)
        except httpx.HTTPError as exc:
            raise LLMError(f"Could not reach the model endpoint: {exc}") from exc
        if response.status_code >= 400:
            raise LLMError(
                f"Model endpoint returned {response.status_code}: {self._reason(response)}"
            )
        try:
            data = response.json()
            content = data["choices"][0]["message"]["content"] or ""
        except (ValueError, KeyError, IndexError) as exc:
            raise LLMError(
                "Model summarization returned an unexpected response."
            ) from exc
        record_usage(data.get("usage"))
        self._record_native_cost(response)
        return content

    @staticmethod
    def _record_native_cost(response: httpx.Response) -> None:
        raw = response.headers.get("X-Gateway-Cost-USD")
        if raw is None:
            return
        try:
            record_cost(float(raw))
        except (TypeError, ValueError):
            return

    @staticmethod
    def _reason(response: httpx.Response) -> str:
        try:
            error = response.json().get("error")
        except ValueError:
            return response.text[:200] or response.reason_phrase
        if isinstance(error, dict):
            return str(error.get("message", error))
        if error:
            return str(error)
        return response.reason_phrase