Ben/category enrichment

agent/pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "json-repair>=0.25.0",
     "shapely>=2.0.0",
     "sqlglot>=25.0.0",
+    "urllib3>=2.7.0",
 ]
 
 [tool.uv.sources]
@@ -57,3 +58,7 @@ select = ["TID"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
+markers = [
+    "real_llm: tests that make live LLM calls",
+    "real_e2e: real end-to-end tests requiring external services",
+]

agent/src/agent/config.py

@@ -19,15 +19,14 @@ class AgentSettings(BaseSettings):
     EMBEDDER_MODEL: str = "nomic-embed-text:latest"
     EMBEDDER_KEY: str = ""
     HYBRID_SEARCH_MAX_TABLES: int = 10
-    MAX_PROFILES_TO_FETCH: int = 3
-    PROFILE_FETCH_CONCURRENCY: int = Field(default=4, gt=0)
     LANGFUSE_SECRET_KEY: str = Field(min_length=1)
     LANGFUSE_PUBLIC_KEY: str = Field(min_length=1)
     LANGFUSE_BASE_URL: str = Field(min_length=1)
+    CACHE_KEY_DELIMITER: str = "#@#"
 
     # ── Jeen Integration ──────────────────────────────────────────────────────
-    JEEN_LLM_CORE_URL: str = ""  # If empty, agent gracefully skips fetching
-    JEEN_API_KEY: str = ""       # If empty, agent gracefully skips fetching
+    JEEN_LLM_CORE_URL: str = "http://schema-modeler.dev161.internal/api/mcp"  # If empty, agent gracefully skips fetching
+    JEEN_API_KEY: str = "mcp_ecd023ab04f849b36aef5d797525365c4c095052e6e07577d065bfe82507ae67"       # If empty, agent gracefully skips fetching
     SKILLS_HOT_RELOAD: bool = False  # If true, bypass Redis cache for skills
     NOMINATIM_USER_AGENT: str = "text2sql-agent/1.0"  # Nominatim acceptable-use identifier
     NOMINATIM_URL: str = "https://nominatim.openstreetmap.org/search"
@@ -54,28 +53,27 @@ class AgentSettings(BaseSettings):
     REDIS_URL: str = "redis://localhost:6379"
 
     # Langfuse prompt names
-    LANGFUSE_PROMPT_EXTRACTOR: str = "text2sql/extractor"
+    LANGFUSE_PROMPT_EXTRACTOR: str = "text2sql/enrichment_extractor"
     LANGFUSE_PROMPT_SCHEMA_EXPLORER: str = "text2sql/schema_explorer"
     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 = (
-        "text2sql/finalizer_sql_explanation"
-    )
+    LANGFUSE_PROMPT_FINALIZER: str = "text2sql/finalizer"
     LANGFUSE_PROMPT_REJECTION_ROUTER: str = "text2sql/rejection_router"
-    LANGFUSE_PROMPT_LOC_EXTRACTOR: str = "text2sql/extractor"
+    LANGFUSE_PROMPT_CATEGORY_ENRICHMENT: str = "text2sql/category_enrichment"
+    LANGFUSE_PROMPT_LOC_EXTRACTOR: str = "text2sql/loc_extractor"
+    LANGFUSE_PROMPT_LOC_EXTRACTOR_INSTRUCTION: str = (
+        "text2sql/location_wkt_instruction"
+    )
+    LANGFUSE_PROMPT_REFINER_STEP1: str = "text2sql/refiner_step1"
+    LANGFUSE_PROMPT_REFINER_STEP2: str = "text2sql/refiner_step2"
     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)
+    MAX_REFINER_ITERATIONS: int = Field(default=20, gt=0)
+    REFINER_SCHEMA_CONTEXT_TABLES: int = Field(default=8, gt=0)
 
     # ── G2-01: Table Scoping ──────────────────────────────────────────────────
     DEFAULT_TABLE_SCOPING_MODE: Literal["strict", "hybrid"] = "hybrid"
 
-    # ── G2-03: Advanced Schema Explorer phases ────────────────────────────────
-    ENABLE_SEMANTIC_TYPING: bool = False   # single batched LLM call — adds id/timestamp/category labels
-    ENABLE_JOIN_GRAPH: bool = False
-    ENABLE_SCHEMA_SUMMARIZATION: bool = False  # generated once at profile-time, not at runtime
     ENABLE_SKILL_INJECTION: bool = False
 
     # ── G2-04: Satisfaction Check ─────────────────────────────────────────────
