add sso and use util package

.dockerignore

@@ -0,0 +1,9 @@
+.venv
+**/__pycache__
+*.pyc
+.git
+.pytest_cache
+node_modules
+dist
+.env
+.env.local

.gitignore

@@ -31,3 +31,10 @@ __pycache__/
 *.py[cod]
 *.pyo
 *.
+
+# Private keys / Secrets
+deploy_key
+
+# Snowflake local config
+snowflake.conf
+*.conf

agent/Dockerfile

@@ -0,0 +1,26 @@
+FROM python:3.12-slim
+
+RUN pip install uv
+RUN apt-get update && apt-get install -y git openssh-client
+
+# Configure git to use ssh for github.com
+RUN git config --global url."git@github.com:".insteadOf "https://github.com/"
+
+WORKDIR /app
+
+# Copy workspace
+COPY core /app/core
+COPY agent /app/agent
+
+WORKDIR /app/agent
+
+# Mount SSH deploy key secret and sync dependencies
+RUN --mount=type=secret,id=deploy_key,target=/root/.ssh/id_rsa,mode=0600 \
+    mkdir -p -m 0700 /root/.ssh && \
+    ssh-keyscan github.com >> /root/.ssh/known_hosts && \
+    uv sync
+
+# Add virtualenv bin to PATH to run uvicorn
+ENV PATH="/app/agent/.venv/bin:$PATH"
+
+CMD ["uvicorn", "agent.main:app", "--host", "0.0.0.0", "--port", "8001"]

agent/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "agent"
+version = "0.1.0"
+description = "Text2SQL Agent Service"
+readme = "README.md"
+requires-python = ">=3.12,<3.13"
+dependencies = [
+    "fastapi>=0.115.6",
+    "uvicorn[standard]>=0.32.1",
+    "langgraph",
+    "langchain-ollama",
+    "core",
+    "esca-sdk",
+    "pydantic-settings>=2.7.0",
+    "trino>=0.328.0",
+    "langfuse>=2.0.0",
+    "langchain",
+    "langchain-openai"
+]
+
+[tool.uv.sources]
+core = { path = "../core" }
+esca-sdk = { git = "https://github.com/elirazpevz/esca.git", subdirectory = "packages/esca-sdk" }
+
+[build-system]
+requires = ["uv_build>=0.11.7,<0.12.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.0.3",
+]

agent/scripts/upload_all_prompts.py

