ai-rpg/backend/app/core/llm.py

"""OpenAI-compatible LLM client with tool calling, streaming, and logging."""
from __future__ import annotations

import json
import time
import uuid
from typing import Any, AsyncIterator, Dict, List, Optional

import httpx
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.settings_service import get_all_settings, cast_setting
from app.logging_setup import get_logger
from app.models import LlmCallLog

log = get_logger("llm")


# vendor-specific end-of-sequence / control tokens that some local models
# emit into the content stream. Strip them so they don't leak to the user.
_EOS_TOKENS = (
    "<eos>",
    "</s>",
    "<|endoftext|>",
    "<|im_end|>",
    "<|end|>",
    "<|eot_id|>",
    "<|eom_id|>",
)


def _clean_model_text(text: str) -> str:
    """Remove vendor-specific end-of-sequence tokens and collapse whitespace.

    Some local models leak control tokens into the visible content stream.
    We strip them so they never reach the user.
    """
    if not text:
        return text
    cleaned = text
    for tok in _EOS_TOKENS:
        cleaned = cleaned.replace(tok, "")
    while "\n\n\n" in cleaned:
        cleaned = cleaned.replace("\n\n\n", "\n\n")
    return cleaned




class LlmResponse:
    """Non-streaming response wrapper."""

    def __init__(self, text: str, tool_calls: List[Dict[str, Any]], usage: Optional[Dict[str, int]]):
        self.text = text
        self.tool_calls = tool_calls
        self.usage = usage or {}


