Spaces:
Sleeping
Sleeping
File size: 821 Bytes
531ff91 c93626e 531ff91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
# 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
|