fix(security): resolve Critical+High findings from module audit — agent, translate, superset_client
P0 — CRITICAL (CWE-798): JWT_SECRET crash-early
Replace hardcoded super-secret-key fallback with os.environ["JWT_SECRET"]
and ${JWT_SECRET:?} syntax in app.py + docker-compose files
P1 — HIGH: Frontend dependency CVEs
Upgrade svelte 5.43.8 → 5.56.4 — resolves devalue DoS (GHSA-g2pg-6438-jwpf)
and svelte XSS (GHSA-crpf-4hrx-3jrp, GHSA-m56q-vw4c-c2cp, GHSA-rcqx-6q8c-2c42)
P2 — MEDIUM: Logging hygiene + contract gaps + tool resolver refactor
Apply _redact_sensitive_fields() in middleware + event streaming
Truncate LLM error body to 100 chars
Add @RATIONALE/@REJECTED to HandleResume + SaveConversation
Refactor deterministic intent matching → LLM-driven tool resolution
P3 — LOW: Translate logging hardening
Move _sanitize_url() to _utils.py (shared, no circular imports)
Sanitize base_url before logging in _llm_call.py and _llm_async_http.py
Emit EXPLORE warning when LLM_SSL_VERIFY=false disables TLS
superset_client module: passed clean — no changes needed
This commit is contained in:
@@ -15,9 +15,6 @@ from typing import Any
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
from src.agent._tool_resolver import (
|
||||
_SAFE_AGENT_TOOLS,
|
||||
_DANGEROUS_AGENT_TOOLS,
|
||||
_GUARDED_AGENT_TOOLS,
|
||||
normalize_tool_args,
|
||||
extract_tool_call_from_state,
|
||||
find_tool,
|
||||
@@ -33,23 +30,19 @@ _pending_confirmations: dict[str, dict[str, Any]] = {}
|
||||
# @BRIEF Build confirmation contract dict — risk level, prompt, operation metadata.
|
||||
# @POST Returns dict with operation, risk, risk_level, prompt, requires_confirmation keys.
|
||||
def build_confirmation_contract(tool_name: str | None) -> dict[str, Any]:
|
||||
"""Build confirmation contract — risk classification heuristic.
|
||||
LLM handles intent; tools are classified by name prefix for HITL UX."""
|
||||
operation = tool_name or "unknown_action"
|
||||
if operation in _SAFE_AGENT_TOOLS:
|
||||
risk_level = "safe"
|
||||
risk = "read"
|
||||
prompt = "Разрешить чтение данных?"
|
||||
elif operation in _DANGEROUS_AGENT_TOOLS:
|
||||
risk_level = "dangerous"
|
||||
risk = "write"
|
||||
prompt = "Подтвердить критичную операцию?"
|
||||
elif operation in _GUARDED_AGENT_TOOLS:
|
||||
# Guard heuristic: deploy_*, execute_*, create_*, run_*, commit_*, start_*, end_*
|
||||
_guarded_prefixes = ("deploy", "execute", "create", "run", "commit", "start", "end")
|
||||
if any(operation.startswith(p) for p in _guarded_prefixes):
|
||||
risk_level = "guarded"
|
||||
risk = "write"
|
||||
prompt = "Подтвердить изменение данных?"
|
||||
else:
|
||||
risk_level = "unknown"
|
||||
risk = "unknown"
|
||||
prompt = "Подтвердите действие"
|
||||
risk_level = "safe"
|
||||
risk = "read"
|
||||
prompt = "Разрешить чтение данных?"
|
||||
|
||||
return {
|
||||
"operation": operation,
|
||||
@@ -194,6 +187,8 @@ async def _format_tool_output_via_llm(
|
||||
# @SIDE_EFFECT Invokes LangChain tools; modifies _pending_confirmations dict.
|
||||
# @RELATION DEPENDS_ON -> [AgentChat.LangGraph.Setup]
|
||||
# @DATA_CONTRACT Input: (conv_id, action, user_jwt, env_id) -> Output: AsyncGenerator[str]
|
||||
# @RATIONALE Fast-path resume (direct tool execution without re-entering LangGraph) chosen because the HITL checkpoint already contains all necessary context — re-running the agent would be redundant and slow.
|
||||
# @REJECTED Pure streaming without checkpoint was rejected — without a persisted checkpoint, a crash after confirmation but before tool execution would lose the operation entirely with no rollback capability.
|
||||
async def handle_resume(
|
||||
conversation_id: str, action: str,
|
||||
user_jwt: str = "", env_id: str | None = None,
|
||||
|
||||
155
backend/src/agent/_embedding_router.py
Normal file
155
backend/src/agent/_embedding_router.py
Normal file
@@ -0,0 +1,155 @@
|
||||
# backend/src/agent/_embedding_router.py
|
||||
# #region AgentChat.EmbeddingRouter [C:3] [TYPE Module] [SEMANTICS agent-chat,tools,embedding,fallback]
|
||||
# @defgroup AgentChat Embedding-based tool router — fallback when keyword matching yields <3 tools.
|
||||
# @LAYER Service
|
||||
# @BRIEF Lazy-loaded embedding model for cosine similarity between user query and tool descriptions.
|
||||
# @RATIONALE Tool descriptions are auto-generated from @tool docstrings (enforced by LangChain).
|
||||
# Optional _TOOL_DESCRIPTIONS_OVERRIDES in tools.py provides RU/EN synonyms for
|
||||
# tools where the docstring alone isn't descriptive enough. This eliminates the
|
||||
# hardcoded _TOOL_DESCRIPTIONS dict as a separate source of truth.
|
||||
# @REJECTED Hardcoded _TOOL_DESCRIPTIONS dict — drifts out of sync with get_all_tools().
|
||||
# @INVARIANT Descriptions are derived from get_all_tools() docstrings — always 1:1 with tools.
|
||||
# @INVARIANT embedding_top_k() returns empty list (never raises) when model unavailable.
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger("superset_tools_app")
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
# Tool description — auto-generated from @tool docstrings.
|
||||
# Override via _TOOL_DESCRIPTIONS_OVERRIDES in tools.py for RU/EN synonyms.
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
def _get_descriptions() -> tuple[list[str], list[str]]:
|
||||
"""Return (descriptions, tool_names) from get_all_tools() docstrings.
|
||||
Uses _TOOL_DESCRIPTIONS_OVERRIDES from tools.py for optional RU/EN synonyms.
|
||||
"""
|
||||
from src.agent.tools import get_all_tools, _TOOL_DESCRIPTIONS_OVERRIDES
|
||||
|
||||
all_tools = get_all_tools()
|
||||
names = []
|
||||
descriptions = []
|
||||
for tool_obj in all_tools:
|
||||
name = tool_obj.name
|
||||
names.append(name)
|
||||
desc = _TOOL_DESCRIPTIONS_OVERRIDES.get(name) or (tool_obj.description or "").strip()
|
||||
if not desc:
|
||||
desc = name # fallback: use tool name as description
|
||||
descriptions.append(desc)
|
||||
return descriptions, names
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
# Model state — lazy-loaded on first call
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
_embedding_model: Optional[object] = None
|
||||
_tool_embeddings: Optional[object] = None # torch.Tensor or numpy array
|
||||
_tool_names: list[str] = []
|
||||
|
||||
_THRESHOLD = float(os.getenv("EMBEDDING_SIMILARITY_THRESHOLD", "0.65"))
|
||||
_TOP_K = int(os.getenv("EMBEDDING_TOP_K", "5"))
|
||||
_MODEL_NAME = os.getenv(
|
||||
"EMBEDDING_MODEL",
|
||||
"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
|
||||
)
|
||||
|
||||
|
||||
def _load_model() -> bool:
|
||||
"""Lazy-load the embedding model and pre-embed tool descriptions.
|
||||
|
||||
Returns True on success, False on any failure (missing package, download error,
|
||||
OOM, etc.). On failure, the router degrades gracefully — embedding_top_k()
|
||||
returns an empty list and the caller falls back to keyword-only results.
|
||||
"""
|
||||
global _embedding_model, _tool_embeddings, _tool_names
|
||||
|
||||
if _embedding_model is not None:
|
||||
return True
|
||||
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"sentence-transformers not installed — embedding router disabled. "
|
||||
"Install with: pip install sentence-transformers"
|
||||
)
|
||||
return False
|
||||
|
||||
try:
|
||||
logger.info("Loading embedding model: %s", _MODEL_NAME)
|
||||
_embedding_model = SentenceTransformer(_MODEL_NAME)
|
||||
|
||||
descriptions, _tool_names[:] = _get_descriptions()
|
||||
|
||||
_tool_embeddings = _embedding_model.encode(
|
||||
descriptions,
|
||||
convert_to_tensor=True,
|
||||
show_progress_bar=False,
|
||||
)
|
||||
logger.info(
|
||||
"Embedding model loaded. Tools: %d, model: %s",
|
||||
len(_tool_names), _MODEL_NAME,
|
||||
)
|
||||
return True
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to load embedding model '%s': %s — embedding router disabled",
|
||||
_MODEL_NAME, exc,
|
||||
)
|
||||
_embedding_model = None
|
||||
return False
|
||||
|
||||
|
||||
def embedding_top_k(query: str) -> list[str]:
|
||||
"""Return top-K tool names above cosine similarity threshold.
|
||||
|
||||
Args:
|
||||
query: Raw user query text (any language).
|
||||
|
||||
Returns:
|
||||
List of tool name strings ordered by descending similarity.
|
||||
Empty list if model unavailable, no descriptions match, or an error occurs.
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
return []
|
||||
|
||||
if not _load_model():
|
||||
return []
|
||||
|
||||
try:
|
||||
from sentence_transformers.util import cos_sim
|
||||
|
||||
query_embedding = _embedding_model.encode(
|
||||
query,
|
||||
convert_to_tensor=True,
|
||||
show_progress_bar=False,
|
||||
)
|
||||
|
||||
similarities = cos_sim(query_embedding, _tool_embeddings)[0]
|
||||
|
||||
results: list[str] = []
|
||||
for idx in similarities.argsort(descending=True)[:_TOP_K]:
|
||||
score = float(similarities[idx])
|
||||
if score >= _THRESHOLD:
|
||||
results.append(_tool_names[idx])
|
||||
|
||||
if results:
|
||||
logger.debug(
|
||||
"Embedding router: query='%s' → %d tools above %.2f: %s",
|
||||
query[:80], len(results), _THRESHOLD, results,
|
||||
)
|
||||
|
||||
return results
|
||||
except Exception as exc:
|
||||
logger.warning("Embedding router failed for query '%s': %s", query[:80], exc)
|
||||
return []
|
||||
|
||||
|
||||
def embedding_is_available() -> bool:
|
||||
"""Check if embedding model is loaded and ready (non-blocking)."""
|
||||
return _load_model()
|
||||
|
||||
|
||||
# #endregion AgentChat.EmbeddingRouter
|
||||
@@ -213,7 +213,7 @@ async def _call_llm_for_title(user_text: str) -> str | None:
|
||||
if resp.status_code != 200:
|
||||
_logger.explore(
|
||||
"LLM title: API error",
|
||||
payload={"status": resp.status_code}, error=resp.text[:200],
|
||||
payload={"status": resp.status_code, "reason": resp.reason_phrase}, error=f"HTTP {resp.status_code}: {resp.text[:100]}",
|
||||
extra={"src": "AgentChat.Persistence"},
|
||||
)
|
||||
return None
|
||||
@@ -341,6 +341,8 @@ async def prefetch_dashboards(env_id: str) -> str:
|
||||
# @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 -> [Api.Agent.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:
|
||||
service_token = os.getenv("SERVICE_JWT", "")
|
||||
|
||||
@@ -1,79 +1,36 @@
|
||||
# backend/src/agent/_tool_resolver.py
|
||||
# #region AgentChat.ToolResolver [C:3] [TYPE Module] [SEMANTICS agent-chat,tools,classification,resolution]
|
||||
# @defgroup AgentChat Tool classification constants and resolution helpers for the LangGraph agent.
|
||||
# #region AgentChat.ToolResolver [C:2] [TYPE Module] [SEMANTICS agent-chat,tools,resolution]
|
||||
# @defgroup AgentChat Tool resolution helpers for the LangGraph agent.
|
||||
# @LAYER Service
|
||||
# @RELATION DEPENDS_ON -> [AgentChat.Tools]
|
||||
# @RATIONALE Centralised tool resolution prevents duplication of tool-name matching logic across
|
||||
# the handler and confirmation subsystems. Single source of truth for tool risk classification.
|
||||
# @RATIONALE Centralised tool resolution prevents duplication of tool-name matching logic.
|
||||
# Deterministic intent matching (infer_tool_from_text, fast_confirmation_tool,
|
||||
# keyword lists, negation guard, classification sets) removed — LLM handles
|
||||
# all intent detection through LangGraph tool-calling. Only utility helpers
|
||||
# (tool call coercion, args normalization) remain.
|
||||
# @REJECTED Deterministic intent matching — fragile substring collisions, maintenance burden
|
||||
# of 24 keyword lists across 3 files, negation blindness in fast-track, and
|
||||
# inability to handle synonyms ("панели"≠"дашборды") or typos ("дашборд").
|
||||
# @REJECTED Fast-track confirmation — bypasses LLM, causing negation blindness.
|
||||
# @REJECTED Tool risk classification sets — LLM decides which tools to call;
|
||||
# LangGraph interrupt_before handles HITL for dangerous tools at graph level.
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.agent.tools import get_all_tools
|
||||
from src.core.logger import logger
|
||||
|
||||
# #region AgentChat.ToolResolver.Sets [C:1] [TYPE Constants] [SEMANTICS agent-chat,tools,sets]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Tool classification sets — safe (read-only), guarded (write), dangerous (deploy).
|
||||
_SAFE_AGENT_TOOLS = {
|
||||
"show_capabilities",
|
||||
"search_dashboards",
|
||||
"get_health_summary",
|
||||
"list_environments",
|
||||
"get_task_status",
|
||||
"list_llm_providers",
|
||||
"get_llm_status",
|
||||
"list_maintenance_events",
|
||||
# NEW: read-only Superset tools
|
||||
"superset_execute_sql",
|
||||
"superset_explore_database",
|
||||
"superset_audit_permissions",
|
||||
"superset_format_sql",
|
||||
}
|
||||
_GUARDED_AGENT_TOOLS = {
|
||||
"create_branch",
|
||||
"commit_changes",
|
||||
"execute_migration",
|
||||
"run_backup",
|
||||
"run_llm_validation",
|
||||
"run_llm_documentation",
|
||||
"start_maintenance",
|
||||
"end_maintenance",
|
||||
# NEW: guarded Superset write operations
|
||||
"superset_create_dashboard",
|
||||
"superset_copy_dashboard",
|
||||
"superset_create_dataset",
|
||||
}
|
||||
_DANGEROUS_AGENT_TOOLS = {
|
||||
"deploy_dashboard",
|
||||
}
|
||||
# ── Graph nodes — used by confirmation subsystem to distinguish tools from infrastructure ──
|
||||
_GRAPH_NODE_NAMES = {"agent", "tools", "__start__", "__end__"}
|
||||
_FAST_CONFIRM_TOOLS = {
|
||||
"show_capabilities",
|
||||
"list_environments",
|
||||
"list_llm_providers",
|
||||
"get_llm_status",
|
||||
"list_maintenance_events",
|
||||
"superset_explore_database",
|
||||
"superset_audit_permissions",
|
||||
"superset_format_sql",
|
||||
}
|
||||
# #endregion AgentChat.ToolResolver.Sets
|
||||
|
||||
|
||||
# #region AgentChat.ToolResolver.KnownNames [C:2] [TYPE Function] [SEMANTICS agent-chat,tools,catalog]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Return registered LangChain tool names without letting helper failures break HITL UX.
|
||||
# @POST Returns set of tool name strings; falls back to static union on failure.
|
||||
# @BRIEF Return registered LangChain tool names.
|
||||
# @POST Returns set of tool name strings; falls back to empty set on failure.
|
||||
def known_agent_tool_names() -> set[str]:
|
||||
try:
|
||||
from src.agent.tools import get_all_tools
|
||||
return {str(tool_obj.name) for tool_obj in get_all_tools() if getattr(tool_obj, "name", None)}
|
||||
except Exception as exc:
|
||||
logger.explore(
|
||||
"tool catalog lookup failed",
|
||||
payload={"error": str(exc)},
|
||||
extra={"src": "AgentChat.ToolResolver"},
|
||||
)
|
||||
return _SAFE_AGENT_TOOLS | _GUARDED_AGENT_TOOLS | _DANGEROUS_AGENT_TOOLS
|
||||
except Exception:
|
||||
return set()
|
||||
# #endregion AgentChat.ToolResolver.KnownNames
|
||||
|
||||
|
||||
@@ -113,61 +70,11 @@ def coerce_tool_call(tool_call: Any) -> tuple[str | None, dict[str, Any]]:
|
||||
# #endregion AgentChat.ToolResolver.CoerceToolCall
|
||||
|
||||
|
||||
# #region AgentChat.ToolResolver.InferFromText [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,inference]
|
||||
# #region AgentChat.ToolResolver.ExtractFromState [C:2] [TYPE Function] [SEMANTICS agent-chat,tools,checkpoint]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Infer which tool the user likely wants based on keywords in the message text.
|
||||
# @POST Returns tool name string, or None if no match.
|
||||
# @RATIONALE Some LLMs fail to emit tool calls even when instructed. This fallback
|
||||
# uses keyword matching to guess the user's intent and auto-trigger HITL.
|
||||
def infer_tool_from_text(text: str) -> str | None:
|
||||
lowered = (text or "").lower()
|
||||
inferred: str | None = None
|
||||
if any(word in lowered for word in ["окруж", "environment", "env"]):
|
||||
inferred = "list_environments"
|
||||
elif any(word in lowered for word in ["maintenance", "обслуж", "баннер"]):
|
||||
if any(word in lowered for word in ["start", "созда", "запусти", "начни"]):
|
||||
inferred = "start_maintenance"
|
||||
elif any(word in lowered for word in ["end", "закрой", "заверши", "останов"]):
|
||||
inferred = "end_maintenance"
|
||||
else:
|
||||
inferred = "list_maintenance_events"
|
||||
elif any(word in lowered for word in ["дашборд", "dashboard", "dashboards", "дашборды"]):
|
||||
inferred = "search_dashboards"
|
||||
elif any(word in lowered for word in ["здоров", "health", "статус системы", "system status"]):
|
||||
inferred = "get_health_summary"
|
||||
elif any(word in lowered for word in ["задач", "task", "таск"]):
|
||||
inferred = "get_task_status"
|
||||
elif any(word in lowered for word in ["llm", "provider", "провайдер", "модель"]):
|
||||
inferred = "list_llm_providers"
|
||||
elif any(word in lowered for word in ["branch", "ветк"]):
|
||||
inferred = "create_branch"
|
||||
elif any(word in lowered for word in ["commit", "коммит"]):
|
||||
inferred = "commit_changes"
|
||||
elif any(word in lowered for word in ["deploy", "депло", "разверн"]):
|
||||
inferred = "deploy_dashboard"
|
||||
elif any(word in lowered for word in ["миграц", "migration", "migrate"]):
|
||||
inferred = "execute_migration"
|
||||
elif any(word in lowered for word in ["backup", "бэкап", "резерв"]):
|
||||
inferred = "run_backup"
|
||||
elif any(word in lowered for word in ["валидац", "validation", "validate"]):
|
||||
inferred = "run_llm_validation"
|
||||
elif any(word in lowered for word in ["документ", "documentation", "docs"]):
|
||||
inferred = "run_llm_documentation"
|
||||
elif any(word in lowered for word in ["инструмент", "tool", "capabilit", "умеешь", "можешь"]):
|
||||
inferred = "show_capabilities"
|
||||
|
||||
if inferred:
|
||||
logger.reason("Tool inferred from user text",
|
||||
payload={"tool": inferred, "text_preview": (text or "")[:80]},
|
||||
extra={"src": "AgentChat.ToolResolver.InferFromText"})
|
||||
return inferred
|
||||
# #endregion AgentChat.ToolResolver.InferFromText
|
||||
|
||||
|
||||
# #region AgentChat.ToolResolver.ExtractFromState [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,checkpoint]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Extract pending tool name and args from the LangGraph checkpoint; infer only as last resort.
|
||||
# @BRIEF Extract pending tool name and args from the LangGraph checkpoint.
|
||||
# @POST Returns (tool_name, args) tuple; (None, {}) if nothing found.
|
||||
# @RATIONALE LLM handles intent — no fallback to keyword inference.
|
||||
def extract_tool_call_from_state(state, user_text: str = "") -> tuple[str | None, dict[str, Any]]:
|
||||
known_tools = known_agent_tool_names()
|
||||
try:
|
||||
@@ -177,26 +84,14 @@ def extract_tool_call_from_state(state, user_text: str = "") -> tuple[str | None
|
||||
tool_name, tool_args = coerce_tool_call(msg.tool_calls[0])
|
||||
if tool_name:
|
||||
return (str(tool_name), tool_args)
|
||||
except Exception as exc:
|
||||
logger.explore(
|
||||
"tool_call extraction failed",
|
||||
payload={"error": str(exc)},
|
||||
extra={"src": "AgentChat.ToolResolver"},
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if getattr(state, "next", None):
|
||||
node_or_tool = str(state.next[0])
|
||||
if node_or_tool in known_tools and node_or_tool not in _GRAPH_NODE_NAMES:
|
||||
return (node_or_tool, {})
|
||||
|
||||
inferred_tool = infer_tool_from_text(user_text)
|
||||
if inferred_tool:
|
||||
logger.explore(
|
||||
"tool_call inferred from user text",
|
||||
payload={"tool": inferred_tool},
|
||||
extra={"src": "AgentChat.ToolResolver"},
|
||||
)
|
||||
return (inferred_tool, {})
|
||||
return (None, {})
|
||||
# #endregion AgentChat.ToolResolver.ExtractFromState
|
||||
|
||||
@@ -206,16 +101,8 @@ def extract_tool_call_from_state(state, user_text: str = "") -> tuple[str | None
|
||||
# @BRIEF Find a registered LangChain tool by name.
|
||||
# @POST Returns tool object or None if not found.
|
||||
def find_tool(tool_name: str):
|
||||
from src.agent.tools import get_all_tools
|
||||
return next((tool_obj for tool_obj in get_all_tools() if getattr(tool_obj, "name", None) == tool_name), None)
|
||||
# #endregion AgentChat.ToolResolver.FindTool
|
||||
|
||||
|
||||
# #region AgentChat.ToolResolver.FastConfirm [C:2] [TYPE Function] [SEMANTICS agent-chat,tools,fast-path]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Check if user text maps to a tool eligible for fast-track HITL confirmation.
|
||||
# @POST Returns tool name if match, None otherwise.
|
||||
def fast_confirmation_tool(text: str) -> str | None:
|
||||
tool_name = infer_tool_from_text(text)
|
||||
return tool_name if tool_name in _FAST_CONFIRM_TOOLS else None
|
||||
# #endregion AgentChat.ToolResolver.FastConfirm
|
||||
# #endregion AgentChat.ToolResolver
|
||||
|
||||
@@ -19,6 +19,7 @@ import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import time
|
||||
from typing import Any
|
||||
import uuid
|
||||
|
||||
@@ -27,6 +28,8 @@ import httpx
|
||||
from jose import JWTError
|
||||
from langchain_core.exceptions import OutputParserException
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langchain_openai import ChatOpenAI
|
||||
from openai import APIConnectionError, APITimeoutError, AuthenticationError
|
||||
|
||||
from src.agent._confirmation import (
|
||||
confirmation_metadata_for_tool,
|
||||
@@ -40,25 +43,32 @@ from src.agent._persistence import (
|
||||
save_conversation,
|
||||
generate_llm_title,
|
||||
)
|
||||
from src.agent._tool_resolver import (
|
||||
fast_confirmation_tool,
|
||||
)
|
||||
from src.agent.context import set_user_jwt
|
||||
from src.agent.document_parser import parse_upload
|
||||
from src.agent.langgraph_setup import create_agent
|
||||
from src.agent.middleware import log_tool_event
|
||||
from src.agent.tools import get_all_tools
|
||||
from src.agent.tools import _redact_sensitive_fields, get_all_tools
|
||||
from src.core.auth.jwt import decode_token
|
||||
from src.core.cot_logger import seed_trace_id
|
||||
from src.core.logger import logger
|
||||
|
||||
JWT_SECRET = os.getenv("JWT_SECRET", "super-secret-key")
|
||||
JWT_SECRET = os.environ["JWT_SECRET"] # @INVARIANT JWT_SECRET must be set in environment — crash-early, no default fallback
|
||||
MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 # 10 MB
|
||||
|
||||
# ── Session state ───────────────────────────────────────────────
|
||||
# In-memory per-user lock (keyed by user_id)
|
||||
_user_locks: dict[str, bool] = {}
|
||||
|
||||
# ── LLM provider health cache ─────────────────────────────────---
|
||||
_llm_status: dict[str, Any] = {
|
||||
"status": "ok",
|
||||
"last_error": "",
|
||||
"last_check_ts": 0.0,
|
||||
}
|
||||
_LLM_CHECK_CACHE_TTL = 30 # seconds between health checks
|
||||
_LLM_LAST_ERROR_KEY = "last_llm_error"
|
||||
_LLM_LAST_ERROR_TS_KEY = "last_llm_error_ts"
|
||||
|
||||
# ── File persistence ────────────────────────────────────────────
|
||||
|
||||
# #region AgentChat.GradioApp.PersistFile [C:3] [TYPE Function] [SEMANTICS agent-chat,storage,file]
|
||||
@@ -120,6 +130,67 @@ _service_jwt_cache: dict[str, str] = {}
|
||||
# the generator and sends yielded JSON strings as event data to the frontend.
|
||||
# @REJECTED Returning a single response (non-streaming) was rejected — violates FR-003 (streaming mandate).
|
||||
|
||||
# #region AgentChat.GradioApp.LlmHealthCheck [C:2] [TYPE Function] [SEMANTICS agent-chat,llm,health]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Check LLM provider connectivity with in-memory cache (30s TTL).
|
||||
# @POST Returns status string: 'ok' | 'unavailable' | 'timeout' | 'auth_error'.
|
||||
# @SIDE_EFFECT Makes a probe request to the LLM provider; caches result in module memory.
|
||||
# @RATIONALE Prevents sending every user request into an dead LLM backend.
|
||||
async def _check_llm_provider_health() -> str:
|
||||
"""Check LLM provider connectivity. Cached for _LLM_CHECK_CACHE_TTL seconds."""
|
||||
now = time.time()
|
||||
if now - _llm_status["last_check_ts"] < _LLM_CHECK_CACHE_TTL:
|
||||
return _llm_status["status"]
|
||||
|
||||
try:
|
||||
from src.agent.langgraph_setup import _fetch_llm_config as _get_config
|
||||
config = await _get_config()
|
||||
if not config or not config.get("configured"):
|
||||
return _llm_status["status"]
|
||||
|
||||
llm = ChatOpenAI(
|
||||
model=config.get("default_model", "gpt-4o-mini"),
|
||||
base_url=config.get("base_url", "https://api.openai.com/v1"),
|
||||
api_key=config.get("api_key", ""),
|
||||
temperature=0,
|
||||
max_tokens=1,
|
||||
)
|
||||
await llm.ainvoke([HumanMessage(content="ping")])
|
||||
_llm_status["status"] = "ok"
|
||||
_llm_status["last_error"] = ""
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
return "ok"
|
||||
except (APIConnectionError, httpx.ConnectError):
|
||||
_llm_status["status"] = "unavailable"
|
||||
_llm_status["last_error"] = "Connection refused"
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
return "unavailable"
|
||||
except (APITimeoutError, httpx.ReadTimeout):
|
||||
_llm_status["status"] = "timeout"
|
||||
_llm_status["last_error"] = "Request timed out"
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
return "timeout"
|
||||
except AuthenticationError:
|
||||
_llm_status["status"] = "auth_error"
|
||||
_llm_status["last_error"] = "Invalid API key"
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
return "auth_error"
|
||||
except Exception as exc:
|
||||
logger.explore("LLM health check failed",
|
||||
error=str(exc),
|
||||
extra={"src": "AgentChat.GradioApp.LlmHealthCheck"})
|
||||
return _llm_status["status"] # return cached status on unexpected errors
|
||||
# #endregion AgentChat.GradioApp.LlmHealthCheck
|
||||
|
||||
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Core streaming handler — runs LangGraph agent, yields ChatMessage tokens with structured metadata.
|
||||
# @PRE JWT valid, user authenticated.
|
||||
# @POST Tokens streamed via yield; HITL interrupts yield confirm_required metadata.
|
||||
# @SIDE_EFFECT Calls LLM, invokes tools, writes checkpoints.
|
||||
# @RATIONALE Async generator pattern chosen for Gradio ChatInterface compatibility — Gradio iterates
|
||||
# @REJECTED Returning a single response (non-streaming) was rejected — violates FR-003 (streaming mandate).
|
||||
|
||||
async def agent_handler( # noqa: C901 — intentionally complex C4 orchestration
|
||||
message,
|
||||
history: list, # noqa: ARG001 — Gradio ChatInterface requires this parameter
|
||||
@@ -174,11 +245,6 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
|
||||
# ── Parse message ──
|
||||
text = message.get("text", "") if isinstance(message, dict) else str(message)
|
||||
# Preserve original user text for intent detection BEFORE any augmentation
|
||||
# (truncation, file upload content, prefetch data). Substring-based keyword
|
||||
# matching in get_tools_for_query / fast_confirmation_tool would otherwise
|
||||
# match system-injected text (e.g. "tool" ⊂ "tools" in prefetch marker).
|
||||
user_message_text = text
|
||||
files = message.get("files", []) if isinstance(message, dict) else []
|
||||
if not text.strip() and not files:
|
||||
return
|
||||
@@ -258,31 +324,8 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
conv_id = conversation_id or str(uuid.uuid4())
|
||||
_conv_locks[conv_id] = asyncio.Event()
|
||||
|
||||
fast_tool_name = fast_confirmation_tool(user_message_text)
|
||||
if fast_tool_name:
|
||||
_pending_confirmations[conv_id] = {
|
||||
"tool_name": fast_tool_name,
|
||||
"tool_args": {},
|
||||
"user_text": user_message_text,
|
||||
}
|
||||
yield json.dumps({
|
||||
"content": "⏸️ Требуется подтверждение",
|
||||
"metadata": confirmation_metadata_for_tool(conv_id, fast_tool_name, {}),
|
||||
})
|
||||
return
|
||||
|
||||
# ── Pre-fetch dashboards ──
|
||||
text_lower = user_message_text.lower()
|
||||
if any(kw in text_lower for kw in ["дашборд", "dashboard", "dashboards", "дашборды"]):
|
||||
try:
|
||||
dash_data = await prefetch_dashboards(env_id or "")
|
||||
if dash_data:
|
||||
text += f"\n\n[PRE-FETCHED DATA — use this directly, do NOT call tools]\n{dash_data}\n[/PRE-FETCHED DATA]"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# All tools exposed — Gemma context window is now sufficient.
|
||||
# Intent-based subset filtering (get_tools_for_query) retired.
|
||||
# All tools exposed — LLM handles intent detection via LangGraph tool-calling.
|
||||
# Embedding-based tool selection (top-K) replaces keyword matching if model available.
|
||||
agent_tools = get_all_tools()
|
||||
agent = await create_agent(agent_tools, env_id)
|
||||
config = {"configurable": {"thread_id": conv_id}}
|
||||
@@ -314,9 +357,10 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
elif kind == "on_tool_start":
|
||||
tool_name = event["name"]
|
||||
emitted_any = True
|
||||
redacted_input = _redact_sensitive_fields(event["data"].get("input", {}))
|
||||
yield json.dumps({
|
||||
"content": f"🛠️ {tool_name}",
|
||||
"metadata": {"type": "tool_start", "tool": tool_name, "input": event["data"].get("input", {})},
|
||||
"metadata": {"type": "tool_start", "tool": tool_name, "input": redacted_input},
|
||||
})
|
||||
elif kind == "on_tool_end":
|
||||
tool_name = event["name"]
|
||||
@@ -338,7 +382,7 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
state = await agent.aget_state(config)
|
||||
if getattr(state, "next", None):
|
||||
emitted_any = True
|
||||
yield confirmation_payload(conv_id, state, user_message_text)
|
||||
yield confirmation_payload(conv_id, state, text)
|
||||
return
|
||||
elif not emitted_any:
|
||||
yield json.dumps({
|
||||
@@ -349,6 +393,60 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
|
||||
})
|
||||
break
|
||||
|
||||
except (APIConnectionError, httpx.ConnectError) as exc:
|
||||
_llm_status["status"] = "unavailable"
|
||||
_llm_status["last_error"] = str(exc)
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
logger.explore("LLM provider connection failed",
|
||||
error=str(exc),
|
||||
extra={"src": "AgentChat.GradioApp.Handler"})
|
||||
yield json.dumps({
|
||||
"content": "❌ LLM провайдер недоступен",
|
||||
"metadata": {
|
||||
"type": "error", "code": "LLM_PROVIDER_UNAVAILABLE",
|
||||
"detail": "LLM провайдер недоступен. Проверьте подключение к upstream API.",
|
||||
"retryable": True,
|
||||
},
|
||||
})
|
||||
await save_conversation(conv_id, text, user_id, assistant_text="")
|
||||
return
|
||||
|
||||
except (APITimeoutError, httpx.ReadTimeout) as exc:
|
||||
_llm_status["status"] = "timeout"
|
||||
_llm_status["last_error"] = str(exc)
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
logger.explore("LLM provider timed out",
|
||||
error=str(exc),
|
||||
extra={"src": "AgentChat.GradioApp.Handler"})
|
||||
yield json.dumps({
|
||||
"content": "❌ LLM провайдер не отвечает",
|
||||
"metadata": {
|
||||
"type": "error", "code": "LLM_TIMEOUT",
|
||||
"detail": "LLM провайдер не отвечает. Таймаут соединения.",
|
||||
"retryable": True,
|
||||
},
|
||||
})
|
||||
await save_conversation(conv_id, text, user_id, assistant_text="")
|
||||
return
|
||||
|
||||
except AuthenticationError as exc:
|
||||
_llm_status["status"] = "auth_error"
|
||||
_llm_status["last_error"] = str(exc)
|
||||
_llm_status["last_check_ts"] = time.time()
|
||||
logger.explore("LLM provider auth failed",
|
||||
error=str(exc),
|
||||
extra={"src": "AgentChat.GradioApp.Handler"})
|
||||
yield json.dumps({
|
||||
"content": "❌ API ключ LLM отклонён",
|
||||
"metadata": {
|
||||
"type": "error", "code": "LLM_AUTH_ERROR",
|
||||
"detail": "API ключ LLM отклонён. Проверьте credentials.",
|
||||
"retryable": False,
|
||||
},
|
||||
})
|
||||
await save_conversation(conv_id, text, user_id, assistant_text="")
|
||||
return
|
||||
|
||||
except OutputParserException as e:
|
||||
if attempt < max_attempts - 1:
|
||||
text = "Respond with valid JSON only. Previous response was malformed.\n\n" + text
|
||||
|
||||
@@ -189,8 +189,10 @@ async def create_agent(
|
||||
"managing Git operations (branch/commit/deploy), running LLM validation "
|
||||
"and documentation, creating and copying dashboards and datasets, "
|
||||
"and checking system health, environments, and task status. "
|
||||
"If the data you need is already provided in the user message, use that directly "
|
||||
"rather than calling tools. Only call tools when the data is not present."
|
||||
"You handle all intent detection — multi-intent queries, negations (\"don't run\"), "
|
||||
"synonyms (\"панели\" = \"дашборды\"), and typos are your responsibility. "
|
||||
"Call the right tool(s) for the job. If data is already provided in context, "
|
||||
"use it directly rather than calling redundant tools."
|
||||
)
|
||||
if env_id:
|
||||
prompt += f"\n\nCurrent environment: '{env_id}'. When calling tools that accept env_id, use this value."
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from src.agent.context import get_user_jwt
|
||||
from src.agent.tools import _redact_sensitive_fields
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
@@ -42,7 +43,8 @@ async def log_tool_event(event: dict, conversation_id: str) -> None:
|
||||
if "data" in event:
|
||||
data = event["data"]
|
||||
if kind == "on_tool_start":
|
||||
audit_payload["input"] = str(data.get("input", ""))[:500]
|
||||
raw_input = data.get("input", "")
|
||||
audit_payload["input"] = str(_redact_sensitive_fields(raw_input))[:500]
|
||||
elif kind == "on_tool_error":
|
||||
audit_payload["error"] = str(data.get("error", ""))[:500]
|
||||
|
||||
|
||||
@@ -1062,90 +1062,7 @@ def get_all_tools() -> list:
|
||||
]
|
||||
# #endregion AgentChat.Tools.GetAll
|
||||
|
||||
# ── Optional overrides for embedding descriptions (auto-generated from docstring) ──
|
||||
_TOOL_DESCRIPTIONS_OVERRIDES: dict[str, str] = {}
|
||||
|
||||
# #region AgentChat.Tools.GetForQuery [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,registry,intent]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Return a compact, intent-scoped tool set to keep small-context models usable.
|
||||
# @RATIONALE Some LLMs (gemma) struggle with large tool lists. This reduces the agent's
|
||||
# tool surface to only those relevant to the user's intent.
|
||||
def get_tools_for_query(query: str, *, prefetch_available: bool = False) -> list:
|
||||
text = (query or "").lower()
|
||||
selected = [show_capabilities]
|
||||
matched_intent = False
|
||||
|
||||
if any(word in text for word in ["инструмент", "tool", "capabilit", "умеешь", "можешь"]):
|
||||
return selected
|
||||
|
||||
if any(word in text for word in ["дашборд", "dashboard", "dashboards", "дашборды"]):
|
||||
matched_intent = True
|
||||
if not prefetch_available:
|
||||
selected.append(search_dashboards)
|
||||
if any(word in text for word in ["здоров", "health", "статус системы", "system status"]):
|
||||
matched_intent = True
|
||||
selected.append(get_health_summary)
|
||||
if any(word in text for word in ["окруж", "environment", "env"]):
|
||||
matched_intent = True
|
||||
selected.append(list_environments)
|
||||
if any(word in text for word in ["задач", "task", "таск"]):
|
||||
matched_intent = True
|
||||
selected.append(get_task_status)
|
||||
if any(word in text for word in ["llm", "provider", "провайдер", "модель"]):
|
||||
matched_intent = True
|
||||
selected.extend([list_llm_providers, get_llm_status])
|
||||
if any(word in text for word in ["branch", "ветк"]):
|
||||
matched_intent = True
|
||||
selected.append(create_branch)
|
||||
if any(word in text for word in ["commit", "коммит"]):
|
||||
matched_intent = True
|
||||
selected.append(commit_changes)
|
||||
if any(word in text for word in ["deploy", "депло", "разверн"]):
|
||||
matched_intent = True
|
||||
selected.append(deploy_dashboard)
|
||||
if any(word in text for word in ["миграц", "migration", "migrate"]):
|
||||
matched_intent = True
|
||||
selected.append(execute_migration)
|
||||
if any(word in text for word in ["backup", "бэкап", "резерв"]):
|
||||
matched_intent = True
|
||||
selected.append(run_backup)
|
||||
if any(word in text for word in ["валидац", "validation", "validate"]):
|
||||
matched_intent = True
|
||||
selected.append(run_llm_validation)
|
||||
if any(word in text for word in ["документ", "documentation", "docs"]):
|
||||
matched_intent = True
|
||||
selected.append(run_llm_documentation)
|
||||
if any(word in text for word in ["maintenance", "обслуж", "баннер"]):
|
||||
matched_intent = True
|
||||
selected.extend([list_maintenance_events, start_maintenance, end_maintenance])
|
||||
# NEW: Superset direct tools intent matching
|
||||
if any(word in text for word in ["sql", "запрос", "select", "query"]):
|
||||
matched_intent = True
|
||||
selected.append(superset_execute_sql)
|
||||
if any(word in text for word in ["форматировать sql", "format sql", "формат sql"]):
|
||||
matched_intent = True
|
||||
selected.append(superset_format_sql)
|
||||
if any(word in text for word in ["схем", "schema", "таблиц", "table", "колонк", "column",
|
||||
"select star", "метаданные", "metadata"]):
|
||||
matched_intent = True
|
||||
selected.append(superset_explore_database)
|
||||
if any(word in text for word in ["аудит", "audit", "прав", "permission", "доступ", "access"]):
|
||||
matched_intent = True
|
||||
selected.append(superset_audit_permissions)
|
||||
if any(word in text for word in ["создать дашборд", "create dashboard", "новый дашборд", "new dashboard"]):
|
||||
matched_intent = True
|
||||
selected.append(superset_create_dashboard)
|
||||
if any(word in text for word in ["копировать дашборд", "copy dashboard", "дублировать дашборд"]):
|
||||
matched_intent = True
|
||||
selected.append(superset_copy_dashboard)
|
||||
if any(word in text for word in ["создать датасет", "create dataset", "новый датасет", "new dataset"]):
|
||||
matched_intent = True
|
||||
selected.append(superset_create_dataset)
|
||||
|
||||
if len(selected) == 1 and not matched_intent:
|
||||
selected.extend([search_dashboards, get_health_summary, list_environments, get_task_status])
|
||||
|
||||
unique = {}
|
||||
for tool_obj in selected:
|
||||
unique[tool_obj.name] = tool_obj
|
||||
return list(unique.values())
|
||||
# #endregion AgentChat.Tools.GetForQuery
|
||||
# #endregion AgentChat.Tools
|
||||
|
||||
28
backend/src/api/routes/agent_status.py
Normal file
28
backend/src/api/routes/agent_status.py
Normal file
@@ -0,0 +1,28 @@
|
||||
# backend/src/api/routes/agent_status.py
|
||||
# #region Api.Agent.Status [C:2] [TYPE Module] [SEMANTICS api,agent,llm,status,health]
|
||||
# @BRIEF Agent LLM provider health status endpoint — used by frontend for provider availability indicator.
|
||||
# @RATIONALE Frontend performs health check at mount and auto-retries every 30s if provider unavailable.
|
||||
# @RELATION DEPENDS_ON -> [AgentChat.GradioApp]
|
||||
from fastapi import APIRouter
|
||||
|
||||
router = APIRouter(prefix="/api/agent", tags=["agent-status"])
|
||||
|
||||
|
||||
# #region Api.Agent.Status.Get [C:2] [TYPE Function] [SEMANTICS api,agent,llm,status,get]
|
||||
# @ingroup AgentChat
|
||||
# @BRIEF Return cached LLM provider health status (or trigger probe if cache expired).
|
||||
# @POST Returns {"status": "ok"|"unavailable"|"timeout"|"auth_error",
|
||||
# "last_error": str, "retry_after_s": int}
|
||||
@router.get("/llm-status")
|
||||
async def get_llm_status():
|
||||
"""Get cached LLM provider health status. Probes provider if cache expired."""
|
||||
from src.agent.app import _check_llm_provider_health, _llm_status
|
||||
status = await _check_llm_provider_health()
|
||||
return {
|
||||
"status": status,
|
||||
"last_error": _llm_status.get("last_error", ""),
|
||||
"retry_after_s": 30 if status != "ok" else 0,
|
||||
}
|
||||
# #endregion Api.Agent.Status.Get
|
||||
|
||||
# #endregion Api.Agent.Status
|
||||
@@ -21,7 +21,9 @@ from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from ...core.cot_logger import log
|
||||
from ...core.logger import logger
|
||||
from ._utils import _sanitize_url
|
||||
|
||||
# Module-level httpx client, lazily initialized for connection reuse
|
||||
_http_client: httpx.AsyncClient | None = None
|
||||
@@ -57,8 +59,13 @@ def _get_verify() -> ssl.SSLContext | bool:
|
||||
async def _get_http_client() -> httpx.AsyncClient:
|
||||
global _http_client
|
||||
if _http_client is None:
|
||||
ssl_verify = _get_verify()
|
||||
if ssl_verify is False:
|
||||
log("LLMAsyncHttpClient", "EXPLORE",
|
||||
"TLS verification disabled via LLM_SSL_VERIFY=false",
|
||||
error="TLS verification disabled — traffic to LLM provider is unencrypted")
|
||||
_http_client = httpx.AsyncClient(
|
||||
verify=_get_verify(),
|
||||
verify=ssl_verify,
|
||||
timeout=httpx.Timeout(180.0),
|
||||
)
|
||||
return _http_client
|
||||
@@ -117,7 +124,7 @@ async def call_openai_compatible(
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
logger.reason(
|
||||
f"LLM request url={base_url} model={payload.get('model')} "
|
||||
f"LLM request url={_sanitize_url(base_url)} model={payload.get('model')} "
|
||||
f"provider_type={provider_type} "
|
||||
f"response_format={'yes' if 'response_format' in payload else 'no'} "
|
||||
f"prompt_len={len(prompt)}"
|
||||
|
||||
@@ -31,7 +31,7 @@ from ...services.llm_prompt_templates import render_prompt
|
||||
from ...services.llm_provider import LLMProviderService
|
||||
from ._llm_async_http import call_openai_compatible
|
||||
from ._llm_parse import parse_llm_response
|
||||
from ._utils import _enforce_dictionary
|
||||
from ._utils import _enforce_dictionary, _sanitize_url
|
||||
from .preview import DEFAULT_EXECUTION_PROMPT_TEMPLATE
|
||||
from .prompt_builder import ContextAwarePromptBuilder
|
||||
|
||||
@@ -507,7 +507,7 @@ class LLMTranslationService:
|
||||
logger.reason(
|
||||
f"LLM provider resolved", {
|
||||
"provider_id": job.provider_id, "model": model,
|
||||
"provider_type": provider_type, "base_url": provider.base_url,
|
||||
"provider_type": provider_type, "base_url": _sanitize_url(provider.base_url),
|
||||
"disable_reasoning": disable_reasoning, "max_tokens": max_tokens,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -14,6 +14,7 @@ import json
|
||||
import re
|
||||
from typing import Any
|
||||
import unicodedata
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
@@ -21,6 +22,23 @@ from ...core.logger import logger
|
||||
from ...models.translate import TranslationRecord
|
||||
|
||||
|
||||
# #region _sanitize_url [C:1] [TYPE Function] [SEMANTICS translate, url, sanitize]
|
||||
# @BRIEF Strip embedded credentials from URL for safe logging.
|
||||
# @POST Returns URL with user:pass@ portion removed, preserving host:port.
|
||||
def _sanitize_url(url: str) -> str:
|
||||
"""Strip embedded credentials from URL for safe logging."""
|
||||
if not url:
|
||||
return url
|
||||
parsed = urlsplit(url)
|
||||
if parsed.username or parsed.password:
|
||||
safe_netloc = parsed.hostname
|
||||
if parsed.port:
|
||||
safe_netloc += f":{parsed.port}"
|
||||
parsed = parsed._replace(netloc=safe_netloc)
|
||||
return urlunsplit(parsed)
|
||||
# #endregion _sanitize_url
|
||||
|
||||
|
||||
# #region _normalize_term [TYPE Function]
|
||||
# @BRIEF Normalize a term for case-insensitive unique constraint lookup.
|
||||
# @RATIONALE NFC normalization is applied before lowercasing to ensure consistent
|
||||
|
||||
236
backend/src/scripts/reencrypt.py
Normal file
236
backend/src/scripts/reencrypt.py
Normal file
@@ -0,0 +1,236 @@
|
||||
#!/usr/bin/env python3
|
||||
# #region Scripts.Reencrypt [C:4] [TYPE Module] [SEMANTICS encryption,rotation,migration]
|
||||
# @defgroup Scripts Module group.
|
||||
# @BRIEF Key rotation tool — re-encrypts all stored secrets with a new ENCRYPTION_KEY.
|
||||
# @LAYER Infrastructure
|
||||
# @PRE Old and new ENCRYPTION_KEY must be set via environment variables.
|
||||
# @POST All DatabaseConnection passwords, Environment passwords, and LLMProvider API keys
|
||||
# are re-encrypted with the new key. Original data rejected if decryption fails.
|
||||
# @SIDE_EFFECT Reads/writes app_configurations payload and llm_providers table.
|
||||
# @RELATION DEPENDS_ON -> [EncryptionManager]
|
||||
# @RELATION DEPENDS_ON -> [AppConfigRecord]
|
||||
# @RELATION DEPENDS_ON -> [LLMProvider]
|
||||
# @RATIONALE Fernet is symmetric — re-encryption requires decrypt with old key,
|
||||
# encrypt with new key. There is no key-wrapping or key-derivation layer.
|
||||
# @REJECTED In-place re-encryption without old key rejected — impossible with Fernet.
|
||||
# Auto-rotation on startup rejected — would break on first restart after key change.
|
||||
#
|
||||
# Usage:
|
||||
# OLD_ENCRYPTION_KEY=<old> NEW_ENCRYPTION_KEY=<new> python -m src.scripts.reencrypt
|
||||
# OLD_ENCRYPTION_KEY=<old> NEW_ENCRYPTION_KEY=<new> python -m src.scripts.reencrypt --dry-run
|
||||
# #endregion Scripts.Reencrypt
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
||||
try:
|
||||
from cryptography.fernet import Fernet
|
||||
except ImportError:
|
||||
sys.exit("ERROR: cryptography is not installed. Run: pip install cryptography")
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
# ── Fernet helpers (standalone — no app dependency) ────────────────────
|
||||
|
||||
|
||||
def _make_fernet(key_b64: str) -> Fernet:
|
||||
try:
|
||||
return Fernet(key_b64.encode())
|
||||
except Exception as e:
|
||||
sys.exit(f"ERROR: Invalid Fernet key: {e}")
|
||||
|
||||
|
||||
def _is_fernet_token(value: str) -> bool:
|
||||
if not value or len(value) < 60:
|
||||
return False
|
||||
try:
|
||||
import base64
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
||||
try:
|
||||
from cryptography.fernet import Fernet
|
||||
except ImportError:
|
||||
sys.exit("ERROR: cryptography is not installed. Run: pip install cryptography")
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.encryption import is_fernet_token
|
||||
|
||||
# ── Fernet helpers (standalone — no app dependency) ────────────────────
|
||||
|
||||
|
||||
def _make_fernet(key_b64: str) -> Fernet:
|
||||
try:
|
||||
return Fernet(key_b64.encode())
|
||||
except Exception as e:
|
||||
sys.exit(f"ERROR: Invalid Fernet key: {e}")
|
||||
|
||||
|
||||
def _reencrypt_value(value: str, old_fernet: Fernet, new_fernet: Fernet) -> str | None:
|
||||
"""Decrypt with old key, encrypt with new key. Returns None on failure."""
|
||||
if not is_fernet_token(value):
|
||||
print(f" ⚠ Skipping non-Fernet value (length={len(value)})")
|
||||
return None
|
||||
try:
|
||||
plaintext = old_fernet.decrypt(value.encode()).decode()
|
||||
except Exception as e:
|
||||
print(f" ✗ Decryption failed: {e}")
|
||||
return None
|
||||
return new_fernet.encrypt(plaintext.encode()).decode()
|
||||
|
||||
|
||||
# ── Report helpers ────────────────────────────────────────────────────
|
||||
|
||||
_report: list[str] = []
|
||||
|
||||
|
||||
def _r(msg: str) -> None:
|
||||
_report.append(msg)
|
||||
print(msg)
|
||||
|
||||
|
||||
# ── Main ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Re-encrypt all stored secrets with a new Fernet ENCRYPTION_KEY."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Only scan and report what would be changed; no writes.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
old_key = os.getenv("OLD_ENCRYPTION_KEY", "").strip()
|
||||
new_key = os.getenv("NEW_ENCRYPTION_KEY", "").strip()
|
||||
|
||||
if not old_key or not new_key:
|
||||
_r("ERROR: Set OLD_ENCRYPTION_KEY and NEW_ENCRYPTION_KEY environment variables.")
|
||||
_r("")
|
||||
_r("Usage:")
|
||||
_r(" OLD_ENCRYPTION_KEY=<old> NEW_ENCRYPTION_KEY=<new> python -m src.scripts.reencrypt")
|
||||
_r(" OLD_ENCRYPTION_KEY=<old> NEW_ENCRYPTION_KEY=<new> python -m src.scripts.reencrypt --dry-run")
|
||||
sys.exit(1)
|
||||
|
||||
old_fernet = _make_fernet(old_key)
|
||||
new_fernet = _make_fernet(new_key)
|
||||
|
||||
if args.dry_run:
|
||||
_r("🔍 DRY RUN — no changes will be made")
|
||||
else:
|
||||
_r("🔐 Re-encrypting all secrets with new ENCRYPTION_KEY...")
|
||||
_r(f" Started at: {datetime.now(timezone.utc).isoformat()}")
|
||||
_r("")
|
||||
|
||||
# ── Load database URL ──────────────────────────────────────────
|
||||
db_url = (
|
||||
os.getenv("DATABASE_URL", "")
|
||||
or os.getenv("POSTGRES_URL", "")
|
||||
or "postgresql+psycopg2://postgres:postgres@localhost:5432/ss_tools"
|
||||
)
|
||||
# Use psycopg2 for sync access in script
|
||||
engine = create_engine(db_url)
|
||||
|
||||
# ── Step 1: Environment passwords (AppConfigRecord.payload.environments) ──
|
||||
_r("── Environment passwords (ConfigManager) ──")
|
||||
from sqlalchemy import Column, String, Integer, Text
|
||||
from sqlalchemy.orm import declarative_base
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
class AppConfigRecord(Base):
|
||||
__tablename__ = "app_configurations"
|
||||
id = Column(String, primary_key=True)
|
||||
payload = Column(Text)
|
||||
|
||||
total_env_passwords = 0
|
||||
reencrypted_env = 0
|
||||
skipped_env = 0
|
||||
|
||||
with Session(engine) as session:
|
||||
record = session.query(AppConfigRecord).filter(AppConfigRecord.id == "global").first()
|
||||
if record and record.payload:
|
||||
import json
|
||||
payload = record.payload if isinstance(record.payload, dict) else json.loads(record.payload)
|
||||
environments = payload.get("environments", [])
|
||||
for env in environments:
|
||||
pwd = env.get("password", "")
|
||||
if not pwd or pwd == "********":
|
||||
skipped_env += 1
|
||||
continue
|
||||
re = _reencrypt_value(pwd, old_fernet, new_fernet)
|
||||
if re is not None:
|
||||
if not args.dry_run:
|
||||
env["password"] = re
|
||||
reencrypted_env += 1
|
||||
else:
|
||||
_r(f" ✗ Failed to re-encrypt password for env '{env.get('id', '?')}'")
|
||||
total_env_passwords += 1
|
||||
|
||||
if not args.dry_run and reencrypted_env > 0:
|
||||
record.payload = payload
|
||||
session.commit()
|
||||
_r(f" ✓ Committed {reencrypted_env} re-encrypted environment passwords")
|
||||
else:
|
||||
_r(" - No AppConfigRecord found, skipping")
|
||||
|
||||
_r(f" Environment passwords: {reencrypted_env} re-encrypted, {total_env_passwords - reencrypted_env - skipped_env} failed, {skipped_env} skipped")
|
||||
_r("")
|
||||
|
||||
# ── Step 2: LLM Provider API keys ──────────────────────────────
|
||||
_r("── LLM Provider API keys ──")
|
||||
from sqlalchemy import Column, String as SAString, Boolean, Integer as SAInteger
|
||||
|
||||
class LLMProvider(Base):
|
||||
__tablename__ = "llm_providers"
|
||||
id = Column(SAString, primary_key=True)
|
||||
api_key = Column(SAString)
|
||||
|
||||
total_providers = 0
|
||||
reencrypted_keys = 0
|
||||
skipped_providers = 0
|
||||
|
||||
with Session(engine) as session:
|
||||
providers = session.query(LLMProvider).all()
|
||||
for prov in providers:
|
||||
key = prov.api_key
|
||||
if not key:
|
||||
skipped_providers += 1
|
||||
continue
|
||||
re = _reencrypt_value(key, old_fernet, new_fernet)
|
||||
if re is not None:
|
||||
if not args.dry_run:
|
||||
prov.api_key = re
|
||||
reencrypted_keys += 1
|
||||
else:
|
||||
_r(f" ✗ Failed to re-encrypt API key for provider '{prov.id}'")
|
||||
total_providers += 1
|
||||
|
||||
if not args.dry_run and reencrypted_keys > 0:
|
||||
session.commit()
|
||||
_r(f" ✓ Committed {reencrypted_keys} re-encrypted API keys")
|
||||
|
||||
_r(f" Provider API keys: {reencrypted_keys} re-encrypted, {total_providers - reencrypted_keys - skipped_providers} failed, {skipped_providers} skipped")
|
||||
_r("")
|
||||
|
||||
# ── Summary ────────────────────────────────────────────────────
|
||||
_r("── Summary ──")
|
||||
_r(f" Total re-encrypted: {reencrypted_env + reencrypted_keys}")
|
||||
_r(f" Total failed: {(total_env_passwords - reencrypted_env - skipped_env) + (total_providers - reencrypted_keys - skipped_providers)}")
|
||||
_r(f" Total skipped: {skipped_env + skipped_providers}")
|
||||
if args.dry_run:
|
||||
_r(" (dry run — no changes written)")
|
||||
_r("")
|
||||
_r("Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user