class LlmClient:
    """Lightweight OpenAI-compatible chat-completions client."""

    def __init__(self, settings_map: Dict[str, Any]):
        self.base_url: str = str(settings_map.get("llm.base_url", "")).rstrip("/")
        self.api_key: str = str(settings_map.get("llm.api_key", "dummy"))
        self.model: str = str(settings_map.get("llm.model", "local-model"))
        self.temperature: float = float(cast_setting("llm.temperature", settings_map.get("llm.temperature", 0.7)))
        self.max_tokens: int = int(cast_setting("llm.max_tokens", settings_map.get("llm.max_tokens", 1024)))
        self.timeout: int = int(cast_setting("llm.request_timeout", settings_map.get("llm.request_timeout", 120)))
        self.streaming: bool = bool(cast_setting("llm.streaming", settings_map.get("llm.streaming", True)))

    @classmethod
    async def from_db(cls, db: AsyncSession) -> "LlmClient":
        s = await get_all_settings(db)
        return cls(s)

    def _headers(self) -> Dict[str, str]:
        h = {"Content-Type": "application/json"}
        if self.api_key and self.api_key != "dummy":
            h["Authorization"] = f"Bearer {self.api_key}"
        return h

    async def chat(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Any = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        purpose: str = "orchestrator",
        user_id: Optional[uuid.UUID] = None,
        session_id: Optional[uuid.UUID] = None,
        db: Optional[AsyncSession] = None,
    ) -> LlmResponse:
        """Non-streaming chat completion with tool support."""
        url = f"{self.base_url}/chat/completions"
        payload: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature if temperature is not None else self.temperature,
            "max_tokens": max_tokens or self.max_tokens,
            "stream": False,
        }
        if tools:
            payload["tools"] = tools
            if tool_choice is not None:
                payload["tool_choice"] = tool_choice
        started = time.monotonic()
        err: Optional[str] = None
        text = ""
        tool_calls: List[Dict[str, Any]] = []
        usage: Dict[str, int] = {}
        try:
            # Use explicit timeout config so connect/read/write/pool timeouts
            # are all visible — a bare `timeout=N` hides WHICH stage failed.
            timeout = httpx.Timeout(
                connect=10.0,           # 10s to establish TCP connection
                read=float(self.timeout),  # full request timeout
                write=10.0,
                pool=5.0,
            )
            async with httpx.AsyncClient(timeout=timeout) as client:
                resp = await client.post(url, json=payload, headers=self._headers())
                resp.raise_for_status()
                data = resp.json()
            choice = (data.get("choices") or [{}])[0]
            msg = choice.get("message", {})
            text = _clean_model_text(msg.get("content") or "")
            tool_calls = msg.get("tool_calls") or []
            usage = data.get("usage") or {}
        except httpx.ConnectError as e:
            err = f"ConnectError: {e}"
            # Surface the URL + cause so the operator can see WHY (DNS, refused, etc.)
            cause = getattr(e, "__cause__", None) or getattr(e, "__context__", None)
            log.error(
                "llm_connect_failed",
                purpose=purpose,
                url=url,
                base_url=self.base_url,
                model=self.model,
                error=err,
                cause=str(cause) if cause else None,
            )
            raise
        except Exception as e:
            err = f"{type(e).__name__}: {e}"
            log.error(
                "llm_call_failed",
                purpose=purpose,
                url=url,
                base_url=self.base_url,
                model=self.model,
                error=err,
            )
            raise
        finally:
            latency_ms = int((time.monotonic() - started) * 1000)
            if db is not None:
                db.add(LlmCallLog(
                    user_id=user_id,
                    session_id=session_id,
                    purpose=purpose,
                    model=self.model,
                    base_url=self.base_url,
                    prompt_messages=messages,
                    tools=tools,
                    response_text=text,
                    tool_calls=tool_calls,
                    prompt_tokens=usage.get("prompt_tokens"),
                    completion_tokens=usage.get("completion_tokens"),
                    total_tokens=usage.get("total_tokens"),
                    latency_ms=latency_ms,
                    error=err,
                ))
                try:
                    await db.commit()
                except Exception:
                    await db.rollback()
        return LlmResponse(text=text, tool_calls=tool_calls, usage=usage)

    async def stream_chat(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Any = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        purpose: str = "orchestrator",
        user_id: Optional[uuid.UUID] = None,
        session_id: Optional[uuid.UUID] = None,
        db: Optional[AsyncSession] = None,
    ) -> AsyncIterator[Dict[str, Any]]:
        """Streaming chat completion. Yields incremental deltas.

        Yields dicts of the form:
          {"type": "delta", "content": "..."} - text delta
          {"type": "tool_calls", "tool_calls": [...]} - final tool calls (if any)
          {"type": "done", "usage": {...}}
          {"type": "error", "error": "..."}
        """
        url = f"{self.base_url}/chat/completions"
        payload: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature if temperature is not None else self.temperature,
            "max_tokens": max_tokens or self.max_tokens,
            "stream": True,
        }
        if tools:
            payload["tools"] = tools
            if tool_choice is not None:
                payload["tool_choice"] = tool_choice

        started = time.monotonic()
        full_text_parts: List[str] = []
        tool_call_accum: Dict[int, Dict[str, Any]] = {}
        usage: Dict[str, int] = {}
        err: Optional[str] = None

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                async with client.stream("POST", url, json=payload, headers=self._headers()) as resp:
                    resp.raise_for_status()
                    async for line in resp.aiter_lines():
                        if not line or not line.startswith("data:"):
                            continue
                        data_str = line[5:].strip()
                        if data_str == "[DONE]":
                            break
                        try:
                            chunk = json.loads(data_str)
                        except json.JSONDecodeError:
                            continue
                        choices = chunk.get("choices") or []
                        if not choices:
                            if chunk.get("usage"):
                                usage = chunk["usage"]
                            continue
                        delta = choices[0].get("delta", {})
                        if delta.get("content"):
                            piece = _clean_model_text(delta["content"])
                            if piece:
                                full_text_parts.append(piece)
                                yield {"type": "delta", "content": piece}
                        if delta.get("tool_calls"):
                            for tc in delta["tool_calls"]:
                                idx = tc.get("index", 0)
                                acc = tool_call_accum.setdefault(idx, {
                                    "id": tc.get("id", ""),
                                    "type": "function",
                                    "function": {"name": "", "arguments": ""},
                                })
                                if tc.get("id"):
                                    acc["id"] = tc["id"]
                                if tc.get("function", {}).get("name"):
                                    acc["function"]["name"] += tc["function"]["name"]
                                if tc.get("function", {}).get("arguments"):
                                    acc["function"]["arguments"] += tc["function"]["arguments"]
                        if chunk.get("usage"):
                            usage = chunk["usage"]
        except Exception as e:
            err = f"{type(e).__name__}: {e}"
            log.error("llm_stream_failed", purpose=purpose, error=err)
            yield {"type": "error", "error": err}
            return

        full_text = "".join(full_text_parts)
        final_tool_calls = [tool_call_accum[i] for i in sorted(tool_call_accum.keys())]
        if final_tool_calls:
            yield {"type": "tool_calls", "tool_calls": final_tool_calls}
        yield {"type": "done", "usage": usage, "full_text": full_text}

        latency_ms = int((time.monotonic() - started) * 1000)
        if db is not None:
            db.add(LlmCallLog(
                user_id=user_id,
                session_id=session_id,
                purpose=purpose,
                model=self.model,
                base_url=self.base_url,
                prompt_messages=messages,
                tools=tools,
                response_text=full_text,
                tool_calls=final_tool_calls,
                prompt_tokens=usage.get("prompt_tokens"),
                completion_tokens=usage.get("completion_tokens"),
                total_tokens=usage.get("total_tokens"),
                latency_ms=latency_ms,
                error=err,
            ))
            try:
                await db.commit()
            except Exception:
                await db.rollback()


def build_tool_schema(name: str, description: str, params: Dict[str, Any]) -> Dict[str, Any]:
    """Helper to build an OpenAI-style tool schema."""
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": description,
            "parameters": params,
        },
    }