@@ -0,0 +1,188 @@
+import os
+from langfuse import Langfuse
+
+# ==============================================================================
+# CONFIGURATION - Set your private Langfuse server credentials here
+# ==============================================================================
+LANGFUSE_PUBLIC_KEY = "pk-lf-..."
+LANGFUSE_SECRET_KEY = "sk-lf-..."
+LANGFUSE_BASE_URL = "http://localhost:3000"  # Replace with your Langfuse URL
+
+def main():
+    print(f"Connecting to Langfuse at {LANGFUSE_BASE_URL}...")
+    langfuse_client = Langfuse(
+        public_key=LANGFUSE_PUBLIC_KEY,
+        secret_key=LANGFUSE_SECRET_KEY,
+        host=LANGFUSE_BASE_URL,
+    )
+
+    prompts_to_upload = [
+        {
+            "name": "text2sql/extractor",
+            "prompt": [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a query enrichment assistant for a text-to-SQL system.\n\n"
+                        "Your job is to read the user's natural-language query and add context that "
+                        "makes ambiguous or implicit terms clearer for downstream processing.\n\n"
+                        "Add enrichment entries for things like:\n"
+                        "  • Abbreviations or acronyms that have a specific meaning "
+                        "(e.g. 'MDA' → 'Magen David Adom')\n"
+                        "  • Relative time expressions that can be resolved to absolute dates "
+                        "(e.g. 'last quarter' → 'Q1 2025, Jan 1 – Mar 31 2025')\n"
+                        "  • Ambiguous proper nouns where context helps "
+                        "(e.g. 'Jordan' used as a country vs. a person's name)\n"
+                        "  • Domain-specific shorthand the downstream system may not know\n\n"
+                        "Do NOT try to identify which database table or column to use — that is handled by a "
+                        "separate schema exploration phase.\n"
+                        "Do NOT add enrichments for terms that are already fully clear from the query.\n"
+                        "If the query needs no enrichment, return an empty enrichments list."
+                    )
+                },
+                {
+                    "role": "user",
+                    "content": "{{user_query}}"
+                }
+            ],
+            "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 is clearly more appropriate: Pick the best match and proceed.\n\n"
+                        "## When to Flag Ambiguity (ONLY these cases)\n\n"
+                        "Set ambiguity_detected=true ONLY IF:\n"
+                        "1. Two or more completely independent tables could each independently and fully answer the user's question, "
+                        "and you have no way to determine which one the user wants (e.g., 'orders' vs 'orders_archive' with no "
+                        "indication of time range, or two fact tables from different business domains that both seem equally relevant).\n"
+                        "2. There is no table in the catalog that can answer the query at all.\n\n"
+                        "Do NOT flag ambiguity for: join decisions, column choices, filter logic, or any decision you can make yourself.\n\n"
+                        "## Output Format\n\n"
+                        "Output MUST be a valid JSON object with the following keys:\n"
+                        "- schema_plan: detailed query plan describing tables, columns, joins, and Esca reference IDs (empty string if ambiguity_detected is true)\n"
+                        "- ambiguity_detected: boolean, true ONLY in the hard-blocker cases described above\n"
+                        "- ambiguity_message: a concise question to ask the user to resolve the blocker (empty string if ambiguity_detected is false)\n"
+                        "- candidate_options: list of strings (table names or options) for the user to choose from (empty list if ambiguity_detected is false)\n"
+                        "Return only the raw JSON, no markdown formatting (no ```json code blocks)."
+                    )
+                },
+                {
+                    "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."
+                },
+                {
+                    "role": "user",
+                    "content": "Plan: {{schema_plan}}\nQuery: {{user_query}}{{feedback_str}}"
+                }
+            ],
+            "type": "chat"
+        },
+        {
+            "name": "text2sql/refiner",
+            "prompt": [
+                {
+                    "role": "system",
+                    "content": "You are a Trino SQL expert. Fix the SQL query based on the database error. Output ONLY the fixed SQL query, nothing else (no backticks, no explanation)."
+                },
+                {
+                    "role": "user",
+                    "content": "SQL: {{sql}}\nError: {{error}}"
+                }
+            ],
+            "type": "chat"
+        },
+        {
+            "name": "text2sql/finalizer_summary",
+            "prompt": [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a helpful data assistant. Summarize the findings for the user nicely. "
+                        "You are given the SQL query that was executed and a preview of the queried data (columns and first few rows) to help you understand the context of the results. "
+                        "Note: If the columns contain single items or aliases like `_col0` with a numeric value, this is the result of an aggregation query (such as `COUNT(*)`). Use this direct result to answer the user's question.\n\n"
+                        "Data Preview:\n{{data_preview}}"
+                    )
+                },
+                {
+                    "role": "user",
+                    "content": "User asked: {{user_query}}\nSQL Query: {{sql_query}}\nData Ref: {{raw_data_ref}}"
+                }
+            ],
+            "type": "chat"
+        },
+        {
+            "name": "text2sql/finalizer_sql_explanation",
+            "prompt": [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a database analyst assistant. Explain the following SQL query in clear, natural language. "
+                        "Describe what fields and tables are queried, any filters, joins, groupings, or aggregations, and what the query accomplishes. "
+                        "Keep the explanation concise and professional."
+                    )
+                },
+                {
+                    "role": "user",
+                    "content": "SQL Query:\n{{sql_query}}"
+                }
+            ],
+            "type": "chat"
+        },
+        {
+            "name": "text2sql/rejection_router",
+            "prompt": [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a sub-agent routing system. Analyze the user feedback on the previous query builder plan and select which phase to route the execution back to:\n"
+                        "- extractor: Choose this if the user is correcting the query intent, entities, constants, adding or mentioning new entities or asking a completely different question.\n"
+                        "- schema_explorer: Choose this if the user points out a wrong table, suggests using a different table/view, or references table schema selection issues.\n"
+                        "- query_builder: Choose this if the user is requesting a minor change in the SQL query itself (like modifying a WHERE clause, sorting, limits, joins, or specific syntax adjustments), but the tables and overall query intent are correct."
+                    )
+                },
+                {
+                    "role": "user",
+                    "content": "User Feedback: {{feedback}}"
+                }
+            ],
+            "type": "chat"
+        }
+    ]
+
+    for p in prompts_to_upload:
+        try:
+            print(f"Uploading prompt: {p['name']}...")
+            langfuse_client.create_prompt(
+                name=p["name"],
+                prompt=p["prompt"],
+                type=p["type"]
+            )
+            print(f"  ✓ Successfully uploaded {p['name']}")
+        except Exception as e:
+            print(f"  ✗ Failed to upload {p['name']}: {e}")
+
+if __name__ == "__main__":
+    main()

agent/src/agent/__init__.py

@@ -0,0 +1 @@
+# Agent package

agent/src/agent/config.py