@@ -95,7 +93,6 @@ class AgentSettings(BaseSettings):
 
     # ── G2-05: Redis Schema Cache ─────────────────────────────────────────────
     SCHEMA_CACHE_TTL: int = Field(default=600, gt=0)    # seconds — DDL content
-    PROFILE_CACHE_TTL: int = Field(default=1800, gt=0)  # seconds — table profile statistics
 
 
 settings = AgentSettings()

agent/src/agent/graph.py

@@ -61,7 +61,9 @@ class InvalidConfigurationException(ValueError):
 # ── G2-01: Config validator node ──────────────────────────────────────────────
 
 
-def validate_config_node(state: AgentState, config: RunnableConfig | None = None) -> dict:
+def validate_config_node(
+    state: AgentState, config: RunnableConfig | None = None
+) -> dict:
     """
     First node after START.  Resolves scoping_mode from state (or falls back
     to the env default) and enforces strict-mode preconditions.
@@ -74,7 +76,9 @@ def validate_config_node(state: AgentState, config: RunnableConfig | None = None
     publish_node_event_sync(thread_id, "validate_config")
 
     runtime_flags = state.get("runtime_flags") or {}
-    mode: str = state.get("scoping_mode") or runtime_flags.get("DEFAULT_TABLE_SCOPING_MODE", settings.DEFAULT_TABLE_SCOPING_MODE)
+    mode: str = state.get("scoping_mode") or runtime_flags.get(
+        "DEFAULT_TABLE_SCOPING_MODE", settings.DEFAULT_TABLE_SCOPING_MODE
+    )
 
     if mode == "strict":
         allowed = state.get("allowed_tables")
@@ -90,10 +94,12 @@ def validate_config_node(state: AgentState, config: RunnableConfig | None = None
 # ── G2-02: HITL escalation node ───────────────────────────────────────────────
 
 
-def hitl_escalation_node(state: AgentState, config: RunnableConfig | None = None) -> dict:
+def hitl_escalation_node(
+    state: AgentState, config: RunnableConfig | None = None
+) -> dict:
     """
     Execution pauses HERE via LangGraph interrupt_before before this node runs.
-    The API consumer then calls graph.update_state() to inject a corrected query 
+    The API consumer then calls graph.update_state() to inject a corrected query
     or provide explicit guidance, rather than just clearing the state.
     After update_state the graph resumes from this node, which immediately
     routes to extractor via its direct edge.
