[KM-512] add Pydantic model the LLM fills via function calling in sql_query, and add same signature for db and tabular

src/models/sql_query.py

@@ -0,0 +1,8 @@
+"""Structured output model for LLM-generated SQL queries."""
+
+from pydantic import BaseModel, Field
+
+
+class SQLQuery(BaseModel):
+    sql: str = Field(description="A single SQL SELECT statement. No markdown, no explanation inline.")
+    reasoning: str = Field(description="One sentence: what this query answers.")

src/query/base.py

@@ -3,6 +3,8 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 
+from sqlalchemy.ext.asyncio import AsyncSession
+
 from src.rag.base import RetrievalResult
 
 
@@ -15,6 +17,7 @@ class QueryResult:
     rows: list[dict]
     row_count: int
     metadata: dict = field(default_factory=dict)
+    # metadata should include "column_types": {"col_name": "dtype"} when available
 
 
 class BaseExecutor(ABC):
@@ -23,5 +26,6 @@ class BaseExecutor(ABC):
         self,
         results: list[RetrievalResult],
         user_id: str,
+        db: AsyncSession,
         limit: int = 100,
     ) -> list[QueryResult]: ...

src/query/executor.py

@@ -2,6 +2,8 @@
 
 import asyncio
 
+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
@@ -16,6 +18,7 @@ class QueryExecutor:
         self,
         results: list[RetrievalResult],
         user_id: str,
+        db: AsyncSession,
         limit: int = 100,
     ) -> list[QueryResult]:
         db_results = [r for r in results if r.source_type == "database"]
@@ -29,8 +32,8 @@ class QueryExecutor:
             return []
 
         batches = await asyncio.gather(
-            db_executor.execute(db_results, user_id, limit) if db_results else _empty(),
-            tabular_executor.execute(tabular_results, user_id, limit) if tabular_results else _empty(),
+            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(),
             return_exceptions=True,
         )
 

src/query/executors/tabular.py

@@ -7,6 +7,8 @@ Flow:
   4. Return QueryResult per document.
 """
 
+from sqlalchemy.ext.asyncio import AsyncSession
+
 from src.middlewares.logging import get_logger
 from src.query.base import BaseExecutor, QueryResult
 from src.rag.base import RetrievalResult
@@ -21,6 +23,7 @@ class TabularExecutor(BaseExecutor):
         self,
         results: list[RetrievalResult],
         user_id: str,
+        db: AsyncSession,
         limit: int = 100,
     ) -> list[QueryResult]:
         # TODO: implement