routes/agents/agent_sql.py

import csv
import json
import os
import re
import socket
import threading
from typing import Any, Dict, List, Optional
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse

import psycopg2
from psycopg2 import extras
from psycopg2 import pool as pgpool
from openai.types.chat import ChatCompletionMessageParam

from routes.agents.clients import get_deepseek_client
from routes.agents.tools import DEEPSEEK_MODEL, DEEPSEEK_SQL_MODEL, SQL_SYSTEM_PROMPT


def _create_sql_input(context: str, user_question: str) -> str:
    if not context.strip():
        return user_question
    return (
        "Use this retrieved business context first when generating SQL, then answer the user question.\n"
        "----------------\n"
        "START CONTEXT\n"
        f"{context}\n"
        "END CONTEXT\n"
        "----------------\n"
        f"QUESTION: {user_question}"
    )

# ============================================================================
# COUNTRY MAP
# ============================================================================

COUNTRY_MAP: dict[str, str] = {}
_COUNTRY_CSV_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..", "..", "public", "data", "masterdata", "m_country.csv")
)
try:
    with open(_COUNTRY_CSV_PATH, encoding="utf-8") as _f:
        for _row in csv.DictReader(_f):
            _code = _row.get("country_code", "").strip().upper()
            for _key in ("country", "country_key", "country_code"):
                _val = _row.get(_key, "").strip()
                if _val:
                    COUNTRY_MAP[_val.lower()] = _code
except Exception as _e:
    print(f"Warning: Could not load country mapping: {_e}")

# ============================================================================
# SQL SYSTEM PROMPT
# ============================================================================

