busya
590b09f587
Replace src.core.auth.jwt import with lightweight _jwt_decoder.py that uses jose.jwt directly with JWT_SECRET env var. Avoids pulling AuthConfig → AUTH_DATABASE_URL/SECRET_KEY validators → SQLAlchemy models into agent. Removed: AUTH_SECRET_KEY, AUTH_DATABASE_URL from agent compose env. Removed: src/core/auth/ copies from Dockerfile.agent COPY section. Added: src/agent/_jwt_decoder.py — stateless JWT validation.
462 lines
19 KiB
Python
462 lines
19 KiB
Python
# backend/src/agent/_persistence.py
|
|
# #region AgentChat.Persistence [C:3] [TYPE Module] [SEMANTICS agent-chat,persistence,save,prefetch,title]
|
|
# @defgroup AgentChat Conversation persistence helpers — save, clean titles, LLM title generation, prefetch.
|
|
# @LAYER Service
|
|
# @RELATION DEPENDS_ON -> [AgentChat.Context]
|
|
# @RATIONALE Persistence logic is extracted to keep the handler under 400 lines and avoid
|
|
# mixing HTTP concerns with streaming logic.
|
|
|
|
import asyncio
|
|
from datetime import datetime
|
|
import os
|
|
import re
|
|
from typing import Any
|
|
import uuid
|
|
|
|
import httpx
|
|
|
|
from src.agent._config import AGENT_PREFETCH_DASHBOARD_LIMIT as _PREFETCH_LIMIT, FASTAPI_URL, SERVICE_JWT as _SERVICE_JWT
|
|
from src.agent._llm_params import add_temperature_if_supported
|
|
from src.core.logger import logger
|
|
|
|
SAVE_API_URL = FASTAPI_URL + "/api/agent/conversations/save"
|
|
TITLE_MAX_LENGTH = 80
|
|
|
|
# ── Rule-based title cleaning ────────────────────────────────────
|
|
|
|
|
|
# #region AgentChat.Persistence.CleanTitle [C:2] [TYPE Function] [SEMANTICS agent-chat,persistence,title]
|
|
# @ingroup AgentChat
|
|
# @BRIEF Clean raw user text into a readable conversation title — strip file markers,
|
|
# pre-fetch blocks, JSON/CSV, URLs; truncate to 80 chars at word boundary.
|
|
# @POST Returns cleaned title string; "Новый диалог" for empty/unparseable input.
|
|
# @RATIONALE Raw user_text often contains file content, pre-fetched data, or pasted JSON/CSV
|
|
# that produces unreadable titles. Rule-based cleaning is instant (no LLM latency)
|
|
# and handles 90% of cases. LLM titling is a best-effort async layer on top.
|
|
# @REJECTED Truncating raw text without cleaning was rejected — produces titles like
|
|
# "--- Uploaded file content --- id,name,status 1,Dashboard A..."
|
|
def clean_title(user_text: str) -> str: # noqa: C901
|
|
if not user_text or not user_text.strip():
|
|
return "Новый диалог"
|
|
|
|
text = user_text.strip()
|
|
|
|
# ── Already a HITL title? Skip cleaning ──
|
|
if text.startswith("✅ ") or text.startswith("⏹️ "):
|
|
return text[:TITLE_MAX_LENGTH]
|
|
|
|
# ── Strip file upload markers (cut before first occurrence) ──
|
|
file_markers = [
|
|
"\n--- Uploaded file content ---",
|
|
"--- Uploaded file content ---", # no leading newline
|
|
"\n[PRE-FETCHED DATA",
|
|
"[PRE-FETCHED DATA", # no leading newline
|
|
"\n[/PRE-FETCHED DATA]",
|
|
"[/PRE-FETCHED DATA]", # no leading newline
|
|
]
|
|
cut_pos = len(text)
|
|
for marker in file_markers:
|
|
pos = text.find(marker)
|
|
if pos != -1 and pos < cut_pos:
|
|
cut_pos = pos
|
|
if cut_pos < len(text):
|
|
text = text[:cut_pos].strip()
|
|
|
|
if not text:
|
|
return "Новый диалог"
|
|
|
|
# ── Take first meaningful segment ──
|
|
# Priority: first sentence ending with .!?, else first line, else full text
|
|
sentence_end = -1
|
|
for m in re.finditer(r"[.!?]\s", text):
|
|
sentence_end = m.start()
|
|
break
|
|
if sentence_end > 3: # don't cut "Привет." into "Привет"
|
|
text = text[: sentence_end + 1].strip()
|
|
elif "\n" in text:
|
|
text = text.split("\n")[0].strip()
|
|
|
|
if not text:
|
|
return "Новый диалог"
|
|
|
|
# ── Detect structured data (JSON/CSV/URL) ──
|
|
if text.startswith("{") or text.startswith("["):
|
|
prefix = "Данные: "
|
|
inner = text[1:57].strip().rstrip(",")
|
|
return prefix + inner + ("…" if len(text) > 60 else "")
|
|
if text.startswith("http://") or text.startswith("https://"):
|
|
try:
|
|
from urllib.parse import urlparse
|
|
|
|
domain = urlparse(text).netloc or "ссылка"
|
|
except Exception:
|
|
domain = "ссылка"
|
|
return domain
|
|
|
|
# ── Detect code (starts with def/class/import) ──
|
|
if any(text.startswith(kw) for kw in ("def ", "class ", "import ", "from ")):
|
|
first_line = text.split("\n")[0].strip()
|
|
return first_line[:TITLE_MAX_LENGTH]
|
|
|
|
# ── Hard truncate at word boundary ──
|
|
if len(text) > TITLE_MAX_LENGTH:
|
|
cut = text.rfind(" ", 0, TITLE_MAX_LENGTH)
|
|
if cut == -1:
|
|
cut = TITLE_MAX_LENGTH - 1
|
|
text = text[:cut].rstrip(".,;:!?") + "…"
|
|
|
|
# ── Fallback for empty/whitespace ──
|
|
if not text.strip():
|
|
return "Новый диалог"
|
|
|
|
return text
|
|
|
|
|
|
# #endregion AgentChat.Persistence.CleanTitle
|
|
|
|
|
|
# #region AgentChat.Persistence.DetectState [C:2] [TYPE Function] [SEMANTICS agent-chat,persistence,error-detection]
|
|
# @ingroup AgentChat
|
|
# @BRIEF Detect error/cancelled state from assistant message text.
|
|
# @POST Returns "error", "cancelled", or None.
|
|
def detect_message_state(text: str) -> str | None:
|
|
t = text.lower() if text else ""
|
|
error_markers = ["недоступен", "unavailable", "ошибка", "error", "произошла", "try again"]
|
|
cancel_markers = ["отменен", "cancelled", "отклонен", "denied"]
|
|
if any(m in t for m in cancel_markers):
|
|
return "cancelled"
|
|
if any(m in t for m in error_markers):
|
|
return "error"
|
|
return None
|
|
|
|
|
|
# #endregion AgentChat.Persistence.DetectState
|
|
|
|
|
|
# #region AgentChat.Persistence.ExtractUserId [C:2] [TYPE Function] [SEMANTICS agent-chat,persistence,auth]
|
|
# @ingroup AgentChat
|
|
# @BRIEF Extract user ID from JWT payload.
|
|
# @POST Returns user_id string or "unknown".
|
|
def extract_user_id(jwt_str: str) -> str:
|
|
try:
|
|
from src.agent._jwt_decoder import decode_token
|
|
|
|
payload = decode_token(jwt_str)
|
|
return payload.get("sub", payload.get("user_id", "unknown"))
|
|
except Exception:
|
|
return "unknown"
|
|
|
|
|
|
# #endregion AgentChat.Persistence.ExtractUserId
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# LLM title generation (best-effort, async)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
# Per-conversation lock to prevent concurrent title generation
|
|
_title_locks: dict[str, asyncio.Lock] = {}
|
|
|
|
|
|
async def _get_llm_config() -> dict[str, Any] | None:
|
|
"""Fetch LLM provider config from FastAPI for title generation."""
|
|
try:
|
|
fastapi_url = os.getenv("FASTAPI_URL", "http://localhost:8000")
|
|
service_token = os.getenv("SERVICE_JWT", "")
|
|
headers = {"Content-Type": "application/json"}
|
|
if service_token:
|
|
headers["Authorization"] = f"Bearer {service_token}"
|
|
|
|
async with httpx.AsyncClient(timeout=5) as client:
|
|
resp = await client.get(f"{fastapi_url}/api/agent/llm-config", headers=headers)
|
|
if resp.status_code == 200:
|
|
return resp.json()
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
async def _call_llm_for_title(user_text: str) -> str | None:
|
|
"""Call LLM to generate a 3-5 word Russian title. Returns title or None on failure."""
|
|
from src.core.logger import logger as _logger
|
|
|
|
try:
|
|
config = await _get_llm_config()
|
|
if not config or not config.get("configured"):
|
|
_logger.explore("LLM title: no provider configured", extra={"src": "AgentChat.Persistence"})
|
|
return None
|
|
|
|
clean_text = clean_title(user_text)[:200]
|
|
# Skip LLM for trivial titles
|
|
if not clean_text or clean_text in ("Новый диалог",):
|
|
return None
|
|
|
|
prompt = f"Сгенерируй заголовок из 3-5 слов на русском для диалога. Только заголовок, без кавычек и пояснений.\n\nДиалог: {clean_text}"
|
|
|
|
api_key = config.get("api_key", "")
|
|
base_url = config.get("base_url", "")
|
|
model = config.get("default_model", "gpt-4o-mini")
|
|
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 15,
|
|
}
|
|
add_temperature_if_supported(payload, model=model)
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
# Normalise base_url: strip trailing /v1 to avoid double /v1
|
|
base = base_url.rstrip("/")
|
|
if base.endswith("/v1"):
|
|
base = base[:-3]
|
|
api_url = base + "/v1/chat/completions"
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
resp = await client.post(api_url, json=payload, headers=headers)
|
|
if resp.status_code != 200:
|
|
_logger.explore(
|
|
"LLM title: API error",
|
|
payload={"status": resp.status_code, "reason": resp.reason_phrase},
|
|
error=f"HTTP {resp.status_code}: {resp.text[:100]}",
|
|
extra={"src": "AgentChat.Persistence"},
|
|
)
|
|
return None
|
|
data = resp.json()
|
|
title = data.get("choices", [{}])[0].get("message", {}).get("content", "")
|
|
if title:
|
|
# Strip markdown/formatting
|
|
title = re.sub(r'[*_`#"\']', "", title).strip()
|
|
title = title[:100] # safety cap
|
|
if title:
|
|
return title
|
|
except Exception as e:
|
|
_logger.explore("LLM title generation failed", error=str(e), extra={"src": "AgentChat.Persistence"})
|
|
return None
|
|
|
|
|
|
# #region AgentChat.Persistence.GenerateLlmTitle [C:3] [TYPE Function] [SEMANTICS agent-chat,persistence,llm,title]
|
|
# @ingroup AgentChat
|
|
# @BRIEF Async best-effort LLM title generation — patches conversation title via REST.
|
|
# @PRE conversation_id exists. LLM provider configured (best-effort, skips otherwise).
|
|
# @POST Conversation title updated via PATCH if LLM succeeds; no-op on failure.
|
|
# @SIDE_EFFECT HTTP PATCH to FastAPI save endpoint; may call external LLM API.
|
|
# @RATIONALE LLM-generated titles are more readable than rule-based ones (e.g. "CSV-файл:
|
|
# Анализ дашбордов" vs "Проанализируй CSV"). This is a non-blocking best-effort
|
|
# layer — the rule-based title is already saved, LLM just improves it.
|
|
async def generate_llm_title(conv_id: str, user_text: str) -> None:
|
|
if not conv_id or not user_text:
|
|
return
|
|
|
|
# Per-conversation dedup lock
|
|
lock = _title_locks.setdefault(conv_id, asyncio.Lock())
|
|
if lock.locked():
|
|
return # already generating for this conversation
|
|
async with lock:
|
|
title = await _call_llm_for_title(user_text)
|
|
if not title:
|
|
return
|
|
|
|
# Patch the title via the same save endpoint
|
|
try:
|
|
headers = {"Content-Type": "application/json"}
|
|
if _SERVICE_JWT:
|
|
headers["Authorization"] = f"Bearer {_SERVICE_JWT}"
|
|
payload = {
|
|
"conversation_id": conv_id,
|
|
"title": title,
|
|
"user_id": "admin",
|
|
"messages": [],
|
|
}
|
|
async with httpx.AsyncClient(timeout=5) as client:
|
|
await client.post(SAVE_API_URL, json=payload, headers=headers)
|
|
logger.reflect(
|
|
"LLM title updated",
|
|
payload={"conv_id": conv_id, "title": title[:40]},
|
|
extra={"src": "AgentChat.Persistence"},
|
|
)
|
|
except Exception as e:
|
|
logger.explore(
|
|
"LLM title save failed",
|
|
payload={"conv_id": conv_id},
|
|
error=str(e),
|
|
extra={"src": "AgentChat.Persistence"},
|
|
)
|
|
finally:
|
|
_title_locks.pop(conv_id, None)
|
|
|
|
|
|
# #endregion AgentChat.Persistence.GenerateLlmTitle
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Prefetch dashboards
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
# #region AgentChat.Persistence.PrefetchDashboards [C:3] [TYPE Function] [SEMANTICS agent-chat,persistence,prefetch]
|
|
# @ingroup AgentChat
|
|
# @BRIEF Pre-fetch dashboard data so the LLM has it in context without needing to call a tool.
|
|
# @POST Returns formatted dashboard list string, or empty string on failure.
|
|
# @SIDE_EFFECT Makes HTTP GET to FastAPI /api/dashboards.
|
|
# @RATIONALE Some LLMs (gemma) don't call tools even when instructed. Pre-fetch ensures
|
|
# dashboard data is available in context without requiring a tool call.
|
|
async def prefetch_dashboards(env_id: str) -> str:
|
|
try:
|
|
from src.agent.tools import FASTAPI_URL, _dual_auth_headers
|
|
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
resp = await client.get(
|
|
f"{FASTAPI_URL}/api/dashboards",
|
|
params={"q": "", "env_id": env_id or ""},
|
|
headers=_dual_auth_headers(),
|
|
)
|
|
if resp.status_code != 200:
|
|
return ""
|
|
data = resp.json()
|
|
dashboards = data.get("dashboards", [])
|
|
if not dashboards:
|
|
return "No dashboards found."
|
|
limit = _PREFETCH_LIMIT
|
|
total = len(dashboards)
|
|
lines = []
|
|
for db in dashboards[:limit]:
|
|
title = db.get("title", "Untitled")
|
|
dashboard_id = db.get("id") or db.get("dashboard_id")
|
|
modified = (db.get("last_modified", "") or "")[:10]
|
|
if modified:
|
|
lines.append(f"- {title} (id: {dashboard_id or 'n/a'}, modified: {modified})")
|
|
else:
|
|
lines.append(f"- {title} (id: {dashboard_id or 'n/a'})")
|
|
suffix = ""
|
|
if total > limit:
|
|
suffix = f"\n... {total - limit} more dashboards omitted. Ask for a narrower search if needed."
|
|
return f"Available dashboards in environment '{env_id or 'default'}' ({total} total):\n" + "\n".join(lines) + suffix
|
|
except Exception as e:
|
|
logger.explore(
|
|
"Prefetch dashboards failed",
|
|
payload={"env_id": env_id},
|
|
error=str(e),
|
|
extra={"src": "AgentChat.Persistence.PrefetchDashboards"},
|
|
)
|
|
return ""
|
|
|
|
|
|
# #endregion AgentChat.Persistence.PrefetchDashboards
|
|
|
|
|
|
# #region AgentChat.Persistence.PrefetchDatabases [C:2] [TYPE Function] [SEMANTICS agent-chat,persistence,prefetch,databases]
|
|
# @ingroup AgentChat
|
|
# @BRIEF Pre-fetch all databases for an environment into a formatted string for runtime context.
|
|
# @SIDE_EFFECT HTTP GET to FastAPI /api/agent/superset/databases.
|
|
# @RATIONALE Позволяет LLM видеть database_id для каждой БД, не вызывая отдельный инструмент.
|
|
async def prefetch_databases(env_id: str) -> str:
|
|
try:
|
|
from src.agent.tools import FASTAPI_URL, _dual_auth_headers
|
|
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
resp = await client.get(
|
|
f"{FASTAPI_URL}/api/agent/superset/databases",
|
|
params={"environment_id": env_id or ""},
|
|
headers=_dual_auth_headers(),
|
|
)
|
|
if resp.status_code != 200:
|
|
return ""
|
|
databases = resp.json()
|
|
if not databases:
|
|
return "No databases found."
|
|
lines = ["Available databases (use database_id for SQL tools):"]
|
|
for db in databases:
|
|
db_id = db.get("id", "?")
|
|
db_name = db.get("database_name", db.get("name", "?"))
|
|
db_engine = db.get("backend", db.get("engine", ""))
|
|
if db_engine:
|
|
lines.append(f" • DB #{db_id}: {db_name} ({db_engine})")
|
|
else:
|
|
lines.append(f" • DB #{db_id}: {db_name}")
|
|
return "\n".join(lines)
|
|
except Exception as e:
|
|
logger.explore(
|
|
"Prefetch databases failed",
|
|
payload={"env_id": env_id},
|
|
error=str(e),
|
|
extra={"src": "AgentChat.Persistence.PrefetchDatabases"},
|
|
)
|
|
return ""
|
|
|
|
|
|
# #endregion AgentChat.Persistence.PrefetchDatabases
|
|
|
|
|
|
# #region AgentChat.Persistence.SaveConversation [C:4] [TYPE Function] [SEMANTICS agent-chat,persistence,save]
|
|
# @ingroup AgentChat
|
|
# @BRIEF Save conversation to DB via FastAPI REST. Cleans title via clean_title().
|
|
# @PRE conversation_id is valid. FASTAPI_URL reachable.
|
|
# @POST Conversation + messages saved via POST /api/agent/conversations/save.
|
|
# @SIDE_EFFECT HTTP POST to FastAPI; writes to AgentConversation and AgentMessage tables.
|
|
# @DATA_CONTRACT Input: (conv_id, user_text, user_id, assistant_text) -> Output: None (side-effect only)
|
|
# @RELATION DISPATCHES -> [AgentChat.Api.Conversations]
|
|
# @RATIONALE Anonymous Gradio sessions (anon_ prefix) default to "admin" because the agent runs in an internal network behind an auth proxy — all users within the network are trusted.
|
|
# @REJECTED Requiring explicit authentication for Gradio was rejected — the agent is designed for internal-network use where the auth proxy handles auth; adding a separate auth layer would create unnecessary friction and duplicate the proxy's responsibility.
|
|
async def save_conversation(conv_id: str, user_text: str, user_id: str = "admin", assistant_text: str = "") -> None:
|
|
try:
|
|
headers = {"Content-Type": "application/json"}
|
|
if _SERVICE_JWT:
|
|
headers["Authorization"] = f"Bearer {_SERVICE_JWT}"
|
|
|
|
# Normalize user_id: anonymous Gradio sessions use "anon_" prefix
|
|
if not user_id or user_id.startswith("anon_"):
|
|
user_id = "admin"
|
|
|
|
messages: list[dict[str, Any]] = [
|
|
{
|
|
"id": str(uuid.uuid4()),
|
|
"conversation_id": conv_id,
|
|
"role": "user",
|
|
"text": user_text.strip(),
|
|
"state": None,
|
|
"created_at": datetime.utcnow().isoformat(),
|
|
},
|
|
]
|
|
if assistant_text.strip():
|
|
assistant_state = detect_message_state(assistant_text)
|
|
messages.append(
|
|
{
|
|
"id": str(uuid.uuid4()),
|
|
"conversation_id": conv_id,
|
|
"role": "assistant",
|
|
"text": assistant_text.strip(),
|
|
"state": assistant_state,
|
|
"created_at": datetime.utcnow().isoformat(),
|
|
}
|
|
)
|
|
|
|
title = clean_title(user_text)
|
|
payload = {
|
|
"conversation_id": conv_id,
|
|
"title": title,
|
|
"user_id": user_id,
|
|
"messages": messages,
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
await client.post(SAVE_API_URL, json=payload, headers=headers)
|
|
logger.reflect(
|
|
"Conversation saved",
|
|
payload={"conv_id": conv_id, "msg_count": len(messages), "title": title[:40]},
|
|
extra={"src": "AgentChat.Persistence.SaveConversation"},
|
|
)
|
|
except Exception as e:
|
|
logger.explore(
|
|
"Failed to save conversation",
|
|
payload={"conv_id": conv_id},
|
|
error=str(e),
|
|
extra={"src": "AgentChat.Persistence"},
|
|
)
|
|
|
|
|
|
# #endregion AgentChat.Persistence.SaveConversation
|
|
# #endregion AgentChat.Persistence
|