337 lines
13 KiB
Python
337 lines
13 KiB
Python
"""ESPD AI chat with tools (doc edit, common fields, RAG, GOST check)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from app.services.espd.html_util import html_to_plain, plain_to_html
|
|
from app.services.espd.models import CONTENT_CONTEXT_KEYS, DOC_TYPE_CONFIG, DocumentType
|
|
from app.services.llm import (
|
|
build_context_usage,
|
|
openrouter_chat,
|
|
serialize_tool_calls,
|
|
_extract_json,
|
|
)
|
|
|
|
ESPD_SYSTEM = """Ты помощник по оформлению документов ЕСПД (спецификация, руководства, ПМИ и др.).
|
|
Используй tools для предложений правок. Не применяй изменения сам.
|
|
Если нужно уточнение — ask_user.
|
|
Для проверки по ГОСТам используй gost_check / rag_search.
|
|
HTML документа редактируй аккуратно, сохраняя структуру заголовков.
|
|
"""
|
|
|
|
ESPD_TOOLS: list[dict[str, Any]] = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "edit_doc_html",
|
|
"description": "Заменить HTML-содержимое выбранного типа документа ЕСПД.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"doc_type": {
|
|
"type": "string",
|
|
"enum": [dt.value for dt in DocumentType if dt != DocumentType.SPEC],
|
|
},
|
|
"html": {"type": "string"},
|
|
"mode": {
|
|
"type": "string",
|
|
"enum": ["replace", "append"],
|
|
"description": "replace — полный текст; append — дописать в конец",
|
|
},
|
|
},
|
|
"required": ["doc_type", "html"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "set_common_fields",
|
|
"description": "Изменить общие поля основной надписи ЕСПД.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"decimal": {"type": "string"},
|
|
"progname": {"type": "string"},
|
|
"author": {"type": "string"},
|
|
"checker": {"type": "string"},
|
|
"ncontrol": {"type": "string"},
|
|
"approver": {"type": "string"},
|
|
"year": {"type": "string"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"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": {
|
|
"doc_type": {"type": "string"},
|
|
"focus": {"type": "string", "description": "На что обратить внимание"},
|
|
"use_current_html": {"type": "boolean"},
|
|
"text": {"type": "string"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"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 {}
|
|
|
|
|
|
def espd_tool_calls_to_edits(
|
|
tool_calls: list[dict[str, Any]],
|
|
) -> tuple[list[dict[str, Any]], str | None, list[dict[str, Any]]]:
|
|
"""Returns (proposal_edits, ask_question, executable_tool_requests)."""
|
|
edits: list[dict[str, Any]] = []
|
|
exec_reqs: list[dict[str, Any]] = []
|
|
question: str | None = None
|
|
for tc in tool_calls:
|
|
fn = tc.get("function") or {}
|
|
name = fn.get("name") or ""
|
|
args = _parse_args(fn.get("arguments"))
|
|
if name == "ask_user":
|
|
question = str(args.get("question") or "").strip() or question
|
|
elif name == "edit_doc_html":
|
|
edits.append(
|
|
{
|
|
"op": "edit_doc_html",
|
|
"doc_type": args.get("doc_type"),
|
|
"html": args.get("html") or "",
|
|
"mode": args.get("mode") or "replace",
|
|
}
|
|
)
|
|
elif name == "set_common_fields":
|
|
fields = {k: v for k, v in args.items() if k in {
|
|
"decimal", "progname", "author", "checker", "ncontrol", "approver", "year"
|
|
} and v is not None}
|
|
if fields:
|
|
edits.append({"op": "set_common_fields", "fields": fields})
|
|
elif name in ("rag_search", "gost_check"):
|
|
exec_reqs.append({"name": name, "arguments": args, "tool_call_id": tc.get("id")})
|
|
return edits, question, exec_reqs
|
|
|
|
|
|
async def chat_espd(
|
|
message: str,
|
|
*,
|
|
history: list[dict[str, str]] | None = None,
|
|
common: dict[str, Any] | None = None,
|
|
doc_contents: dict[str, str] | None = None,
|
|
active_doc_type: str | None = None,
|
|
project_id: int | None = None,
|
|
rag_search_fn=None,
|
|
gost_check_fn=None,
|
|
) -> dict[str, Any]:
|
|
"""ESPD assistant with multi-round tools. rag_search_fn/gost_check_fn optional callables."""
|
|
hist = list(history or [])
|
|
if not hist or hist[-1].get("content") != message:
|
|
hist.append({"role": "user", "content": message})
|
|
|
|
docs_summary = {
|
|
k: (html_to_plain(v)[:2000] if v else "") for k, v in (doc_contents or {}).items()
|
|
}
|
|
catalog = [
|
|
{"id": dt.value, "name": cfg["name"], "content_key": CONTENT_CONTEXT_KEYS.get(dt, "")}
|
|
for dt, cfg in DOC_TYPE_CONFIG.items()
|
|
if dt != DocumentType.SPEC
|
|
]
|
|
snapshot = (
|
|
f"Активный документ: {active_doc_type or '—'}\n"
|
|
f"Общие поля: {json.dumps(common or {}, ensure_ascii=False)}\n"
|
|
f"Типы документов: {json.dumps(catalog, ensure_ascii=False)}\n"
|
|
f"Тексты (plain, усечённо): {json.dumps(docs_summary, ensure_ascii=False)}"
|
|
)
|
|
messages: list[dict[str, Any]] = [
|
|
{"role": "system", "content": ESPD_SYSTEM},
|
|
{"role": "user", "content": snapshot},
|
|
{"role": "assistant", "content": "Контекст ЕСПД принят. Жду команды."},
|
|
]
|
|
for m in hist[-20:]:
|
|
if m.get("content"):
|
|
messages.append({"role": m["role"], "content": m["content"]})
|
|
|
|
all_edits: list[dict[str, Any]] = []
|
|
all_tools: list[dict[str, Any]] = []
|
|
question: str | None = None
|
|
final_content = ""
|
|
rag_notes: list[str] = []
|
|
|
|
for round_no in range(1, 5):
|
|
msg = await openrouter_chat(
|
|
messages,
|
|
temperature=0.2,
|
|
max_tokens=3500,
|
|
tools=ESPD_TOOLS,
|
|
tool_choice="auto",
|
|
)
|
|
content = (msg.get("content") or "").strip()
|
|
tool_calls = msg.get("tool_calls") or []
|
|
if tool_calls:
|
|
all_tools.extend(serialize_tool_calls(tool_calls, round_no=round_no))
|
|
edits, q, exec_reqs = espd_tool_calls_to_edits(tool_calls)
|
|
all_edits.extend(edits)
|
|
if q:
|
|
question = q
|
|
final_content = content
|
|
break
|
|
|
|
messages.append({"role": "assistant", "content": content or "", "tool_calls": tool_calls})
|
|
for req in exec_reqs:
|
|
result_payload: dict[str, Any]
|
|
if req["name"] == "rag_search" and rag_search_fn:
|
|
result_payload = await _maybe_await(
|
|
rag_search_fn,
|
|
req["arguments"].get("query") or "",
|
|
int(req["arguments"].get("top_k") or 5),
|
|
bool(req["arguments"].get("gost_only")),
|
|
project_id,
|
|
)
|
|
rag_notes.append(json.dumps(result_payload, ensure_ascii=False)[:4000])
|
|
elif req["name"] == "gost_check" and gost_check_fn:
|
|
html = ""
|
|
dt = req["arguments"].get("doc_type") or active_doc_type
|
|
if req["arguments"].get("use_current_html", True) and dt:
|
|
html = (doc_contents or {}).get(dt, "")
|
|
text = req["arguments"].get("text") or html_to_plain(html)
|
|
result_payload = await _maybe_await(
|
|
gost_check_fn,
|
|
text,
|
|
req["arguments"].get("focus") or "",
|
|
project_id,
|
|
)
|
|
rag_notes.append(json.dumps(result_payload, ensure_ascii=False)[:6000])
|
|
if result_payload.get("reply"):
|
|
final_content = str(result_payload["reply"])
|
|
else:
|
|
result_payload = {
|
|
"status": "accepted_as_proposal" if req["name"] not in ("rag_search", "gost_check") else "unavailable",
|
|
"tool": req["name"],
|
|
}
|
|
# also acknowledge propose tools
|
|
messages.append(
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": req.get("tool_call_id") or f"call_{round_no}_{req['name']}",
|
|
"content": json.dumps(result_payload, ensure_ascii=False),
|
|
}
|
|
)
|
|
# acknowledge edit proposals without exec
|
|
for tc in tool_calls:
|
|
fn = (tc.get("function") or {}).get("name")
|
|
if fn in ("edit_doc_html", "set_common_fields"):
|
|
messages.append(
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": tc.get("id") or f"call_{round_no}_{fn}",
|
|
"content": json.dumps(
|
|
{"status": "accepted_as_proposal", "tool": fn},
|
|
ensure_ascii=False,
|
|
),
|
|
}
|
|
)
|
|
if exec_reqs and not edits:
|
|
continue
|
|
if edits and not exec_reqs:
|
|
# one more round for natural reply optional
|
|
continue
|
|
continue
|
|
|
|
if content:
|
|
parsed = _extract_json(content)
|
|
if isinstance(parsed, dict) and parsed.get("reply"):
|
|
final_content = str(parsed["reply"])
|
|
else:
|
|
final_content = content
|
|
break
|
|
|
|
reply = question or final_content
|
|
if not reply:
|
|
if all_edits:
|
|
reply = f"Предложено правок: {len(all_edits)}. Подтвердите применение."
|
|
elif rag_notes:
|
|
reply = "Результаты поиска по базе знаний получены (см. инструменты)."
|
|
else:
|
|
reply = "Нет предложенных изменений."
|
|
|
|
return {
|
|
"reply": reply,
|
|
"edits": all_edits,
|
|
"tools": all_tools,
|
|
"context": build_context_usage(messages),
|
|
}
|
|
|
|
|
|
async def _maybe_await(fn, *args):
|
|
result = fn(*args)
|
|
if hasattr(result, "__await__"):
|
|
return await result
|
|
return result
|
|
|
|
|
|
def apply_espd_edits(state, edits: list[dict[str, Any]]):
|
|
"""Apply proposal edits onto EspdState (mutates)."""
|
|
for edit in edits:
|
|
op = edit.get("op")
|
|
if op == "edit_doc_html":
|
|
dt = edit.get("doc_type")
|
|
if not dt:
|
|
continue
|
|
html = edit.get("html") or ""
|
|
mode = edit.get("mode") or "replace"
|
|
if mode == "append":
|
|
prev = state.get_doc_html(dt)
|
|
html = (prev or "") + html
|
|
state.set_doc_html(dt, html)
|
|
elif op == "set_common_fields":
|
|
fields = edit.get("fields") or {}
|
|
for k, v in fields.items():
|
|
if hasattr(state.common, k) and v is not None:
|
|
setattr(state.common, k, str(v))
|
|
# sync legacy
|
|
if state.get_doc_html("ro"):
|
|
state.ro_ai_content = html_to_plain(state.get_doc_html("ro"))
|
|
return state
|