prazy1208/text2sql

1"""2Gen-SQL Agent (Step 7 — SQL synthesis).3Inputs: rephrased_question, business_insights, few_shot_examples, selected_tables, selected_columns;4optional relationships (same FK row dicts as Table Agent — from metadata_store JSON) for JOIN hints.5Output: generated_sql (single read-only statement) plus reasoning_summary (often empty when model returns SQL only).6 7Uses the dedicated Gen-SQL model via llm_client (agent_name=gen_sql).8Dialect: PostgreSQL (Supabase).9"""10 11from __future__ import annotations12 13import json14import logging15import re16 17from backend.services.llm_client import AGENT_GEN_SQL, chat_completion18 19logger = logging.getLogger(__name__)20 21GEN_SQL_SYSTEM = """You are the SQL Generation Agent. Follow the user message exactly for process and output shape.22Dialect: PostgreSQL (Supabase). Emit a single read-only statement: SELECT or WITH … SELECT only (no INSERT, UPDATE, DELETE, DDL).23Use PostgreSQL syntax and built-ins (e.g. date_trunc, INTERVAL '1 month', CURRENT_DATE, standard string/date types).24Output must match what the user message asks for (typically SQL only, no prose)."""25 26GEN_SQL_AGENT_PROMPT = """You are the SQL Generation Agent in a Natural Language to SQL system.27 28Your task is to generate a correct, efficient, and executable SQL query.29 30--------------------------------------------------31INPUTS32--------------------------------------------------33 34ANALYTICAL QUESTION:35{rephrased_question}36 37SELECTED TABLES:38{selected_tables}39 40SELECTED COLUMNS:41{selected_columns}42 43TABLE RELATIONSHIPS:44{relationships_block}45 46FEW-SHOT PATTERNS:47{few_shot_block}48 49BUSINESS RULES:50{business_rules_block}51 52--------------------------------------------------53SQL GENERATION PROCESS (FOLLOW STRICTLY)54--------------------------------------------------55 56You MUST construct the SQL query following the logical order of SQL execution:57 581. FROM / JOIN592. WHERE603. GROUP BY614. HAVING625. SELECT636. DISTINCT647. ORDER BY658. LIMIT / OFFSET66 67--------------------------------------------------68 69Step 1: FROM / JOIN70- Start with base tables71- If multiple tables are present:72  - Use the provided relationships to JOIN tables73  - Do NOT assume joins74  - Do NOT skip intermediate tables75 76--------------------------------------------------77 78Step 2: WHERE79- Apply filtering conditions (dates, categories, thresholds, etc.)80- Use business rules where applicable81 82--------------------------------------------------83 84Step 3: GROUP BY85- Apply only if aggregation is required86 87--------------------------------------------------88 89Step 4: HAVING90- Apply only if filtering on aggregated results is needed91 92--------------------------------------------------93 94Step 5: SELECT95- Include required columns96- Apply aggregation functions where needed97- Use aliases for readability98 99--------------------------------------------------100 101Step 6: DISTINCT102- Apply only if the query explicitly requires unique values103 104--------------------------------------------------105 106Step 7: ORDER BY107- Apply when sorting or ranking is needed108 109--------------------------------------------------110 111Step 8: LIMIT / OFFSET112- Apply LIMIT to restrict output size113- Use LIMIT 100 by default unless specified otherwise114 115--------------------------------------------------116IMPORTANT NOTE117--------------------------------------------------118 119- Not all steps are required for every query120- Only include clauses that are necessary based on the analytical question121- Do NOT force unused clauses into the query122 123--------------------------------------------------124CRITICAL RULES125--------------------------------------------------126 127- Use ONLY provided tables and columns128- Do NOT invent tables or columns129- Do NOT include columns not selected earlier130- Ensure joins are correct and complete using relationships131- Use only SELECT queries (NO INSERT, UPDATE, DELETE, DROP)132- Ensure compatibility with PostgreSQL (Supabase runs PostgreSQL). Use standard PostgreSQL syntax and functions.133- Prefer simple, clear, and efficient queries134 135--------------------------------------------------136OUTPUT FORMAT137--------------------------------------------------138 139Return ONLY the SQL query.140Do NOT include explanations or extra text.141"""142 143 144def _format_business_rules(insights: list[str] | None) -> str:145    if not insights:146        return "(none)"147    lines = [f"- {str(x).strip()}" for x in insights if x and str(x).strip()]148    return "\n".join(lines) if lines else "(none)"149 150 151def _format_few_shot(examples: list[dict] | None) -> str:152    if not examples:153        return "(none)"154    blocks: list[str] = []155    for i, ex in enumerate(examples, start=1):156        if not isinstance(ex, dict):157            continue158        eid = ex.get("id", "—")159        qt = str(ex.get("query_type") or "").strip()160        qn = str(ex.get("question_text") or "").strip()161        sql = str(ex.get("sql_query") or "").strip()162        blocks.append(163            f"### Example {i} (id={eid})\n"164            f"- query_type: {qt or '—'}\n"165            f"- question_text: {qn or '—'}\n"166            f"- sql_query:\n{sql or '(none)'}"167        )168    return "\n\n".join(blocks) if blocks else "(none)"169 170 171def _format_selected_tables(tables: list[str] | None) -> str:172    if not tables:173        return "(none — do not invent tables)"174    return "\n".join(f"- `{t}`" for t in tables if t and str(t).strip())175 176 177def _format_selected_columns(selected: dict[str, list[str]] | None) -> str:178    if not selected:179        return "(none — do not invent columns)"180    lines: list[str] = []181    for table_fqn, cols in sorted(selected.items()):182        if not table_fqn or not str(table_fqn).strip():183            continue184        if not isinstance(cols, list):185            continue186        names = [str(c).strip() for c in cols if c and str(c).strip()]187        if not names:188            lines.append(f"- `{table_fqn}`: (no columns listed)")189        else:190            lines.append(f"- `{table_fqn}`: {', '.join(names)}")191    return "\n".join(lines) if lines else "(none — do not invent columns)"192 193 194def _format_relationships_block(relationships: list[dict] | None) -> str:195    """Same shape as Table/Column agents: rows with relationship_text from domain table_relationships."""196    if not relationships:197        return (198            "(None provided. Infer joins only from selected columns, few-shot patterns, and keys such as *_id. "199            "Do not invent foreign keys.)"200        )201    lines: list[str] = []202    for r in relationships:203        rt = (r.get("relationship_text") or "").strip()204        if rt:205            lines.append(f"- {rt}")206    return "\n".join(lines) if lines else "(none)"207 208 209def _parse_gen_sql_response(response_text: str) -> tuple[str, str] | None:210    """211    Return (generated_sql, reasoning_summary).212    Supports: JSON legacy {{generated_sql, reasoning_summary}}, fenced ```sql```, or plain SQL starting WITH/SELECT/EXPLAIN.213    """214    text = (response_text or "").strip()215    if not text:216        return None217 218    # ```sql ... ``` or ``` ... ```219    if "```" in text:220        if "```json" in text.lower():221            text = re.sub(r"^.*?```json\s*", "", text, flags=re.IGNORECASE | re.DOTALL)222        else:223            text = re.sub(r"^.*?```(?:sql|postgresql)?\s*", "", text, flags=re.IGNORECASE | re.DOTALL)224        text = re.sub(r"\s*```\s*$", "", text, flags=re.DOTALL).strip()225 226    # Legacy JSON body227    if text.lstrip().startswith("{"):228        try:229            data = json.loads(text)230            if isinstance(data, dict) and "generated_sql" in data:231                sql = str(data.get("generated_sql") or "").strip()232                reason = str(data.get("reasoning_summary") or "").strip()233                return sql, reason234        except (json.JSONDecodeError, TypeError):235            pass236 237    # Plain SQL (prompt: return ONLY the query)238    if re.match(r"^\s*(WITH|SELECT|EXPLAIN)\b", text, re.IGNORECASE):239        return text.strip(), ""240 241    logger.warning("Gen-SQL response is neither valid JSON nor SQL-looking: %r", text[:300])242    return None243 244 245def run_gen_sql(246    use_case: str,247    rephrased_question: str,248    business_insights: list[str] | None,249    few_shot_examples: list[dict] | None,250    selected_tables: list[str] | None,251    selected_columns: dict[str, list[str]] | None,252    *,253    relationships: list[dict] | None = None,254) -> dict:255    """256    Synthesize one SQL query from pipeline context.257 258    relationships: optional FK rows for the active domain (from list_relationships_from_metadata; same as Table Agent).259 260    Returns:261        { "generated_sql": str, "reasoning_summary": str }262    """263    _ = use_case  # not in current user-facing prompt; keep signature for callers264    rq = (rephrased_question or "").strip()265    if not rq:266        return {267            "generated_sql": "",268            "reasoning_summary": "Missing rephrased_question; cannot generate SQL.",269        }270 271    if not selected_tables:272        return {273            "generated_sql": "",274            "reasoning_summary": "No tables selected upstream; cannot generate SQL.",275        }276 277    user_msg = GEN_SQL_AGENT_PROMPT.format(278        rephrased_question=rq,279        selected_tables=_format_selected_tables(selected_tables),280        selected_columns=_format_selected_columns(selected_columns),281        relationships_block=_format_relationships_block(relationships),282        few_shot_block=_format_few_shot(few_shot_examples),283        business_rules_block=_format_business_rules(business_insights),284    )285 286    messages = [287        {"role": "system", "content": GEN_SQL_SYSTEM},288        {"role": "user", "content": user_msg},289    ]290 291    try:292        raw = chat_completion(messages, agent_name=AGENT_GEN_SQL)293    except Exception as e:294        logger.exception("Gen-SQL Agent LLM call failed")295        return {296            "generated_sql": "",297            "reasoning_summary": f"LLM error: {e}",298        }299 300    parsed = _parse_gen_sql_response(raw)301    if parsed is None:302        return {303            "generated_sql": "",304            "reasoning_summary": "Could not extract SQL from model response.",305        }306 307    sql, reason = parsed308    return {309        "generated_sql": sql,310        "reasoning_summary": reason,311    }312