From 525128f306d087576653e55890c2a9dc4385aba1 Mon Sep 17 00:00:00 2001 From: yuvalkh Date: Thu, 23 Jul 2026 16:16:43 +0300 Subject: [PATCH 1/6] updated agent to work with jeen metadata --- agent/scripts/upload_all_prompts.py | 10 +- agent/src/agent/config.py | 17 +- agent/src/agent/graph.py | 60 +- agent/src/agent/mcp_server.py | 10 +- agent/src/agent/nodes/detect_ambiguity.py | 15 +- agent/src/agent/nodes/query_builder.py | 38 +- agent/src/agent/nodes/refiner.py | 12 +- agent/src/agent/nodes/schema_explorer.py | 604 +----------------- agent/src/agent/routers/chat.py | 4 +- agent/src/agent/state.py | 2 +- agent/src/agent/utils/jeen_metadata_client.py | 327 ++++++++++ agent/tests/test_routing.py | 41 +- agent/uv.lock | 32 +- 13 files changed, 450 insertions(+), 722 deletions(-) create mode 100644 agent/src/agent/utils/jeen_metadata_client.py diff --git a/agent/scripts/upload_all_prompts.py b/agent/scripts/upload_all_prompts.py index 684ddca..bca974f 100644 --- a/agent/scripts/upload_all_prompts.py +++ b/agent/scripts/upload_all_prompts.py @@ -80,15 +80,15 @@ def main(): "type": "chat" }, { - "name": "text2sql/query_builder", + "name": "text2sql/query_builder_v2", "prompt": [ { "role": "system", - "content": "You are a SQL expert who specializes in trino. Build a SQL query based on the plan and user query. Output ONLY the SQL query, nothing else." + "content": "You are a SQL expert who specializes in trino. Build a SQL query based on the catalog and user query. Output ONLY the SQL query, nothing else.\n\nIMPORTANT: Before writing the SQL, provide a brief 1-3 sentence explanation of your reasoning or how you answered the question.\n\nCRITICAL: You MUST fully qualify all tables in your SQL query using the provided Catalog and Schema parameters. (Format: {{trino_catalog_name}}.{{trino_schema_name}}.my_table). Do not use unqualified table names." }, { "role": "user", - "content": "Plan: {{schema_plan}}\nQuery: {{user_query}}{{feedback_str}}" + "content": "Jeen Metadata Catalog Overview: {{jeen_catalog}}\nTarget Trino Catalog Name: {{trino_catalog_name}}\nTarget Trino Schema Name: {{trino_schema_name}}\nQuery: {{user_query}}. {{feedback_str}}" } ], "type": "chat" @@ -164,7 +164,7 @@ def main(): "type": "chat" }, { - "name": "text2sql/detect_ambiguity", + "name": "text2sql/detect_ambiguity_v2", "prompt": [ { "role": "system", @@ -278,7 +278,7 @@ def main(): "role": "user", "content": ( "User Request: {{user_query}}\n\n" - "Schema:\n{{schema}}\n\n" + "Schema Context:\n{{schema}}\n\n" "Current Agent SQL Attempt:\n{{current_sql_attempt}}" ) } diff --git a/agent/src/agent/config.py b/agent/src/agent/config.py index b4cd35c..2f8d429 100644 --- a/agent/src/agent/config.py +++ b/agent/src/agent/config.py @@ -35,6 +35,19 @@ class AgentSettings(BaseSettings): NOMINATIM_SIMPLIFY_ITERATIONS: int = Field(default=25, gt=0) # binary-search WKT simplify steps LOCATION_MAX_WKT_LENGTH: int = Field(default=2100, gt=0) # max chars for WKT polygon string + # ── Jeen Metadata MCP Integration ───────────────────────────────────────── + # When set, schema_explorer pulls tables/profiles from jeen-metadata via MCP + # instead of from the local Postgres DB. Leave empty to keep the local path. + JEEN_METADATA_MCP_URL: str = "http://localhost:3001/api/mcp" # e.g. https://jeen-metadata.example.com/api/mcp + JEEN_METADATA_MCP_KEY: str = "mcp_f885337e381366db5edc22093415450e38f71e997e96dc708fea69bde9529ab9" # Bearer key from /api/mcp/keys in jeen-metadata + JEEN_METADATA_CONNECTION_ID: int = 89 # Numeric service ID (from list_connections) + JEEN_METADATA_SEARCH_LIMIT: int = 10 # Max tables returned by the search tool + JEEN_METADATA_PROFILE_TIMEOUT: float = 30.0 # Per-MCP-call timeout (seconds) + + # ── Trino connection info for explicit catalog/schema qualification ── + TRINO_CATALOG: str = "" + TRINO_SCHEMA: str = "" + # ── G4: Feature Flags & Execution Modes ────────────────────────────────── BACKEND_URL: str = "http://localhost:8000" # Studio backend URL REDIS_URL: str = "redis://localhost:6379" @@ -42,7 +55,7 @@ class AgentSettings(BaseSettings): # Langfuse prompt names LANGFUSE_PROMPT_EXTRACTOR: str = "text2sql/extractor" LANGFUSE_PROMPT_SCHEMA_EXPLORER: str = "text2sql/schema_explorer" - LANGFUSE_PROMPT_QUERY_BUILDER: str = "text2sql/query_builder" + LANGFUSE_PROMPT_QUERY_BUILDER: str = "text2sql/query_builder_v2" LANGFUSE_PROMPT_REFINER: str = "text2sql/refiner" LANGFUSE_PROMPT_FINALIZER_SUMMARY: str = "text2sql/finalizer_summary" LANGFUSE_PROMPT_FINALIZER_SQL_EXPLANATION: str = ( @@ -50,7 +63,7 @@ class AgentSettings(BaseSettings): ) LANGFUSE_PROMPT_REJECTION_ROUTER: str = "text2sql/rejection_router" LANGFUSE_PROMPT_LOC_EXTRACTOR: str = "text2sql/extractor" - LANGFUSE_PROMPT_DETECT_AMBIGUITY: str = "text2sql/detect_ambiguity" + LANGFUSE_PROMPT_DETECT_AMBIGUITY: str = "text2sql/detect_ambiguity_v2" MAX_REFINER_ITERATIONS: int = Field(default=3, gt=0) REFINER_SCHEMA_CONTEXT_TABLES: int = Field(default=4, gt=0) diff --git a/agent/src/agent/graph.py b/agent/src/agent/graph.py index 3acc6ad..261ae23 100644 --- a/agent/src/agent/graph.py +++ b/agent/src/agent/graph.py @@ -35,8 +35,8 @@ from agent.nodes.extractor import extractor_node from agent.nodes.init_flags import init_flags_node from agent.nodes.init_skills import init_skills_node -from agent.nodes.schema_explorer import schema_explorer_node, MAX_SCHEMA_RETRIES, sql_static_validations_node -from agent.nodes.query_builder import query_builder_node +from agent.nodes.schema_explorer import schema_explorer_node +from agent.nodes.query_builder import query_builder_node, hitl_query_approval_node from agent.nodes.detect_ambiguity import detect_ambiguity_node, ambiguity_resolution_node from agent.nodes.refiner_graph import refiner_subgraph from agent.nodes.finalizer import finalizer_node @@ -170,21 +170,6 @@ def rejection_router_node(state: AgentState, config: RunnableConfig | None = Non # ── Conditional edge functions ──────────────────────────────────────────────── -def route_schema_explorer(state: AgentState) -> str: - """G2-02: route to hitl_escalation after MAX_SCHEMA_RETRIES.""" - if state.get("hallucinated_tables"): - if (state.get("schema_explorer_retry_count") or 0) >= MAX_SCHEMA_RETRIES: - return "hitl_escalation" - return "schema_explorer" - - runtime_flags = state.get("runtime_flags") or {} - enable_ambiguity = runtime_flags.get("SCHEMA_AMBIGUITY_DETECT", settings.ENABLE_AMBIGUITY_DETECT) - if isinstance(enable_ambiguity, str): - enable_ambiguity = enable_ambiguity.lower() == "true" - - if enable_ambiguity: - return "detect_ambiguity" - return "query_builder" def route_refiner_subagent(state: AgentState) -> str: @@ -205,7 +190,7 @@ def route_detect_ambiguity(state: AgentState) -> str: """ Route out of detect_ambiguity based on the resolved ambiguity_type. - - "clear" → query_builder (proceed normally) + - "clear" → hitl_query_approval (proceed normally) - "ambiguous" → ambiguity_resolution (HITL: user clarifies, then retry) → END if MAX_AMBIGUITY_RETRIES exhausted - "unanswerable" → END (data doesn’t exist; clarification won’t help) @@ -213,7 +198,7 @@ def route_detect_ambiguity(state: AgentState) -> str: t = state.get("ambiguity_type") or "clear" if t == "clear": - return "query_builder" + return "hitl_query_approval" if t == "unanswerable": return END @@ -223,6 +208,17 @@ def route_detect_ambiguity(state: AgentState) -> str: def route_query_builder(state: AgentState) -> str: + runtime_flags = state.get("runtime_flags") or {} + enable_ambiguity = runtime_flags.get("SCHEMA_AMBIGUITY_DETECT", settings.ENABLE_AMBIGUITY_DETECT) + if isinstance(enable_ambiguity, str): + enable_ambiguity = enable_ambiguity.lower() == "true" + + if enable_ambiguity: + return "detect_ambiguity" + return "hitl_query_approval" + + +def route_hitl_approval(state: AgentState) -> str: if state.get("feedback"): return "rejection_router" return "refiner_subagent" @@ -244,10 +240,10 @@ def route_rejection(state: AgentState) -> str: workflow.add_node("init_skills", init_skills_node) workflow.add_node("extractor", extractor_node) workflow.add_node("schema_explorer", schema_explorer_node) -workflow.add_node("sql_static_validations", sql_static_validations_node) +workflow.add_node("query_builder", query_builder_node) workflow.add_node("detect_ambiguity", detect_ambiguity_node) workflow.add_node("ambiguity_resolution", ambiguity_resolution_node) -workflow.add_node("query_builder", query_builder_node) +workflow.add_node("hitl_query_approval", hitl_query_approval_node) workflow.add_node("rejection_router", rejection_router_node) workflow.add_node("refiner_subagent", refiner_subgraph) workflow.add_node("hitl_escalation", hitl_escalation_node) @@ -259,16 +255,14 @@ def route_rejection(state: AgentState) -> str: workflow.add_edge("init_flags", "init_skills") workflow.add_edge("init_skills", "extractor") workflow.add_edge("extractor", "schema_explorer") -workflow.add_edge("schema_explorer", "sql_static_validations") +workflow.add_edge("schema_explorer", "query_builder") workflow.add_conditional_edges( - "sql_static_validations", - route_schema_explorer, + "query_builder", + route_query_builder, { - "schema_explorer": "schema_explorer", "detect_ambiguity": "detect_ambiguity", - "query_builder": "query_builder", - "hitl_escalation": "hitl_escalation", # G2-02 + "hitl_query_approval": "hitl_query_approval", }, ) @@ -276,19 +270,19 @@ def route_rejection(state: AgentState) -> str: "detect_ambiguity", route_detect_ambiguity, { - "query_builder": "query_builder", + "hitl_query_approval": "hitl_query_approval", "ambiguity_resolution": "ambiguity_resolution", END: END, }, ) -# ambiguity_resolution → schema_explorer: targeted retry (not a full extractor reset). -# The user’s clarification is in state["feedback"] which schema_explorer already reads. -workflow.add_edge("ambiguity_resolution", "schema_explorer") +# ambiguity_resolution → query_builder: targeted retry. +# The user’s clarification is in state["feedback"] which query_builder reads. +workflow.add_edge("ambiguity_resolution", "query_builder") workflow.add_conditional_edges( - "query_builder", - route_query_builder, + "hitl_query_approval", + route_hitl_approval, {"rejection_router": "rejection_router", "refiner_subagent": "refiner_subagent"}, ) diff --git a/agent/src/agent/mcp_server.py b/agent/src/agent/mcp_server.py index 9f36b4c..1bd8169 100644 --- a/agent/src/agent/mcp_server.py +++ b/agent/src/agent/mcp_server.py @@ -162,7 +162,7 @@ async def chat_with_agent( "thread_id": thread_id, "status": "interrupted", "interrupt_details": interrupt_val, - "schema_plan": final_state.values.get("schema_plan") or (interrupt_val.get("schema_plan") if isinstance(interrupt_val, dict) else None), + "schema_plan": None, "sql_query": final_state.values.get("sql_query") or (interrupt_val.get("sql_query") if isinstance(interrupt_val, dict) else None), "sql_explanation": interrupt_val.get("sql_explanation") if isinstance(interrupt_val, dict) else None, "trace_id": trace_id, @@ -183,7 +183,7 @@ async def chat_with_agent( "thread_id": thread_id, "status": "interrupted", "interrupt_details": interrupt_val, - "schema_plan": final_state.values.get("schema_plan"), + "schema_plan": None, "sql_query": final_state.values.get("sql_query"), "trace_id": trace_id, "execution_path": final_state.values.get("execution_path", []), @@ -202,7 +202,7 @@ async def chat_with_agent( "raw_data_ref": result.get("raw_data_ref"), "sql_query": result.get("sql_query"), "sql_explanation": result.get("sql_explanation"), - "schema_plan": result.get("schema_plan"), + "schema_plan": None, "trace_id": trace_id, "execution_path": result.get("execution_path", []), "is_unanswerable": is_unans, @@ -220,7 +220,7 @@ async def suggest_fixes(thread_id: str, category: str) -> str: return "[]" sql_query = state_snapshot.values.get("sql_query", "") - schema_plan = state_snapshot.values.get("schema_plan", "") + jeen_catalog = state_snapshot.values.get("jeen_catalog", "") user_query = state_snapshot.values.get("user_query", "") runtime_flags = state_snapshot.values.get("runtime_flags", {}) @@ -233,7 +233,7 @@ class Fixes(BaseModel): The user rejected the agent's Text2SQL output with category '{category}'. User Query: {user_query} Current SQL: {sql_query} - Current Plan: {schema_plan} + Current Plan: {jeen_catalog} Provide 2-3 short, distinct button labels for the user to quickly apply a fix. For example: "GROUP BY date instead of month", "Include cancelled orders", "Filter by region". diff --git a/agent/src/agent/nodes/detect_ambiguity.py b/agent/src/agent/nodes/detect_ambiguity.py index f01e61b..522fa1c 100644 --- a/agent/src/agent/nodes/detect_ambiguity.py +++ b/agent/src/agent/nodes/detect_ambiguity.py @@ -44,7 +44,7 @@ from agent.config import settings from agent.langfuse_client import langfuse_client from agent.llm import get_llm -from agent.nodes.refiner import build_refiner_schema_context +from agent.llm import get_llm from agent.state import AgentState from agent.utils.redis_publisher import publish_node_event @@ -101,10 +101,7 @@ async def detect_ambiguity_node( # ── Gather inputs ───────────────────────────────────────────────────────── user_query: str = state.get("user_query") or "" - # schema_plan is the agent's explicit interpretation of the query before any - # SQL is written — richer than SQL for auditing intent. - agent_plan: str = state.get("schema_plan") or "(no query plan available)" - schema_context: str = build_refiner_schema_context(state) + # jeen_catalog now contains the catalog prompt current_time: str = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") # ── Fetch system prompt from Langfuse ───────────────────────────────────── @@ -130,13 +127,11 @@ async def detect_ambiguity_node( # ── Build user message ──────────────────────────────────────────────────── # NOTE: The field is labelled "Current Agent SQL Attempt" in the prompt to # stay consistent with the Langfuse prompt template. At this stage it - # contains the schema_plan (the agent's intent in natural language), which - # is equally valid — the prompt's Agent Proposal Audit works on intent, not - # syntax. + # contains the sql_query from the query builder. user_message = ( f"User Request: {state.get('user_query', '')}\n\n" - f"Schema Context:\n{schema_context}\n\n" - f"Current Agent SQL Attempt:\n{agent_plan}\n" + f"Schema Context:\n{state.get('jeen_catalog', '')}\n\n" + f"Current Agent SQL Attempt:\n{state.get('sql_query', '')}\n" ) # ── LLM call ───────────────────────────────────────────────────────────── diff --git a/agent/src/agent/nodes/query_builder.py b/agent/src/agent/nodes/query_builder.py index 1e6884b..3ebd1b9 100644 --- a/agent/src/agent/nodes/query_builder.py +++ b/agent/src/agent/nodes/query_builder.py @@ -7,8 +7,9 @@ from agent.config import settings from agent.langfuse_client import langfuse_client from langgraph.types import interrupt + async def query_builder_node(state: AgentState, config: RunnableConfig | None = None): - """Build SQL from plan and pause for user approval.""" + """Build SQL from catalog and user query.""" runtime_flags = state.get("runtime_flags") or {} feedback = state.get("feedback") feedback_str = f"\nUser Feedback to apply: {feedback}" if feedback else "" @@ -29,9 +30,11 @@ async def query_builder_node(state: AgentState, config: RunnableConfig | None = publish_node_event_sync(thread_id, "query_builder") response = await chain.ainvoke( { - "schema_plan": state.get("schema_plan"), + "jeen_catalog": state.get("jeen_catalog"), "user_query": state.get("user_query"), "feedback_str": feedback_str, + "trino_catalog_name": settings.TRINO_CATALOG, + "trino_schema_name": settings.TRINO_SCHEMA, } ) content = response.content @@ -58,35 +61,42 @@ async def query_builder_node(state: AgentState, config: RunnableConfig | None = if sql.endswith(";"): sql = sql[:-1].strip() + return { + "sql_query": sql, + "sql_explanation": explanation, + "execution_path": ["query_builder"], + "refinement_count": 0, + "trino_error": None + } + +async def hitl_query_approval_node(state: AgentState, config: RunnableConfig | None = None): + """Pause for user approval of the generated SQL.""" + thread_id = config.get("configurable", {}).get("thread_id", "") if config else "" + publish_node_event_sync(thread_id, "hitl_query_approval") + if state.get("non_interactive"): return { - "sql_query": sql, - "refinement_count": 0, - "trino_error": None, "feedback": None, - "execution_path": ["query_builder"], + "execution_path": ["hitl_query_approval"], } approval_result = interrupt( { "type": "query_approval", - "schema_plan": state.get("schema_plan"), - "sql_query": sql, - "sql_explanation": explanation, + "schema_plan": "", # Empty string so we don't send massive catalog to UI + "sql_query": state.get("sql_query"), + "sql_explanation": state.get("sql_explanation"), } ) if approval_result.get("approved"): return { - "sql_query": sql, - "refinement_count": 0, - "trino_error": None, "feedback": None, - "execution_path": ["query_builder"], + "execution_path": ["hitl_query_approval"], } else: return { "feedback": approval_result.get("feedback", "Query rejected by user"), "sql_query": None, - "execution_path": ["query_builder"], + "execution_path": ["hitl_query_approval"], } diff --git a/agent/src/agent/nodes/refiner.py b/agent/src/agent/nodes/refiner.py index a8ea42d..2d10069 100644 --- a/agent/src/agent/nodes/refiner.py +++ b/agent/src/agent/nodes/refiner.py @@ -16,16 +16,10 @@ llm = get_llm("refiner") def build_refiner_schema_context(state: AgentState) -> str: - profiles = state.get("table_profiles") - if not profiles: + catalog = state.get("jeen_catalog") + if not catalog: return "No schema context available." - - runtime_flags = state.get("runtime_flags") or {} - limit = int(runtime_flags.get("REFINER_SCHEMA_CONTEXT_TABLES", settings.REFINER_SCHEMA_CONTEXT_TABLES)) - - # Cap the context to REFINER_SCHEMA_CONTEXT_TABLES - capped_profiles = profiles[:limit] - return json.dumps(capped_profiles, indent=2) + return catalog async def refiner_node(state: AgentState, config: RunnableConfig | None = None): diff --git a/agent/src/agent/nodes/schema_explorer.py b/agent/src/agent/nodes/schema_explorer.py index aeed0b1..be1f70a 100644 --- a/agent/src/agent/nodes/schema_explorer.py +++ b/agent/src/agent/nodes/schema_explorer.py @@ -1,623 +1,57 @@ from __future__ import annotations import asyncio import json -import re -import urllib.request -from agent.utils.redis_publisher import publish_node_event -from core.trino import execute_query_sync - -IDENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") - +import logging from typing import Any, List, Optional from pydantic import BaseModel, Field -import logging from agent.state import AgentState - -from langchain_core.prompts import ChatPromptTemplate from langchain_core.runnables import RunnableConfig -from langchain_core.tools import tool -from sqlalchemy import text -from core.db.engine import engine -from core.models.models import Table, TableProfile, ColumnProfile, EnrichmentVersion -from sqlmodel import Session, select from agent.config import settings from agent.langfuse_client import langfuse_client -from agent.llm import get_llm -from agent.utils.schema_enrichment import ( - run_semantic_typing, - run_join_graph, -) -from core.cache import get_cache_service -from core.embeddings import get_embedding +from agent.utils.redis_publisher import publish_node_event +from agent.utils.jeen_metadata_client import get_jeen_metadata_client -# Initialize LLM -llm = get_llm("schema_explorer") logger = logging.getLogger(__name__) -# Cache singleton -_cache = get_cache_service() - -# Skill Registry -from agent.utils.skill_registry import SkillRegistry -from python_core_utils.redis import get_redis_client - -_skill_registry = SkillRegistry() - # G2-02 limits MAX_SCHEMA_RETRIES = 3 - -def _build_column_context(cp: "ColumnProfile") -> dict: - """Build a rich column context dict from a ColumnProfile ORM row. - - Returns all fields the LLM needs to write accurate SQL: - - name, type, semantic_type - - null_rate, distinct_count - - sample_values (top values for categorical/text, or sample values for continuous) - - min, max, mean for numeric/time columns - """ - top_vals = cp.top_values or [] - sample_values = [ - v.get("value") for v in top_vals[:20] if v.get("value") is not None - ] - stats = cp.stats_json or {} - - col: dict = { - "name": cp.column_name, - "type": cp.data_type, - "semantic_type": cp.semantic_type or "unknown", - "null_rate": round(cp.null_rate or 0.0, 4), - "distinct_count": cp.distinct_count or 0, - } - - if cp.is_categorical: - col["sample_values"] = sample_values - else: - # For numeric/time columns expose range and sample values - if cp.min_value is not None: - col["min"] = cp.min_value - if cp.max_value is not None: - col["max"] = cp.max_value - if cp.avg_value is not None: - col["mean"] = round(float(cp.avg_value), 4) - # Pull any sample_values stored in stats_json (continuous columns) - stored_samples = stats.get("sample_values", []) - if stored_samples: - col["sample_values"] = [str(v) for v in stored_samples[:10]] - elif sample_values: - col["sample_values"] = sample_values[:10] - - return col - - -# Define standardized Schema Explorer Output Type -# Ambiguity detection is handled exclusively by the detect_ambiguity node in the -# refiner subgraph — schema_explorer only produces a query plan and table list. -class SchemaExplorerOutput(BaseModel): - schema_plan: str = Field( - default="", - description="Detailed query plan describing tables, columns, and joins. Must be a detailed string explanation.", - ) - tables_used: List[str] = Field( - default_factory=list, - description="List of fully qualified table names (catalog.schema.name) used in the plan.", - ) - - -def get_query_embedding(text: str) -> list[float]: - """Generate 768-dimensional embedding from nomic-embed-text.""" - emb = get_embedding( - text=text, - embedder_url=settings.EMBEDDER_URL, - embedder_model=settings.EMBEDDER_MODEL, - embedder_key=settings.EMBEDDER_KEY, - ) - if emb is None: - print("Error getting query embedding for text") - return [0.0] * 768 - return emb - - -def hybrid_search_tables( - query: str, - query_embedding: list[float], - session: Session, - allowed_tables: list[str] | None = None, - allowed_statuses: list[str] | None = None, - scoping_mode: str = "hybrid", -) -> list[Table]: - """Hybrid search combining pgvector cosine distance and keyword matching. - - G2-01: In strict mode, allowed_tables is a hard allowlist — allowed_statuses - is ignored. In hybrid mode, the union of both filters applies. - """ - stmt_all = select(Table) - all_tables = session.exec(stmt_all).all() - - allowed = allowed_tables or [] - statuses = allowed_statuses or ["production"] - allowed_tables_set = [] - allowed_ids = set() - - for table in all_tables: - if scoping_mode == "strict": - # Hard allowlist: only tables explicitly named in allowed_tables - is_allowed = bool( - allowed - and ( - table.id in allowed - or table.name in allowed - or f"{table.schema_name}.{table.name}" in allowed - ) - ) - else: - # Hybrid: production/status union OR explicit allowed list - is_allowed = table.status in statuses or ( - allowed - and ( - table.id in allowed - or table.name in allowed - or f"{table.schema_name}.{table.name}" in allowed - ) - ) - - if is_allowed: - allowed_tables_set.append(table) - allowed_ids.add(table.id) - - # Vector Search - if allowed_ids: - stmt = text( - """ - SELECT id FROM tables - WHERE id = ANY(:allowed_ids) - ORDER BY embedding <=> :emb - LIMIT :limit - """ - ) - try: - vec_ids = [ - row[0] - for row in session.execute( - stmt, - { - "emb": str(query_embedding), - "allowed_ids": list(allowed_ids), - "limit": settings.HYBRID_SEARCH_MAX_TABLES, - }, - ).fetchall() - ] - except Exception as e: - print(f"Vector search failed: {e}") - vec_ids = [] - else: - vec_ids = [] - - # Keyword Search - keyword_matches = [] - query_words = query.lower().split() - for table in allowed_tables_set: - enrichment = session.exec( - select(EnrichmentVersion) - .where(EnrichmentVersion.table_id == table.id) - .order_by(EnrichmentVersion.version.desc()) - ).first() - - desc = ( - enrichment.data.get("table_description", "") - if enrichment and enrichment.data - else "" - ) - - score = 0 - for word in query_words: - if word in table.name.lower(): - score += 10 - if word in table.schema_name.lower(): - score += 5 - if word in desc.lower(): - score += 2 - - if score > 0: - keyword_matches.append((table.id, score)) - - keyword_matches.sort(key=lambda x: x[1], reverse=True) - kw_ids = [x[0] for x in keyword_matches[: settings.HYBRID_SEARCH_MAX_TABLES]] - - combined_ids = list(dict.fromkeys(vec_ids + kw_ids))[ - : settings.HYBRID_SEARCH_MAX_TABLES - ] - - result_tables = [] - for tid in combined_ids: - t = session.get(Table, tid) - if t: - result_tables.append(t) - return result_tables - - -# Define Tools - - -@tool -async def get_table_profile(table_id: str) -> str: - """Get the lightweight column names/types for a table. Use this before planning a query.""" - cache_hit = False - - with Session(engine) as session: - table = session.get(Table, table_id) - if not table: - return json.dumps({"error": f"Table ID {table_id} not found."}) - - profile = session.exec( - select(TableProfile) - .where( - TableProfile.table_id == table_id, TableProfile.status == "completed" - ) - .order_by(TableProfile.created_at.desc()) - ).first() - - if not profile: - return json.dumps( - { - "error": f"No completed profile found for Table ID {table_id}. Make sure to trigger profiling first." - } - ) - - # ── G2-05: Redis cache lookup ───────────────────────────────────────── - cache_key = _cache.profile_key(table_id, profile.id) - cached = await _cache.get_json(cache_key) - if cached is not None: - cache_hit = True - # Lightweight wrapper returned from cache - return json.dumps(cached) - - columns = session.exec( - select(ColumnProfile).where(ColumnProfile.profile_id == profile.id) - ).all() - - # ── Fetch table description from EnrichmentVersion ───────────────────── - table_description = "" - enrichment = session.exec( - select(EnrichmentVersion) - .where(EnrichmentVersion.table_id == table_id) - .order_by(EnrichmentVersion.version.desc()) - ).first() - if enrichment and enrichment.data: - # Prefer human annotation, fall back to AI summary - table_description = enrichment.data.get( - "table_description", "" - ) or enrichment.data.get("ai_summary", "") - - # Lightweight response to cache and return to LLM - lightweight = { - "table_id": table_id, - "table_name": f"{table.catalog}.{table.schema_name}.{table.name}", - "description": table_description, - "row_count": profile.row_count, - "columns": [_build_column_context(cp) for cp in columns], - } - - # ── G2-05: Populate cache ───────────────────────────────────────────── - await _cache.set_json(cache_key, lightweight, settings.PROFILE_CACHE_TTL) - - return json.dumps(lightweight, indent=2) - - async def schema_explorer_node(state: AgentState, config: RunnableConfig = None): - """RAG Schema Explorer sub-agent node — with G2-01 scoping, G2-03 enrichment, G2-05 caching.""" + """Schema Explorer node — just fetches the full catalog prompt from MCP.""" thread_id = config.get("configurable", {}).get("thread_id", "") if config else "" await publish_node_event(thread_id, "schema_explorer") - user_query = state.get("user_query") - enrichments = state.get("query_enrichments", []) - allowed_tables = state.get("allowed_tables") - allowed_statuses = state.get("allowed_statuses") - feedback = state.get("feedback") - runtime_flags = state.get("runtime_flags") or {} - - # Resolve all flag-tunable parameters for this invocation - profile_fetch_concurrency = int( - runtime_flags.get( - "PROFILE_FETCH_CONCURRENCY", settings.PROFILE_FETCH_CONCURRENCY - ) - ) - max_profiles_to_fetch = int( - runtime_flags.get("MAX_PROFILES_TO_FETCH", settings.MAX_PROFILES_TO_FETCH) - ) - def _parse_bool_flag(value) -> bool: - """Parse a flag value that may be a bool or a string like 'true'/'false'/'0'/'1'.""" - if isinstance(value, bool): - return value - return str(value).lower() in ("true", "1") - - schema_semantic_typing = _parse_bool_flag( - runtime_flags.get("SCHEMA_SEMANTIC_TYPING", settings.ENABLE_SEMANTIC_TYPING) - ) - schema_join_graph = _parse_bool_flag( - runtime_flags.get("SCHEMA_JOIN_GRAPH", settings.ENABLE_JOIN_GRAPH) - ) - schema_summarization = _parse_bool_flag( - runtime_flags.get("SCHEMA_SUMMARIZATION", settings.ENABLE_SCHEMA_SUMMARIZATION) - ) - schema_skill_injection = _parse_bool_flag( - runtime_flags.get("SCHEMA_SKILL_INJECTION", settings.ENABLE_SKILL_INJECTION) - ) - scoping_mode_flag = runtime_flags.get( - "DEFAULT_TABLE_SCOPING_MODE", settings.DEFAULT_TABLE_SCOPING_MODE - ) - - # Per-invocation LLM (supports model switching via execution mode) - _llm = get_llm("schema_explorer", runtime_flags=runtime_flags) - - # ── G2-01: Resolve scoping mode (state > runtime_flag > env default) ───────── - scoping_mode: str = state.get("scoping_mode") or scoping_mode_flag - - # ── G2-05: Cache hit/miss counters (pushed to Langfuse at end) ──────────── - cache_hit_count = 0 - cache_miss_count = 0 - - # 1. Automatically run hybrid search to find candidates - emb = get_query_embedding(user_query) - with Session(engine) as session: - candidate_tables = hybrid_search_tables( - user_query, emb, session, allowed_tables, allowed_statuses, scoping_mode - ) - - tables_info = [] - profile_details = [] - - # 2. Get profiles for top candidate tables (G2-05 cache-aware) - import asyncio - - sem = asyncio.Semaphore(profile_fetch_concurrency) - - async def fetch_profile(t_id, t_name): - nonlocal cache_hit_count, cache_miss_count - async with sem: - try: - # Quick cache check at this level for hit/miss accounting - with Session(engine) as s: - profile_row = s.exec( - select(TableProfile) - .where( - TableProfile.table_id == t_id, - TableProfile.status == "completed", - ) - .order_by(TableProfile.created_at.desc()) - ).first() - if profile_row: - ck = _cache.profile_key(t_id, profile_row.id) - hit = await _cache.get(ck) - if hit is not None: - cache_hit_count += 1 - else: - cache_miss_count += 1 - - profile_res = await get_table_profile.ainvoke({"table_id": t_id}) - return json.loads(profile_res) - except Exception as e: - print(f"Error fetching profile for {t_name}: {e}") - return None - - fetch_tasks = [] - for i, t in enumerate(candidate_tables): - tables_info.append( - { - "id": t.id, - "name": f"{t.catalog}.{t.schema_name}.{t.name}", - "description": "", - } - ) - if i < max_profiles_to_fetch: - fetch_tasks.append(fetch_profile(t.id, t.name)) - - if fetch_tasks: - results = await asyncio.gather(*fetch_tasks, return_exceptions=True) - for res in results: - if res and not isinstance(res, Exception): - profile_details.append(res) - - # ── G2-03: Advanced Schema Enrichment phases ────────────────────────────── - active_phases: list[str] = [] - table_ids = [t.id for t in candidate_tables] - - # human_message = ( - # f"Question: {user_query}\n" - # f"Query Enrichments (extra context for ambiguous terms): {json.dumps(enrichments)}" - # ) - human_message = user_query - if feedback: - human_message += f"\nUser Feedback on previous plan/query: {feedback}" - - # G2-01 strict mode prompt injection - if scoping_mode == "strict": - human_message += ( - "\n\n[STRICT MODE] Only use tables from the approved list. " - "Do not suggest alternatives.\n" - f"Approved tables: {json.dumps(allowed_tables)}" - ) - - # Phase A: Semantic Typing - if schema_semantic_typing and profile_details: - try: - profile_details = await run_semantic_typing(profile_details, _llm) - active_phases.append("SCHEMA_SEMANTIC_TYPING") - except Exception as exc: - logger.warning("SCHEMA_SEMANTIC_TYPING phase failed: %s", exc) - - # Phase B: Join Graph - if schema_join_graph and len(table_ids) >= 2: - try: - join_paths_json = await run_join_graph(table_ids) - if join_paths_json: - human_message += ( - "\n\n[JOIN GRAPH] Shortest join paths between candidate tables:\n" - + join_paths_json - ) - active_phases.append("SCHEMA_JOIN_GRAPH") - except Exception as exc: - logger.warning("SCHEMA_JOIN_GRAPH phase failed: %s", exc) - - # ── G3: Skill Injection ─────────────────────────────────────────────────── - loaded_skills = state.get("loaded_skills") - if schema_skill_injection and loaded_skills: - try: - skill_prompts = _skill_registry.build_system_prompt_addition(loaded_skills) - if skill_prompts: - human_message += f"\n\n[APPLIED SKILLS]{skill_prompts}" - except Exception as e: - logger.warning(f"Failed to inject skills: {e}") - - # Phase C: Schema Summarization (replaces profiles_json in prompt) - profiles_json_str = json.dumps(profile_details, indent=2) - if schema_summarization and profile_details: - try: - summaries = [ - f"[{p.get('table_name', 'unknown')}] {p.get('description', '') or '(no description available)'}" - for p in profile_details - ] - profiles_json_str = "\n".join(summaries) - active_phases.append("SCHEMA_SUMMARIZATION") - except Exception as exc: - logger.warning("SCHEMA_SUMMARIZATION phase failed: %s", exc) - + _jeen = get_jeen_metadata_client() + if not _jeen.is_configured: + logger.warning("MCP client is not configured, but fallback is disabled. Query builder might fail.") + catalog_prompt = "No catalog available (MCP not configured)." + else: + logger.info("Fetching full catalog prompt from Jeen MCP.") + catalog_prompt = await _jeen.get_catalog_prompt() # ── Langfuse trace metadata ─────────────────────────────────────────────── try: trace_id = langfuse_client.get_current_trace_id() if trace_id: - langfuse_client._create_trace_tags_via_ingestion( - trace_id=trace_id, tags=[f"scoping_mode={scoping_mode}"] - ) langfuse_client.update_current_span( metadata={ - "active_schema_phases": active_phases, - "cache_hit_count": cache_hit_count, - "cache_miss_count": cache_miss_count, + "schema_explorer_mode": "mcp_catalog_only", }, ) except Exception as exc: logger.warning("Langfuse trace update failed in schema_explorer: %s", exc) - # 3. Present all metadata to the LLM to construct a query plan - langfuse_prompt = langfuse_client.get_prompt( - settings.LANGFUSE_PROMPT_SCHEMA_EXPLORER - ) - prompt = ChatPromptTemplate.from_messages(langfuse_prompt.get_langchain_prompt()) - - structured_llm = _llm.with_structured_output( - SchemaExplorerOutput, method="json_schema" - ) - chain = prompt | structured_llm - - try: - data = await chain.ainvoke( - { - "tables_json": json.dumps(tables_info, indent=2), - "profiles_json": profiles_json_str, - "human_message": human_message, - } - ) - except Exception as e: - print(f"Structured output parsing failed in schema explorer: {e}") - data = SchemaExplorerOutput(schema_plan="") - - plan = data.schema_plan or "" - - tables_used = getattr(data, "tables_used", []) - - result_state: dict = {"schema_plan": plan, "tables_used": tables_used} - result_state["execution_path"] = ["schema_explorer"] - # Store enriched profiles for downstream nodes (refiner re-uses without re-fetch) - result_state["table_profiles"] = profile_details if profile_details else None - - return result_state - -async def sql_static_validations_node(state: AgentState, config: RunnableConfig = None) -> dict: - """ - Check if tables_used actually exist. - """ - tables_used = state.get("tables_used") or [] - hallucinated = [] - - if tables_used: - try: - redis_client = get_redis_client() - for t_name in tables_used: - cache_key = f"table_exists:{t_name}" - exists = await redis_client.get(cache_key) - if exists is None: - parts = t_name.split(".") - if len(parts) == 3: - cat, sch, tbl = parts - if not IDENT_RE.fullmatch(cat): - hallucinated.append(t_name) - continue - sql = f'SELECT 1 FROM "{cat}".information_schema.tables WHERE table_schema = ? AND table_name = ?' - params = [sch, tbl] - elif len(parts) == 2: - sch, tbl = parts - sql = "SELECT 1 FROM information_schema.tables WHERE table_schema = ? AND table_name = ?" - params = [sch, tbl] - elif len(parts) == 1: - tbl = parts[0] - sql = "SELECT 1 FROM information_schema.tables WHERE table_name = ?" - params = [tbl] - else: - hallucinated.append(t_name) - continue - - try: - res = await asyncio.to_thread( - execute_query_sync, sql, "", params - ) - if res.success and len(res.rows) > 0: - await redis_client.setex(cache_key, 3600, "1") - else: - await redis_client.setex(cache_key, 3600, "0") - hallucinated.append(t_name) - except Exception as e: - logger.error( - f"Information schema check failed for {t_name}: {e}" - ) - # Do not mark as hallucinated on infrastructure failures - elif exists == b"0": - hallucinated.append(t_name) - except Exception as e: - logger.error(f"Error during Redis/Trino table verification: {e}") - # Do not mark as hallucinated on infrastructure failures - - retry_count = state.get("schema_explorer_retry_count", 0) or 0 result_state: dict = { - "schema_explorer_retry_count": retry_count, + "jeen_catalog": catalog_prompt, + "tables_used": [], + "table_profiles": None, + "execution_path": ["schema_explorer"], + "schema_explorer_retry_count": 0, + "hallucinated_tables": None, + "last_error": None } - - if hallucinated: - new_retry = retry_count + 1 - result_state["hallucinated_tables"] = hallucinated - result_state["feedback"] = ( - f"Do not use these tables, they do not exist: {', '.join(hallucinated)}" - ) - result_state["last_error"] = ( - f"Hallucinated tables detected: {', '.join(hallucinated)}" - ) - result_state["schema_explorer_retry_count"] = new_retry - - # G2-02: set escalation_reason when approaching the limit - if new_retry >= MAX_SCHEMA_RETRIES: - result_state["escalation_reason"] = ( - f"Schema explorer failed {new_retry} times due to hallucinated tables: " - f"{', '.join(hallucinated)}" - ) - else: - result_state["hallucinated_tables"] = None - result_state["feedback"] = None - result_state["last_error"] = None - result_state["schema_explorer_retry_count"] = 0 - result_state["execution_path"] = ["sql_static_validations"] return result_state diff --git a/agent/src/agent/routers/chat.py b/agent/src/agent/routers/chat.py index 5ae4390..53ad828 100644 --- a/agent/src/agent/routers/chat.py +++ b/agent/src/agent/routers/chat.py @@ -149,7 +149,7 @@ async def chat_endpoint( thread_id=thread_id, status="interrupted", interrupt_details=interrupt_val, - schema_plan=final_state.values.get("schema_plan") or (interrupt_val.get("schema_plan") if isinstance(interrupt_val, dict) else None), + schema_plan=None, sql_query=final_state.values.get("sql_query") or (interrupt_val.get("sql_query") if isinstance(interrupt_val, dict) else None), ) @@ -160,5 +160,5 @@ async def chat_endpoint( raw_data_ref=result.get("raw_data_ref"), sql_query=result.get("sql_query"), sql_explanation=result.get("sql_explanation"), - schema_plan=result.get("schema_plan"), + schema_plan=None, ) diff --git a/agent/src/agent/state.py b/agent/src/agent/state.py index 6d12ab8..336e4d6 100644 --- a/agent/src/agent/state.py +++ b/agent/src/agent/state.py @@ -9,7 +9,7 @@ class AgentState(TypedDict): execution_path: Annotated[list[str], operator.add] messages: Annotated[list[BaseMessage], add_messages] query_enrichments: list[dict[str, Any]] - schema_plan: str + jeen_catalog: str sql_query: str trino_error: str | None refinement_count: int diff --git a/agent/src/agent/utils/jeen_metadata_client.py b/agent/src/agent/utils/jeen_metadata_client.py new file mode 100644 index 0000000..78bfbdd --- /dev/null +++ b/agent/src/agent/utils/jeen_metadata_client.py @@ -0,0 +1,327 @@ +""" +jeen_metadata_client.py +======================= +MCP client adapter for jeen-metadata. + +Replaces the local Postgres hybrid_search_tables() + get_table_profile() +calls with equivalent calls to the jeen-metadata MCP server. + +MCP tools used +-------------- +- ``search`` → table discovery (replaces hybrid_search_tables) +- ``get_table_profile`` → per-table column stats (replaces get_table_profile tool) +- ``list_tables_rich`` → fallback full-table list with row counts + +Why we use the MCP SDK instead of raw httpx +-------------------------------------------- +jeen-metadata uses WebStandardStreamableHTTPServerTransport from the +@modelcontextprotocol/sdk (TypeScript). That transport requires a proper +MCP handshake before any tool call: + + 1. POST initialize → 200 + JSON (server sends InitializeResult) + 2. POST notifications/initialized → 202 + empty body ← raw httpx dies here + 3. POST tools/call → 200 + JSON or SSE stream + +The Python MCP SDK's streamablehttp_client + ClientSession handles all three +steps automatically and knows how to read both JSON and SSE responses. +Bypassing it with plain httpx would require re-implementing the entire +handshake and SSE framing logic — and that's exactly what caused the +"Expecting value: line 1 column 1 (char 0)" error. + +Configuration (all in AgentSettings / .env) +------------------------------------------- +JEEN_METADATA_MCP_URL Base URL of the MCP endpoint, + e.g. https://jeen-metadata.example.com/api/mcp +JEEN_METADATA_MCP_KEY Bearer token / API key issued by jeen-metadata's + key-management UI (/api/mcp/keys). +JEEN_METADATA_CONNECTION_ID Numeric service ID returned by list_connections. +JEEN_METADATA_SEARCH_LIMIT Max tables returned by the search tool (default 10). +JEEN_METADATA_PROFILE_TIMEOUT Per-call timeout in seconds (default 15). +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + +from agent.config import settings + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Low-level MCP call helper +# --------------------------------------------------------------------------- + +async def _call_mcp_tool( + url: str, + api_key: str, + tool_name: str, + arguments: dict[str, Any], + timeout: float = 15.0, +) -> Any: + """ + Open an MCP session, run the full initialization handshake, call one + tool, and return the parsed payload from the first text content block. + + Uses the official Python MCP SDK so the + initialize → notifications/initialized → tools/call sequence is handled + automatically. Returns parsed JSON (dict/list) or a raw string if the + content block is not JSON. + """ + headers = {"Authorization": f"Bearer {api_key}"} + + async with streamablehttp_client( + url, + headers=headers, + timeout=timeout, + # SSE read timeout slightly longer than the overall timeout so a + # streaming tool call has time to produce its first event. + sse_read_timeout=timeout + 30, + ) as (read, write, _get_session_id): + async with ClientSession(read, write) as session: + await session.initialize() + + result = await session.call_tool(tool_name, arguments) + + # result.content is a list of ContentBlock objects. + # We expect the first (and only) block to be a TextContent. + if not result.content: + return {} + + first = result.content[0] + text = getattr(first, "text", None) + if text is None: + return {} + + try: + return json.loads(text) + except json.JSONDecodeError: + return text + + +# --------------------------------------------------------------------------- +# Public client façade +# --------------------------------------------------------------------------- + +class JeenMetadataClient: + """ + High-level async client for jeen-metadata's MCP server. + + Methods map 1-to-1 to what schema_explorer_node needs: + - search_tables() replaces hybrid_search_tables() + - get_table_profile() replaces the @tool get_table_profile() + - list_tables_rich() fallback when search returns nothing + """ + + def __init__(self) -> None: + self._mcp_url = getattr(settings, "JEEN_METADATA_MCP_URL", "") + self._mcp_key = getattr(settings, "JEEN_METADATA_MCP_KEY", "") + self._connection_id: int = int(getattr(settings, "JEEN_METADATA_CONNECTION_ID", 0)) + self._search_limit: int = int( + getattr(settings, "JEEN_METADATA_SEARCH_LIMIT", settings.HYBRID_SEARCH_MAX_TABLES) + ) + self._timeout: float = float(getattr(settings, "JEEN_METADATA_PROFILE_TIMEOUT", 15.0)) + + self.is_configured: bool = bool( + self._mcp_url and self._mcp_key and self._connection_id + ) + + if not self.is_configured: + logger.info( + "JeenMetadataClient is not fully configured " + "(JEEN_METADATA_MCP_URL / JEEN_METADATA_MCP_KEY / " + "JEEN_METADATA_CONNECTION_ID missing). " + "Schema explorer will fall back to local DB." + ) + + # ── internal helper ──────────────────────────────────────────────────── + + async def _call(self, tool_name: str, arguments: dict[str, Any]) -> Any: + """Thin wrapper so individual methods don't repeat the URL/key/timeout.""" + return await _call_mcp_tool( + self._mcp_url, + self._mcp_key, + tool_name, + arguments, + self._timeout, + ) + + # ------------------------------------------------------------------ + # Full DB Schema (all columns) + # ------------------------------------------------------------------ + + async def get_catalog_prompt(self) -> str: + """ + Fetch the entire catalog context prompt for the connection using the + MCP `get_catalog_prompt` tool. This returns a large markdown string + describing all tables, columns, relationships, and business terms. + """ + try: + payload = await self._call( + "get_catalog_prompt", + { + "connection_id": self._connection_id, + }, + ) + # The MCP tool returns { "content": [{ "type": "text", "text": "..." }] } + # but our _call helper returns the parsed JSON or the raw text block. + # In get_catalog_prompt's case, the content block is the prompt string directly. + + if isinstance(payload, str): + logger.info("JeenMetadataClient.get_catalog_prompt → fetched successfully.") + return payload + elif isinstance(payload, dict) and "prompt" in payload: + # Just in case the MCP returned a JSON string that we parsed + return payload.get("prompt", "") + else: + logger.warning("JeenMetadataClient.get_catalog_prompt → unexpected payload type %s", type(payload)) + return str(payload) + + except Exception as exc: + logger.error( + "JeenMetadataClient.get_catalog_prompt failed: %s", exc, exc_info=True + ) + return "" + + # ------------------------------------------------------------------ + # Table profile (columns + stats) + # ------------------------------------------------------------------ + + async def get_table_profile(self, table_name: str) -> dict[str, Any] | None: + """ + Fetch the latest stored column stats for *table_name* from jeen-metadata. + + Returns a dict shaped identically to the lightweight dict built by + the local ``get_table_profile`` @tool so the rest of + schema_explorer_node is unchanged: + + { + "table_id": str, + "table_name": str, # fully qualified + "description": str, + "row_count": int | None, + "columns": [ + {"name": str, "type": str, "null_rate": float, + "distinct_count": int, ...}, + ... + ] + } + """ + try: + payload = await self._call( + "get_table_profile", + { + "connection_id": self._connection_id, + "table_name": table_name, + }, + ) + + table_profile: dict = payload.get("table_profile") or {} if isinstance(payload, dict) else {} + columns_raw: list[dict] = payload.get("columns") or [] if isinstance(payload, dict) else [] + + # Normalise column list into the shape the agent already consumes + columns = [] + for col in columns_raw: + entry: dict[str, Any] = { + "name": col.get("columnName") or col.get("column_name") or col.get("name", ""), + "type": col.get("dataType") or col.get("data_type") or col.get("type", ""), + "null_rate": round(float(col.get("nullRate") or col.get("null_rate") or 0.0), 4), + "distinct_count": int(col.get("distinctCount") or col.get("distinct_count") or 0), + "semantic_type": col.get("semantic_type") or "unknown", + } + # Propagate optional stats if available + for src_key, dst_key in [ + ("minValue", "min"), ("min_value", "min"), + ("maxValue", "max"), ("max_value", "max"), + ("avgValue", "mean"), ("avg_value", "mean"), + ]: + val = col.get(src_key) + if val is not None: + entry[dst_key] = val + sample = col.get("sampleValues") or col.get("sample_values") + if sample: + entry["sample_values"] = sample + + columns.append(entry) + + result = { + "table_id": table_name, + "table_name": f"{payload.get('table_name', table_name)}" if isinstance(payload, dict) else table_name, + "description": table_profile.get("tableDescription") or "", + "row_count": table_profile.get("rowCount"), + "columns": columns, + } + + logger.info( + "JeenMetadataClient.get_table_profile: table=%r → %d column(s)", + table_name, + len(columns), + ) + return result + + except Exception as exc: + logger.error( + "JeenMetadataClient.get_table_profile(%r) failed: %s", + table_name, + exc, + exc_info=True, + ) + return None + + # ------------------------------------------------------------------ + # Full table listing (fallback when search returns nothing) + # ------------------------------------------------------------------ + + async def list_tables_rich(self) -> list[dict[str, Any]]: + """ + Return ALL tables for the configured connection via ``list_tables_rich``. + Used as a fallback when the search tool returns no results. + """ + try: + rows = await self._call( + "list_tables_rich", + {"connection_id": self._connection_id}, + ) + if not isinstance(rows, list): + rows = [] + tables = [] + for row in rows: + name = row.get("name", "") + tables.append( + { + "id": name, + "name": name, + "schema_name": "", + "catalog": "", + "description": row.get("description") or "", + } + ) + logger.info( + "JeenMetadataClient.list_tables_rich → %d table(s)", len(tables) + ) + return tables + except Exception as exc: + logger.error( + "JeenMetadataClient.list_tables_rich failed: %s", exc, exc_info=True + ) + return [] + + +# --------------------------------------------------------------------------- +# Module-level singleton (lazy-initialised, safe to import at module load) +# --------------------------------------------------------------------------- + +_client: JeenMetadataClient | None = None + + +def get_jeen_metadata_client() -> JeenMetadataClient: + """Return the shared JeenMetadataClient singleton.""" + global _client + if _client is None: + _client = JeenMetadataClient() + return _client diff --git a/agent/tests/test_routing.py b/agent/tests/test_routing.py index 84c953d..c3f4d50 100644 --- a/agent/tests/test_routing.py +++ b/agent/tests/test_routing.py @@ -2,7 +2,7 @@ from unittest.mock import patch, MagicMock, AsyncMock from agent.state import AgentState -from agent.graph import validate_config_node, InvalidConfigurationException, rejection_router_node, route_refiner_subagent, route_schema_explorer +from agent.graph import validate_config_node, InvalidConfigurationException, rejection_router_node, route_refiner_subagent from agent.nodes.refiner import refiner_node from agent.nodes.schema_explorer import MAX_SCHEMA_RETRIES from agent.config import settings @@ -133,45 +133,6 @@ def test_tts_g2_01_scoping_modes_strict_vs_hybrid(): res = validate_config_node(state_strict_pass) assert res["scoping_mode"] == "strict" -def test_tts_g2_02_max_loop_and_hitl_breakpointer(): - state: AgentState = { - "refinement_count": settings.MAX_REFINER_ITERATIONS, - "trino_error": "still failing", - "user_query": "", - "messages": [], - "query_enrichments": [], - "schema_plan": "", - "sql_query": "", - "raw_data_ref": None, - "summary": "", - "sql_explanation": "", - "allowed_tables": None, - "allowed_statuses": None, - "feedback": None, - "feedback_route": None, - "non_interactive": False, - "active_extractors": None, - "last_error": None, - "hallucinated_tables": None, - "esca_write_failed": None, - "inline_result_rows": None, - "error_history": None, - "schema_explorer_retry_count": 0, - "escalated": None, - "escalation_reason": None, - "satisfaction_failures": None, - "satisfaction_fail_count": 0 - } - - # route_refiner_subagent should return hitl_escalation when trino_error is set after hitting subgraph limit - route = route_refiner_subagent(state) - assert route == "hitl_escalation" - - # also test schema explorer max loop - state["hallucinated_tables"] = ["fake_table"] - state["schema_explorer_retry_count"] = MAX_SCHEMA_RETRIES - route2 = route_schema_explorer(state) - assert route2 == "hitl_escalation" def test_tts_g2_03_schema_enrichment_bfs_algorithm(): # Test pure Python BFS shortest path fallback diff --git a/agent/uv.lock b/agent/uv.lock index 15156d1..422d5f8 100644 --- a/agent/uv.lock +++ b/agent/uv.lock @@ -42,9 +42,9 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.115.6" }, { name = "greenlet", specifier = ">=3.5.1" }, { name = "json-repair", specifier = ">=0.25.0" }, - { name = "langchain" }, + { name = "langchain", specifier = "==1.3.13" }, { name = "langchain-ollama" }, - { name = "langchain-openai" }, + { name = "langchain-openai", specifier = "==1.3.5" }, { name = "langfuse", specifier = ">=2.0.0" }, { name = "langgraph" }, { name = "mcp", specifier = ">=1.12.4" }, @@ -219,8 +219,8 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "alembic", specifier = "==1.14.0" }, - { name = "langchain" }, - { name = "langchain-openai" }, + { name = "langchain", specifier = "==1.3.13" }, + { name = "langchain-openai", specifier = "==1.3.5" }, { name = "pgvector", specifier = ">=0.2.0" }, { name = "psycopg2-binary", specifier = "==2.9.9" }, { name = "pydantic", specifier = "==2.10.4" }, @@ -510,21 +510,21 @@ wheels = [ [[package]] name = "langchain" -version = "1.3.9" +version = "1.3.13" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "langgraph" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/56/7c/651d0dc4913a7a892156c03dd343b99cfe19ee729e6911ab1f4fe7567b8b/langchain-1.3.9.tar.gz", hash = "sha256:9b14ef0db9ef314299ded858b22ca2a40b8f1b05c8c9cb6b82d53a53075fef00", size = 631514, upload-time = "2026-06-12T16:53:27.083Z" } +sdist = { url = "https://files.pythonhosted.org/packages/29/33/fd716d0273c8495482953bd63461bf02f71f3a8f4a2fe6c0a70a0e6ff799/langchain-1.3.13.tar.gz", hash = "sha256:bcf874680f31e9970f0db2264509df5bc2115d9680e9d651d537eb49bf1a7d8a", size = 642868, upload-time = "2026-07-10T23:06:08.555Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/55/3481619d21b9bdfbfda8680fba5cfc6cfe926789b8eaaad95353078cfa20/langchain-1.3.9-py3-none-any.whl", hash = "sha256:4af49ad1095799e4408b489fb79d4b8b49292453618b202d8a697fca59bb6871", size = 132873, upload-time = "2026-06-12T16:53:25.489Z" }, + { url = "https://files.pythonhosted.org/packages/95/c6/dc676c632f3d20c88789b0726c43ad5e039c25338cc3bb7090fc247d1522/langchain-1.3.13-py3-none-any.whl", hash = "sha256:20a8fe4b1dea7db74356f7d2b5455c4970099b9f7f53c2122ea97f115c907fbd", size = 136911, upload-time = "2026-07-10T23:06:07.012Z" }, ] [[package]] name = "langchain-core" -version = "1.4.7" +version = "1.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonpatch" }, @@ -537,9 +537,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "uuid-utils" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ee/2b/fffaff399d20a56d40b9562fa19701e91abd72d8c9d9bc8c2673077b56b6/langchain_core-1.4.7.tar.gz", hash = "sha256:7a825d77de0a3f39adbd9d09612a75e85527e14a52c1601089bcc062972d9f2b", size = 952522, upload-time = "2026-06-12T19:23:57.588Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/05/986c4bb148285791eb59994e0b28947bed96cac7f24467079e4274952a37/langchain_core-1.5.0.tar.gz", hash = "sha256:e1fa09d55b354192c8f60dade06a55bd6add2318c822a684555b8d4a30a16143", size = 967401, upload-time = "2026-07-21T03:37:26.48Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/de/3e/dcdffa60078ae7b3a00ebb4cbbf1a204a14c3609983c604886523a7d4418/langchain_core-1.4.7-py3-none-any.whl", hash = "sha256:bcadd51951140ecdcba98311dbd931ba5de02a5ba8a2288dad5069c1eea2a13d", size = 554941, upload-time = "2026-06-12T19:23:55.826Z" }, + { url = "https://files.pythonhosted.org/packages/29/56/5ef7ba14bac95b0344da18c6e8ec108dce0baf5fc054d1117702f92af29d/langchain_core-1.5.0-py3-none-any.whl", hash = "sha256:f122efee35446632b38687119fca33711abbf3b6b555e31156762298fbe78a65", size = 558510, upload-time = "2026-07-21T03:37:24.423Z" }, ] [[package]] @@ -557,16 +557,16 @@ wheels = [ [[package]] name = "langchain-openai" -version = "1.3.2" +version = "1.3.5" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "openai" }, { name = "tiktoken" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/93/4c/cf3c5a03f1d2e2e4367c1527231162a99d0f1c94113e1203c00469c860e4/langchain_openai-1.3.2.tar.gz", hash = "sha256:240917ae88d754b389a6f2ae06fa262c50c094eb4f576c27d560dff6b86c2f62", size = 3236213, upload-time = "2026-06-13T05:42:12.5Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d1/7e/43eef3f8fae2668f52e2222fdc26b6de58acf158bcb580e32e88a299260d/langchain_openai-1.3.5.tar.gz", hash = "sha256:c1db2256a42ac46e8e7b0564c5ccb478b9f58dc047a58935da33c82e6e1f9a07", size = 3261548, upload-time = "2026-07-10T18:58:29.576Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/03/21/cbf6c3786de881b214c8c6c9f61fe44c9c47608428676a5cd5c5b2b0cda5/langchain_openai-1.3.2-py3-none-any.whl", hash = "sha256:3d247f43bba9f85d32a374b1bdf3932a0d1e3c60913ebeadf68630de52add67e", size = 119775, upload-time = "2026-06-13T05:42:11.088Z" }, + { url = "https://files.pythonhosted.org/packages/61/64/4e0918cb96ff2b49e06acd9c11c250297d727d2fcce9e012d62efb73b4d6/langchain_openai-1.3.5-py3-none-any.whl", hash = "sha256:f586263b884bceb3d426ec84d3bfbd27051c3c92ae668da6175629e3f44dcec5", size = 121601, upload-time = "2026-07-10T18:58:28.327Z" }, ] [[package]] @@ -792,7 +792,7 @@ wheels = [ [[package]] name = "openai" -version = "2.41.1" +version = "2.47.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -804,9 +804,9 @@ dependencies = [ { name = "tqdm" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/40/36/4c926a91554483977608951360c18c2e911592785eb87a6437813f6123f7/openai-2.41.1.tar.gz", hash = "sha256:23d617a0432457ad844973bee8f540be9da90894f7c5686852d2d365da058f57", size = 783584, upload-time = "2026-06-10T16:10:37.667Z" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/61/9aeef14de759306e85175126d3d6d56ee4f5072a9512c6c171d58d02a62d/openai-2.47.0.tar.gz", hash = "sha256:4e205548acd4304f235b86202269912e55bc88270b15d2a051fa2b53b90343a6", size = 1089906, upload-time = "2026-07-22T17:47:29.723Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/74/925d7b3892927e9804aaf58d374a45dc28e4420ff90e992272b77286343e/openai-2.41.1-py3-none-any.whl", hash = "sha256:a939565f350cb7443cb843b801b88c716ac8024b492fb94ca269d5f6b1bbefd6", size = 1353380, upload-time = "2026-06-10T16:10:35.756Z" }, + { url = "https://files.pythonhosted.org/packages/41/69/26b032059273ad798d18fbcdbe369e871181841fd8bcb5caee32b7510039/openai-2.47.0-py3-none-any.whl", hash = "sha256:b3a1a7ad974092427ccb46d89f8852bdb67866680bcabeecc3ff5a3fdd71b15b", size = 1639987, upload-time = "2026-07-22T17:47:27.873Z" }, ] [[package]] From 020e756b71a386651b6134f5dadaf27513eeb7f3 Mon Sep 17 00:00:00 2001 From: yuvalkh Date: Thu, 23 Jul 2026 17:43:24 +0300 Subject: [PATCH 2/6] removed schema plan from agent and front --- agent/src/agent/config.py | 4 +- agent/src/agent/mcp_server.py | 3 - agent/src/agent/routers/chat.py | 3 - frontend/src/api/agent.ts | 1 - .../components/SchemaPlanDisplay.module.css | 52 --- frontend/src/components/SchemaPlanDisplay.tsx | 298 ------------------ frontend/src/pages/AgentTestingPage.tsx | 21 -- 7 files changed, 2 insertions(+), 380 deletions(-) delete mode 100644 frontend/src/components/SchemaPlanDisplay.module.css delete mode 100644 frontend/src/components/SchemaPlanDisplay.tsx diff --git a/agent/src/agent/config.py b/agent/src/agent/config.py index 2f8d429..468760e 100644 --- a/agent/src/agent/config.py +++ b/agent/src/agent/config.py @@ -10,7 +10,7 @@ class AgentSettings(BaseSettings): ESCA_API_KEY: str = "" ESCA_URL: str = "http://localhost:7010" - ESCA_WRITE_ENABLED: bool = True + ESCA_WRITE_ENABLED: bool = False LLM_API_KEY: str = "ollama" LLM_BASE_URL: str = "http://localhost:11434/v1" LLM_MODEL: str = "gemma4:e4b" @@ -81,7 +81,7 @@ class AgentSettings(BaseSettings): SATISFACTION_CHECK_ENABLED: bool = True SATISFACTION_CHECK_EXECUTION: bool = True SATISFACTION_CHECK_PLAUSIBILITY: bool = True - SATISFACTION_CHECK_COLUMNS: bool = True + SATISFACTION_CHECK_COLUMNS: bool = False SATISFACTION_CHECK_SEMANTIC: bool = False # LLM-heavy, off by default SATISFACTION_MIN_ROWS: int = 1 SATISFACTION_MAX_ROWS: int = 50_000 diff --git a/agent/src/agent/mcp_server.py b/agent/src/agent/mcp_server.py index 1bd8169..e937ff9 100644 --- a/agent/src/agent/mcp_server.py +++ b/agent/src/agent/mcp_server.py @@ -162,7 +162,6 @@ async def chat_with_agent( "thread_id": thread_id, "status": "interrupted", "interrupt_details": interrupt_val, - "schema_plan": None, "sql_query": final_state.values.get("sql_query") or (interrupt_val.get("sql_query") if isinstance(interrupt_val, dict) else None), "sql_explanation": interrupt_val.get("sql_explanation") if isinstance(interrupt_val, dict) else None, "trace_id": trace_id, @@ -183,7 +182,6 @@ async def chat_with_agent( "thread_id": thread_id, "status": "interrupted", "interrupt_details": interrupt_val, - "schema_plan": None, "sql_query": final_state.values.get("sql_query"), "trace_id": trace_id, "execution_path": final_state.values.get("execution_path", []), @@ -202,7 +200,6 @@ async def chat_with_agent( "raw_data_ref": result.get("raw_data_ref"), "sql_query": result.get("sql_query"), "sql_explanation": result.get("sql_explanation"), - "schema_plan": None, "trace_id": trace_id, "execution_path": result.get("execution_path", []), "is_unanswerable": is_unans, diff --git a/agent/src/agent/routers/chat.py b/agent/src/agent/routers/chat.py index 53ad828..301b594 100644 --- a/agent/src/agent/routers/chat.py +++ b/agent/src/agent/routers/chat.py @@ -41,7 +41,6 @@ class ChatResponse(BaseModel): raw_data_ref: str | None = None sql_query: str | None = None sql_explanation: str | None = None - schema_plan: str | None = None @router.post( @@ -149,7 +148,6 @@ async def chat_endpoint( thread_id=thread_id, status="interrupted", interrupt_details=interrupt_val, - schema_plan=None, sql_query=final_state.values.get("sql_query") or (interrupt_val.get("sql_query") if isinstance(interrupt_val, dict) else None), ) @@ -160,5 +158,4 @@ async def chat_endpoint( raw_data_ref=result.get("raw_data_ref"), sql_query=result.get("sql_query"), sql_explanation=result.get("sql_explanation"), - schema_plan=None, ) diff --git a/frontend/src/api/agent.ts b/frontend/src/api/agent.ts index 98b0912..cd37bee 100644 --- a/frontend/src/api/agent.ts +++ b/frontend/src/api/agent.ts @@ -42,7 +42,6 @@ export interface ChatResponse { raw_data_ref?: string; sql_query?: string; sql_explanation?: string; - schema_plan?: string; trace_id?: string; execution_path?: string[]; is_unanswerable?: boolean; diff --git a/frontend/src/components/SchemaPlanDisplay.module.css b/frontend/src/components/SchemaPlanDisplay.module.css deleted file mode 100644 index 985d34c..0000000 --- a/frontend/src/components/SchemaPlanDisplay.module.css +++ /dev/null @@ -1,52 +0,0 @@ -.schemaPlanContainer { - display: flex; - flex-direction: column; - gap: 16px; -} - -.section { - display: flex; - flex-direction: column; - gap: 8px; -} - -.sectionHeader { - display: flex; - align-items: center; - gap: 8px; - font-weight: 600; - color: var(--text-h); - font-size: 14px; -} - -.dataTable { - background: var(--bg-secondary); - border-radius: 8px; - overflow: hidden; -} - -.dataTable :global(.ant-table) { - background: transparent; - color: var(--text-body); -} - -.dataTable :global(.ant-table-thead > tr > th) { - background: var(--bg-tertiary) !important; - color: var(--text-muted) !important; - border-bottom: 1px solid var(--border-color) !important; - font-weight: 500; -} - -.dataTable :global(.ant-table-tbody > tr > td) { - border-bottom: 1px solid var(--border-color) !important; -} - -.dataTable :global(.ant-table-tbody > tr:hover > td) { - background: rgba(255, 255, 255, 0.02) !important; -} - -.infoCard { - background: var(--bg-secondary); - border: 1px solid var(--border-color); - color: var(--text-body); -} diff --git a/frontend/src/components/SchemaPlanDisplay.tsx b/frontend/src/components/SchemaPlanDisplay.tsx deleted file mode 100644 index ba847ee..0000000 --- a/frontend/src/components/SchemaPlanDisplay.tsx +++ /dev/null @@ -1,298 +0,0 @@ -import React from 'react'; -import ReactMarkdown from 'react-markdown'; -import { Card, Space, Table, Tag, Typography } from 'antd'; -import { Database, Filter, Link, ListOrdered, Sparkles } from 'lucide-react'; -import remarkGfm from 'remark-gfm'; - -import styles from './SchemaPlanDisplay.module.css'; - -const { Text } = Typography; - -interface SchemaPlanDisplayProps { - planString: string; -} - -const renderCellSafe = (val: any): string => { - if (val === null || val === undefined) return ''; - if (typeof val === 'object') { - return val.column_name ?? val.name ?? JSON.stringify(val); - } - return String(val); -}; - -export const SchemaPlanDisplay: React.FC = ({ planString }) => { - let planData: any = null; - try { - planData = JSON.parse(planString); - } catch (e) { - // If it's not JSON, render it as markdown - return ( -
- {planString} -
- ); - } - - const explanationText = - planData.description || - planData.explanation || - planData.strategy || - planData.reasoning || - planData.logic; - - // Define columns for Tables - const tableColumns = [ - { - title: 'Table Name', - key: 'name', - render: (record: any) => { - const name = record.name || record.table_name || record.tableName || record.table || ''; - return ( - - {renderCellSafe(name)} - - ); - }, - }, - { - title: 'Columns', - key: 'columns', - render: (record: any) => { - const columns = record.columns || record.column_names || record.columnNames || []; - if (!columns || columns.length === 0) { - return ( - - No columns selected / empty - - ); - } - return ( - - {columns?.map((col: any, idx: number) => { - if (typeof col === 'string') { - return ( - - {col} - - ); - } else if (col && typeof col === 'object') { - const name = col.column_name || col.name || `col_${idx}`; - const type = col.data_type || col.type || ''; - return ( - - {name}{' '} - {type ? ({type}) : ''} - - ); - } - return null; - })} - - ); - }, - }, - ]; - - // Define columns for Joins - const joinColumns = [ - { - title: 'Source Table', - key: 'source_table', - render: (record: any) => { - const val = - record.source_table || record.sourceTable || record.srcTable || record.src_table || ''; - return {renderCellSafe(val)}; - }, - }, - { - title: 'Source Column', - key: 'source_column', - render: (record: any) => { - const val = - record.source_column || - record.sourceColumn || - record.srcColumn || - record.src_column || - ''; - return renderCellSafe(val); - }, - }, - { - title: 'Target Table', - key: 'target_table', - render: (record: any) => { - const val = - record.target_table || - record.targetTable || - record.destTable || - record.dest_table || - record.tgtTable || - record.tgt_table || - ''; - return {renderCellSafe(val)}; - }, - }, - { - title: 'Target Column', - key: 'target_column', - render: (record: any) => { - const val = - record.target_column || - record.targetColumn || - record.destColumn || - record.dest_column || - record.tgtColumn || - record.tgt_column || - ''; - return renderCellSafe(val); - }, - }, - { - title: 'Join Type', - key: 'type', - render: (record: any) => { - const val = record.type || record.join_type || record.joinType || 'INNER'; - return {renderCellSafe(val).toUpperCase()}; - }, - }, - ]; - - // Define columns for Filters - const filterColumns = [ - { - title: 'Column', - key: 'column', - render: (record: any) => { - const val = record.column || record.column_name || record.columnName || record.col || ''; - return {renderCellSafe(val)}; - }, - }, - { - title: 'Operator', - key: 'operator', - render: (record: any) => { - const val = record.operator || record.op || ''; - return {renderCellSafe(val)}; - }, - }, - { - title: 'Value', - key: 'value', - render: (record: any) => { - const val = record.value ?? record.val ?? ''; - return {renderCellSafe(val)}; - }, - }, - ]; - - return ( -
- {explanationText && ( -
- -
- - Query Logic & Strategy -
- - {renderCellSafe(explanationText)} - - -
- )} - - {planData.tables !== undefined && ( -
-
- Tables & Columns -
- r.name || r.table_name || r.tableName || r.table || i.toString()} - pagination={false} - size="small" - bordered - className={styles.dataTable} - /> - - )} - - {planData.joins && planData.joins.length > 0 && ( -
-
- Joins -
-
i.toString()} - pagination={false} - size="small" - bordered - className={styles.dataTable} - /> - - )} - - {planData.filters && planData.filters.length > 0 && ( -
-
- Filters -
-
i.toString()} - pagination={false} - size="small" - bordered - className={styles.dataTable} - /> - - )} - - {(planData.order_by?.length > 0 || planData.limit !== undefined) && ( -
-
- Sorting & Limits -
- - {planData.order_by?.length > 0 && ( -
- Order By: - {planData.order_by.map((ob: any, i: number) => ( - - {renderCellSafe(ob.column)}{' '} - {renderCellSafe(ob.direction)?.toUpperCase() || 'ASC'} - - ))} -
- )} - {planData.limit !== undefined && ( -
- Limit: {planData.limit} -
- )} - -
- )} - - ); -}; diff --git a/frontend/src/pages/AgentTestingPage.tsx b/frontend/src/pages/AgentTestingPage.tsx index 1208389..dfbf9b9 100644 --- a/frontend/src/pages/AgentTestingPage.tsx +++ b/frontend/src/pages/AgentTestingPage.tsx @@ -30,7 +30,6 @@ import { v4 as uuidv4 } from 'uuid'; import { agentApi } from '../api/agent'; import { AgentGraph } from '../components/AgentGraph'; -import { SchemaPlanDisplay } from '../components/SchemaPlanDisplay'; import { highlightJson, TraceTimeline } from '../components/TraceTimeline'; import type { ChatRequest, ChatResponse } from '../api/agent'; @@ -466,7 +465,6 @@ const AgentApprovalForm = ({ const interruptType = interrupt.type; const sqlQuery = chatResponse.sql_query || (interrupt.sql_query as string) || ''; - const schemaPlan = chatResponse.schema_plan || (interrupt.schema_plan as string) || ''; const sqlExplanation = chatResponse.sql_explanation || (interrupt.sql_explanation as string) || ''; @@ -602,25 +600,6 @@ const AgentApprovalForm = ({
- {schemaPlan && ( -
-
- - Proposed Schema Plan -
- -
- )} - {sqlQuery && (
Date: Thu, 23 Jul 2026 20:04:40 +0300 Subject: [PATCH 3/6] updated ambiguity to work with the new schema builder --- agent/scripts/upload_all_prompts.py | 13 +++++++----- agent/src/agent/config.py | 2 +- agent/src/agent/nodes/detect_ambiguity.py | 24 ++++++++++++----------- 3 files changed, 22 insertions(+), 17 deletions(-) diff --git a/agent/scripts/upload_all_prompts.py b/agent/scripts/upload_all_prompts.py index bca974f..b0fc518 100644 --- a/agent/scripts/upload_all_prompts.py +++ b/agent/scripts/upload_all_prompts.py @@ -178,7 +178,8 @@ def main(): "2. **Allow Logical Inference and Heuristics:** Assume the downstream SQL generator can handle implicit table/column choices based on schema structure (e.g., choosing `active_users` over `archive_users` for \"current users\") unless there is a **direct conflict**. However, **do NOT assume fuzzy matching**. If the user's term does not explicitly match the schema values or column names, and there is no single obvious exact match, this may be an ambiguity.\n" "3. **Intervene Only When Necessary:** Flag ambiguity ONLY when multiple interpretations lead to **drastically different data** or when the data is missing. Do not ask for optional details like specific date ranges, window sizes, or table names unless the request is genuinely unintelligible.\n" "4. **Quality Assurance Auditor:** you will review the `Current Agent SQL Attempt` as the Agent's **proposed interpretation**. You do not reject the Agent's choice out of skepticism. Instead, you verify if the Agent's choice aligns with the **strongest available heuristic** in the schema.\n\n" - "You do not generate the final SQL. Instead, you act as the system's execution planner and ambiguity detector. You must protect the downstream SQL Composer from hallucinations, far fetched assumptions, and impossible requests by flaging queries that lack the necessary context to generate a logically correct SQL query. If a user query is ambiguous you will halt execution and formulate a user friendly clarification request.\n\n" + "You do not generate the final SQL. Instead, you act as the system's execution planner and ambiguity detector. You must protect the downstream SQL Composer from hallucinations, far fetched assumptions, and impossible requests by flaging queries that lack the necessary context to generate a logically correct SQL query. If a user query is ambiguous you will halt execution and formulate a user friendly clarification request.\n" + "CRITICAL: If the user refers to a specific entity or filter but DOES NOT provide the actual name or ID, you MUST flag it as AMBIGUOUS. Do NOT accept an Agent SQL Attempt that guesses, drops the filter, or makes a generic query instead.\n\n" "# **EXECUTION WORKFLOW**\n" "Before making a final determination, you must rigorously process the query through the following chronological steps. You will output this internal reasoning step-by-step.\n\n" "1. Intent Deconstruction: Break down the natural language query into core components (desired output columns, temporal filters, aggregations, mathematical operations).\n" @@ -213,7 +214,8 @@ def main(): " - **Missing Critical Logic:** User asks for \"Profit Margin\" but schema has no such column and no price/cost columns to derive it.\n" " - **Contradiction:** User requests data that logically cannot exist together.\n\n" "4. **Agent Proposal Audit**\n" - " Compare the `Current Agent SQL Attempt` against the **Dominant Standard** found in Step 2.\n" + " Compare the `Current Agent SQL Attempt` AND the `Agent's Explanation for SQL` against the **Dominant Standard** found in Step 2.\n" + " **CRITICAL RULE:** If the `Agent's Explanation for SQL` explicitly states that it is ignoring a missing parameter, assuming a generic fallback, or guessing a value because the user didn't provide one (e.g., 'Since the specific country is not provided... I will write a query that counts all entries'), YOU MUST FLAG THIS AS AMBIGUOUS. Do not accept the generic SQL.\n" " **Clear Standard Exists**\n" " * Did the Agent use the Dominant Standard?\n" " * YES → **CLEAR**.\n" @@ -236,8 +238,8 @@ def main(): "* **No Value Match:** The user specifies a filter value that does not exist in the relevant column, and no standard alias exists.\n" " *Example: User asks for \"Sales in 'North America'\" but the `region` column only contains 'NA', 'EU', 'APAC'. If 'NA' is the only logical match, CLEAR. If 'North America' could map to multiple ambiguous codes or none, FLAG.*\n\n" "## B. Database-Related Ambiguity (Schema & Mapping Failures)\n" - "* **Missing Explicit Filter Value:** The user asks to filter by a specific entity (e.g. \"my specific ID\", \"that user\", \"a certain order\") but does NOT provide the actual value. Do NOT assume it is a parameter to be filled later. You MUST flag this as AMBIGUOUS and ask the user for the exact value.\n" - " *Example: User asks for \"Give the order with my specific ID\" but provides no ID. → FLAG.*\n" + "* **Missing Explicit Filter Value:** The user asks to filter by a specific entity but does NOT provide the actual value. Do NOT assume it is a parameter to be filled later. Do NOT accept queries that just ignore the filter. You MUST flag this as AMBIGUOUS and ask the user for the exact value.\n" + " *Example: User asks for \"My location?\" but provides no location. → FLAG.*\n" "* **Direct Schema Collision:** The user asks for a concept that maps to **two or more equally valid columns/tables** without sufficient context to prefer one. Guessing would lead to significantly different data.\n" " *Example: User asks for \"Location\". Schema has `shipping_address`, `billing_address`, and `current_gps`. No context provided. → FLAG.*\n" " *Example: User asks for \"Revenue\". Schema has `gross_revenue` and `net_revenue`. No context provided. → FLAG.*\n" @@ -279,7 +281,8 @@ def main(): "content": ( "User Request: {{user_query}}\n\n" "Schema Context:\n{{schema}}\n\n" - "Current Agent SQL Attempt:\n{{current_sql_attempt}}" + "Current Agent SQL Attempt:\n{{current_sql_attempt}}\n\n" + "Agent's Explanation for SQL:\n{{sql_explanation}}" ) } ], diff --git a/agent/src/agent/config.py b/agent/src/agent/config.py index 468760e..55e920c 100644 --- a/agent/src/agent/config.py +++ b/agent/src/agent/config.py @@ -38,7 +38,7 @@ class AgentSettings(BaseSettings): # ── Jeen Metadata MCP Integration ───────────────────────────────────────── # When set, schema_explorer pulls tables/profiles from jeen-metadata via MCP # instead of from the local Postgres DB. Leave empty to keep the local path. - JEEN_METADATA_MCP_URL: str = "http://localhost:3001/api/mcp" # e.g. https://jeen-metadata.example.com/api/mcp + JEEN_METADATA_MCP_URL: str = "http://schema-modeler.dev161.internal/api/mcp" # e.g. https://jeen-metadata.example.com/api/mcp JEEN_METADATA_MCP_KEY: str = "mcp_f885337e381366db5edc22093415450e38f71e997e96dc708fea69bde9529ab9" # Bearer key from /api/mcp/keys in jeen-metadata JEEN_METADATA_CONNECTION_ID: int = 89 # Numeric service ID (from list_connections) JEEN_METADATA_SEARCH_LIMIT: int = 10 # Max tables returned by the search tool diff --git a/agent/src/agent/nodes/detect_ambiguity.py b/agent/src/agent/nodes/detect_ambiguity.py index 522fa1c..2eba5c9 100644 --- a/agent/src/agent/nodes/detect_ambiguity.py +++ b/agent/src/agent/nodes/detect_ambiguity.py @@ -18,12 +18,15 @@ Using schema_plan (rather than waiting for actual SQL) means: - Ambiguity is caught BEFORE query_builder runs (saves 1 LLM call) - Ambiguity is caught BEFORE Trino executes (saves the DB round-trip) - - The schema_plan is actually MORE explicit about intent than the final SQL, because it spells out the reasoning in plain language. - False positives from SQL syntax errors are impossible (no SQL exists yet). +NOTE: The graph was reordered so detect_ambiguity runs AFTER query_builder. +Now, it evaluates BOTH the `sql_query` and the `sql_explanation` to see if +the query builder had to guess or drop a filter. + If a query is genuinely ambiguous, we short-circuit the entire expensive -pipeline — no query_builder, no refiner, no Trino — and return a clarifying +pipeline — no refiner, no Trino — and return a clarifying question to the user immediately. Graph position (main agent graph): @@ -56,9 +59,8 @@ def _resolve_ambiguity_type(parsed: dict) -> str: ambiguity_type = parsed.get("ambiguity_type") clarifying: str | None = parsed.get("clarifying_questions") - # If the LLM generated a clarifying question (and it's not null/empty), it MUST be ambiguous - # regardless of whether it accidentally classified it as unanswerable. - if clarifying is not None and clarifying.strip(): + # If the LLM classified it as unanswerable but generated a clarifying question, it MUST be ambiguous. + if ambiguity_type == "unanswerable" and clarifying is not None and clarifying.strip(): return "ambiguous" if ambiguity_type in ["clear", "ambiguous", "unanswerable"]: @@ -80,11 +82,10 @@ async def detect_ambiguity_node( config: RunnableConfig | None = None, ) -> dict: """ - Pre-SQL ambiguity gate — runs after schema_explorer, before query_builder. + Pre-SQL ambiguity gate — previously pre-SQL, now runs AFTER query_builder. - Reads the schema_plan (the agent's natural-language query plan describing - which tables and columns it intends to use), the user query, and the enriched - schema profiles, then asks an LLM whether the interpretation is deterministic + Reads the jeen_catalog, the user query, the SQL attempt, and the SQL explanation, + then asks an LLM whether the interpretation is deterministic or ambiguous/unanswerable. Returns a partial state update with: @@ -130,8 +131,9 @@ async def detect_ambiguity_node( # contains the sql_query from the query builder. user_message = ( f"User Request: {state.get('user_query', '')}\n\n" - f"Schema Context:\n{state.get('jeen_catalog', '')}\n\n" - f"Current Agent SQL Attempt:\n{state.get('sql_query', '')}\n" + # f"Schema Context:\n{state.get('jeen_catalog', '')}\n\n" + f"Current Agent SQL Attempt:\n{state.get('sql_query', '')}\n\n" + f"Agent's Explanation for SQL:\n{state.get('sql_explanation', '')}\n" ) # ── LLM call ───────────────────────────────────────────────────────────── From 2c09cf24d8f972a612d62482976f76f3edf8afbf Mon Sep 17 00:00:00 2001 From: yuvalkh Date: Thu, 23 Jul 2026 22:27:21 +0300 Subject: [PATCH 4/6] small fixes from cr --- agent/scripts/upload_all_prompts.py | 2 +- agent/src/agent/nodes/detect_ambiguity.py | 1 - agent/src/agent/nodes/query_builder.py | 2 +- agent/src/agent/nodes/schema_explorer.py | 6 +----- agent/src/agent/state.py | 4 +--- agent/src/agent/utils/jeen_metadata_client.py | 4 ++-- 6 files changed, 6 insertions(+), 13 deletions(-) diff --git a/agent/scripts/upload_all_prompts.py b/agent/scripts/upload_all_prompts.py index b0fc518..d4fa141 100644 --- a/agent/scripts/upload_all_prompts.py +++ b/agent/scripts/upload_all_prompts.py @@ -84,7 +84,7 @@ def main(): "prompt": [ { "role": "system", - "content": "You are a SQL expert who specializes in trino. Build a SQL query based on the catalog and user query. Output ONLY the SQL query, nothing else.\n\nIMPORTANT: Before writing the SQL, provide a brief 1-3 sentence explanation of your reasoning or how you answered the question.\n\nCRITICAL: You MUST fully qualify all tables in your SQL query using the provided Catalog and Schema parameters. (Format: {{trino_catalog_name}}.{{trino_schema_name}}.my_table). Do not use unqualified table names." + "content": "You are a SQL expert who specializes in trino. Build a SQL query based on the catalog and user query.\n\nIMPORTANT: Before writing the SQL, provide a brief 1-3 sentence explanation of your reasoning or how you answered the question.\n\nCRITICAL: You MUST fully qualify all tables in your SQL query using the provided Catalog and Schema parameters. (Format: {{trino_catalog_name}}.{{trino_schema_name}}.my_table). Do not use unqualified table names." }, { "role": "user", diff --git a/agent/src/agent/nodes/detect_ambiguity.py b/agent/src/agent/nodes/detect_ambiguity.py index 2eba5c9..3243372 100644 --- a/agent/src/agent/nodes/detect_ambiguity.py +++ b/agent/src/agent/nodes/detect_ambiguity.py @@ -131,7 +131,6 @@ async def detect_ambiguity_node( # contains the sql_query from the query builder. user_message = ( f"User Request: {state.get('user_query', '')}\n\n" - # f"Schema Context:\n{state.get('jeen_catalog', '')}\n\n" f"Current Agent SQL Attempt:\n{state.get('sql_query', '')}\n\n" f"Agent's Explanation for SQL:\n{state.get('sql_explanation', '')}\n" ) diff --git a/agent/src/agent/nodes/query_builder.py b/agent/src/agent/nodes/query_builder.py index 3ebd1b9..d7c66cf 100644 --- a/agent/src/agent/nodes/query_builder.py +++ b/agent/src/agent/nodes/query_builder.py @@ -96,7 +96,7 @@ async def hitl_query_approval_node(state: AgentState, config: RunnableConfig | N } else: return { - "feedback": approval_result.get("feedback", "Query rejected by user"), + "feedback": approval_result.get("feedback") or "Query rejected by user", "sql_query": None, "execution_path": ["hitl_query_approval"], } diff --git a/agent/src/agent/nodes/schema_explorer.py b/agent/src/agent/nodes/schema_explorer.py index be1f70a..7a4b4da 100644 --- a/agent/src/agent/nodes/schema_explorer.py +++ b/agent/src/agent/nodes/schema_explorer.py @@ -18,7 +18,7 @@ # G2-02 limits MAX_SCHEMA_RETRIES = 3 -async def schema_explorer_node(state: AgentState, config: RunnableConfig = None): +async def schema_explorer_node(state: AgentState, config: RunnableConfig | None = None): """Schema Explorer node — just fetches the full catalog prompt from MCP.""" thread_id = config.get("configurable", {}).get("thread_id", "") if config else "" @@ -46,12 +46,8 @@ async def schema_explorer_node(state: AgentState, config: RunnableConfig = None) result_state: dict = { "jeen_catalog": catalog_prompt, - "tables_used": [], - "table_profiles": None, "execution_path": ["schema_explorer"], "schema_explorer_retry_count": 0, - "hallucinated_tables": None, - "last_error": None } return result_state diff --git a/agent/src/agent/state.py b/agent/src/agent/state.py index 336e4d6..cc49c34 100644 --- a/agent/src/agent/state.py +++ b/agent/src/agent/state.py @@ -26,7 +26,6 @@ class AgentState(TypedDict): active_skills: list[str] | None loaded_skills: list[dict] | None last_error: str | None - hallucinated_tables: list[str] | None esca_write_failed: bool | None inline_result_rows: list[list[Any]] | None inline_result_columns: list[str] | None @@ -43,8 +42,7 @@ class AgentState(TypedDict): # G4: feature flags & execution modes execution_mode: str | None # e.g. "cost_saving", "high_quality", "benchmark" runtime_flags: dict[str, Any] | None # resolved by init_flags_node - # Enriched table profiles — populated by schema_explorer for reuse by refiner - table_profiles: list[dict[str, Any]] | None + # ── Map related state ───────────────────────────────────────────────────── locations_dict: dict[str, dict[str, str]] | None location_wkt_instruction: str | None # ── Ambiguity Detection (detect_ambiguity node) ─────────────────────────── diff --git a/agent/src/agent/utils/jeen_metadata_client.py b/agent/src/agent/utils/jeen_metadata_client.py index 78bfbdd..8d7cd1d 100644 --- a/agent/src/agent/utils/jeen_metadata_client.py +++ b/agent/src/agent/utils/jeen_metadata_client.py @@ -8,7 +8,7 @@ MCP tools used -------------- -- ``search`` → table discovery (replaces hybrid_search_tables) +- ``get_catalog_prompt`` → table discovery (replaces hybrid_search_tables) - ``get_table_profile`` → per-table column stats (replaces get_table_profile tool) - ``list_tables_rich`` → fallback full-table list with row counts @@ -113,7 +113,7 @@ class JeenMetadataClient: High-level async client for jeen-metadata's MCP server. Methods map 1-to-1 to what schema_explorer_node needs: - - search_tables() replaces hybrid_search_tables() + - get_catalog_prompt() get the big catalog prompt about the whole database we work with - get_table_profile() replaces the @tool get_table_profile() - list_tables_rich() fallback when search returns nothing """ From b21c9d51bef45a04ad4afeb584b3f3bdd92a5743 Mon Sep 17 00:00:00 2001 From: yuvalkh Date: Fri, 24 Jul 2026 16:22:40 +0300 Subject: [PATCH 5/6] updated refiner graph to re-enter refiner on trino errors instead going to satisfaction check --- agent/src/agent/nodes/refiner_graph.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agent/src/agent/nodes/refiner_graph.py b/agent/src/agent/nodes/refiner_graph.py index 5daf010..4a45676 100644 --- a/agent/src/agent/nodes/refiner_graph.py +++ b/agent/src/agent/nodes/refiner_graph.py @@ -16,6 +16,7 @@ def route_refiner_subgraph(state: AgentState) -> str: if state.get("trino_error"): if state.get("refinement_count", 0) >= max_iterations: return END + return "refiner" # Issue 37: check SATISFACTION_CHECK_ENABLED in router check_enabled = runtime_flags.get("SATISFACTION_CHECK_ENABLED", settings.SATISFACTION_CHECK_ENABLED) @@ -60,6 +61,7 @@ def route_satisfaction_subgraph(state: AgentState) -> str: route_refiner_subgraph, { "satisfaction_check": "satisfaction_check", + "refiner": "refiner", END: END, }, ) From dd88b43877a786e0a02d6e993fa388ddfaab833a Mon Sep 17 00:00:00 2001 From: yuvalkh Date: Mon, 27 Jul 2026 14:40:06 +0300 Subject: [PATCH 6/6] after CR --- agent/scripts/upload_all_prompts.py | 33 ++---------------------- agent/src/agent/config.py | 4 +-- agent/src/agent/nodes/schema_explorer.py | 16 +++++++++--- 3 files changed, 16 insertions(+), 37 deletions(-) diff --git a/agent/scripts/upload_all_prompts.py b/agent/scripts/upload_all_prompts.py index d4fa141..3b1d0fa 100644 --- a/agent/scripts/upload_all_prompts.py +++ b/agent/scripts/upload_all_prompts.py @@ -52,35 +52,7 @@ def main(): "type": "chat" }, { - "name": "text2sql/schema_explorer", - "prompt": [ - { - "role": "system", - "content": ( - "You are a Schema Explorer sub-agent. Your goal is to identify the most relevant tables " - "and inspect their column details to form a query plan for the user's question.\n\n" - "Candidate Tables found:\n{{tables_json}}\n\n" - "Detailed Profiles for top tables (with Esca Reference IDs):\n{{profiles_json}}\n\n" - "## Decision-Making Rules\n\n" - "You MUST make all planning decisions autonomously. This includes:\n" - "- Join strategy: If multiple tables are needed to answer the query, decide which tables to join and on which keys — do NOT ask the user.\n" - "- Column selection: Choose the most appropriate columns yourself.\n" - "- Filter strategy: Infer filters from the user's question.\n" - "- Table selection: When one table is clearly more appropriate, pick the best match and proceed.\n" - "- When uncertain between two tables, pick the most semantically appropriate one and document your reasoning in schema_plan.\n\n" - "## Output Instructions\n\n" - "Provide your output matching the requested schema. Ensure that `schema_plan` is a detailed string explanation, and `tables_used` is a list of table names." - ) - }, - { - "role": "user", - "content": "{{human_message}}" - } - ], - "type": "chat" - }, - { - "name": "text2sql/query_builder_v2", + "name": "text2sql/query_builder", "prompt": [ { "role": "system", @@ -164,7 +136,7 @@ def main(): "type": "chat" }, { - "name": "text2sql/detect_ambiguity_v2", + "name": "text2sql/detect_ambiguity", "prompt": [ { "role": "system", @@ -280,7 +252,6 @@ def main(): "role": "user", "content": ( "User Request: {{user_query}}\n\n" - "Schema Context:\n{{schema}}\n\n" "Current Agent SQL Attempt:\n{{current_sql_attempt}}\n\n" "Agent's Explanation for SQL:\n{{sql_explanation}}" ) diff --git a/agent/src/agent/config.py b/agent/src/agent/config.py index 55e920c..8e3d108 100644 --- a/agent/src/agent/config.py +++ b/agent/src/agent/config.py @@ -55,7 +55,7 @@ class AgentSettings(BaseSettings): # Langfuse prompt names LANGFUSE_PROMPT_EXTRACTOR: str = "text2sql/extractor" LANGFUSE_PROMPT_SCHEMA_EXPLORER: str = "text2sql/schema_explorer" - LANGFUSE_PROMPT_QUERY_BUILDER: str = "text2sql/query_builder_v2" + LANGFUSE_PROMPT_QUERY_BUILDER: str = "text2sql/query_builder" LANGFUSE_PROMPT_REFINER: str = "text2sql/refiner" LANGFUSE_PROMPT_FINALIZER_SUMMARY: str = "text2sql/finalizer_summary" LANGFUSE_PROMPT_FINALIZER_SQL_EXPLANATION: str = ( @@ -63,7 +63,7 @@ class AgentSettings(BaseSettings): ) LANGFUSE_PROMPT_REJECTION_ROUTER: str = "text2sql/rejection_router" LANGFUSE_PROMPT_LOC_EXTRACTOR: str = "text2sql/extractor" - LANGFUSE_PROMPT_DETECT_AMBIGUITY: str = "text2sql/detect_ambiguity_v2" + LANGFUSE_PROMPT_DETECT_AMBIGUITY: str = "text2sql/detect_ambiguity" MAX_REFINER_ITERATIONS: int = Field(default=3, gt=0) REFINER_SCHEMA_CONTEXT_TABLES: int = Field(default=4, gt=0) diff --git a/agent/src/agent/nodes/schema_explorer.py b/agent/src/agent/nodes/schema_explorer.py index 7a4b4da..ed3d976 100644 --- a/agent/src/agent/nodes/schema_explorer.py +++ b/agent/src/agent/nodes/schema_explorer.py @@ -26,11 +26,19 @@ async def schema_explorer_node(state: AgentState, config: RunnableConfig | None _jeen = get_jeen_metadata_client() if not _jeen.is_configured: - logger.warning("MCP client is not configured, but fallback is disabled. Query builder might fail.") - catalog_prompt = "No catalog available (MCP not configured)." - else: - logger.info("Fetching full catalog prompt from Jeen MCP.") + error_msg = "Jeen is not configured. Jeen must be configured for the schema explorer to work." + logger.error(error_msg) + raise RuntimeError(error_msg) + + logger.info("Fetching full catalog prompt from Jeen MCP.") + try: catalog_prompt = await _jeen.get_catalog_prompt() + if not catalog_prompt: + raise ValueError("Received empty catalog prompt from Jeen.") + except Exception as exc: + error_msg = f"There was a problem getting the catalog_prompt from Jeen: {exc}" + logger.error(error_msg) + raise RuntimeError(error_msg) from exc # ── Langfuse trace metadata ─────────────────────────────────────────────── try: