sales-trainer/backend/app/llm.py

"""OpenAI-compatible LLM client (OpenAI / DeepSeek / custom base URL).

Mirrors the MiroFish provider-agnostic pattern. Credentials live in .env only.
"""
from __future__ import annotations

import json
import re
from typing import Any

from openai import OpenAI

from .config import Config


class LLMError(Exception):
    pass


def _strip_thinking_trace(text: str) -> str:
    """Remove ReACT-style chain-of-thought / fences, keep the final JSON text."""
    for fence in ("```json", "```"):
        idx = text.rfind(fence)
        if idx != -1:
            after = text[idx:].lstrip()
            lang_len = after.find("\n")
            body = after[lang_len:] if lang_len != -1 else after
            end = body.rfind("```")
            if end != -1:
                body = body[:end]
            body = body.strip()
            if body:
                return body
    for marker in ("\n\n[", "\n\n{"):
        idx = text.rfind(marker)
        if idx != -1:
            candidate = text[idx:].strip()
            if candidate and candidate[0] in "{[":
                return candidate
    return text


class LLMClient:
    def __init__(
        self,
        *,
        base_url: str | None = None,
        api_key: str | None = None,
        model: str | None = None,
    ) -> None:
        self.base_url = base_url or Config.LLM_BASE_URL
        self.api_key = api_key or Config.LLM_API_KEY
        self.model = model or Config.LLM_MODEL
        if not self.api_key:
            raise LLMError("LLM_API_KEY is not configured in .env")
        if not self.base_url:
            raise LLMError("LLM_BASE_URL is not configured (unknown provider)")
        self.client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    def complete(
        self,
        system_prompt: str,
        user_prompt: str,
        *,
        temperature: float = 0.5,
        max_tokens: int = 3000,
    ) -> str:
        try:
            resp = self.client.chat.completions.create(
                model=self.model,
                temperature=temperature,
                max_tokens=max_tokens,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
            )
        except Exception as exc:  # network/auth/provider
            raise LLMError(f"LLM call failed: {exc}") from exc
        text = (resp.choices[0].message.content or "").strip()
        if not text:
            raise LLMError("LLM returned empty response")
        return text

    def complete_json(
        self,
        system_prompt: str,
        user_prompt: str,
        *,
        temperature: float = 0.2,
        max_tokens: int = 6000,
    ) -> dict[str, Any]:
        text = self.complete(
            system_prompt,
            user_prompt,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        text = _strip_thinking_trace(text)
        try:
            return json.loads(text)
        except json.JSONDecodeError as exc:
            # Last-ditch: strip leading text before the first { or [
            match = re.search(r"[{\[].*[}\]]", text, re.DOTALL)
            if match:
                try:
                    return json.loads(match.group(0))
                except json.JSONDecodeError:
                    pass
            raise LLMError(f"LLM returned invalid JSON: {exc}") from exc

    def complete_conversation(
        self,
        messages: list[dict[str, str]],
        *,
        temperature: float = 0.6,
        max_tokens: int = 1200,
    ) -> str:
        # Normalize internal role labels (customer/seller) to the roles an OpenAI-compatible
        # chat endpoint accepts: system/user/assistant. customer=assistant (the persona/LLM),
        # seller=user (the trainee). Anything else maps to a safe default.
        api_messages = []
        for m in messages:
            role = (m.get("role") or "").lower()
            if role == "system":
                mapped = "system"
            elif role in ("customer", "assistant"):
                mapped = "assistant"
            elif role in ("seller", "user"):
                mapped = "user"
            else:
                mapped = "user"
            api_messages.append({"role": mapped, "content": m.get("text") or m.get("content") or ""})
        try:
            resp = self.client.chat.completions.create(
                model=self.model,
                temperature=temperature,
                max_tokens=max_tokens,
                messages=api_messages,
            )
        except Exception as exc:
            raise LLMError(f"LLM call failed: {exc}") from exc
        text = (resp.choices[0].message.content or "").strip()
        if not text:
            raise LLMError("LLM returned empty response")
        return text