_SQL_GEN_PROMPT = """
You are an expert in generating PostgreSQL queries. You are an expert in correcting buggy SQL statements for PostgreSQL. You will be given SQL statements that may have errors and you will rewrite them into valid, executable PostgreSQL. Here is the database schema:

Tables:
- m_synonym(from TEXT PRIMARY KEY, to TEXT)
- m_country(country_key VARCHAR PRIMARY KEY, country_code VARCHAR, country VARCHAR, region VARCHAR, subregion VARCHAR)
- m_customer_shipto(customer_shipto_key VARCHAR PRIMARY KEY, customer_shipto VARCHAR)
- m_customer_soldto(customer_soldto_key VARCHAR PRIMARY KEY, customer_soldto VARCHAR, customer_group VARCHAR)
- m_product(product_article_key VARCHAR PRIMARY KEY, product_article VARCHAR, product_group_key VARCHAR, product_group VARCHAR)
- fact_inventory(country_key VARCHAR, product_article_key VARCHAR, inventory_volume NUMERIC, source VARCHAR)
- fact_orders(date DATE, country_key VARCHAR, document_number VARCHAR PRIMARY KEY, customer_soldto_key VARCHAR, customer_shipto_key VARCHAR, product_article_key VARCHAR, po_number VARCHAR, order_volume NUMERIC, order_eur NUMERIC, order_lc NUMERIC, order_usd NUMERIC, source VARCHAR)
- fact_overdues(country_key VARCHAR, customer_soldto_key VARCHAR, customer_shipto_key VARCHAR, product_article_key VARCHAR, document_number VARCHAR PRIMARY KEY, due_date DATE, overdue_eur NUMERIC, overdue_lc NUMERIC, overdue_usd NUMERIC, source VARCHAR)
- fact_sales(date DATE, country_key VARCHAR, customer_soldto_key VARCHAR, customer_shipto_key VARCHAR, product_article_key VARCHAR, sales_volume NUMERIC, sales_eur NUMERIC, sales_lc NUMERIC, sales_usd NUMERIC, source VARCHAR)

1) Target database is PostgreSQL. Do NOT use T-SQL features such as square brackets [], TOP, NVARCHAR, or GO.
2) Prefer sargable date ranges to leverage indexes, e.g. instead of EXTRACT(YEAR FROM date) = 2025, use date >= '2025-01-01' AND date < '2026-01-01'.
3) Keep identifiers unquoted and lowercase unless quoting is required. Our schemas use lowercase table/column names.
4) Use LIMIT for pagination/row limiting; ORDER BY only columns present in the SELECT scope.
5) Do not invent new columns or change the intention of the original query; fix only syntax/semantic issues.
6) If a CTE or alias is referenced incorrectly, correct the aliasing while preserving intent.
7) Return only the corrected PostgreSQL SQL in the structured output field correctedSQL.

Generate syntactically correct SQL for PostgreSQL using only these tables and columns. Do not use T-SQL features. Return only the SQL statement.

Here are example user inputs and the expected outputs:

User input 1:
WITH YTD AS (
    SELECT product_article_key,
           SUM(sales_eur) AS YTD
    FROM fact_sales
    WHERE date >= DATE_TRUNC('year', CURRENT_DATE)
      AND date <= CURRENT_DATE
    GROUP BY product_article_key
    ORDER BY YTD DESC
    LIMIT 10
),
PYTD AS (
    SELECT product_article_key,
           SUM(sales_eur) AS PYTD
    FROM fact_sales
    WHERE date >= DATE_TRUNC('year', CURRENT_DATE - INTERVAL '1 year')
      AND date <= (CURRENT_DATE - INTERVAL '1 year')
    GROUP BY product_article_key
)
SELECT 
    mp.product_article,
    c.YTD,
    p.PYTD,
    (c.YTD - COALESCE(p.PYTD, 0)) AS diff
FROM YTD c
LEFT JOIN PYTD p
  ON c.product_article_key = p.product_article_key
LEFT JOIN m_product mp
  ON c.product_article_key = mp.product_article_key
ORDER BY c.YTD DESC;


User input 2:
WITH sales_ytd AS (
    SELECT SUM(sales_eur) AS ytd
    FROM fact_sales
    WHERE country_key = 'MY'
      AND date >= DATE_TRUNC('year', CURRENT_DATE)
      AND date <= CURRENT_DATE
),
sales_pytd AS (
    SELECT SUM(sales_eur) AS pytd
    FROM fact_sales
    WHERE country_key = 'MY'
      AND date >= DATE_TRUNC('year', CURRENT_DATE) - INTERVAL '1 year'
      AND date < DATE_TRUNC('year', CURRENT_DATE)
),
orders AS (
    SELECT SUM(order_eur) AS orders
    FROM fact_orders
    WHERE country_key = 'MY'
),
overdues AS (
    SELECT SUM(overdue_eur) AS overdues
    FROM fact_overdues
    WHERE country_key = 'MY'
),
inventory AS (
    SELECT SUM(inventory_volume) AS inventory
    FROM fact_inventory
    WHERE country_key = 'MY'
)
SELECT
    s.ytd,
    p.pytd,
    (s.ytd - COALESCE(p.pytd, 0)) AS diff,
    o.orders,
    od.overdues,
    i.inventory
FROM sales_ytd s
CROSS JOIN sales_pytd p
CROSS JOIN orders o
CROSS JOIN overdues od
CROSS JOIN inventory i;


User input 3:
SELECT
    mp.product_article,
    SUM(fi.inventory_volume) AS inventory
FROM fact_inventory fi
JOIN m_product mp ON fi.product_article_key = mp.product_article_key
GROUP BY mp.product_article
ORDER BY inventory DESC
LIMIT 10;

User input 4:
WITH monthly_sales AS (
    SELECT
        DATE_TRUNC('month', date) AS month,
        EXTRACT(MONTH FROM date) AS month_num,
        TO_CHAR(date, 'Mon') AS month_name,
        EXTRACT(YEAR FROM date) AS year,
        SUM(sales_eur) AS total_sales_eur
    FROM fact_sales
    GROUP BY DATE_TRUNC('month', date),
             EXTRACT(MONTH FROM date),
             TO_CHAR(date, 'Mon'),
             EXTRACT(YEAR FROM date)
),
filtered_sales AS (
    SELECT *
    FROM monthly_sales
    WHERE year IN (EXTRACT(YEAR FROM CURRENT_DATE),
                   EXTRACT(YEAR FROM CURRENT_DATE) - 1)
),
pivoted AS (
    SELECT
        fs.month_num,
        fs.month_name AS "MMM",
        MAX(CASE WHEN fs.year = EXTRACT(YEAR FROM CURRENT_DATE) - 1 THEN fs.total_sales_eur END) AS PY,
        MAX(CASE WHEN fs.year = EXTRACT(YEAR FROM CURRENT_DATE) THEN fs.total_sales_eur END) AS YTD
    FROM filtered_sales fs
    GROUP BY fs.month_num, fs.month_name
)
SELECT "MMM",
       COALESCE(PY::text, '') AS PY,
       COALESCE(YTD::text, '') AS YTD,
       COALESCE((YTD - PY)::text, '') AS diff
FROM (
    SELECT "MMM", PY, YTD, (YTD - PY) AS diff, month_num
    FROM pivoted
    ORDER BY month_num
) sub
UNION ALL
SELECT 'Total' AS "MMM",
       SUM(PY)::text AS PY,
       SUM(YTD)::text AS YTD,
       (SUM(YTD) - SUM(PY))::text AS diff
FROM pivoted;


"""

