chrysopedia/backend/pipeline/llm_client.py

"""Synchronous LLM client with primary/fallback endpoint logic.

Uses the OpenAI-compatible API (works with Ollama, vLLM, OpenWebUI, etc.).
Celery tasks run synchronously, so this uses ``openai.OpenAI`` (not Async).

Supports two modalities:
- **chat**: Standard JSON mode with ``response_format: {"type": "json_object"}``
- **thinking**: For reasoning models that emit ``<think>...</think>`` blocks
  before their answer. Skips ``response_format``, appends JSON instructions to
  the system prompt, and strips think tags from the response.
"""

from __future__ import annotations

import logging
import re
from typing import TypeVar

import openai
from pydantic import BaseModel

from config import Settings

logger = logging.getLogger(__name__)

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

# ── Think-tag stripping ──────────────────────────────────────────────────────

_THINK_PATTERN = re.compile(r"<think>.*?</think>", re.DOTALL)


def strip_think_tags(text: str) -> str:
    """Remove ``<think>...</think>`` blocks from LLM output.

    Thinking/reasoning models often prefix their JSON with a reasoning trace
    wrapped in ``<think>`` tags. This strips all such blocks (including
    multiline and multiple occurrences) and returns the cleaned text.

    Handles:
    - Single ``<think>...</think>`` block
    - Multiple blocks in one response
    - Multiline content inside think tags
    - Responses with no think tags (passthrough)
    - Empty input (passthrough)
    """
    if not text:
        return text
    cleaned = _THINK_PATTERN.sub("", text)
    return cleaned.strip()


class LLMClient:
    """Sync LLM client that tries a primary endpoint and falls back on failure."""

    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self._primary = openai.OpenAI(
            base_url=settings.llm_api_url,
            api_key=settings.llm_api_key,
        )
        self._fallback = openai.OpenAI(
            base_url=settings.llm_fallback_url,
            api_key=settings.llm_api_key,
        )

    # ── Core completion ──────────────────────────────────────────────────

    def complete(
        self,
        system_prompt: str,
        user_prompt: str,
        response_model: type[BaseModel] | None = None,
        modality: str = "chat",
        model_override: str | None = None,
    ) -> str:
        """Send a chat completion request, falling back on connection/timeout errors.

        Parameters
        ----------
        system_prompt:
            System message content.
        user_prompt:
            User message content.
        response_model:
            If provided and modality is "chat", ``response_format`` is set to
            ``{"type": "json_object"}``. For "thinking" modality, JSON
            instructions are appended to the system prompt instead.
        modality:
            Either "chat" (default) or "thinking". Thinking modality skips
            response_format and strips ``<think>`` tags from output.
        model_override:
            Model name to use instead of the default. If None, uses the
            configured default for the endpoint.

        Returns
        -------
        str
            Raw completion text from the model (think tags stripped if thinking).
        """
        kwargs: dict = {}
        effective_system = system_prompt

        if modality == "thinking":
            # Thinking models often don't support response_format: json_object.
            # Instead, append explicit JSON instructions to the system prompt.
            if response_model is not None:
                json_schema_hint = (
                    "\n\nYou MUST respond with ONLY valid JSON. "
                    "No markdown code fences, no explanation, no preamble — "
                    "just the raw JSON object."
                )
                effective_system = system_prompt + json_schema_hint
        else:
            # Chat modality — use standard JSON mode
            if response_model is not None:
                kwargs["response_format"] = {"type": "json_object"}

        messages = [
            {"role": "system", "content": effective_system},
            {"role": "user", "content": user_prompt},
        ]

        primary_model = model_override or self.settings.llm_model
        fallback_model = self.settings.llm_fallback_model

        logger.info(
            "LLM request: model=%s, modality=%s, response_model=%s",
            primary_model,
            modality,
            response_model.__name__ if response_model else None,
        )

        # --- Try primary endpoint ---
        try:
            response = self._primary.chat.completions.create(
                model=primary_model,
                messages=messages,
                max_tokens=self.settings.llm_max_tokens,
                **kwargs,
            )
            raw = response.choices[0].message.content or ""
            usage = getattr(response, "usage", None)
            if usage:
                logger.info(
                    "LLM response: prompt_tokens=%s, completion_tokens=%s, total=%s, content_len=%d, finish=%s",
                    usage.prompt_tokens, usage.completion_tokens, usage.total_tokens,
                    len(raw), response.choices[0].finish_reason,
                )
            if modality == "thinking":
                raw = strip_think_tags(raw)
            return raw

        except (openai.APIConnectionError, openai.APITimeoutError) as exc:
            logger.warning(
                "Primary LLM endpoint failed (%s: %s), trying fallback at %s",
                type(exc).__name__,
                exc,
                self.settings.llm_fallback_url,
            )

        # --- Try fallback endpoint ---
        try:
            response = self._fallback.chat.completions.create(
                model=fallback_model,
                messages=messages,
                max_tokens=self.settings.llm_max_tokens,
                **kwargs,
            )
            raw = response.choices[0].message.content or ""
            usage = getattr(response, "usage", None)
            if usage:
                logger.info(
                    "LLM response (fallback): prompt_tokens=%s, completion_tokens=%s, total=%s, content_len=%d, finish=%s",
                    usage.prompt_tokens, usage.completion_tokens, usage.total_tokens,
                    len(raw), response.choices[0].finish_reason,
                )
            if modality == "thinking":
                raw = strip_think_tags(raw)
            return raw

        except (openai.APIConnectionError, openai.APITimeoutError, openai.APIError) as exc:
            logger.error(
                "Fallback LLM endpoint also failed (%s: %s). Giving up.",
                type(exc).__name__,
                exc,
            )
            raise

    # ── Response parsing ─────────────────────────────────────────────────

    def parse_response(self, text: str, model: type[T]) -> T:
        """Parse raw LLM output as JSON and validate against a Pydantic model.

        Parameters
        ----------
        text:
            Raw JSON string from the LLM.
        model:
            Pydantic model class to validate against.

        Returns
        -------
        T
            Validated Pydantic model instance.

        Raises
        ------
        pydantic.ValidationError
            If the JSON doesn't match the schema.
        ValueError
            If the text is not valid JSON.
        """
        try:
            return model.model_validate_json(text)
        except Exception:
            logger.error(
                "Failed to parse LLM response as %s. Response text: %.500s",
                model.__name__,
                text,
            )
            raise