added global RaG

2026-09-09 16:23:28 +03:00
parent 959813c11d
commit 4e16d98cab
14 changed files with 941 additions and 256 deletions
@@ -0,0 +1,180 @@
+"""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),
+    }