File size: 1,703 Bytes
ba550a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """QueryExecutor — dispatches retrieval results to the appropriate executor by source_type."""
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_executor import db_executor
from src.query.executors.tabular import tabular_executor
from src.rag.base import RetrievalResult
logger = get_logger("query_executor")
class QueryExecutor:
async def execute(
self,
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"]
tabular_results = [
r for r in results
if r.source_type == "document"
and r.metadata.get("data", {}).get("file_type") in ("csv", "xlsx")
]
async def _empty() -> list[QueryResult]:
return []
batches = await asyncio.gather(
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,
)
query_results: list[QueryResult] = []
for batch in batches:
if isinstance(batch, Exception):
logger.error("executor failed", error=str(batch))
continue
query_results.extend(batch)
logger.info("query execution complete", total=len(query_results))
return query_results
query_executor = QueryExecutor()
|