# ============================================================================
# DATA-SOURCE ACCESS CONTROL
# ============================================================================
#
# Business data lives in the fact_* tables. A user's profiles."Source access"
# lists which data sources they may query (e.g. "Sales, Inventory, Overdue,
# Order"); the SQL agent refuses to run a query that touches any fact table
# outside that list. The m_* master/dimension tables hold only reference data
# (country, product, customer names) and are always joinable.

ALL_FACT_TABLES = {"fact_sales", "fact_inventory", "fact_overdues", "fact_orders"}

# Normalised source label -> fact table. Both singular and plural spellings are
# accepted because the profiles value is free text entered by an admin.
_SOURCE_LABEL_TO_TABLE = {
    "sales": "fact_sales",
    "inventory": "fact_inventory",
    "inventories": "fact_inventory",
    "overdue": "fact_overdues",
    "overdues": "fact_overdues",
    "order": "fact_orders",
    "orders": "fact_orders",
}

# Human-readable label per fact table, for the access-denied message and the
# generation prompt.
_TABLE_TO_SOURCE_LABEL = {
    "fact_sales": "Sales",
    "fact_inventory": "Inventory",
    "fact_overdues": "Overdue",
    "fact_orders": "Order",
}


class SQLAccessError(Exception):
    """Raised when a generated query touches a fact table the user cannot access."""


def parse_source_access(source_access: Optional[str]) -> set[str]:
    """Turn a profiles."Source access" string into a set of allowed fact tables."""
    allowed: set[str] = set()
    if not source_access:
        return allowed
    for token in re.split(r"[,;/|]", source_access):
        table = _SOURCE_LABEL_TO_TABLE.get(token.strip().lower())
        if table:
            allowed.add(table)
    return allowed


# Comment strippers so a fact-table name mentioned inside a comment is neither a
# false positive for the access check nor a way to smuggle intent past it.
_LINE_COMMENT_RE = re.compile(r"--[^\n]*")
_BLOCK_COMMENT_RE = re.compile(r"/\*.*?\*/", re.DOTALL)


def referenced_fact_tables(sql: str) -> set[str]:
    """Return the fact tables referenced by name in a SQL statement.

    Postgres requires the real table name in FROM/JOIN even when aliased, so a
    literal word-boundary scan of the known fact-table names is a sound way to
    detect which restricted tables a query reads.
    """
    stripped = _BLOCK_COMMENT_RE.sub(" ", sql)
    stripped = _LINE_COMMENT_RE.sub(" ", stripped)
    found: set[str] = set()
    for table in ALL_FACT_TABLES:
        if re.search(rf"\b{re.escape(table)}\b", stripped, re.IGNORECASE):
            found.add(table)
    return found


def enforce_table_access(sql: str, allowed_tables: set[str]) -> None:
    """Raise SQLAccessError if `sql` reads any fact table not in `allowed_tables`."""
    forbidden = referenced_fact_tables(sql) - allowed_tables
    if forbidden:
        labels = ", ".join(sorted(_TABLE_TO_SOURCE_LABEL.get(t, t) for t in forbidden))
        raise SQLAccessError(
            f"You do not have access to the following data source(s): {labels}. "
            "Contact an administrator to request access."
        )


DEFAULT_RESULT_LIMIT = 500

# Server-side cap on how long a generated query may run before Postgres kills it.
STATEMENT_TIMEOUT_MS = int(os.getenv("SQL_STATEMENT_TIMEOUT_MS", "30000"))

# Max pooled connections; excess concurrent SQL queries fail fast with a pool
# error (surfaced as execution_error) rather than piling onto the database.
POOL_MAX_CONNECTIONS = int(os.getenv("SQL_POOL_MAX_CONNECTIONS", "5"))


