ai-rpg/app/core/llm.py

"""LLM client — OpenAI-compatible API wrapper with retry, logging, and streaming.

Usage:
    client = LlmClient.from_settings(settings_dict)
    resp = await client.complete(
        stage="orchestrator_phase1",
        messages=[{"role": "system", "content": "..."}, ...],
        tools=[...],            # optional
        temperature=0.7,
        max_tokens=2048,
        stream=False,           # if True, returns an async iterator of deltas
        user_id=..., world_id=..., step_id=...,
    )
"""

from __future__ import annotations

import asyncio
import json
import time
import uuid
from collections.abc import AsyncIterator
from typing import Any

import httpx
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.logging import get_logger
from app.models import LlmCallLog

_logger = get_logger(__name__)


# Patterns for parsing tool calls emitted as text (some local models don't
# use the OpenAI function-calling format and instead emit calls as text).
import re as _re

_TOOL_CALL_PATTERNS = [
    # call:name{json_args}  or  call:name(json_args)
    _re.compile(r"call:(\w+)\s*[\{\(]([^}\)]*)[\}\)]"),
    # <tool_call>name{args}</tool_call>  or  <tool_call>\n{"name": ..., "arguments": ...}\n</tool_call>
    _re.compile(r"<tool_call>\s*(\w+)\s*\{([^}]*)\}\s*</tool_call>"),
    # name{"key": "value", ...}  (function call style)
    _re.compile(r"\b(\w+)\s*\(\s*(\{[^}]*\})\s*\)"),
]


def _parse_text_tool_calls(content: str) -> list[dict]:
    """Parse tool calls emitted as text by some local models.

    Handles patterns like:
    - call:calc{"expression": "2+2"}
    - <tool_call>calc{"expression": "2+2"}</tool_call>
    - calc({"expression": "2+2"})

    Returns a list of OpenAI-format tool_call dicts.
    """
    if not content:
        return []
    calls: list[dict] = []
    for pattern in _TOOL_CALL_PATTERNS:
        for match in pattern.finditer(content):
            name = match.group(1)
            args_str = match.group(2).strip()
            if not args_str:
                args = {}
            else:
                try:
                    args = json.loads(args_str)
                except json.JSONDecodeError:
                    # Try fixing common issues: single quotes, unquoted keys
                    try:
                        fixed = args_str.replace("'", '"')
                        # Add quotes around bare keys
                        fixed = _re.sub(r"(\w+)\s*:", r'"\1":', fixed)
                        args = json.loads(fixed)
                    except json.JSONDecodeError:
                        args = {"_raw": args_str}
            calls.append({
                "id": f"parsed_{len(calls)}",
                "type": "function",
                "function": {"name": name, "arguments": json.dumps(args)},
            })
    return calls


class LLMError(Exception):
    """Base LLM error."""

    def __init__(self, code: str, message: str, status: str = "api_error"):
        super().__init__(message)
        self.code = code
        self.status = status


class LLMTimeoutError(LLMError):
    def __init__(self, message: str = "LLM call timed out"):
        super().__init__("llm_timeout", message, status="timeout")


class LLMUnavailableError(LLMError):
    def __init__(self, message: str = "LLM provider unavailable"):
        super().__init__("llm_unavailable", message, status="api_error")


class LLMResponseError(LLMError):
    def __init__(self, message: str, code: str = "parse_error"):
        super().__init__(code, message, status="parse_error")


