app/services/agent_sql.py

# app/services/agent_sql.py
from typing import Dict, Any, Tuple
from app.tools.sql_tool import SQLTool
from app.tools.llm_sqlgen import SQLGenTool
from app.tools.llm_answer import AnswerLLM

class SQLAgent:
    def __init__(self, sql: SQLTool, sqlgen: SQLGenTool, answer_llm: AnswerLLM):
        self.sql = sql
        self.sqlgen = sqlgen
        self.answer_llm = answer_llm

    def ask(self, user_id: str | None, question: str) -> Tuple[str, Dict[str, Any]]:
        sql = self.sqlgen.generate_sql(question)
        result = self.sql.execute_sql_readonly(sql)
        answer = self.answer_llm.generate(question, sql, result["columns"], result["rows"])
        meta = {
            "sql": sql,
            "rowcount": result["rowcount"],
            "columns": result["columns"],
        }
        return answer, meta