@@ -0,0 +1,31 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+class AgentSettings(BaseSettings):
+    model_config = SettingsConfigDict(env_file=".env", extra="ignore")
+
+    ESCA_API_KEY: str = ""
+    ESCA_URL: str = "http://localhost:7010"
+    LLM_API_KEY: str = "ollama"
+    LLM_BASE_URL: str = "http://localhost:11434/v1"
+    LLM_MODEL: str = "gemma4:e4b"
+    EMBEDDER_URL: str = "http://localhost:11434"
+    EMBEDDER_MODEL: str = "nomic-embed-text:latest"
+    HYBRID_SEARCH_MAX_TABLES: int = 10
+    MAX_PROFILES_TO_FETCH: int = 3
+
+    LANGFUSE_SECRET_KEY: str = Field(min_length=1)
+    LANGFUSE_PUBLIC_KEY: str = Field(min_length=1)
+    LANGFUSE_BASE_URL: str = Field(min_length=1)
+
+    # 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_REFINER: str = "text2sql/refiner"
+    LANGFUSE_PROMPT_FINALIZER_SUMMARY: str = "text2sql/finalizer_summary"
+    LANGFUSE_PROMPT_FINALIZER_SQL_EXPLANATION: str = "text2sql/finalizer_sql_explanation"
+    LANGFUSE_PROMPT_REJECTION_ROUTER: str = "text2sql/rejection_router"
+
+
+settings = AgentSettings()

agent/src/agent/graph.py

@@ -0,0 +1,82 @@
+from langgraph.graph import StateGraph, START, END
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_openai import ChatOpenAI
+from langchain_core.prompts import ChatPromptTemplate
+from agent.state import AgentState
+from agent.nodes.extractor import extractor_node
+from agent.nodes.schema_explorer import schema_explorer_node
+from agent.nodes.query_builder import query_builder_node
+from agent.nodes.refiner import refiner_node
+from agent.nodes.finalizer import finalizer_node
+from agent.config import settings
+from agent.langfuse_client import langfuse_client
+from pydantic import BaseModel, Field
+from typing import Literal
+
+# Initialize LLM for rejection routing classification
+llm = ChatOpenAI(model=settings.LLM_MODEL, base_url=settings.LLM_BASE_URL, api_key=settings.LLM_API_KEY, temperature=0)
+
+class RejectionRoute(BaseModel):
+    route: Literal["extractor", "schema_explorer", "query_builder"] = Field(
+        description="The phase to route the execution back to based on the user feedback."
+    )
+
+def rejection_router_node(state: AgentState):
+    """Classify user rejection feedback and select the appropriate phase to return to."""
+    feedback = state.get("feedback")
+    if not feedback:
+        return {"feedback_route": "query_builder"}
+
+    langfuse_prompt = langfuse_client.get_prompt(settings.LANGFUSE_PROMPT_REJECTION_ROUTER)
+    prompt = ChatPromptTemplate.from_messages(langfuse_prompt.get_langchain_prompt())
+
+    structured_llm = llm.with_structured_output(RejectionRoute, method="json_schema")
+    chain = prompt | structured_llm
+    try:
+        response = chain.invoke({"feedback": feedback})
+        route = response.route
+    except Exception as e:
+        print(f"Structured output parsing failed: {e}")
+        route = "extractor"  # Fallback
+
+    return {"feedback_route": route}
+
+
+def route_refiner(state: AgentState):
+    if state.get("trino_error") and state.get("refinement_count", 0) < 3:
+        return "refiner"
+    return "finalizer"
+
+def route_query_builder(state: AgentState):
+    if state.get("feedback"):
+        return "rejection_router"
+    return "refiner"
+
+def route_rejection(state: AgentState):
+    route = state.get("feedback_route")
+    if route in ["extractor", "schema_explorer", "query_builder"]:
+        return route
+    return "extractor"
+
+workflow = StateGraph(AgentState)
+
+workflow.add_node("extractor", extractor_node)
+workflow.add_node("schema_explorer", schema_explorer_node)
+workflow.add_node("query_builder", query_builder_node)
+workflow.add_node("rejection_router", rejection_router_node)
+workflow.add_node("refiner", refiner_node)
+workflow.add_node("finalizer", finalizer_node)
+
+workflow.add_edge(START, "extractor")
+workflow.add_edge("extractor", "schema_explorer")
+workflow.add_edge("schema_explorer", "query_builder")
+
+workflow.add_conditional_edges("query_builder", route_query_builder, {"rejection_router": "rejection_router", "refiner": "refiner"})
+workflow.add_conditional_edges("rejection_router", route_rejection, {"extractor": "extractor", "schema_explorer": "schema_explorer", "query_builder": "query_builder"})
+
+workflow.add_conditional_edges("refiner", route_refiner, {"refiner": "refiner", "finalizer": "finalizer"})
+workflow.add_edge("finalizer", END)
+
+memory = MemorySaver()
+agent_graph = workflow.compile(checkpointer=memory)
+

agent/src/agent/langfuse_client.py

@@ -0,0 +1,8 @@
+from langfuse import Langfuse
+from agent.config import settings
+
+langfuse_client = Langfuse(
+    public_key=settings.LANGFUSE_PUBLIC_KEY,
+    secret_key=settings.LANGFUSE_SECRET_KEY,
+    host=settings.LANGFUSE_BASE_URL,
+)