233 lines
8.3 KiB
Python
233 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from app.core.config import get_settings
|
|
|
|
|
|
SYSTEM_PROMPT = """Ты помощник инженера по редактированию ГОСТ-таблиц (перечень, спецификация, ведомость).
|
|
Пользователь даёт команды на русском. Учитывай предыдущие сообщения диалога.
|
|
Текущая выбранная таблица указана в сообщениях — работай только с ней, если не сказано иное.
|
|
Ты отвечаешь ТОЛЬКО валидным JSON без markdown:
|
|
{
|
|
"reply": "краткий ответ пользователю",
|
|
"edits": [
|
|
{"op": "update_row", "row_id": 123, "fields": {"note": "..."}},
|
|
{"op": "update_row", "row_index": 5, "fields": {"designation": "..."}},
|
|
{"op": "add_row", "row_index": 10, "fields": {...}},
|
|
{"op": "delete_row", "row_id": 123}
|
|
]
|
|
}
|
|
Правила:
|
|
- Меняй только то, о чём просят (с учётом контекста диалога).
|
|
- Используй row_id из актуального снимка таблицы, если есть.
|
|
- Не выдумывай поля вне списка колонок.
|
|
- Если правок нет — edits: [].
|
|
"""
|
|
|
|
TABLE_TYPE_LABELS = {
|
|
"perechen": "Перечень элементов",
|
|
"specification_pcb": "Спецификация печатной платы",
|
|
"specification": "Спецификация",
|
|
"vedomost": "Ведомость покупных изделий",
|
|
"simple_list": "Бланк заказа",
|
|
}
|
|
|
|
def estimate_tokens(text: str) -> int:
|
|
"""Rough token estimate (works for mixed RU/EN)."""
|
|
if not text:
|
|
return 0
|
|
return max(1, (len(text) + 3) // 4)
|
|
|
|
|
|
def _extract_json(text: str) -> dict[str, Any]:
|
|
text = text.strip()
|
|
if text.startswith("```"):
|
|
text = re.sub(r"^```(?:json)?\s*", "", text)
|
|
text = re.sub(r"\s*```$", "", text)
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError:
|
|
m = re.search(r"\{.*\}", text, re.S)
|
|
if m:
|
|
return json.loads(m.group(0))
|
|
return {"reply": text, "edits": []}
|
|
|
|
|
|
def compress_rows(rows: list[dict[str, Any]], max_rows: int) -> list[dict[str, Any]]:
|
|
slim = []
|
|
for r in rows[:max_rows]:
|
|
item = {"row_id": r.get("id"), "row_index": r.get("row_index")}
|
|
for k, v in r.items():
|
|
if k in ("id", "project_id", "is_auto_generated", "stretch", "overflow_fields"):
|
|
continue
|
|
if v not in ("", None, False, 0) or k in ("is_header", "is_empty"):
|
|
item[k] = v
|
|
slim.append(item)
|
|
return slim
|
|
|
|
|
|
def build_context_usage(
|
|
messages: list[dict[str, str]],
|
|
*,
|
|
limit: int | None = None,
|
|
) -> dict[str, Any]:
|
|
settings = get_settings()
|
|
limit = limit or settings.llm_context_limit
|
|
used = sum(estimate_tokens(m.get("content") or "") for m in messages)
|
|
# role labels / framing overhead
|
|
used += len(messages) * 4
|
|
percent = min(100.0, round(100.0 * used / limit, 1)) if limit else 0.0
|
|
return {"used": used, "limit": limit, "percent": percent}
|
|
|
|
|
|
def build_llm_messages(
|
|
table_type: str,
|
|
rows: list[dict[str, Any]],
|
|
history: list[dict[str, str]],
|
|
) -> list[dict[str, str]]:
|
|
"""history: prior+current turns as {role, content}, roles user|assistant."""
|
|
settings = get_settings()
|
|
snapshot = compress_rows(rows, settings.llm_max_context_rows)
|
|
table_label = TABLE_TYPE_LABELS.get(table_type, table_type)
|
|
system = (
|
|
f"{SYSTEM_PROMPT}\n"
|
|
f"Текущая выбранная таблица: «{table_label}» (код: {table_type})."
|
|
)
|
|
snapshot_msg = (
|
|
f"Текущая выбранная таблица: «{table_label}» ({table_type}).\n"
|
|
f"Актуальный снимок строк (до {settings.llm_max_context_rows}). "
|
|
f"Ориентируйся на row_id отсюда:\n"
|
|
f"{json.dumps(snapshot, ensure_ascii=False)}"
|
|
)
|
|
|
|
messages: list[dict[str, str]] = [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": snapshot_msg},
|
|
{
|
|
"role": "assistant",
|
|
"content": json.dumps(
|
|
{
|
|
"reply": f"Снимок таблицы «{table_label}» принят. Жду команды.",
|
|
"edits": [],
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
]
|
|
|
|
# Keep dialogue; drop oldest if over budget (keep system+snapshot trio + newest)
|
|
dialog = [{"role": m["role"], "content": m["content"]} for m in history if m.get("content")]
|
|
limit = settings.llm_context_limit
|
|
reserved = sum(estimate_tokens(m["content"]) for m in messages) + 500
|
|
budget = max(2000, limit - reserved)
|
|
|
|
trimmed: list[dict[str, str]] = []
|
|
used = 0
|
|
for m in reversed(dialog):
|
|
cost = estimate_tokens(m["content"]) + 4
|
|
if used + cost > budget and trimmed:
|
|
break
|
|
trimmed.append(m)
|
|
used += cost
|
|
trimmed.reverse()
|
|
messages.extend(trimmed)
|
|
return messages
|
|
|
|
|
|
async def chat_edit_table(
|
|
message: str,
|
|
table_type: str,
|
|
rows: list[dict[str, Any]],
|
|
history: list[dict[str, str]] | None = None,
|
|
) -> dict[str, Any]:
|
|
settings = get_settings()
|
|
if not settings.openrouter_api_key:
|
|
empty_msgs = build_llm_messages(table_type, rows, history or [{"role": "user", "content": message}])
|
|
return {
|
|
"reply": "OpenRouter API key не настроен (OPENROUTER_API_KEY).",
|
|
"edits": [],
|
|
"context": build_context_usage(empty_msgs),
|
|
}
|
|
|
|
hist = list(history or [])
|
|
if not hist or hist[-1].get("role") != "user" or hist[-1].get("content") != message:
|
|
hist.append({"role": "user", "content": message})
|
|
|
|
messages = build_llm_messages(table_type, rows, hist)
|
|
context = build_context_usage(messages)
|
|
|
|
payload = {
|
|
"model": settings.openrouter_model,
|
|
"messages": messages,
|
|
"temperature": 0.2,
|
|
"max_tokens": 2000,
|
|
}
|
|
headers = {
|
|
"Authorization": f"Bearer {settings.openrouter_api_key}",
|
|
"Content-Type": "application/json",
|
|
"HTTP-Referer": "https://gostgenerator.local",
|
|
"X-Title": "GostGenerator Web",
|
|
}
|
|
url = settings.openrouter_base_url.rstrip("/") + "/chat/completions"
|
|
|
|
proxy = (settings.openrouter_proxy or "").strip() or None
|
|
if proxy and "://" not in proxy:
|
|
proxy = f"socks5://{proxy}"
|
|
|
|
client_kwargs: dict[str, Any] = {"timeout": 90.0}
|
|
if proxy:
|
|
client_kwargs["proxy"] = proxy
|
|
|
|
async with httpx.AsyncClient(**client_kwargs) as client:
|
|
resp = await client.post(url, headers=headers, json=payload)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
content = data["choices"][0]["message"]["content"]
|
|
parsed = _extract_json(content)
|
|
edits = parsed.get("edits") or []
|
|
norm_edits = []
|
|
for e in edits:
|
|
if not isinstance(e, dict):
|
|
continue
|
|
norm_edits.append(
|
|
{
|
|
"op": e.get("op", "update_row"),
|
|
"row_id": e.get("row_id"),
|
|
"row_index": e.get("row_index"),
|
|
"fields": e.get("fields") or {},
|
|
}
|
|
)
|
|
|
|
# include assistant reply in usage estimate
|
|
context = build_context_usage(messages + [{"role": "assistant", "content": content}])
|
|
usage = data.get("usage") or {}
|
|
if usage.get("prompt_tokens"):
|
|
prompt_t = int(usage["prompt_tokens"])
|
|
completion_t = int(usage.get("completion_tokens") or 0)
|
|
total_t = int(usage.get("total_tokens") or (prompt_t + completion_t))
|
|
limit = settings.llm_context_limit
|
|
context = {
|
|
"used": total_t,
|
|
"limit": limit,
|
|
"percent": min(100.0, round(100.0 * total_t / limit, 1)) if limit else 0.0,
|
|
"prompt_tokens": prompt_t,
|
|
"completion_tokens": completion_t,
|
|
}
|
|
|
|
return {"reply": parsed.get("reply") or "", "edits": norm_edits, "context": context}
|
|
|
|
|
|
def estimate_chat_context(
|
|
table_type: str,
|
|
rows: list[dict[str, Any]],
|
|
history: list[dict[str, str]],
|
|
) -> dict[str, Any]:
|
|
messages = build_llm_messages(table_type, rows, history)
|
|
return build_context_usage(messages)
|