[KM-512] connect query executor to user question. add logging for db_executor

src/query/base.py

@@ -27,5 +27,6 @@ class BaseExecutor(ABC):
         results: list[RetrievalResult],
         user_id: str,
         db: AsyncSession,
+        question: str,
         limit: int = 100,
     ) -> list[QueryResult]: ...

src/query/executors/db_executor.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import sqlglot
 import sqlglot.expressions as exp
+import tiktoken
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_openai import AzureChatOpenAI
 from sqlalchemy import text
@@ -35,6 +36,8 @@ from src.utils.db_credential_encryption import decrypt_credentials_dict
 
 logger = get_logger("db_executor")
 
+_enc = tiktoken.get_encoding("cl100k_base")
+
 _SUPPORTED_DB_TYPES = {"postgres", "supabase", "mysql"}
 _MAX_RETRIES = 3
 _MAX_LIMIT = 500
@@ -43,13 +46,15 @@ _SQL_SYSTEM_PROMPT = """\
 You are a SQL data analyst working with a user's database.
 Generate a single SQL SELECT statement that answers the user's question.
 
+Database dialect: {dialect}
+
 Rules:
 - ONLY reference tables and columns listed in the schema below. Do not invent names.
 - Always include a LIMIT clause (max {limit}).
 - Do not use DELETE, UPDATE, INSERT, DROP, TRUNCATE, ALTER, CREATE, or any DDL.
 - Prefer explicit JOINs over subqueries when combining tables.
 - For aggregations, always alias the result column (e.g. COUNT(*) AS order_count).
-- For date filtering, use standard SQL date functions appropriate for the dialect.
+- For date filtering, use dialect-appropriate functions ({dialect} syntax).
 
 Schema:
 {schema}
@@ -81,6 +86,7 @@ class DbExecutor(BaseExecutor):
         results: list[RetrievalResult],
         user_id: str,
         db: AsyncSession,
+        question: str,
         limit: int = 100,
     ) -> list[QueryResult]:
         db_results = [r for r in results if r.source_type == "database"]
@@ -99,7 +105,7 @@ class DbExecutor(BaseExecutor):
         query_results: list[QueryResult] = []
         for client_id, client_results in by_client.items():
             try:
-                qr = await self._execute_for_client(client_id, client_results, user_id, db, limit)
+                qr = await self._execute_for_client(client_id, client_results, user_id, db, question, limit)
                 if qr:
                     query_results.append(qr)
             except Exception as e:
@@ -117,6 +123,7 @@ class DbExecutor(BaseExecutor):
         results: list[RetrievalResult],
         user_id: str,
         db: AsyncSession,
+        question: str,
         limit: int,
     ) -> QueryResult | None:
         client = await database_client_service.get(db, client_id)
@@ -143,17 +150,30 @@ class DbExecutor(BaseExecutor):
             return None
 
         schema_ctx = self._build_schema_context(full_schema)
-        question = self._extract_question(results)
         capped_limit = min(limit, _MAX_LIMIT)
+        dialect = client.db_type
 
         # SQL generation with retry
         validated_sql: str | None = None
         prev_error: str = ""
+        prev_reasoning: str = ""
         for attempt in range(_MAX_RETRIES):
-            error_section = f"Previous attempt failed: {prev_error}\nFix the issue above." if prev_error else ""
+            if prev_error:
+                error_section = (
+                    f"Previous attempt reasoning: {prev_reasoning}\n"
+                    f"Previous attempt failed: {prev_error}\n"
+                    "Fix the issue above."
+                )
+            else:
+                error_section = ""
             try:
+                prompt_text = schema_ctx + error_section + question
+                input_tokens = len(_enc.encode(prompt_text))
+                logger.info("sql generation input tokens", attempt=attempt + 1, tokens=input_tokens)
+
                 result: SQLQuery = await self._chain.ainvoke({
                     "schema": schema_ctx,
+                    "dialect": dialect,
                     "limit": capped_limit,
                     "error_section": error_section,
                     "question": question,
@@ -162,10 +182,19 @@ class DbExecutor(BaseExecutor):
                 validation_error = self._validate(sql, full_schema, capped_limit)
                 if validation_error:
                     prev_error = validation_error
+                    prev_reasoning = result.reasoning
                     logger.warning("sql validation failed", attempt=attempt + 1, error=validation_error)
                     continue
-                validated_sql = sql
-                logger.info("sql generated", attempt=attempt + 1, reasoning=result.reasoning)
+                validated_sql = self._enforce_limit(sql, capped_limit)
+                output_tokens = len(_enc.encode(result.sql)) + len(_enc.encode(result.reasoning))
+                logger.info(
+                    "sql generated",
+                    attempt=attempt + 1,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    total_tokens=input_tokens + output_tokens,
+                    reasoning=result.reasoning,
+                )
                 break
             except Exception as e:
                 prev_error = str(e)
@@ -272,13 +301,6 @@ class DbExecutor(BaseExecutor):
             lines.append("")
         return "\n".join(lines).strip()
 
-    def _extract_question(self, results: list[RetrievalResult]) -> str:
-        # The search_query rewritten by the orchestrator is not in RetrievalResult —
-        # the content field carries schema descriptions. Return a generic fallback;
-        # callers that have the original question should pass it explicitly.
-        # TODO: thread the original user question through to execute() when wiring into the agent.
-        return "Answer the user's data question using the schema provided."
-
     # ------------------------------------------------------------------
     # Guardrails
     # ------------------------------------------------------------------

src/query/query_executor.py

@@ -6,7 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from src.middlewares.logging import get_logger
 from src.query.base import QueryResult
-from src.query.executors.db import db_executor
+from src.query.executors.db_executor import db_executor
 from src.query.executors.tabular import tabular_executor
 from src.rag.base import RetrievalResult
 
@@ -19,6 +19,7 @@ class QueryExecutor:
         results: list[RetrievalResult],
         user_id: str,
         db: AsyncSession,
+        question: str,
         limit: int = 100,
     ) -> list[QueryResult]:
         db_results = [r for r in results if r.source_type == "database"]
@@ -32,8 +33,8 @@ class QueryExecutor:
             return []
 
         batches = await asyncio.gather(
-            db_executor.execute(db_results, user_id, db, limit) if db_results else _empty(),
-            tabular_executor.execute(tabular_results, user_id, db, limit) if tabular_results else _empty(),
+            db_executor.execute(db_results, user_id, db, question, limit) if db_results else _empty(),
+            tabular_executor.execute(tabular_results, user_id, db, question, limit) if tabular_results else _empty(),
             return_exceptions=True,
         )