223 lines
9.7 KiB
Python
223 lines
9.7 KiB
Python
# backend/src/agent/langgraph_setup.py
|
|
# #region AgentChat.LangGraph.Setup [C:4] [TYPE Module] [SEMANTICS agent-chat,langgraph,agent]
|
|
# @defgroup AgentChat LangGraph agent setup: create_react_agent with PostgresSaver.
|
|
# @PRE LLM provider configured via backend API /api/agent/llm-config.
|
|
# @POST Compiled StateGraph ready for astream_events().
|
|
# @SIDE_EFFECT Initializes checkpointer and message history tables on first call.
|
|
# @RELATION DEPENDS_ON -> [AgentChat.Tools]
|
|
# @RATIONALE LangGraph create_react_agent provides built-in tool calling + checkpointing + interrupt/resume.
|
|
# @REJECTED Using only environment variables for LLM config was rejected — FastAPI API-based config allows runtime switching without restart.
|
|
# RunnableWithMessageHistory wrapper is NOT used — PostgresSaver handles history natively.
|
|
|
|
# ── Monkey-patch: OpenAI SDK for Pydantic BaseModel classes ──
|
|
# LangChain BaseTool objects carry an ``args_schema`` field that is a Pydantic
|
|
# BaseModel *class* reference (not an instance). When the OpenAI SDK recursively
|
|
# transforms the request body, it hits ``isinstance(data, pydantic.BaseModel)``
|
|
# which is True for both instances AND classes. It then calls ``model_dump()``
|
|
# on the class, which fails with:
|
|
#
|
|
# PydanticSerializationError: Unable to serialize unknown type: ModelMetaclass
|
|
#
|
|
# The fix: skip model_dump for classes, only dump instances.
|
|
import inspect as _inspect
|
|
import os
|
|
|
|
import httpx
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
from langgraph.prebuilt import create_react_agent
|
|
import openai._utils._transform as _openai_transform
|
|
import psycopg
|
|
from psycopg.rows import dict_row
|
|
import pydantic as _pydantic
|
|
import pydantic_core as _pydantic_core
|
|
|
|
from src.agent._config import AGENT_CONFIRM_TOOLS, AGENT_INTERRUPT_BEFORE as _INTERRUPT_BEFORE, FASTAPI_URL
|
|
from src.agent._llm_params import chat_openai_kwargs
|
|
from src.core.logger import logger
|
|
|
|
_original_transform = _openai_transform._async_transform_recursive
|
|
|
|
async def _patched_transform(data, *, annotation, inner_type=None):
|
|
if isinstance(data, _pydantic.BaseModel):
|
|
if _inspect.isclass(data):
|
|
# BaseModel CLASS (not instance) — skip model_dump
|
|
return data
|
|
# BaseModel INSTANCE — intercept PydanticSerializationError
|
|
try:
|
|
return await _original_transform(data, annotation=annotation, inner_type=inner_type)
|
|
except _pydantic_core.PydanticSerializationError:
|
|
print(f"[PATCH] Caught PydanticSerializationError on {type(data).__name__}")
|
|
# Fallback: dump with exclude of type-ref fields
|
|
serializable = {}
|
|
for field_name in data.model_fields_set:
|
|
val = getattr(data, field_name)
|
|
if isinstance(val, type) and issubclass(val, _pydantic.BaseModel):
|
|
serializable[field_name] = val.model_json_schema()
|
|
else:
|
|
serializable[field_name] = val
|
|
return serializable
|
|
return await _original_transform(data, annotation=annotation, inner_type=inner_type)
|
|
|
|
_openai_transform._async_transform_recursive = _patched_transform
|
|
|
|
# ── Postgres checkpointer (FR-004/FR-012/FR-027) ──
|
|
_CHECKPOINTER: AsyncPostgresSaver | None = None
|
|
_CHECKPOINTER_INIT = False
|
|
_CHECKPOINTER_CONN = None
|
|
|
|
|
|
async def init_checkpointer() -> None:
|
|
"""Initialize the PostgreSQL checkpointer (FR-004/FR-012/FR-027).
|
|
|
|
Called once at agent startup. Creates a persistent psycopg async connection
|
|
and passes it to AsyncPostgresSaver. Runs .setup() to create checkpoint tables.
|
|
Connection stays open for the lifetime of the agent process.
|
|
"""
|
|
global _CHECKPOINTER, _CHECKPOINTER_INIT, _CHECKPOINTER_CONN
|
|
if _CHECKPOINTER_INIT:
|
|
return
|
|
db_url = os.getenv("DATABASE_URL")
|
|
# Convert SQLAlchemy-style URL to psycopg format
|
|
pg_url = db_url.replace("postgresql+psycopg2://", "postgres://").replace("postgresql://", "postgres://")
|
|
_CHECKPOINTER_CONN = await psycopg.AsyncConnection.connect(pg_url, autocommit=True, row_factory=dict_row)
|
|
_CHECKPOINTER = AsyncPostgresSaver(_CHECKPOINTER_CONN)
|
|
await _CHECKPOINTER.setup()
|
|
_CHECKPOINTER_INIT = True
|
|
|
|
# ── LLM config (no cache — fetched on each create_agent call) ──
|
|
_llm_config: dict | None = None
|
|
|
|
|
|
def configure_from_api(llm_config: dict) -> None:
|
|
"""Update LLM config from FastAPI response. Called at startup."""
|
|
global _llm_config
|
|
_llm_config = llm_config
|
|
|
|
|
|
async def _fetch_llm_config() -> dict | None:
|
|
"""Fetch LLM config from FastAPI.
|
|
|
|
Called on every create_agent() to pick up Admin UI changes immediately.
|
|
Falls back to cached config if fetch fails.
|
|
"""
|
|
global _llm_config
|
|
try:
|
|
fastapi_url = FASTAPI_URL
|
|
async with httpx.AsyncClient(timeout=5) as client:
|
|
resp = await client.get(f"{fastapi_url}/api/agent/llm-config")
|
|
if resp.status_code == 200:
|
|
config = resp.json()
|
|
if config.get("configured"):
|
|
_llm_config = config
|
|
return config
|
|
except Exception as e:
|
|
logger.explore("Failed to fetch LLM config from FastAPI", error=str(e),
|
|
extra={"src": "AgentChat.LangGraph.Setup"})
|
|
return _llm_config
|
|
|
|
|
|
def _interrupt_before_from_env() -> list[str]:
|
|
"""Return LangGraph node names that must pause for HITL confirmation."""
|
|
if AGENT_CONFIRM_TOOLS:
|
|
return ["tools"]
|
|
raw = _INTERRUPT_BEFORE
|
|
if not raw:
|
|
return []
|
|
return [name.strip() for name in raw.split(",") if name.strip()]
|
|
|
|
|
|
async def create_agent(
|
|
tools: list,
|
|
env_id: str | None = None,
|
|
interrupt_before: list[str] | None = None,
|
|
):
|
|
"""Create the LangGraph agent with PostgreSQL checkpointer and message history.
|
|
|
|
LLM configuration source:
|
|
llm_config from FastAPI /api/agent/llm-config (fetched on every call).
|
|
If backend has no configured provider, agent raises an error.
|
|
|
|
Returns a compiled StateGraph ready for astream_events().
|
|
interrupt_before is set from AGENT_CONFIRM_TOOLS (or AGENT_INTERRUPT_BEFORE env var)
|
|
to enable HITL guardrails — when pending tools are detected, the graph pauses before
|
|
executing them and yields confirm_required metadata to the frontend.
|
|
Checkpointer is AsyncPostgresSaver (survives container restarts).
|
|
"""
|
|
# Fetch fresh LLM config from FastAPI on every call
|
|
config = await _fetch_llm_config()
|
|
|
|
if config and config.get("configured"):
|
|
api_key = config["api_key"]
|
|
base_url = config.get("base_url")
|
|
model = config.get("default_model")
|
|
config_source = "FastAPI"
|
|
else:
|
|
raise RuntimeError(
|
|
"No LLM provider configured in backend. "
|
|
"Configure one via Settings → AI Providers in the web UI."
|
|
)
|
|
|
|
logger.reason(
|
|
"Creating LangGraph agent",
|
|
payload={"model": model, "config_source": config_source, "tools_count": len(tools), "env_id": env_id},
|
|
extra={"src": "AgentChat.LangGraph.Setup"},
|
|
)
|
|
|
|
llm = ChatOpenAI(**chat_openai_kwargs(
|
|
model=model,
|
|
base_url=base_url,
|
|
api_key=api_key,
|
|
max_tokens=2048,
|
|
))
|
|
|
|
# System prompt — env_id injected deterministically, not in user message
|
|
prompt = (
|
|
"You are a Superset Tools assistant. You have access to tools for searching "
|
|
"dashboards, managing maintenance, running migrations and backups, "
|
|
"executing SQL and exploring databases, auditing permissions, "
|
|
"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. "
|
|
"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. "
|
|
"For maintenance requests, use the RUNTIME CONTEXT current datetime when the user says "
|
|
"\"start\", \"run\", \"now\", \"запусти\", or \"сейчас\" without an explicit start time. "
|
|
"Convert user durations into end_time. Do not ask for ISO datetime in that case. "
|
|
"If a user asks for dashboard maintenance, resolve the dashboard from provided context or tools, "
|
|
"then infer affected tables when possible; ask for table names only after resolution fails."
|
|
)
|
|
if env_id:
|
|
prompt += f"\n\nCurrent environment: '{env_id}'. When calling tools that accept env_id, use this value."
|
|
|
|
# Checkpointer — AsyncPostgresSaver. Fallback to InMemory if Postgres init failed
|
|
if _CHECKPOINTER is not None:
|
|
checkpointer = _CHECKPOINTER
|
|
else:
|
|
checkpointer = InMemorySaver()
|
|
logger.explore(
|
|
"Postgres checkpointer unavailable, falling back to InMemorySaver",
|
|
error="_CHECKPOINTER is None — checkpoints will be lost on restart",
|
|
extra={"src": "AgentChat.LangGraph.Setup"},
|
|
)
|
|
|
|
graph = create_react_agent(
|
|
model=llm,
|
|
tools=tools,
|
|
prompt=prompt,
|
|
version="v2",
|
|
checkpointer=checkpointer,
|
|
interrupt_before=_interrupt_before_from_env() if interrupt_before is None else interrupt_before,
|
|
)
|
|
|
|
logger.reflect(
|
|
"LangGraph agent created",
|
|
payload={"model": model, "checkpointer_type": type(checkpointer).__name__, "tools_count": len(tools)},
|
|
extra={"src": "AgentChat.LangGraph.Setup"},
|
|
)
|
|
return graph
|
|
# #endregion AgentChat.LangGraph.Setup
|