Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
35 changes: 24 additions & 11 deletions agent/src/agent/mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,16 +64,21 @@ async def chat_with_agent(
}
)

result = await agent_graph.ainvoke(Command(
update={
"last_error": None,
"trino_error": None,
"escalated": None,
"escalation_reason": None,
"refinement_count": 0,
},
resume=resume_value
), config=config)
is_node_interrupt = False
if getattr(state_snapshot, "interrupts", None):
is_node_interrupt = True
elif getattr(state_snapshot, "tasks", None):
if any(getattr(task, "interrupts", None) for task in state_snapshot.tasks):
is_node_interrupt = True

# Resume the graph execution depending on how it was interrupted
if is_node_interrupt:
# Interrupted by the `interrupt()` function inside a node
result = await agent_graph.ainvoke(Command(resume=resume_value), config=config)
else:
# Interrupted by an `interrupt_before` breakpoint
await agent_graph.aupdate_state(config, resume_value)
result = await agent_graph.ainvoke(None, config=config)
else:
if not query:
return json.dumps({"error": "Query is required for new chat session."})
Expand Down Expand Up @@ -142,8 +147,16 @@ async def chat_with_agent(
final_state = await agent_graph.aget_state(config)

# Check if interrupted by `interrupt()` function
if final_state.interrupts:
interrupt_val = None
if getattr(final_state, "interrupts", None):
interrupt_val = final_state.interrupts[-1].value
elif getattr(final_state, "tasks", None):
for task in final_state.tasks:
if getattr(task, "interrupts", None):
interrupt_val = task.interrupts[-1].value
break

if interrupt_val is not None:
return json.dumps(
{
"thread_id": thread_id,
Expand Down
2 changes: 2 additions & 0 deletions agent/src/agent/nodes/query_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
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."""
runtime_flags = state.get("runtime_flags") or {}
Expand Down
5 changes: 3 additions & 2 deletions agent/src/agent/nodes/refiner.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ async def refiner_node(state: AgentState, config: RunnableConfig | None = None):
"error": trino_error,
"schema_context": schema_context,
"error_history": json.dumps(error_history),
"user_query": state.get("user_query", ""),
}
)
new_sql = clean_sql(response.content)
Expand Down Expand Up @@ -140,8 +141,8 @@ async def refiner_node(state: AgentState, config: RunnableConfig | None = None):
langfuse_client.update_current_span(
level="ERROR", status_message=f"ESCA write failed: {e}"
)
else:
logging.error(f"ESCA write failed: {e}")

logging.error(f"ESCA write failed: {e}")
raise RuntimeError(f"Failed to write query result to ESCA: {e}")

return {
Expand Down
3 changes: 2 additions & 1 deletion agent/src/agent/nodes/satisfaction_check.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,8 @@ async def satisfaction_check_node(state: AgentState, config: RunnableConfig | No
failures.append(f"[CHECK_A] Execution failed: {state['trino_error']}")

# ── Check B: Row Plausibility ─────────────────────────────────────────────
if _f(runtime_flags, "SATISFACTION_CHECK_PLAUSIBILITY", settings.SATISFACTION_CHECK_PLAUSIBILITY):
# Only run plausibility if there wasn't a hard execution failure
if not state.get("trino_error") and _f(runtime_flags, "SATISFACTION_CHECK_PLAUSIBILITY", settings.SATISFACTION_CHECK_PLAUSIBILITY):
n = len(rows)
min_rows = _f(runtime_flags, "SATISFACTION_MIN_ROWS", settings.SATISFACTION_MIN_ROWS)
max_rows = _f(runtime_flags, "SATISFACTION_MAX_ROWS", settings.SATISFACTION_MAX_ROWS)
Expand Down
37 changes: 29 additions & 8 deletions agent/src/agent/nodes/schema_explorer.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,9 +135,22 @@ def _build_column_context(cp: "ColumnProfile") -> dict:
return col


class TableJoin(BaseModel):
source_table: str = Field(description="Fully qualified name of the source table (catalog.schema.name)")
target_table: str = Field(description="Fully qualified name of the target table (catalog.schema.name)")
join_condition: str = Field(description="SQL condition for the join (e.g., A.id = B.a_id)")

class SchemaPlan(BaseModel):
tables: List[str] = Field(description="List of fully qualified table names (catalog.schema.name) required for the query")
columns: List[str] = Field(description="Columns to select")
joins: List[TableJoin] = Field(default_factory=list, description="How to join the tables")
filters: List[str] = Field(default_factory=list, description="Any WHERE clause filters")

# Define standardized Schema Explorer Output Type
class SchemaExplorerOutput(BaseModel):
schema_plan: Optional[Any] = Field(
reasoning: str = Field(default_factory=str, description="Reasoning for the schema plan - explain your logic.")

schema_plan: Optional[SchemaPlan] = Field(
default=None,
description="Detailed query plan describing tables, columns, and joins.",
)
Expand All @@ -156,6 +169,10 @@ class SchemaExplorerOutput(BaseModel):
default_factory=list,
description="List of fully qualified table names (catalog.schema.name) used in the plan.",
)
error: Optional[str] = Field(
default=None,
description="Error message if any"
)
Comment thread
elirazpevz marked this conversation as resolved.


def get_query_embedding(text: str) -> list[float]:
Expand Down Expand Up @@ -189,7 +206,7 @@ def hybrid_search_tables(
all_tables = session.exec(stmt_all).all()

allowed = allowed_tables or []
statuses = allowed_statuses or ["production"]
statuses = allowed_statuses if allowed_statuses is not None else ["production"]
allowed_tables_set = []
allowed_ids = set()

Expand Down Expand Up @@ -430,7 +447,7 @@ def _parse_bool_flag(value) -> bool:

sem = asyncio.Semaphore(profile_fetch_concurrency)

async def fetch_profile(t_id, t_name):
async def fetch_profile(t_id, t_name, t_fqn):
nonlocal cache_hit_count, cache_miss_count
async with sem:
try:
Expand All @@ -453,22 +470,25 @@ async def fetch_profile(t_id, t_name):
cache_miss_count += 1

profile_res = await get_table_profile.ainvoke({"table_id": t_id})
return json.loads(profile_res)
data = json.loads(profile_res)
data["fully_qualified_name"] = t_fqn
return data
except Exception as e:
print(f"Error fetching profile for {t_name}: {e}")
return None

fetch_tasks = []
for i, t in enumerate(candidate_tables):
fqn = f"{t.catalog}.{t.schema_name}.{t.name}"
tables_info.append(
{
"id": t.id,
"name": f"{t.catalog}.{t.schema_name}.{t.name}",
"name": fqn,
"description": "",
}
)
if i < max_profiles_to_fetch:
fetch_tasks.append(fetch_profile(t.id, t.name))
fetch_tasks.append(fetch_profile(t.id, t.name, fqn))

if fetch_tasks:
results = await asyncio.gather(*fetch_tasks, return_exceptions=True)
Expand Down Expand Up @@ -532,7 +552,7 @@ async def fetch_profile(t_id, t_name):
if schema_summarization and profile_details:
try:
summaries = [
f"[{p.get('table_name', 'unknown')}] {p.get('description', '') or '(no description available)'}"
f"[{p.get('fully_qualified_name', p.get('table_name', 'unknown'))}] {p.get('description', '') or '(no description available)'}"
for p in profile_details
]
profiles_json_str = "\n".join(summaries)
Expand Down Expand Up @@ -593,13 +613,14 @@ async def fetch_profile(t_id, t_name):
ambiguity_detected=False,
ambiguity_message="",
candidate_options=[],
error=str(e),
)

data = await _resolve_ambiguity(data, chain, tables_info, profiles_json_str, human_message, state)

plan = data.schema_plan
if plan is not None and not isinstance(plan, str):
plan = json.dumps(plan)
plan = plan.model_dump_json() if hasattr(plan, "model_dump_json") else json.dumps(plan)
elif plan is None:
plan = ""

Expand Down
54 changes: 54 additions & 0 deletions agent/tests/test_scoping_mode.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
import pytest
from unittest.mock import patch, MagicMock

from agent.nodes.schema_explorer import hybrid_search_tables
from core.models.models import Table

def test_hybrid_search_tables_strict_mode():
session_mock = MagicMock()

table1 = Table(id="t1", name="table_1", schema_name="public", catalog="cat", status="production")
table2 = Table(id="t2", name="table_2", schema_name="public", catalog="cat", status="production")
table3 = Table(id="t3", name="table_3", schema_name="public", catalog="cat", status="deprecated")

session_mock.exec.return_value.all.return_value = [table1, table2, table3]

# Mock execute for vector search
session_mock.execute.return_value.fetchall.return_value = [("t1",)]

# Mock session.get for final return
def mock_get(cls, id):
return {"t1": table1, "t2": table2, "t3": table3}.get(id)
session_mock.get.side_effect = mock_get

# Mock enrichment version for keyword search
session_mock.exec.return_value.first.return_value = None

# Test strict mode allows ONLY t1
results = hybrid_search_tables(
query="table",
query_embedding=[0.0],
session=session_mock,
allowed_tables=["t1"],
allowed_statuses=["production"],
scoping_mode="strict"
)

assert len(results) == 1
assert results[0].id == "t1"

# Test hybrid mode allows t1 and t2 (because status="production")
session_mock.execute.return_value.fetchall.return_value = [("t1",), ("t2",)]
results_hybrid = hybrid_search_tables(
query="table",
query_embedding=[0.0],
session=session_mock,
allowed_tables=["t1"],
allowed_statuses=["production"],
scoping_mode="hybrid"
)

assert len(results_hybrid) == 2
ids = {r.id for r in results_hybrid}
assert "t1" in ids
assert "t2" in ids
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
"""Add unique constraint to tables

Revision ID: 70e8a34ff877
Revises: f9a3d1c8e205
Create Date: 2026-07-07 11:17:43.459879

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import sqlmodel


# revision identifiers, used by Alembic.
revision: str = '70e8a34ff877'
down_revision: Union[str, None] = 'f9a3d1c8e205'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint('uq_table_profiles_table_id', 'table_profiles', type_='unique')
op.drop_index('ix_table_profiles_table_id', table_name='table_profiles')
op.create_index(op.f('ix_table_profiles_table_id'), 'table_profiles', ['table_id'], unique=True)
op.drop_column('table_profiles', 'is_partial')
op.create_unique_constraint('uq_table_fqn', 'tables', ['catalog', 'schema_name', 'name'])
# ### end Alembic commands ###


def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint('uq_table_fqn', 'tables', type_='unique')
op.add_column('table_profiles', sa.Column('is_partial', sa.BOOLEAN(), server_default=sa.text('false'), autoincrement=False, nullable=False))
op.drop_index(op.f('ix_table_profiles_table_id'), table_name='table_profiles')
op.create_index('ix_table_profiles_table_id', 'table_profiles', ['table_id'], unique=False)
op.create_unique_constraint('uq_table_profiles_table_id', 'table_profiles', ['table_id'], postgresql_nulls_not_distinct=False)
# ### end Alembic commands ###
1 change: 1 addition & 0 deletions backend/app/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ class Settings(BaseSettings):
LANGFUSE_PUBLIC_KEY: str = ""
LANGFUSE_SECRET_KEY: str = ""
LANGFUSE_HOST: str = "https://cloud.langfuse.com"
LANGFUSE_REQUEST_TIMEOUT: float = 30.0
CORS_ORIGINS: list[str] = [
"http://localhost:5173",
"http://localhost:3000",
Expand Down
19 changes: 16 additions & 3 deletions backend/app/routers/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,17 @@ class ChatResponse(BaseModel):
trace_id: str | None = None
execution_path: list[str] | None = None

class TraceSpan(BaseModel):
span_name: str
start_time: str | None = None
duration_ms: int = 0
input_tokens: int = 0
output_tokens: int = 0
model: str = "N/A"
status: str
input_preview: str = ""
output_preview: str = ""
Comment thread
elirazpevz marked this conversation as resolved.


# ---------------------------------------------------------------------------
# Helpers
Expand Down Expand Up @@ -187,6 +198,8 @@ async def chat(request: ChatRequest) -> ChatResponse:
"allowed_tables": request.allowed_tables,
"allowed_statuses": request.allowed_statuses,
"extractors": request.extractors,
"active_skills": request.active_skills,
"execution_mode": request.execution_mode,
"hitl_enabled": request.hitl_enabled,
}

Expand Down Expand Up @@ -254,13 +267,13 @@ async def event_generator():
return StreamingResponse(event_generator(), media_type="text/event-stream")


@router.get("/traces/{trace_id}")
@router.get("/traces/{trace_id}", response_model=list[TraceSpan])
async def get_trace_timeline(trace_id: str):
"""Fetch trace from Langfuse and normalize observations for frontend timeline."""
auth = (settings.LANGFUSE_PUBLIC_KEY, settings.LANGFUSE_SECRET_KEY)
url = f"{settings.LANGFUSE_HOST}/api/public/traces/{trace_id}"

async with httpx.AsyncClient() as client:
async with httpx.AsyncClient(timeout=settings.LANGFUSE_REQUEST_TIMEOUT) as client:
resp = await client.get(url, auth=auth)
if resp.status_code != 200:
if resp.status_code == 404:
Expand Down Expand Up @@ -290,7 +303,7 @@ async def get_trace_timeline(trace_id: str):

timeline.append(
{
"span_name": obs.get("name") or obs.get("type"),
"span_name": obs.get("name") or obs.get("type") or "Unknown Span",
"start_time": start_time_str,
"duration_ms": duration_ms,
"input_tokens": obs.get("promptTokens", 0),
Expand Down
24 changes: 23 additions & 1 deletion backend/app/routers/tables.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
UserScope,
)
from fastapi import APIRouter, Depends, Header, HTTPException, Query
from sqlalchemy.exc import IntegrityError
from sqlmodel import Session, col, select

from app.config import settings
Expand Down Expand Up @@ -195,6 +196,20 @@ def create_table(payload: TableCreate, session: Session = Depends(get_session)):
text_to_embed = f"Table name: {name}\nSchema: {schema_name}\nDescription: {description}\nColumns: {', '.join([c.get('name', '') for c in om_columns])}"
embedding = get_embedding(text_to_embed)

# Check for duplicate table
existing = session.exec(
select(Table).where(
Table.catalog == catalog_name,
Table.schema_name == schema_name,
Table.name == name
)
).first()
if existing:
raise HTTPException(
status_code=409,
detail=f"Table '{catalog_name}.{schema_name}.{name}' already exists."
)

# Create the table
table = Table(
name=name,
Expand All @@ -207,7 +222,14 @@ def create_table(payload: TableCreate, session: Session = Depends(get_session)):
embedding=embedding,
)
session.add(table)
session.commit()
try:
session.commit()
except IntegrityError:
session.rollback()
raise HTTPException(
status_code=409,
detail=f"Table '{catalog_name}.{schema_name}.{name}' already exists."
)
Comment on lines +225 to +232

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Ensure atomic creation of the table and its enrichment data.

Using session.commit() here commits the table to the database before the associated EnrichmentVersion is created. If the subsequent commit on line 257 fails, it leaves the database in an inconsistent state (a table without its enrichment data).

Using session.flush() executes the insert, populates table.id, and validates constraints (raising IntegrityError if necessary), but defers the transaction commit until the enrichment data is also added, making the entire operation atomic.

♻️ Proposed fix to ensure atomicity
     try:
-        session.commit()
+        session.flush()
     except IntegrityError:
         session.rollback()
         raise HTTPException(
             status_code=409,
             detail=f"Table '{catalog_name}.{schema_name}.{name}' already exists."
         )
📝 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
try:
session.commit()
except IntegrityError:
session.rollback()
raise HTTPException(
status_code=409,
detail=f"Table '{catalog_name}.{schema_name}.{name}' already exists."
)
try:
session.flush()
except IntegrityError:
session.rollback()
raise HTTPException(
status_code=409,
detail=f"Table '{catalog_name}.{schema_name}.{name}' already exists."
)
🤖 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 `@backend/app/routers/tables.py` around lines 225 - 232, Replace the early
session.commit() in the table-creation flow with session.flush() so the table
insert, generated table.id, and constraint validation occur without finalizing
the transaction. Preserve the existing IntegrityError rollback and conflict
response, and leave the final commit after EnrichmentVersion creation as the
sole transaction commit.


# Extract columns in the format expected by the frontend
def parse_columns(cols):
Expand Down
4 changes: 4 additions & 0 deletions backend/dump_schema.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
import json
from app.main import app
with open("../frontend/openapi.json", "w") as f:
json.dump(app.openapi(), f, indent=2)
Comment on lines +3 to +4

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '--- backend/dump_schema.py ---\n'
cat -n backend/dump_schema.py

printf '\n--- frontend/package.json ---\n'
cat -n frontend/package.json

printf '\n--- path resolution probe ---\n'
python3 - <<'PY'
from pathlib import Path
cwd = Path.cwd()
script_dir = Path('backend').resolve()
print('cwd:', cwd)
print('script_dir:', script_dir)
print('open("../frontend/openapi.json") from repo root ->', (cwd / '../frontend/openapi.json').resolve())
print('expected repo frontend path ->', (cwd / 'frontend/openapi.json').resolve())
print('from backend dir ->', (script_dir / '../frontend/openapi.json').resolve())
PY

Repository: StavPonte11/text2sql-onboarding

Length of output: 260


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '--- backend/dump_schema.py ---'
cat -n backend/dump_schema.py

echo
echo '--- frontend/package.json ---'
cat -n frontend/package.json

echo
echo '--- repo path probe ---'
python3 - <<'PY'
from pathlib import Path
cwd = Path.cwd()
print('cwd:', cwd)
print('repo frontend path:', (cwd / 'frontend/openapi.json').resolve())
print('relative "../frontend/openapi.json" from repo root:', (cwd / '../frontend/openapi.json').resolve())
PY

Repository: StavPonte11/text2sql-onboarding

Length of output: 3440


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '--- references to dump_schema.py ---'
rg -n "dump_schema\.py|openapi\.json" -g '!**/node_modules/**' -g '!**/.git/**' .

echo
echo '--- backend tree ---'
git ls-files backend frontend | sed -n '1,200p'

Repository: StavPonte11/text2sql-onboarding

Length of output: 6192


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '--- references to backend/dump_schema.py in docs and scripts ---'
rg -n "dump_schema\.py|openapi\.json|generate-types|schema\.json|schema\.d.ts" README.md backend/README.md frontend/README.md backend frontend -g '!**/node_modules/**' -g '!**/.git/**'

echo
echo '--- backend README (if present) ---'
if [ -f backend/README.md ]; then
  cat -n backend/README.md
fi

echo
echo '--- top-level README (if present) ---'
if [ -f README.md ]; then
  cat -n README.md
fi

Repository: StavPonte11/text2sql-onboarding

Length of output: 570


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '--- backend/Dockerfile ---'
cat -n backend/Dockerfile

echo
echo '--- backend/pyproject.toml ---'
cat -n backend/pyproject.toml

echo
echo '--- backend/e2e_test.py (if it mentions schema generation) ---'
rg -n "dump_schema|openapi.json|frontend/openapi.json|schema.d.ts" backend -g '!**/.git/**'

Repository: StavPonte11/text2sql-onboarding

Length of output: 4975


Anchor openapi.json to __file__
open("../frontend/openapi.json", "w") only works when the process starts in backend/; from the repo root it resolves outside the repo. Use a script-relative path here.

🤖 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 `@backend/dump_schema.py` around lines 3 - 4, Update the output path in the
schema-dumping flow around app.openapi() to resolve openapi.json relative to
dump_schema.py’s __file__, ensuring it targets the repository’s frontend
directory regardless of the process working directory. Preserve the existing
JSON serialization and formatting.

Loading