class LlmClient:
    """OpenAI-compatible LLM client with retry, logging, and streaming."""

    def __init__(
        self,
        api_url: str,
        api_key: str,
        model: str,
        timeout: float = 60.0,
        max_retries: int = 3,
    ):
        self._api_url = api_url.rstrip("/")
        self._api_key = api_key
        self._model = model
        self._timeout = timeout
        self._max_retries = max_retries

    # ------------------------------------------------------------------ #
    # Construction
    # ------------------------------------------------------------------ #
    @classmethod
    def from_settings(cls, settings: dict[str, Any]) -> "LlmClient":
        return cls(
            api_url=settings.get("llm.api_url", "http://localhost:11434/v1"),
            api_key=settings.get("llm.api_key", ""),
            model=settings.get("llm.model", "qwen2.5-7b-instruct"),
            timeout=float(settings.get("llm.timeout_seconds", 60)),
        )

    # ------------------------------------------------------------------ #
    # Non-streaming call
    # ------------------------------------------------------------------ #
    async def complete(
        self,
        *,
        stage: str,
        messages: list[dict[str, Any]],
        tools: list[dict] | None = None,
        tool_choice: Any = None,
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_tokens: int = 2048,
        user_id: uuid.UUID | None = None,
        world_id: uuid.UUID | None = None,
        step_id: uuid.UUID | None = None,
        session: AsyncSession | None = None,
        stream: bool = False,
    ) -> dict[str, Any]:
        """Make a non-streaming chat completion call.

        Returns a dict with keys:
        - `message`: assistant message (with `content` and optional `tool_calls`)
        - `finish_reason`: stop | length | tool_calls
        - `prompt_tokens`, `completion_tokens`, `latency_ms`
        - `log_id`: id of the LlmCallLog row if `session` provided
        """
        if stream:
            raise ValueError("Use stream_complete() for streaming calls")

        payload: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens,
            "stream": False,
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = tool_choice or "auto"

        start = time.monotonic()
        last_exc: Exception | None = None
        for attempt in range(self._max_retries):
            try:
                async with httpx.AsyncClient(timeout=self._timeout) as client:
                    resp = await client.post(
                        f"{self._api_url}/chat/completions",
                        headers={
                            "Authorization": f"Bearer {self._api_key}",
                            "Content-Type": "application/json",
                        },
                        json=payload,
                    )
                if resp.status_code >= 500:
                    raise LLMUnavailableError(
                        f"LLM provider returned HTTP {resp.status_code}: {resp.text[:300]}"
                    )
                if resp.status_code == 429:
                    raise LLMUnavailableError("LLM provider rate-limited (HTTP 429)")
                if resp.status_code >= 400:
                    # Try to extract error message from JSON body
                    err_body = resp.text[:500]
                    try:
                        err_json = resp.json()
                        if "error" in err_json:
                            err_msg = err_json["error"].get("message", err_body)
                        else:
                            err_msg = err_body
                    except Exception:  # noqa: BLE001
                        err_msg = err_body
                    raise LLMResponseError(
                        f"LLM provider returned HTTP {resp.status_code}: {err_msg}",
                        code="api_error",
                    )
                # Parse JSON response — if this fails, the URL is likely wrong
                # (pointing at an HTML page instead of an OpenAI-compatible API)
                try:
                    data = resp.json()
                except Exception as e:  # noqa: BLE001
                    raise LLMResponseError(
                        f"LLM provider returned non-JSON response (check that "
                        f"api_url points to an OpenAI-compatible endpoint). "
                        f"First 200 chars: {resp.text[:200]}",
                        code="parse_error",
                    ) from e
                break
            except (httpx.TimeoutException, asyncio.TimeoutError) as e:
                last_exc = LLMTimeoutError(str(e))
                _logger.warning(
                    "llm_timeout", stage=stage, attempt=attempt + 1, error=str(e)
                )
            except (httpx.ConnectError, httpx.NetworkError) as e:
                last_exc = LLMUnavailableError(str(e))
                _logger.warning(
                    "llm_connection_error", stage=stage, attempt=attempt + 1, error=str(e)
                )
            except LLMError as e:
                last_exc = e
                _logger.warning(
                    "llm_error", stage=stage, attempt=attempt + 1, error=str(e)
                )
            # exponential backoff
            await asyncio.sleep(min(2**attempt, 4))
        else:
            # All retries exhausted
            if session is not None:
                await self._write_log_safely(
                    session=session,
                    stage=stage,
                    messages=messages,
                    tools=tools,
                    response_message={},
                    tool_calls=None,
                    prompt_tokens=None,
                    completion_tokens=None,
                    latency_ms=int((time.monotonic() - start) * 1000),
                    temperature=temperature,
                    status=last_exc.status if isinstance(last_exc, LLMError) else "api_error",
                    error_message=str(last_exc) if last_exc else "unknown",
                    user_id=user_id,
                    world_id=world_id,
                    step_id=step_id,
                )
            assert last_exc is not None
            raise last_exc

        latency_ms = int((time.monotonic() - start) * 1000)
        choice = data["choices"][0]
        msg = choice.get("message", {})
        finish_reason = choice.get("finish_reason", "stop")
        usage = data.get("usage", {})

        # If the model didn't return tool_calls in the OpenAI format but DID
        # emit them as text (some local models use "call:name{args}" or
        # "<tool_call>name{args}</tool_call>"), try to parse them out.
        if tools and not msg.get("tool_calls"):
            content = msg.get("content", "") or ""
            parsed = _parse_text_tool_calls(content)
            if parsed:
                msg = dict(msg)  # don't mutate the original
                msg["tool_calls"] = parsed
                if finish_reason == "stop":
                    finish_reason = "tool_calls"

        log_id: uuid.UUID | None = None
        if session is not None:
            log_id = await self._write_log_safely(
                session=session,
                stage=stage,
                messages=messages,
                tools=tools,
                response_message=msg,
                tool_calls=msg.get("tool_calls"),
                prompt_tokens=usage.get("prompt_tokens"),
                completion_tokens=usage.get("completion_tokens"),
                latency_ms=latency_ms,
                temperature=temperature,
                status="ok",
                error_message=None,
                user_id=user_id,
                world_id=world_id,
                step_id=step_id,
            )

        return {
            "message": msg,
            "finish_reason": finish_reason,
            "prompt_tokens": usage.get("prompt_tokens"),
            "completion_tokens": usage.get("completion_tokens"),
            "latency_ms": latency_ms,
            "log_id": log_id,
        }

    # ------------------------------------------------------------------ #
    # Streaming call
    # ------------------------------------------------------------------ #
    async def stream_complete(
        self,
        *,
        stage: str,
        messages: list[dict[str, Any]],
        tools: list[dict] | None = None,
        tool_choice: Any = None,
        temperature: float = 0.85,
        top_p: float = 0.95,
        max_tokens: int = 2048,
        user_id: uuid.UUID | None = None,
        world_id: uuid.UUID | None = None,
        step_id: uuid.UUID | None = None,
        session: AsyncSession | None = None,
    ) -> AsyncIterator[dict[str, Any]]:
        """Stream chat completion. Yields dicts with keys:
        - `delta`: {content?, tool_calls?}
        - `finish_reason`: present only on the final chunk
        After the iterator is exhausted, the call is logged to `llm_call_logs`.
        """
        payload: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens,
            "stream": True,
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = tool_choice or "auto"

        start = time.monotonic()
        full_content_parts: list[str] = []
        full_tool_calls: list[dict] = []
        finish_reason: str | None = None
        usage: dict[str, Any] = {}
        status = "ok"
        error_message: str | None = None

        try:
            async with httpx.AsyncClient(timeout=self._timeout) as client:
                async with client.stream(
                    "POST",
                    f"{self._api_url}/chat/completions",
                    headers={
                        "Authorization": f"Bearer {self._api_key}",
                        "Content-Type": "application/json",
                    },
                    json=payload,
                ) as resp:
                    if resp.status_code >= 400:
                        body = await resp.aread()
                        raise LLMResponseError(
                            f"LLM provider returned {resp.status_code}: {body.decode('utf-8', 'ignore')[:500]}",
                            code="api_error",
                        )
                    async for line in resp.aiter_lines():
                        if not line:
                            continue
                        if line.startswith("data: "):
                            line = line[6:]
                        if line.strip() == "[DONE]":
                            break
                        try:
                            chunk = json.loads(line)
                        except json.JSONDecodeError:
                            continue
                        if not chunk.get("choices"):
                            if chunk.get("usage"):
                                usage = chunk["usage"]
                            continue
                        choice = chunk["choices"][0]
                        delta = choice.get("delta", {})
                        if delta.get("content"):
                            full_content_parts.append(delta["content"])
                        if delta.get("tool_calls"):
                            full_tool_calls.extend(delta["tool_calls"])
                        if choice.get("finish_reason"):
                            finish_reason = choice["finish_reason"]
                        yield {"delta": delta, "finish_reason": finish_reason}
        except Exception as e:
            status = "api_error" if not isinstance(e, LLMTimeoutError) else "timeout"
            error_message = str(e)
            _logger.warning("llm_stream_error", stage=stage, error=error_message)
            raise
        finally:
            latency_ms = int((time.monotonic() - start) * 1000)
            if session is not None:
                full_content = "".join(full_content_parts)
                await self._write_log_safely(
                    session=session,
                    stage=stage,
                    messages=messages,
                    tools=tools,
                    response_message={
                        "role": "assistant",
                        "content": full_content,
                        "tool_calls": full_tool_calls or None,
                    },
                    tool_calls=full_tool_calls or None,
                    prompt_tokens=usage.get("prompt_tokens"),
                    completion_tokens=usage.get("completion_tokens"),
                    latency_ms=latency_ms,
                    temperature=temperature,
                    status=status,
                    error_message=error_message,
                    user_id=user_id,
                    world_id=world_id,
                    step_id=step_id,
                )

    # ------------------------------------------------------------------ #
    # Safe logging (separate transaction)
    # ------------------------------------------------------------------ #
    async def _write_log_safely(
        self,
        *,
        session: AsyncSession,
        stage: str,
        messages: list[dict[str, Any]],
        tools: list[dict] | None,
        response_message: dict[str, Any],
        tool_calls: list | None,
        prompt_tokens: int | None,
        completion_tokens: int | None,
        latency_ms: int,
        temperature: float,
        status: str,
        error_message: str | None,
        user_id: uuid.UUID | None,
        world_id: uuid.UUID | None,
        step_id: uuid.UUID | None,
    ) -> uuid.UUID | None:
        """Insert an LlmCallLog row in a nested transaction so it survives rollback.

        Errors here are logged but never raised — logging is best-effort.
        """
        try:
            async with session.begin_nested():
                log = LlmCallLog(
                    user_id=user_id,
                    world_id=world_id,
                    step_id=step_id,
                    stage=stage,
                    model=self._model,
                    request_messages=messages,
                    request_tools=tools,
                    response_message=response_message,
                    tool_calls=tool_calls,
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    latency_ms=latency_ms,
                    temperature=temperature,
                    status=status,
                    error_message=error_message,
                )
                session.add(log)
                await session.flush()
                log_id = log.id
            await session.commit()
            return log_id
        except Exception as e:  # noqa: BLE001
            _logger.error("llm_log_write_failed", stage=stage, error=str(e))
            try:
                await session.rollback()
            except Exception:
                pass
            return None


