prazy1208/text2sql

1"""Routes: POST /session, POST /query (full pipeline incl. Gen-SQL), GET /use-cases, GET /domain-info."""2 3import json4import logging5import re6import time7import uuid8from typing import Any, Literal9 10from fastapi import APIRouter, Header, HTTPException, Query, Response11from pydantic import BaseModel, Field12 13from backend.agents.column_agent import run_column_agent14from backend.agents.few_shot_agent import run_few_shot_agent15from backend.agents.gen_sql_agent import run_gen_sql16from backend.agents.intent_agent import run_intent17from backend.agents.table_agent import run_table_agent18from backend.api.db import (19    create_session,20    delete_session_if_owned,21    get_all_chat_messages_ordered,22    get_latest_pending_intent,23    get_latest_rejected_intent_rephrase,24    get_recent_chat_messages,25    get_session_memory,26    get_session_pipeline_turns,27    get_session_row,28    insert_chat_message,29    insert_column_agent_output,30    insert_few_shot_agent_output,31    insert_gen_sql_agent_output,32    insert_intent_output,33    insert_intent_review,34    insert_table_agent_output,35    merge_session_memory,36    list_sessions_for_client,37    session_exists,38    set_session_title_if_unset,39    update_intent_review_status,40    update_session_use_case,41)42from backend.config import METADATA_STORE_DIR, USE_CASE_TO_SCHEMA, USE_CASES43from backend.services.relationship_retrieval import (44    filter_relationships_for_selected_tables,45    list_relationships_from_metadata,46)47from backend.services.sql_validator import validate_generated_sql48 49logger = logging.getLogger(__name__)50router = APIRouter(tags=["query"])51 52 53class QueryRequest(BaseModel):54    message: str = Field(..., min_length=1)55    use_case: str = Field(...)56    session_id: str | None = Field(None)57    message_type: Literal["new_query", "intent_confirmation", "intent_correction"] = Field("new_query")58    confirmation: Literal["yes", "no"] | None = Field(None)59    client_id: str | None = Field(60        None,61        description="Optional UUID for anonymous chat lists; applied when a new session is created",62    )63 64 65class QueryResponse(BaseModel):66    session_id: str67    rephrased_question: str68    keywords: list[str]69    business_insights: list[str]70    few_shot_examples: list[dict] = Field(default_factory=list)71    selected_tables: list[str] = Field(default_factory=list)72    selected_columns: dict[str, list[str]] = Field(default_factory=dict)73    generated_sql: str = Field(default="", description="Synthesized SQL from Gen-SQL Agent")74    resolved_question: str = ""75    intent_confidence: int = 076    needs_confirmation: bool = False77    clarification_question: str = ""78    conversation_state: str = "completed"79    pending_intent_id: int | None = None80    error: str | None = None81 82 83class SessionResponse(BaseModel):84    session_id: str85 86 87class SessionListItem(BaseModel):88    session_id: str89    title: str | None = None90    use_case: str | None = None91    created_at: str | None = None92    updated_at: str | None = None93 94 95class ChatMessageItem(BaseModel):96    id: int97    role: str98    content: str99    message_type: str100    created_at: str | None = None101 102 103class PipelineTurnItem(BaseModel):104    """Structured pipeline snapshot for one intent_agent_output row (UI replay)."""105 106    model_config = {"extra": "ignore"}107 108    intent_output_id: int109    user_input: str110    rephrased_question: str = ""111    resolved_question: str = ""112    keywords: list[str] = Field(default_factory=list)113    business_insights: list[str] = Field(default_factory=list)114    intent_confidence: int = 0115    selected_tables: list[str] = Field(default_factory=list)116    selected_columns: dict[str, Any] = Field(default_factory=dict)117    few_shot_examples: list[Any] = Field(default_factory=list)118    generated_sql: str = ""119    conversation_state: str = "completed"120    error: str | None = None121 122 123def _normalize_optional_uuid(value: str | None, *, field: str) -> str | None:124    if value is None or not str(value).strip():125        return None126    try:127        return str(uuid.UUID(str(value).strip()))128    except ValueError as e:129        raise HTTPException(status_code=400, detail=f"{field} must be a valid UUID") from e130 131 132def derive_session_title(user_text: str, max_chars: int = 60) -> str:133    """Short label from the first user question (single assign in DB; never overwritten there)."""134    s = " ".join((user_text or "").split())135    if not s:136        return "New chat"137    if len(s) <= max_chars:138        return s139    return s[: max_chars - 1].rstrip() + "…"140 141 142_MAX_PIPELINE_COMPLETION_CHAT_CHARS = 200_000143 144 145def _persist_pipeline_completion_assistant(146    session_id: str,147    generated_sql: str,148    err: str | None,149) -> None:150    """Persist the main assistant turn so GET /sessions/.../messages can reload the thread."""151    parts = ["Here is what I found."]152    e = (err or "").strip()153    if e:154        parts.extend(["", e])155    gs = (generated_sql or "").strip()156    if gs:157        parts.extend(["", "Generated SQL:", gs])158    content = "\n".join(parts)159    if len(content) > _MAX_PIPELINE_COMPLETION_CHAT_CHARS:160        content = content[: _MAX_PIPELINE_COMPLETION_CHAT_CHARS - 40] + "\n… (truncated for storage)"161    try:162        insert_chat_message(163            session_id=session_id,164            role="assistant",165            message_type="pipeline_completed",166            content=content,167        )168    except Exception:169        logger.exception("Failed to persist pipeline completion assistant message")170 171 172def _empty_response(session_id: str, error: str) -> QueryResponse:173    return QueryResponse(174        session_id=session_id,175        rephrased_question="",176        resolved_question="",177        keywords=[],178        business_insights=[],179        few_shot_examples=[],180        selected_tables=[],181        selected_columns={},182        generated_sql="",183        intent_confidence=0,184        needs_confirmation=False,185        clarification_question="",186        conversation_state="error",187        pending_intent_id=None,188        error=error,189    )190 191 192INTENT_CONFIRM_THRESHOLD = 95193 194_OPEN_INVITE_RE = re.compile(195    r"(?is)\b("196    r"do\s+you\s+have\b"197    r"|what\s+(would|do)\s+you\b"198    r"|please\s+(describe|tell\s+me|share|provide)\b"199    r"|could\s+you\s+(please\s+)?(tell|describe|clarify)\b"200    r"|what\s+kind\s+of\b"201    r"|how\s+can\s+i\s+help\b"202    r"|what\s+are\s+you\s+looking\s+for\b"203    r"|tell\s+me\s+more\s+about\s+what\b"204    r"|any\s+analytical\s+question\b"205    r"|topic\s+you\s+would\s+like\s+to\s+explore\b"206    r")\b"207)208 209 210def _is_open_analytical_invitation(text: str) -> bool:211    t = (text or "").strip()212    if not t:213        return False214    return bool(_OPEN_INVITE_RE.search(t))215 216 217def _is_trivial_user_message(msg: str) -> bool:218    m = (msg or "").strip().lower()219    if not m:220        return True221    if len(m) <= 3 and m.isalpha():222        return True223    trivial = {224        "hi",225        "hello",226        "hey",227        "yo",228        "sup",229        "thanks",230        "thank you",231        "thankyou",232        "ok",233        "okay",234        "k",235        "bye",236        "goodbye",237    }238    if m in trivial:239        return True240    if m.startswith("thank") and len(m) < 40:241        return True242    if m in {"good morning", "good afternoon", "good evening", "gm", "gn"}:243        return True244    return False245 246 247def _greeting_assistant_reply(user_message: str) -> str:248    """Short reply for trivial turns; no Yes/No — user should type their next question."""249    m = (user_message or "").strip().lower()250    if m.startswith("thank"):251        return "You're welcome! Whenever you're ready, describe the analysis you want."252    if m in {"bye", "goodbye"}:253        return "Goodbye! Come back anytime you have an analytical question."254    if m in {"hi", "hello", "hey", "yo", "sup", "gm", "gn"} or m in {255        "good morning",256        "good afternoon",257        "good evening",258    }:259        return "Hello! What would you like to analyze?"260    if m in {"ok", "okay", "k"}:261        return "Great — tell me what you'd like to explore in the data."262    return "Hi! When you're ready, describe what you'd like to analyze."263 264 265def _effective_intent_confidence(266    model_confidence: int,267    user_message: str,268    rephrased: str,269    keywords: list[str],270) -> int:271    c = max(0, min(100, int(model_confidence)))272    if _is_trivial_user_message(user_message):273        return min(c, 55)274    kw = keywords or []275    core = (rephrased or "").strip()276    if not kw and len(core) < 24:277        return min(c, 60)278    return c279 280 281def _strip_internal_memory_keys(summary: dict) -> dict:282    skip = {"pending_confirm_kind", "pending_intent_output_id"}283    return {k: v for k, v in summary.items() if k not in skip}284 285 286@router.get("/use-cases")287def get_use_cases() -> list[str]:288    return USE_CASES289 290 291_DOMAIN_META: dict[str, dict] = {292    "healthcare": {293        "display_name": "Healthcare",294        "description": "Patient records, visits, diagnoses, billing, and insurance claims",295        "icon": "heart",296        "example_questions": [297            "How many patients registered in each month of 2024?",298            "What is the average billing amount by insurance type?",299            "Which departments have the most patient visits?",300        ],301    },302    "retail": {303        "display_name": "Retail",304        "description": "Customers, products, orders, inventory, shipments, and promotions",305        "icon": "cart",306        "example_questions": [307            "What are the top 10 best-selling products by revenue?",308            "How has order volume changed month over month?",309            "Which stores have the highest inventory turnover?",310        ],311    },312    "finance": {313        "display_name": "Finance",314        "description": "Accounts, transactions, loans, credit cards, and fraud alerts",315        "icon": "chart",316        "example_questions": [317            "What is the total transaction volume by branch?",318            "Which accounts have the highest average balance?",319            "How many fraud alerts were raised per month?",320        ],321    },322}323 324 325@router.get("/domain-info")326def get_domain_info() -> dict:327    """Return domain descriptions, table lists, and curated example questions for the UI."""328    domains = []329    for use_case in USE_CASES:330        schema = USE_CASE_TO_SCHEMA[use_case]331        meta = _DOMAIN_META.get(use_case, {})332        metadata_file = METADATA_STORE_DIR / f"{schema}_metadata.json"333        tables: list[str] = []334        if metadata_file.is_file():335            try:336                with open(metadata_file, "r", encoding="utf-8") as f:337                    table_list = json.load(f)338                tables = [t["table_name"] for t in table_list if "table_name" in t]339            except (json.JSONDecodeError, KeyError):340                pass341        domains.append({342            "name": use_case,343            "display_name": meta.get("display_name", use_case.capitalize()),344            "description": meta.get("description", ""),345            "icon": meta.get("icon", "database"),346            "tables": tables,347            "table_count": len(tables),348            "example_questions": meta.get("example_questions", []),349        })350    return {"domains": domains}351 352 353@router.post("/session", response_model=SessionResponse)354def create_fresh_session(355    client_id: str | None = Query(356        default=None,357        description="Optional UUID; scopes session for GET /sessions",358    ),359    x_client_id: str | None = Header(default=None, alias="X-Client-Id"),360) -> SessionResponse:361    raw = (x_client_id or client_id or "").strip() or None362    cid = _normalize_optional_uuid(raw, field="client_id") if raw else None363    try:364        return SessionResponse(session_id=create_session(client_id=cid))365    except Exception as e:366        logger.exception("Failed to create session")367        raise HTTPException(status_code=500, detail=f"Failed to create session: {e}") from e368 369 370@router.get("/sessions", response_model=list[SessionListItem])371def list_sessions(372    client_id: str = Query(..., min_length=1, description="UUID previously sent when creating sessions"),373    limit: int = Query(100, ge=1, le=500),374) -> list[SessionListItem]:375    cid = _normalize_optional_uuid(client_id, field="client_id")376    if not cid:377        raise HTTPException(status_code=400, detail="client_id is required")378    rows = list_sessions_for_client(cid, limit=limit)379    return [SessionListItem(**r) for r in rows]380 381 382def _assert_session_client_access(383    session_id: str,384    client_id: str | None,385    x_client_id: str | None,386) -> None:387    if not session_exists(session_id):388        raise HTTPException(status_code=404, detail="Unknown session_id")389    row = get_session_row(session_id)390    if not row:391        raise HTTPException(status_code=404, detail="Unknown session_id")392    scoped = row.get("client_id")393    if scoped:394        raw = (x_client_id or client_id or "").strip()395        effective = _normalize_optional_uuid(raw, field="client_id") if raw else None396        if effective != scoped:397            raise HTTPException(398                status_code=403,399                detail="client_id does not match this session",400            )401 402 403@router.get("/sessions/{session_id}/messages", response_model=list[ChatMessageItem])404def list_session_messages(405    session_id: str,406    client_id: str | None = Query(407        default=None,408        description="Must match session.client_id when the session is scoped",409    ),410    x_client_id: str | None = Header(default=None, alias="X-Client-Id"),411) -> list[ChatMessageItem]:412    _assert_session_client_access(session_id, client_id, x_client_id)413    msgs = get_all_chat_messages_ordered(session_id)414    return [ChatMessageItem(**m) for m in msgs]415 416 417@router.get("/sessions/{session_id}/pipeline-turns", response_model=list[PipelineTurnItem])418def list_session_pipeline_turns(419    session_id: str,420    client_id: str | None = Query(421        default=None,422        description="Must match session.client_id when the session is scoped",423    ),424    x_client_id: str | None = Header(default=None, alias="X-Client-Id"),425) -> list[PipelineTurnItem]:426    """Structured rows from intent/table/column/few-shot/gen-sql tables for rich replay."""427    _assert_session_client_access(session_id, client_id, x_client_id)428    rows = get_session_pipeline_turns(session_id)429    return [PipelineTurnItem(**r) for r in rows]430 431 432@router.delete("/sessions/{session_id}", status_code=204)433def delete_session(434    session_id: str,435    client_id: str | None = Query(436        default=None,437        description="Must match session.client_id (same as listing)",438    ),439    x_client_id: str | None = Header(default=None, alias="X-Client-Id"),440) -> Response:441    """Remove a session and all related app data (FK cascade). Scoped to owning client_id."""442    if not session_exists(session_id):443        raise HTTPException(status_code=404, detail="Unknown session_id")444    raw = (x_client_id or client_id or "").strip()445    if not raw:446        raise HTTPException(447            status_code=400,448            detail="client_id query param or X-Client-Id header is required to delete a session",449        )450    cid = _normalize_optional_uuid(raw, field="client_id")451    if not delete_session_if_owned(session_id, cid):452        raise HTTPException(453            status_code=403,454            detail="Cannot delete this session (wrong client or session has no client_id)",455        )456    return Response(status_code=204)457 458 459@router.post("/query", response_model=QueryResponse)460def post_query(body: QueryRequest) -> QueryResponse:461    """Full POST /query wall time (intent → agents → Gen-SQL), all branches."""462    t0 = time.perf_counter()463    outcome: QueryResponse | None = None464    try:465        outcome = _post_query_core(body)466        return outcome467    finally:468        elapsed_s = time.perf_counter() - t0469        elapsed_ms = elapsed_s * 1000.0470        if outcome is not None:471            logger.info(472                "POST /query pipeline duration %.0f ms (%.2f s) session_id=%s message_type=%s conversation_state=%s",473                elapsed_ms,474                elapsed_s,475                outcome.session_id,476                body.message_type,477                outcome.conversation_state,478            )479        else:480            logger.info(481                "POST /query duration %.0f ms (%.2f s) message_type=%s (handler exited without QueryResponse — e.g. HTTPException)",482                elapsed_ms,483                elapsed_s,484                body.message_type,485            )486 487 488def _post_query_core(body: QueryRequest) -> QueryResponse:489    use_case = body.use_case.strip().lower()490    if use_case not in USE_CASES:491        raise HTTPException(status_code=400, detail=f"use_case must be one of: {USE_CASES}")492 493    if body.session_id:494        if not session_exists(body.session_id):495            raise HTTPException(status_code=400, detail="Invalid or unknown session_id")496        session_id = body.session_id497    else:498        try:499            raw_cid = (body.client_id or "").strip() or None500            cid = _normalize_optional_uuid(raw_cid, field="client_id") if raw_cid else None501            session_id = create_session(client_id=cid)502        except Exception as e:503            logger.exception("Failed to create session")504            return _empty_response("", str(e))505 506    logger.info("Query request: use_case=%s", use_case)507    user_message = body.message.strip()508    if not user_message:509        return _empty_response(session_id, "message cannot be empty")510 511    # Persist incoming user turn for chat history.512    try:513        insert_chat_message(514            session_id=session_id,515            role="user",516            message_type=body.message_type,517            content=user_message,518        )519    except Exception:520        logger.exception("Failed to persist user chat message")521 522    try:523        update_session_use_case(session_id, use_case)524        if body.message_type in ("new_query", "intent_correction"):525            set_session_title_if_unset(session_id, derive_session_title(user_message))526    except Exception:527        logger.exception("Failed to update session title/use_case")528 529    # Handle yes/no confirmation action before running a new intent.530    if body.message_type == "intent_confirmation":531        if body.confirmation not in {"yes", "no"}:532            return _empty_response(session_id, "confirmation must be 'yes' or 'no' for intent_confirmation")533        pending = get_latest_pending_intent(session_id)534        if not pending:535            return _empty_response(session_id, "No pending intent found for confirmation")536        session_mem = get_session_memory(session_id)537        confirm_kind = (session_mem or {}).get("pending_confirm_kind") or "intent_confirm"538 539        if confirm_kind == "open_invite" and body.confirmation == "yes":540            update_intent_review_status(pending["intent_output_id"], "rejected")541            try:542                merge_session_memory(543                    session_id,544                    {"pending_confirm_kind": None, "pending_intent_output_id": None},545                )546            except Exception:547                logger.exception("Failed to clear session memory after open-invite yes")548            follow = (549                "Please type your analytical question in the message box "550                "(for example: total sales by region last quarter)."551            )552            try:553                insert_chat_message(554                    session_id=session_id,555                    role="assistant",556                    message_type="intent_invite_accepted",557                    content=follow,558                )559            except Exception:560                logger.exception("Failed to persist assistant chat message")561            rq = (pending["rephrased_question"] or "").strip()562            return QueryResponse(563                session_id=session_id,564                rephrased_question=rq,565                resolved_question=rq,566                keywords=pending["keywords"] or [],567                business_insights=pending["business_insights"] or [],568                few_shot_examples=[],569                selected_tables=[],570                selected_columns={},571                generated_sql="",572                intent_confidence=int(pending["confidence_score"] or 0),573                needs_confirmation=False,574                clarification_question=follow,575                conversation_state="waiting_analytical_query",576                pending_intent_id=None,577                error=None,578            )579 580        if confirm_kind == "open_invite" and body.confirmation == "no":581            update_intent_review_status(pending["intent_output_id"], "rejected")582            try:583                merge_session_memory(584                    session_id,585                    {"pending_confirm_kind": None, "pending_intent_output_id": None},586                )587            except Exception:588                logger.exception("Failed to clear session memory after open-invite no")589            closing = "Ok, thank you!"590            try:591                insert_chat_message(592                    session_id=session_id,593                    role="assistant",594                    message_type="intent_conversation_closed",595                    content=closing,596                )597            except Exception:598                logger.exception("Failed to persist assistant chat message")599            rq = (pending["rephrased_question"] or "").strip()600            return QueryResponse(601                session_id=session_id,602                rephrased_question=rq,603                resolved_question=rq,604                keywords=pending["keywords"] or [],605                business_insights=pending["business_insights"] or [],606                few_shot_examples=[],607                selected_tables=[],608                selected_columns={},609                generated_sql="",610                intent_confidence=int(pending["confidence_score"] or 0),611                needs_confirmation=False,612                clarification_question=closing,613                conversation_state="conversation_ended",614                pending_intent_id=None,615                error=None,616            )617 618        if body.confirmation == "yes":619            try:620                merge_session_memory(621                    session_id,622                    {"pending_confirm_kind": None, "pending_intent_output_id": None},623                )624            except Exception:625                logger.exception("Failed to clear session memory after intent confirm yes")626            update_intent_review_status(pending["intent_output_id"], "confirmed")627            try:628                insert_chat_message(629                    session_id=session_id,630                    role="assistant",631                    message_type="intent_confirmation_result",632                    content="Thanks, I understood your intent. Proceeding to SQL generation.",633                )634            except Exception:635                logger.exception("Failed to persist confirmation assistant chat message")636            intent = {637                "rephrased_question": pending["rephrased_question"],638                "resolved_question": pending["rephrased_question"],639                "keywords": pending["keywords"] or [],640                "business_insights": pending["business_insights"] or [],641                "confidence_score": pending["confidence_score"],642                "clarification_question": "",643            }644            intent_output_id = pending["intent_output_id"]645        else:646            update_intent_review_status(pending["intent_output_id"], "rejected")647            try:648                merge_session_memory(649                    session_id,650                    {"pending_confirm_kind": None, "pending_intent_output_id": None},651                )652            except Exception:653                logger.exception("Failed to clear session memory after intent reject")654            try:655                insert_chat_message(656                    session_id=session_id,657                    role="assistant",658                    message_type="intent_rejected",659                    content="Thanks, please provide a corrected question.",660                )661            except Exception:662                logger.exception("Failed to persist rejection assistant chat message")663            return QueryResponse(664                session_id=session_id,665                rephrased_question=pending["rephrased_question"],666                resolved_question=pending["rephrased_question"],667                keywords=pending["keywords"] or [],668                business_insights=pending["business_insights"] or [],669                few_shot_examples=[],670                selected_tables=[],671                selected_columns={},672                generated_sql="",673                intent_confidence=int(pending["confidence_score"] or 0),674                needs_confirmation=False,675                clarification_question="Please provide a corrected question.",676                conversation_state="waiting_user_rephrase",677                pending_intent_id=pending["intent_output_id"],678                error=None,679            )680    else:681        # Build conversation context and run intent.682        try:683            recent_messages = get_recent_chat_messages(session_id, limit=8)684            session_summary = get_session_memory(session_id)685        except Exception:686            logger.exception("Failed to read chat context; continuing without it")687            recent_messages = []688            session_summary = {}689        if body.message_type in ("new_query", "intent_correction"):690            try:691                merge_session_memory(692                    session_id,693                    {"pending_confirm_kind": None, "pending_intent_output_id": None},694                )695            except Exception:696                logger.exception("Failed to reset pending confirmation flags in session memory")697 698    schema_name = USE_CASE_TO_SCHEMA[use_case]699    relationships: list[dict] = []700    try:701        relationships = list_relationships_from_metadata(schema_name)702    except Exception as e:703        logger.warning(704            "Could not load FK relationships from metadata for %s: %s",705            schema_name,706            e,707        )708 709    if body.message_type != "intent_confirmation":710        # Greeting / small talk only: respond in kind and wait for a real question (no Yes/No).711        if body.message_type == "new_query" and _is_trivial_user_message(user_message):712            rephrased = user_message.strip()713            resolved_question = rephrased714            keywords: list[str] = []715            business_insights: list[str] = []716            confidence_score = 50717            try:718                intent_output_id = insert_intent_output(719                    session_id=session_id,720                    use_case=use_case,721                    user_input=user_message,722                    rephrased_question=rephrased,723                    keywords=keywords,724                    business_insights=business_insights,725                )726                insert_intent_review(727                    intent_output_id=intent_output_id,728                    confidence_score=confidence_score,729                    confirmation_required=False,730                    confirmation_status="confirmed",731                )732            except Exception as e:733                logger.exception("Failed to persist greeting intent output")734                return QueryResponse(735                    session_id=session_id,736                    rephrased_question=rephrased,737                    resolved_question=resolved_question,738                    keywords=keywords,739                    business_insights=business_insights,740                    few_shot_examples=[],741                    selected_tables=[],742                    selected_columns={},743                    generated_sql="",744                    intent_confidence=confidence_score,745                    needs_confirmation=False,746                    clarification_question="",747                    conversation_state="error",748                    pending_intent_id=None,749                    error=f"Save failed: {e}",750                )751            reply = _greeting_assistant_reply(user_message)752            try:753                insert_chat_message(754                    session_id=session_id,755                    role="assistant",756                    message_type="greeting_ack",757                    content=reply,758                )759            except Exception:760                logger.exception("Failed to persist greeting assistant message")761            return QueryResponse(762                session_id=session_id,763                rephrased_question=rephrased,764                resolved_question=resolved_question,765                keywords=keywords,766                business_insights=business_insights,767                few_shot_examples=[],768                selected_tables=[],769                selected_columns={},770                generated_sql="",771                intent_confidence=confidence_score,772                needs_confirmation=False,773                clarification_question=reply,774                conversation_state="waiting_analytical_query",775                pending_intent_id=None,776                error=None,777            )778 779        pending_for_intent: str | None = None780        if body.message_type == "intent_correction":781            try:782                pending_for_intent = get_latest_rejected_intent_rephrase(session_id)783            except Exception:784                logger.exception("Failed to read rejected intent for correction context")785                pending_for_intent = None786        summary_for_intent: dict = {}787        if isinstance(session_summary, dict):788            summary_for_intent = _strip_internal_memory_keys(session_summary)789        try:790            intent = run_intent(791                user_message,792                use_case,793                recent_messages=recent_messages,794                last_confirmed_intent=None,795                pending_intent=pending_for_intent,796                session_summary=summary_for_intent,797            )798        except Exception as e:799            logger.exception("Intent Agent failed")800            return _empty_response(session_id, str(e))801 802        rephrased = intent.get("rephrased_question", "")803        resolved_question = intent.get("resolved_question") or rephrased804        keywords = intent.get("keywords") or []805        business_insights = intent.get("business_insights") or []806        _ics = intent.get("confidence_score")807        model_confidence = 50 if _ics is None else max(0, min(100, int(_ics)))808        confidence_score = _effective_intent_confidence(809            model_confidence, user_message, rephrased, keywords810        )811        intent["confidence_score"] = confidence_score812        clarification_question = (intent.get("clarification_question") or "").strip()813        needs_confirmation = confidence_score < INTENT_CONFIRM_THRESHOLD814 815        try:816            intent_output_id = insert_intent_output(817                session_id=session_id,818                use_case=use_case,819                user_input=user_message,820                rephrased_question=rephrased,821                keywords=keywords,822                business_insights=business_insights,823            )824            insert_intent_review(825                intent_output_id=intent_output_id,826                confidence_score=confidence_score,827                confirmation_required=needs_confirmation,828                confirmation_status="pending" if needs_confirmation else "confirmed",829            )830        except Exception as e:831            logger.exception("Failed to persist intent output")832            return QueryResponse(833                session_id=session_id,834                rephrased_question=rephrased,835                resolved_question=resolved_question,836                keywords=keywords,837                business_insights=business_insights,838                few_shot_examples=[],839                selected_tables=[],840                selected_columns={},841                generated_sql="",842                intent_confidence=confidence_score,843                needs_confirmation=needs_confirmation,844                clarification_question=clarification_question,845                conversation_state="error",846                pending_intent_id=None,847                error=f"Save failed: {e}",848            )849 850        if needs_confirmation:851            anchor = (rephrased or resolved_question or user_message).strip() or "—"852            prompt_text = (853                f'Did I understand correctly? You want: "{anchor}". Please answer Yes or No.'854            )855            open_invite = _is_open_analytical_invitation(clarification_question)856            confirm_kind = "open_invite" if open_invite else "intent_confirm"857            try:858                merge_session_memory(859                    session_id,860                    {861                        "pending_confirm_kind": confirm_kind,862                        "pending_intent_output_id": intent_output_id,863                    },864                )865            except Exception:866                logger.exception("Failed to persist pending confirmation kind in session memory")867            try:868                insert_chat_message(869                    session_id=session_id,870                    role="assistant",871                    message_type="intent_confirmation_prompt",872                    content=prompt_text,873                )874            except Exception:875                logger.exception("Failed to persist assistant confirmation prompt")876            return QueryResponse(877                session_id=session_id,878                rephrased_question=rephrased,879                resolved_question=resolved_question,880                keywords=keywords,881                business_insights=business_insights,882                few_shot_examples=[],883                selected_tables=[],884                selected_columns={},885                generated_sql="",886                intent_confidence=confidence_score,887                needs_confirmation=True,888                clarification_question=prompt_text,889                conversation_state="waiting_intent_confirmation",890                pending_intent_id=intent_output_id,891                error=None,892            )893    # Common fields for downstream pipeline path894    rephrased = intent.get("rephrased_question", "")895    resolved_question = intent.get("resolved_question") or rephrased896    keywords = intent.get("keywords") or []897    business_insights = intent.get("business_insights") or []898    _ics = intent.get("confidence_score")899    confidence_score = 50 if _ics is None else max(0, min(100, int(_ics)))900    clarification_question = (intent.get("clarification_question") or "").strip()901 902    few_shot_examples: list[dict] = []903    few_shot_error: str | None = None904 905    selected_tables: list[str] = []906    table_error: str | None = None907    try:908        table_out = run_table_agent(909            use_case,910            rephrased,911            keywords,912            relationships=relationships,913        )914        selected_tables = table_out.get("selected_tables") or []915    except Exception as e:916        logger.exception("Table Agent failed")917        table_error = f"Table Agent failed: {e}"918 919    table_agent_output_id: int | None = None920    try:921        table_agent_output_id = insert_table_agent_output(intent_output_id, selected_tables)922    except Exception as e:923        logger.exception("Failed to persist table agent output")924        persist_msg = f"Table selection save failed: {e}"925        table_error = f"{table_error}; {persist_msg}" if table_error else persist_msg926 927    selected_columns: dict[str, list[str]] = {}928    column_error: str | None = None929    if table_agent_output_id is not None:930        if selected_tables:931            try:932                rel_for_columns = filter_relationships_for_selected_tables(933                    schema_name,934                    relationships,935                    selected_tables,936                )937                col_out = run_column_agent(938                    use_case,939                    rephrased,940                    keywords,941                    selected_tables,942                    relationships=rel_for_columns,943                )944                selected_columns = col_out.get("selected_columns") or {}945            except Exception as e:946                logger.exception("Column Agent failed")947                column_error = f"Column Agent failed: {e}"948                selected_columns = {}949        try:950            insert_column_agent_output(table_agent_output_id, selected_columns)951        except Exception as e:952            logger.exception("Failed to persist column agent output")953            persist_col = f"Column selection save failed: {e}"954            column_error = f"{column_error}; {persist_col}" if column_error else persist_col955 956    try:957        fs_out = run_few_shot_agent(rephrased, keywords, business_insights)958        few_shot_examples = fs_out.get("few_shot_examples") or []959    except Exception as e:960        logger.exception("Few-Shot Agent failed")961        few_shot_examples = []962        few_shot_error = f"Few-Shot Agent failed: {e}"963 964    try:965        insert_few_shot_agent_output(intent_output_id, few_shot_examples)966    except Exception as e:967        logger.exception("Failed to persist few-shot agent output")968        persist_fs = f"Few-shot save failed: {e}"969        few_shot_error = f"{few_shot_error}; {persist_fs}" if few_shot_error else persist_fs970 971    generated_sql = ""972    reasoning_summary = ""973    gen_sql_error: str | None = None974 975    try:976        gen_out = run_gen_sql(977            use_case,978            rephrased,979            business_insights,980            few_shot_examples,981            selected_tables,982            selected_columns,983            relationships=relationships,984        )985        generated_sql = (gen_out.get("generated_sql") or "").strip()986        reasoning_summary = (gen_out.get("reasoning_summary") or "").strip()987    except Exception as e:988        logger.exception("Gen-SQL Agent failed")989        gen_sql_error = f"Gen-SQL Agent failed: {e}"990 991    try:992        validation = validate_generated_sql(993            generated_sql,994            selected_tables=selected_tables or None,995            selected_columns=selected_columns or None,996        )997    except Exception as e:998        logger.exception("SQL validator failed")999        validation = {1000            "validation_passed": False,1001            "validation_error_codes": "VALIDATOR_EXCEPTION",1002            "validation_error_message": str(e),1003            "blocked_keywords": "",1004            "is_single_statement": False,1005            "is_select_only": False,1006        }1007 1008    if not validation.get("validation_passed"):1009        msg = (validation.get("validation_error_message") or "").strip() or "SQL validation failed"1010        codes = (validation.get("validation_error_codes") or "").strip()1011        blocked = (validation.get("blocked_keywords") or "").strip()1012        parts = [f"SQL validation: {msg}"]1013        if codes:1014            parts.append(f"({codes})")1015        if blocked:1016            parts.append(f"[blocked: {blocked}]")1017        val_err = " ".join(parts)1018        gen_sql_error = f"{gen_sql_error}; {val_err}" if gen_sql_error else val_err1019 1020    err = few_shot_error1021    if table_error:1022        err = f"{err}; {table_error}" if err else table_error1023    if column_error:1024        err = f"{err}; {column_error}" if err else column_error1025    if gen_sql_error:1026        err = f"{err}; {gen_sql_error}" if err else gen_sql_error1027 1028    try:1029        insert_gen_sql_agent_output(1030            intent_output_id,1031            generated_sql,1032            reasoning_summary if reasoning_summary else None,1033            bool(validation.get("validation_passed")),1034            str(validation.get("validation_error_codes") or ""),1035            str(validation.get("validation_error_message") or ""),1036            str(validation.get("blocked_keywords") or ""),1037            bool(validation.get("is_single_statement")),1038            bool(validation.get("is_select_only")),1039        )1040    except Exception as e:1041        logger.exception("Failed to persist gen-sql agent output")1042        persist_gen = f"Gen-SQL save failed: {e}"1043        err = f"{err}; {persist_gen}" if err else persist_gen1044 1045    _persist_pipeline_completion_assistant(session_id, generated_sql, err)1046 1047    return QueryResponse(1048        session_id=session_id,1049        rephrased_question=rephrased,1050        resolved_question=resolved_question,1051        keywords=keywords,1052        business_insights=business_insights,1053        few_shot_examples=few_shot_examples,1054        selected_tables=selected_tables,1055        selected_columns=selected_columns,1056        generated_sql=generated_sql,1057        intent_confidence=confidence_score,1058        needs_confirmation=False,1059        clarification_question="",1060        conversation_state="completed",1061        pending_intent_id=None,1062        error=err,1063    )1064