class SQLAgent:
    # Process-wide connection pool, created lazily on the first query.
    # ThreadedConnectionPool because queries run in FastAPI's threadpool.
    _pool: Optional[pgpool.ThreadedConnectionPool] = None
    _pool_lock = threading.Lock()

    def __init__(self):
        self.llm = get_deepseek_client()

    # ------------------------------------------------------------------
    # DB connectivity helpers
    # ------------------------------------------------------------------

    @staticmethod
    def is_configured() -> bool:
        return bool(
            os.getenv("SUPABASE_DB_POOLER_URL")
            or os.getenv("SQL_DATABASE_URL")
            or os.getenv("SUPABASE_DB_URL")
            or os.getenv("DATABASE_URL")
            or os.getenv("SUPABASE_DATABASE_URL")
        )

    @staticmethod
    def _get_connection_url() -> str:
        url = (
            os.getenv("SUPABASE_DB_POOLER_URL")
            or os.getenv("SQL_DATABASE_URL")
            or os.getenv("SUPABASE_DB_URL")
            or os.getenv("DATABASE_URL")
            or os.getenv("SUPABASE_DATABASE_URL")
        )
        if not url:
            raise RuntimeError(
                "Missing database connection URL. Set SUPABASE_DB_URL or DATABASE_URL in .env."
            )
        return url

    @staticmethod
    def _is_network_unreachable(exc: Exception) -> bool:
        msg = str(exc).lower()
        return "network is unreachable" in msg or "no route to host" in msg or "errno 101" in msg

    @staticmethod
    def _inject_hostaddr_ipv4(connection_url: str) -> Optional[str]:
        parsed = urlparse(connection_url)
        host = parsed.hostname
        if not host:
            return None
        try:
            socket.inet_aton(host)
            return connection_url
        except OSError:
            pass
        port = parsed.port or 5432
        try:
            ipv4_info = socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM)
        except socket.gaierror:
            return None
        if not ipv4_info:
            return None
        ipv4 = str(ipv4_info[0][4][0])
        params = dict(parse_qsl(parsed.query, keep_blank_values=True))
        params.setdefault("sslmode", "require")
        params["hostaddr"] = ipv4
        return urlunparse(parsed._replace(query=urlencode(params)))

    # ------------------------------------------------------------------
    # SQL safety
    # ------------------------------------------------------------------

    @staticmethod
    def _sanitize_read_only_sql(sql_query: str) -> str:
        sql = (sql_query or "").strip()
        if not sql:
            raise ValueError("SQL query is empty.")
        forbidden = (
            " insert ", " update ", " delete ", " drop ", " alter ", " truncate ",
            " create ", " grant ", " revoke ", " comment ", " do ", " call ", " copy ",
        )
        padded = f" {sql.lower()} "
        if any(token in padded for token in forbidden):
            raise ValueError("Only read-only SQL is allowed.")
        statements = [p.strip() for p in sql.split(";") if p.strip()]
        if len(statements) > 1:
            raise ValueError("Multiple SQL statements are not allowed.")
        first_kw = statements[0].split(None, 1)[0].lower() if statements else ""
        if first_kw not in {"select", "with", "show", "explain"}:
            raise ValueError("Only SELECT/CTE/SHOW/EXPLAIN queries are allowed.")
        return statements[0]

    # ------------------------------------------------------------------
    # Execution
    # ------------------------------------------------------------------

    @classmethod
    def _get_pool(cls) -> pgpool.ThreadedConnectionPool:
        if cls._pool is None:
            with cls._pool_lock:
                if cls._pool is None:
                    connection_url = cls._get_connection_url()
                    try:
                        cls._pool = pgpool.ThreadedConnectionPool(
                            1, POOL_MAX_CONNECTIONS, connection_url, connect_timeout=10
                        )
                    except psycopg2.OperationalError as exc:
                        if not cls._is_network_unreachable(exc):
                            raise
                        ipv4_url = cls._inject_hostaddr_ipv4(connection_url)
                        if not ipv4_url or ipv4_url == connection_url:
                            raise
                        cls._pool = pgpool.ThreadedConnectionPool(
                            1, POOL_MAX_CONNECTIONS, ipv4_url, connect_timeout=10
                        )
        return cls._pool

    @staticmethod
    def _run_query(conn, sanitized: str, result_limit: int) -> list[dict]:
        # Enforce read-only at the database: the text-based sanitizer is a first
        # filter, but Postgres itself rejects any write in this session (e.g. a
        # data-modifying CTE that slips past keyword checks).
        conn.set_session(readonly=True)
        cursor = conn.cursor(cursor_factory=extras.RealDictCursor)
        try:
            # SET (not a startup option) so transaction-mode poolers accept it.
            cursor.execute(f"SET statement_timeout = {STATEMENT_TIMEOUT_MS}")
            cursor.execute(sanitized)
            if cursor.description is None:
                return []
            return [dict(row) for row in cursor.fetchmany(result_limit)]
        finally:
            cursor.close()
            if not conn.closed:
                conn.rollback()  # end the read transaction before returning to the pool

    def execute_sql(self, sql_query: str, result_limit: int = DEFAULT_RESULT_LIMIT) -> list[dict]:
        sanitized = self._sanitize_read_only_sql(sql_query)
        # One retry with a fresh connection: pooled connections can be closed
        # server-side (e.g. by the Supabase pooler) after sitting idle.
        for attempt in range(2):
            db_pool = self._get_pool()
            conn = db_pool.getconn()
            try:
                return self._run_query(conn, sanitized, result_limit)
            except (psycopg2.OperationalError, psycopg2.InterfaceError):
                # Only retry when the connection itself broke; a server-side
                # cancel (e.g. statement_timeout) is also an OperationalError
                # but must not re-run the query.
                broken = bool(conn.closed)
                db_pool.putconn(conn, close=True)
                conn = None
                if not broken or attempt == 1:
                    raise
            finally:
                if conn is not None:
                    db_pool.putconn(conn)
        raise RuntimeError("unreachable")  # loop always returns or raises

    # ------------------------------------------------------------------
    # Country name normalisation
    # ------------------------------------------------------------------

    @staticmethod
    def map_country_names_to_codes(text: str) -> str:
        if not COUNTRY_MAP:
            return text
        pattern = r"\b(" + "|".join(sorted(map(re.escape, COUNTRY_MAP.keys()), key=len, reverse=True)) + r")\b"

        def _replace(match: re.Match) -> str:
            code = COUNTRY_MAP.get(match.group(0).lower())
            return code if code else match.group(0)

        return re.sub(pattern, _replace, text, flags=re.IGNORECASE)

    # ------------------------------------------------------------------
    # SQL generation
    # ------------------------------------------------------------------

    @staticmethod
    def _access_prompt(allowed_tables: Optional[set[str]]) -> str:
        """Extra system guidance so the LLM only writes queries the user may run.

        None means unrestricted (no access control applied). An empty set means
        the user has no data-source access at all.
        """
        if allowed_tables is None:
            return ""
        if not allowed_tables:
            return (
                "\nACCESS CONTROL: This user has no data-source access. Do not query any "
                "fact_* table (fact_sales, fact_inventory, fact_overdues, fact_orders)."
            )
        allowed_list = ", ".join(sorted(allowed_tables))
        forbidden_list = ", ".join(sorted(ALL_FACT_TABLES - allowed_tables)) or "none"
        return (
            f"\nACCESS CONTROL: This user may only query these fact tables: {allowed_list}. "
            f"Never reference these fact tables: {forbidden_list}. "
            "The m_* master/dimension tables are always allowed for joins."
        )

    def generate_sql_query(
        self, user_prompt: str, context: str = "", allowed_tables: Optional[set[str]] = None
    ) -> tuple[str, Optional[dict]]:
        sql_input = _create_sql_input(context, user_prompt)
        preprocessed = self.map_country_names_to_codes(sql_input)

        messages: list[ChatCompletionMessageParam] = [
            {"role": "system", "content": _SQL_GEN_PROMPT + self._access_prompt(allowed_tables)},
            {"role": "user", "content": preprocessed},
        ]

        response = self.llm.chat.completions.create(
            model=DEEPSEEK_SQL_MODEL,
            messages=messages,
            temperature=0,
            max_tokens=4096,
            timeout=30.0,
        )

        token_usage = None
        if hasattr(response, "usage") and response.usage:
            u = response.usage
            token_usage = {
                "prompt_tokens": getattr(u, "prompt_tokens", None),
                "completion_tokens": getattr(u, "completion_tokens", None),
                "total_tokens": getattr(u, "total_tokens", None),
            }

        raw_sql = (response.choices[0].message.content or "").strip()
        code_block = re.compile(r"```(?:sql|json)?\n([\s\S]*?)```", re.IGNORECASE)
        matches = code_block.findall(raw_sql)
        sql = "\n".join(m.strip() for m in matches) if matches else raw_sql.replace("```", "").strip()

        # Strip any LLM preamble: find the first line that opens a SQL statement
        sql_start = re.search(r"(?im)^(WITH|SELECT|SHOW|EXPLAIN)\b", sql)
        if sql_start:
            sql = sql[sql_start.start():]

        # LLM may return {"correctedSQL": "..."} — extract the SQL string
        try:
            parsed = json.loads(sql)
            if isinstance(parsed, dict) and "correctedSQL" in parsed:
                sql = parsed["correctedSQL"].strip()
        except (json.JSONDecodeError, ValueError):
            pass

        # Wrap UNION parts that contain a LIMIT inside parentheses
        if re.search(r"UNION( ALL)?", sql, re.IGNORECASE):
            parts = re.split(r"(UNION(?: ALL)?)", sql, flags=re.IGNORECASE)
            wrapped: list[str] = []
            for part in parts:
                part = part.strip()
                if part.upper().startswith("SELECT") and "LIMIT" in part.upper():
                    part = part.rstrip(";")
                    wrapped.append(f"({part})")
                else:
                    wrapped.append(part)
            sql = "\n".join(wrapped)

        return sql, token_usage

    # ------------------------------------------------------------------
    # Result synthesis
    # ------------------------------------------------------------------

    def synthesize_results(
        self,
        query: str,
        sql_query: str,
        results: List[Dict[str, Any]],
        conversation_history: list[dict] | None = None,
    ) -> str:
        messages: list[ChatCompletionMessageParam] = [{"role": "system", "content": SQL_SYSTEM_PROMPT}]
        if conversation_history:
            messages.extend(conversation_history[-4:])  # last 2 exchanges
        messages.append({
            "role": "user",
            "content": (
                f"Question: {query}\n\n"
                f"SQL Query:\n{sql_query}\n\n"
                f"Query Results (first {len(results)} rows):\n{results}\n\n"
                "Provide a clear, business-friendly summary of these results, "
                "then end your response with ALL query results rendered as a markdown table (with proper column headers)."
            ),
        })

        response = self.llm.chat.completions.create(
            model=DEEPSEEK_MODEL,
            messages=messages,
            temperature=0.3,
        )
        return response.choices[0].message.content.strip()

    # ------------------------------------------------------------------
    # Public entry point
    # ------------------------------------------------------------------

    def process_query(
        self,
        query: str,
        context: str = "",
        conversation_history: list[dict] | None = None,
        allowed_tables: Optional[set[str]] = None,
    ) -> tuple[str, Dict[str, Any]]:
        """Generate and run SQL for `query`.

        `allowed_tables` is the set of fact tables this user may read (see
        parse_source_access). None disables access control (internal callers);
        an empty set denies every fact table. Access is enforced both by
        steering generation and by validating the final SQL before execution,
        so a jailbroken prompt still cannot reach a forbidden table.
        """
        if not self.is_configured():
            raise RuntimeError(
                "SQL execution is not configured. Set SUPABASE_DB_URL or DATABASE_URL."
            )

        sql_query, token_usage = self.generate_sql_query(query, context, allowed_tables=allowed_tables)

        execution_error: Optional[str] = None
        access_denied = False
        results: List[Dict[str, Any]] = []
        try:
            if allowed_tables is not None:
                enforce_table_access(sql_query, allowed_tables)
            results = self.execute_sql(sql_query)
        except SQLAccessError as exc:
            access_denied = True
            execution_error = str(exc)
        except Exception as exc:
            execution_error = str(exc)

        if access_denied:
            # A permission problem, not a malformed query — surface the reason
            # directly rather than framing it as an execution failure.
            answer = execution_error
        elif execution_error:
            answer = f"SQL query was generated but execution failed: {execution_error}"
        else:
            answer = self.synthesize_results(query, sql_query, results, conversation_history=conversation_history)

        data: Dict[str, Any] = {
            "sql_query": sql_query,
            "results": results,
            "execution_error": execution_error,
            "token_usage": token_usage,
        }
        return answer, data