@@ -118,7 +124,18 @@ def hitl_escalation_node(state: AgentState, config: RunnableConfig | None = None
     except Exception:
         pass
 
-    return {"escalated": True, "execution_path": ["hitl_escalation"]}
+    return {
+        "escalated": True,
+        "execution_path": ["hitl_escalation"],
+        # Clear out error and escalation state so the resumed run starts fresh
+        "escalation_reason": None,
+        "rejection_category": None,
+        "satisfaction_failures": None,
+        "satisfaction_fail_count": 0,
+        "trino_error": None,
+        "error_history": [],
+        "refinement_count": 0,
+    }
 
 
 # ── Rejection router ──────────────────────────────────────────────────────────
@@ -157,13 +174,13 @@ def rejection_router_node(state: AgentState, config: RunnableConfig | None = Non
         response = chain.invoke({"feedback": feedback})
         route = response.route
     except Exception as e:
-        logger.error(f"Structured output parsing failed: {e}")
-        raise RuntimeError(f"Rejection router failed to parse structured output: {e}")
+        logger.error(f"Rejection router structured output failed: {e}")
+        raise RuntimeError(f"Failed to route user rejection feedback: {e}") from e
     return {
         "feedback_route": route,
         "raw_data_ref": None,
         "trino_error": None,
-        "execution_path": ["rejection_router"]
+        "execution_path": ["rejection_router"],
     }
 
 

agent/src/agent/langfuse_client.py

@@ -1,8 +1,50 @@
+import warnings
+import urllib3
 from langfuse import Langfuse
+from opentelemetry import trace as otel_trace_api
 from agent.config import settings
 
+# Suppress unverified HTTPS warnings for dev internal endpoints
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+warnings.filterwarnings("ignore", category=urllib3.exceptions.InsecureRequestWarning)
+
 langfuse_client = Langfuse(
     public_key=settings.LANGFUSE_PUBLIC_KEY,
     secret_key=settings.LANGFUSE_SECRET_KEY,
     host=settings.LANGFUSE_BASE_URL,
 )
+
+# Patch update_current_span and get_current_trace_id to safely no-op when running without an active OpenTelemetry span context
+_orig_update_current_span = langfuse_client.update_current_span
+_orig_get_current_trace_id = langfuse_client.get_current_trace_id
+
+
+def _is_valid_langfuse_span(span) -> bool:
+    if span is otel_trace_api.INVALID_SPAN:
+        return False
+    context = span.get_span_context()
+    if not context.is_valid or context.is_remote:
+        return False
+    if not span.is_recording():
+        return False
+    return True
+
+
+def _safe_update_current_span(*args, **kwargs):
+    current_span = otel_trace_api.get_current_span()
+    if not _is_valid_langfuse_span(current_span):
+        return
+    return _orig_update_current_span(*args, **kwargs)
+
+
+def _safe_get_current_trace_id(*args, **kwargs):
+    current_span = otel_trace_api.get_current_span()
+    if not _is_valid_langfuse_span(current_span):
+        return None
+    return _orig_get_current_trace_id(*args, **kwargs)
+
+
+langfuse_client.update_current_span = _safe_update_current_span
+langfuse_client.get_current_trace_id = _safe_get_current_trace_id
+
+

agent/src/agent/llm.py

@@ -69,3 +69,4 @@ def get_llm(
         temperature=temperature,
         timeout=300.0,
     )
+

agent/src/agent/nodes/detect_ambiguity.py

@@ -7,7 +7,7 @@
 The ambiguity-detection prompt requires three inputs:
 
   1. The original user request   — always available
-  2. The database schema         — available from state["table_profiles"] after
+  2. The database schema         — available from state["jeen_catalog"] after
                                    schema_explorer runs
   3. The agent's current plan    — the schema_plan produced by schema_explorer
                                    describes WHICH tables and columns the agent

agent/src/agent/nodes/extractor.py

@@ -151,7 +151,6 @@ def extractor_node(state: AgentState, config: RunnableConfig | None = None) -> d
     # whole list and lose TimeExtractor / LLMExtractor results.
     _extractor_specs: list = [
         (TimeExtractor,    {"runtime_flags": runtime_flags}),
-        (LLMExtractor,     {"runtime_flags": runtime_flags}),
         (LocationExtractor, {"runtime_flags": runtime_flags}),
         *[
             (HTTPExtractor, {"url": e["url"], "name": e["name"], "runtime_flags": runtime_flags})

agent/src/agent/nodes/finalizer.py

@@ -1,5 +1,7 @@
 import json
-import asyncio
+import logging
+
+logger = logging.getLogger(__name__)
 from langchain_core.runnables.config import RunnableConfig
 from agent.utils.redis_publisher import publish_node_event
 from agent.state import AgentState
@@ -9,8 +11,6 @@
 from agent.llm import get_llm
 from agent.utils.esca import get_esca_client
 
-from agent.utils.esca import get_esca_client
-
 
 async def get_esca_preview(esca_id: str, limit: int | None = None) -> str:
     """Load data from Esca and return a preview of the columns and the rows."""
@@ -19,7 +19,6 @@ async def get_esca_preview(esca_id: str, limit: int | None = None) -> str:
 
     try:
         async with get_esca_client() as client:
-            # TODO: instead of fetching everything from esca and then chunk, get only the chunk
             data_bytes = await client.load_head(esca_id)
             data = json.loads(data_bytes.decode())
 
@@ -36,53 +35,42 @@ async def get_esca_preview(esca_id: str, limit: int | None = None) -> str:
                 "preview_count": len(preview_rows),
                 "total_rows": total_rows,
             }
-            return json.dumps(preview_info, indent=2)
+            return json.dumps(preview_info, indent=2, default=str)
     except Exception as e:
-        return f"Error retrieving data preview from Esca: {e}"
-
-
-async def get_sql_explanation(sql_query: str | None, llm) -> str:
-    """Ask LLM to explain the SQL query in natural language."""
-    if not sql_query:
-        return "No SQL query was generated."
-
-    langfuse_prompt = langfuse_client.get_prompt(
-        settings.LANGFUSE_PROMPT_FINALIZER_SQL_EXPLANATION
-    )
-    prompt_sql_explanation = ChatPromptTemplate.from_messages(
-        langfuse_prompt.get_langchain_prompt()
-    )
-
-    chain = prompt_sql_explanation | llm
-    response = await chain.ainvoke({"sql_query": sql_query})
-    return response.content
+        logger.error(f"Error retrieving data preview from Esca: {e}")
+        return "Data preview is currently unavailable."
 
 
 async def finalizer_node(state: AgentState, config: RunnableConfig | None = None):
-    """Summarize data."""
+    """Summarize data using the unified Hebrew finalizer prompt."""
     thread_id = config.get("configurable", {}).get("thread_id", "") if config else ""
-    from agent.utils.redis_publisher import publish_node_event
     await publish_node_event(thread_id, "finalizer")
+
     raw_data_ref = state.get("raw_data_ref")
-    esca_write_failed = state.get("esca_write_failed", False)
     inline_result_rows = state.get("inline_result_rows")
+    inline_result_columns = state.get("inline_result_columns")
     runtime_flags = state.get("runtime_flags") or {}
     llm = get_llm("finalizer", runtime_flags=runtime_flags)
 
-    esca_write_enabled = str(runtime_flags.get("ESCA_WRITE_ENABLED", settings.ESCA_WRITE_ENABLED)).lower() == "true"
-
+    esca_write_enabled = (
+        str(
+            runtime_flags.get("ESCA_WRITE_ENABLED", settings.ESCA_WRITE_ENABLED)
+        ).lower()
+        == "true"
+    )
+
     raw_limit = runtime_flags.get("PREVIEW_ROWS_COUNT", settings.PREVIEW_ROWS_COUNT)
     limit = int(raw_limit) if raw_limit is not None else None
-
     preview_str = ""
-    if not esca_write_enabled:
+    if not esca_write_enabled or not raw_data_ref:
         if inline_result_rows is not None:
             preview_rows = inline_result_rows[:limit] if limit is not None else inline_result_rows
-            columns = (
-                list(preview_rows[0].keys())
-                if preview_rows and isinstance(preview_rows[0], dict)
-                else []
-            )
+            if inline_result_columns:
+                columns = inline_result_columns
+            elif preview_rows and isinstance(preview_rows[0], dict):
+                columns = list(preview_rows[0].keys())
+            else:
+                columns = []
             preview_info = {
                 "columns": columns,
                 "preview_rows": preview_rows,
@@ -95,30 +83,24 @@ async def finalizer_node(state: AgentState, config: RunnableConfig | None = None
     else:
         preview_str = await get_esca_preview(raw_data_ref, limit=limit)
 
-    langfuse_prompt_summary = langfuse_client.get_prompt(
-        settings.LANGFUSE_PROMPT_FINALIZER_SUMMARY
-    )
-    prompt_summary = ChatPromptTemplate.from_messages(
-        langfuse_prompt_summary.get_langchain_prompt()
+    prompt_name = settings.LANGFUSE_PROMPT_FINALIZER
+    langfuse_prompt = langfuse_client.get_prompt(prompt_name)
+    prompt_finalizer = ChatPromptTemplate.from_messages(
+        langfuse_prompt.get_langchain_prompt()
     )
 
-    summary_chain = prompt_summary | llm
-
-    summary_task = summary_chain.ainvoke(
+    chain = prompt_finalizer | llm
+    response = await chain.ainvoke(
         {
-            "user_query": state["user_query"],
+            "user_request": state.get("user_query") or "",
             "sql_query": state.get("sql_query") or "",
-            "raw_data_ref": raw_data_ref,
-            "data_preview": preview_str,
+            "sql_translation": state.get("sql_explanation") or "",
+            "sql_results": preview_str,
         }
     )
 
-    sql_task = get_sql_explanation(state.get("sql_query"), llm)
-
-    summary_response, sql_explanation = await asyncio.gather(summary_task, sql_task)
-
     return {
-        "summary": summary_response.content,
-        "sql_explanation": sql_explanation,
+        "summary": response.content,
+        "sql_explanation": state.get("sql_explanation", ""),
         "execution_path": ["finalizer"],
     }