Skip to content

Made agent to work with Jeen-Metadata #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuvalkh merged 6 commits into from
Jul 27, 2026
Merged

Made agent to work with Jeen-Metadata #30

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
525128f
updated agent to work with jeen metadata
yuvalkh Jul 23, 2026
020e756
removed schema plan from agent and front
yuvalkh Jul 23, 2026
56382e7
updated ambiguity to work with the new schema builder
yuvalkh Jul 23, 2026
2c09cf2
small fixes from cr
yuvalkh Jul 23, 2026
b21c9d5
updated refiner graph to re-enter refiner on trino errors instead goi…
yuvalkh Jul 24, 2026
dd88b43
after CR
yuvalkh Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 10 additions & 36 deletions agent/scripts/upload_all_prompts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,44 +51,16 @@ 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",
"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.\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"
Expand Down Expand Up @@ -178,7 +150,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"
Expand Down Expand Up @@ -213,7 +186,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"
Expand All @@ -236,8 +210,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"
Expand Down Expand Up @@ -278,8 +252,8 @@ def main():
"role": "user",
"content": (
"User Request: {{user_query}}\n\n"
"Schema:\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}}"
)
}
],
Expand Down
17 changes: 15 additions & 2 deletions agent/src/agent/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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://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
JEEN_METADATA_PROFILE_TIMEOUT: float = 30.0 # Per-MCP-call timeout (seconds)
Comment on lines +41 to +45

@coderabbitai coderabbitai Bot Jul 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🔴 Critical | ⚡ Quick win

Critical: remove committed MCP credentials and environment-specific defaults; use cleartext HTTP only intentionally.

JEEN_METADATA_MCP_KEY embeds a live-looking bearer token directly in source, and JEEN_METADATA_MCP_URL/JEEN_METADATA_CONNECTION_ID hardcode an internal environment. This leaks a secret into VCS and also contradicts the section comment on Lines 39-40 ("Leave empty to keep the local path") — with these non-empty defaults, JeenMetadataClient.is_configured is always true, so the documented local-DB fallback can never engage without an explicit override. The URL is also plaintext http:// (flagged by static analysis, CWE-319).

Rotate the leaked key immediately and default these to empty, sourcing real values from the environment.

🔒 Proposed fix 
-    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_MCP_URL: str = ""   # e.g. https://jeen-metadata.example.com/api/mcp
+    JEEN_METADATA_MCP_KEY: str = ""   # Bearer key from /api/mcp/keys in jeen-metadata
+    JEEN_METADATA_CONNECTION_ID: int = 0  # Numeric service ID (from list_connections)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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
JEEN_METADATA_PROFILE_TIMEOUT: float = 30.0 # Per-MCP-call timeout (seconds)
JEEN_METADATA_MCP_URL: str = "" # e.g. https://jeen-metadata.example.com/api/mcp
JEEN_METADATA_MCP_KEY: str = "" # Bearer key from /api/mcp/keys in jeen-metadata
JEEN_METADATA_CONNECTION_ID: int = 0 # 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)
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@agent/src/agent/config.py` around lines 41 - 45, Update the configuration
defaults for JEEN_METADATA_MCP_KEY, JEEN_METADATA_MCP_URL, and
JEEN_METADATA_CONNECTION_ID to empty values so JeenMetadataClient.is_configured
remains false unless explicitly configured. Ensure real credentials and
environment-specific connection settings are sourced from environment variables,
remove the committed bearer token, and preserve the intentional cleartext HTTP
behavior only when an explicit URL is supplied.

Source: Linters/SAST tools


# ── 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"
Expand Down Expand Up @@ -68,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
Expand Down
60 changes: 27 additions & 33 deletions agent/src/agent/graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -205,15 +190,15 @@ 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)
"""
t = state.get("ambiguity_type") or "clear"

if t == "clear":
return "query_builder"
return "hitl_query_approval"

if t == "unanswerable":
return END
Expand All @@ -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"
Expand All @@ -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)
Expand All @@ -259,36 +255,34 @@ 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",
},
)

workflow.add_conditional_edges(
"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"},
)

Expand Down
7 changes: 2 additions & 5 deletions agent/src/agent/mcp_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,6 @@ 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),
"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,
Expand All @@ -183,7 +182,6 @@ async def chat_with_agent(
"thread_id": thread_id,
"status": "interrupted",
"interrupt_details": interrupt_val,
"schema_plan": final_state.values.get("schema_plan"),
"sql_query": final_state.values.get("sql_query"),
"trace_id": trace_id,
"execution_path": final_state.values.get("execution_path", []),
Expand All @@ -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": result.get("schema_plan"),
"trace_id": trace_id,
"execution_path": result.get("execution_path", []),
"is_unanswerable": is_unans,
Expand All @@ -220,7 +217,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", {})

Expand All @@ -233,7 +230,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".
Expand Down
Loading