RAG_helper/services/llm_client.py

import asyncio
import logging

import httpx

from config import settings


class LLMUnavailableError(RuntimeError):
    """Внешний LLM недоступен после всех попыток — сигнал для вызывающего кода."""
    pass

logger = logging.getLogger(__name__)


DEFAULT_USER_TEMPLATE = """Вопрос пациента:
{question}

Выдержки из базы знаний операторов:
{sources}

Ответь пациенту в чате по правилам из системного сообщения."""


CHAT_USER_TEMPLATE = """Новая реплика пациента:
{question}

Выдержки из базы знаний операторов (по последней реплике):
{sources}

Ответь пациенту с учётом истории диалога выше и правил из системного сообщения."""


class LLMClient:
    def __init__(
        self,
        api_key: str | None = None,
        model: str | None = None,
        base_url: str | None = None,
    ):
        self.api_key = api_key or settings.deepseek_api_key
        self.model = model or settings.deepseek_model
        self.base_url = (base_url or settings.deepseek_base_url).rstrip("/")

    def _format_sources(self, sources: list[dict]) -> str:
        if not sources:
            return "(источники не найдены)"
        lines = []
        for i, src in enumerate(sources, 1):
            meta = src.get("metadata", {})
            doc_name = meta.get("document_name", "Документ")
            section = meta.get("section", "")
            lines.append(
                f"[{i}] {src['text']}\n"
                f"    (Источник: {doc_name}, раздел: {section})"
            )
        return "\n".join(lines)

    async def answer(
        self,
        question: str,
        sources: list[dict],
        system_prompt: str | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> dict:
        """Generate a patient-facing answer using RAG context.

        Returns dict with 'text' and 'assembled_prompt'.
        """
        effective_system = system_prompt or ""
        effective_temp = temperature if temperature is not None else 0.2
        effective_max_tokens = max_tokens or 1200

        formatted_sources = self._format_sources(sources)
        user_message = DEFAULT_USER_TEMPLATE.format(
            question=question,
            sources=formatted_sources,
        )

        assembled_prompt = f"[SYSTEM]\n{effective_system}\n\n[USER]\n{user_message}"

        url = f"{self.base_url}/chat/completions"
        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": effective_system},
                {"role": "user", "content": user_message},
            ],
            "temperature": effective_temp,
            "max_tokens": effective_max_tokens,
        }

        data = await self._call_with_retry(url, payload)
        content = data["choices"][0]["message"]["content"]
        logger.info("LLM response: %d chars, model=%s, temp=%.2f", len(content), self.model, effective_temp)
        return {"text": content.strip(), "assembled_prompt": assembled_prompt}

    async def chat(
        self,
        question: str,
        sources: list[dict],
        history: list[dict],
        system_prompt: str | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> dict:
        """Generate a patient-facing answer using RAG + conversation history.

        `history` — список предыдущих сообщений треда в формате
        [{"role": "user"|"assistant", "content": str}, ...] (без текущей реплики).

        Returns dict with 'text' and 'assembled_prompt'.
        """
        effective_system = system_prompt or ""
        effective_temp = temperature if temperature is not None else 0.2
        effective_max_tokens = max_tokens or 1200

        formatted_sources = self._format_sources(sources)
        user_message = CHAT_USER_TEMPLATE.format(
            question=question,
            sources=formatted_sources,
        )

        messages: list[dict] = [{"role": "system", "content": effective_system}]
        messages.extend(history)
        messages.append({"role": "user", "content": user_message})

        assembled_prompt_parts = [f"[SYSTEM]\n{effective_system}"]
        for m in history:
            tag = "USER" if m["role"] == "user" else "ASSISTANT"
            assembled_prompt_parts.append(f"[{tag}]\n{m['content']}")
        assembled_prompt_parts.append(f"[USER]\n{user_message}")
        assembled_prompt = "\n\n".join(assembled_prompt_parts)

        url = f"{self.base_url}/chat/completions"
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": effective_temp,
            "max_tokens": effective_max_tokens,
        }

        data = await self._call_with_retry(url, payload)
        content = data["choices"][0]["message"]["content"]
        logger.info("LLM chat response: %d chars, history=%d, model=%s", len(content), len(history), self.model)
        return {"text": content.strip(), "assembled_prompt": assembled_prompt}

    async def _call_with_retry(self, url: str, payload: dict) -> dict:
        """POST к DeepSeek с одним ретраем — модель периодически моргает по сети."""
        last_error: Exception | None = None
        for attempt in range(2):
            try:
                async with httpx.AsyncClient(timeout=60.0) as client:
                    response = await client.post(
                        url,
                        json=payload,
                        headers={
                            "Authorization": f"Bearer {self.api_key}",
                            "Content-Type": "application/json",
                        },
                    )
                    response.raise_for_status()
                    return response.json()
            except Exception as e:
                last_error = e
                logger.warning(
                    "LLM call failed (attempt %d, %s: %s)",
                    attempt + 1, type(e).__name__, e,
                )
                if attempt < 1:
                    await asyncio.sleep(0.5)
        raise LLMUnavailableError(
            f"LLM unavailable after retries: {type(last_error).__name__}: {last_error}"
        ) from last_error