# --------------------------------------------------------------------------- #
# Mock LLM client (for tests)
# --------------------------------------------------------------------------- #
class MockLlmClient:
    """Replay-based mock LLM client. Returns pre-recorded responses per stage."""

    def __init__(self, replay_data: dict[str, list[dict]] | None = None):
        self._replay = replay_data or {}
        self._call_counts: dict[str, int] = {}
        # Allow recording mode
        self.recorded_calls: list[dict[str, Any]] = []

    def set_replay(self, stage: str, responses: list[dict]) -> None:
        self._replay[stage] = responses
        self._call_counts.pop(stage, None)

    async def complete(self, *, stage: str, messages=None, tools=None, **kwargs) -> dict[str, Any]:
        idx = self._call_counts.get(stage, 0)
        responses = self._replay.get(stage, [])
        if idx >= len(responses):
            raise LLMResponseError(
                f"Replay exhausted for stage {stage} (call #{idx + 1})",
                code="replay_exhausted",
            )
        resp = responses[idx]
        self._call_counts[stage] = idx + 1
        self.recorded_calls.append({"stage": stage, "messages": messages, "tools": tools})

        # Mimic the real client's return shape
        return {
            "message": resp.get("message", {"role": "assistant", "content": resp.get("content", "")}),
            "finish_reason": resp.get("finish_reason", "stop"),
            "prompt_tokens": resp.get("prompt_tokens", 0),
            "completion_tokens": resp.get("completion_tokens", 0),
            "latency_ms": 0,
            "log_id": None,
        }

    async def stream_complete(self, *, stage: str, messages=None, tools=None, **kwargs):
        idx = self._call_counts.get(stage, 0)
        responses = self._replay.get(stage, [])
        if idx >= len(responses):
            raise LLMResponseError(
                f"Replay exhausted for stage {stage} (call #{idx + 1})",
                code="replay_exhausted",
            )
        resp = responses[idx]
        self._call_counts[stage] = idx + 1
        content = resp.get("message", {}).get("content", resp.get("content", ""))
        # Yield content in 3 chunks for streaming tests
        chunk_size = max(1, len(content) // 3)
        for i in range(0, len(content), chunk_size):
            yield {"delta": {"content": content[i : i + chunk_size]}, "finish_reason": None}
        yield {"delta": {}, "finish_reason": "stop"}


def get_mock_client() -> MockLlmClient:
    """Convenience factory — used in tests and as a fallback in dev when no LLM configured."""
    return MockLlmClient()