GostGenerator/web/backend/app/services/rag/chat.py

"""Global knowledge-base chat (RAG / GOST check only)."""

from __future__ import annotations

import json
from typing import Any

from app.services.llm import (
    build_context_usage,
    openrouter_chat,
    serialize_tool_calls,
    _extract_json,
)

KNOWLEDGE_SYSTEM = """Ты помощник по нормативной базе (ГОСТы ЕСПД/ЕСКД и загруженные документы).
Используй tools: rag_search для поиска по корпусу, gost_check для проверки текста с цитатами.
Если нужно уточнение — ask_user.
Отвечай по-русски. Таблицы оформляй markdown (| колонка | ... |).
Не выдумывай номера ГОСТов — опирайся на найденные фрагменты.
"""

KNOWLEDGE_TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "rag_search",
            "description": "Поиск по общей базе знаний (ГОСТы и документы).",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string"},
                    "top_k": {"type": "integer"},
                    "gost_only": {"type": "boolean"},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "gost_check",
            "description": "Проверить текст на соответствие нормам из RAG (с цитатами).",
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {"type": "string"},
                    "focus": {"type": "string"},
                },
                "required": ["text"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "ask_user",
            "description": "Уточняющий вопрос.",
            "parameters": {
                "type": "object",
                "properties": {"question": {"type": "string"}},
                "required": ["question"],
            },
        },
    },
]


def _parse_args(raw: Any) -> dict[str, Any]:
    if isinstance(raw, dict):
        return raw
    try:
        data = json.loads(raw or "{}")
        return data if isinstance(data, dict) else {}
    except json.JSONDecodeError:
        return {}


async def _maybe_await(fn, *args):
    result = fn(*args)
    if hasattr(result, "__await__"):
        return await result
    return result


async def chat_knowledge(
    message: str,
    *,
    history: list[dict[str, str]] | None = None,
    project_id: int | None = None,
    rag_search_fn=None,
    gost_check_fn=None,
) -> dict[str, Any]:
    hist = list(history or [])
    if not hist or hist[-1].get("content") != message:
        hist.append({"role": "user", "content": message})

    messages: list[dict[str, Any]] = [
        {"role": "system", "content": KNOWLEDGE_SYSTEM},
    ]
    for m in hist[-24:]:
        if m.get("content"):
            messages.append({"role": m["role"], "content": m["content"]})

    all_tools: list[dict[str, Any]] = []
    question: str | None = None
    final_content = ""

    for round_no in range(1, 5):
        msg = await openrouter_chat(
            messages,
            temperature=0.2,
            max_tokens=3500,
            tools=KNOWLEDGE_TOOLS,
            tool_choice="auto",
        )
        content = (msg.get("content") or "").strip()
        tool_calls = msg.get("tool_calls") or []
        if not tool_calls:
            if content:
                parsed = _extract_json(content)
                if isinstance(parsed, dict) and parsed.get("reply"):
                    final_content = str(parsed["reply"])
                else:
                    final_content = content
            break

        all_tools.extend(serialize_tool_calls(tool_calls, round_no=round_no))
        messages.append({"role": "assistant", "content": content or "", "tool_calls": tool_calls})

        stop_after_ask = False
        for tc in tool_calls:
            fn = tc.get("function") or {}
            name = fn.get("name") or ""
            args = _parse_args(fn.get("arguments"))
            tool_id = tc.get("id") or f"call_{round_no}_{name}"

            if name == "ask_user":
                question = str(args.get("question") or "").strip() or question
                result_payload = {"status": "asked", "question": question}
                stop_after_ask = True
            elif name == "rag_search" and rag_search_fn:
                result_payload = await _maybe_await(
                    rag_search_fn,
                    args.get("query") or "",
                    int(args.get("top_k") or 5),
                    bool(args.get("gost_only", True)),
                    project_id,
                )
            elif name == "gost_check" and gost_check_fn:
                result_payload = await _maybe_await(
                    gost_check_fn,
                    args.get("text") or "",
                    args.get("focus") or "",
                    project_id,
                )
                if isinstance(result_payload, dict) and result_payload.get("reply"):
                    final_content = str(result_payload["reply"])
            else:
                result_payload = {"status": "unavailable", "tool": name}

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_id,
                    "content": json.dumps(result_payload, ensure_ascii=False),
                }
            )

        if stop_after_ask:
            final_content = content
            break

    reply = question or final_content or "Нет ответа. Попробуйте уточнить запрос."
    return {
        "reply": reply,
        "edits": [],
        "tools": all_tools,
        "context": build_context_usage(messages),
    }