RAG_helper/services/llm_client.py

import logging
from pathlib import Path

import httpx

from config import settings

logger = logging.getLogger(__name__)

SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent.parent / "prompts" / "system_prompt.md"


def _load_system_prompt() -> str:
    try:
        return SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip()
    except FileNotFoundError:
        logger.warning("System prompt file not found at %s — using empty prompt", SYSTEM_PROMPT_PATH)
        return ""


DEFAULT_SYSTEM_PROMPT = _load_system_prompt()

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 DEFAULT_SYSTEM_PROMPT
        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,
        }

        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()
            data = response.json()

        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 DEFAULT_SYSTEM_PROMPT
        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,
        }

        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()
            data = response.json()

        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}