Spaces:
Runtime error
Runtime error
File size: 5,010 Bytes
ba71fca | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | from agents.base import BaseAgent
import sqlite3
import pandas as pd
from datetime import datetime
import re
class CoffeeDatabaseAgent(BaseAgent):
def __init__(self, db_path: str, llm_agent):
self.db_path = db_path
self.llm_agent = llm_agent
def check_connection(self) -> bool:
try:
conn = sqlite3.connect(self.db_path)
conn.execute("SELECT 1;")
conn.close()
return True
except Exception as e:
print(f"[DatabaseAgent] DB Connection Error: {e}")
return False
def get_schema_overview(self) -> str:
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
tables = cursor.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
if not tables:
return "(No tables found)"
lines = []
for (table,) in tables:
columns = cursor.execute(f"PRAGMA table_info({table})").fetchall()
col_list = [col[1] for col in columns]
col_meta = []
for col in col_list:
if any(kw in col.lower() for kw in ["status", "type", "category", "state", "flag"]):
vals = cursor.execute(f"SELECT DISTINCT {col} FROM {table} LIMIT 21").fetchall()
vals = [str(v[0]) for v in vals if v[0] is not None]
if 1 < len(vals) <= 20:
col_meta.append(f"{col}: [{', '.join(vals)}]")
if col_meta:
lines.append(f"{table}({', '.join(col_list)}) # " + "; ".join(col_meta))
else:
lines.append(f"{table}({', '.join(col_list)})")
conn.close()
return "\n".join(lines)
except Exception as e:
return f"(Schema error: {e})"
def llm_to_sql(self, command: str, schema: str) -> str:
prompt = (
"You are a database specialist agent. Given the schema and user request, generate a single, safe SQL SELECT statement. "
"If you cannot answer with available tables/columns, reply NO_SQL.\n"
f"Database schema overview:\n{schema}\n"
f"User request: {command}\nSQL:"
)
response = self.llm_agent.invoke([{"role": "user", "content": prompt}])
if hasattr(response, 'content'):
sql = response.content
else:
sql = str(response)
match = re.search(r"```sql\s*(.*?)```", sql, re.DOTALL | re.IGNORECASE)
if match:
sql = match.group(1).strip()
elif sql.startswith("```") and sql.endswith("```"):
sql = sql[3:-3].strip()
elif sql.upper().startswith("SQL:"):
sql = sql[4:].strip()
return sql
def query_database(self, sql: str, max_rows: int = 15) -> pd.DataFrame:
if sql.upper() == "NO_SQL":
raise ValueError("Cannot answer this question with the given schema.")
conn = sqlite3.connect(self.db_path)
try:
df = pd.read_sql_query(sql, conn)
finally:
conn.close()
return df.head(max_rows)
def process(self, state: dict) -> dict:
command = ""
for a in state.get("assigned_agents", []):
if a["agent"] == "coffee_db_agent":
command = a["command"]
break
logs = []
now = datetime.now().isoformat(timespec='seconds')
logs.append(f"[{now}] [DatabaseAgent] Received command: '{command}'")
if not self.check_connection():
log = f"[{now}] [DatabaseAgent][ERROR] Cannot connect to database ({self.db_path})"
return {"result": "❌ Database connection failed.", "logs": [log]}
schema = self.get_schema_overview()
logs.append(f"[{now}] [DatabaseAgent] Schema overview: {schema}")
try:
sql = self.llm_to_sql(command, schema)
logs.append(f"[{now}] [DatabaseAgent] LLM generated SQL: {sql}")
except Exception as e:
log = f"[{now}] [DatabaseAgent][ERROR] LLM to SQL failed: {e}"
return {"result": f"❌ [LLM Error]: {e}", "logs": logs + [log]}
try:
df = self.query_database(sql)
if df.empty:
result = "No data found for your request."
else:
result = f"Query result:\n{df.to_string(index=False)}"
logs.append(f"[{now}] [DatabaseAgent] Query executed successfully. Rows: {len(df)}")
except Exception as e:
result = f"❌ [DB Error]: {e}\n(SQL: {sql})"
logs.append(f"[{now}] [DatabaseAgent][ERROR] Query execution failed: {e}")
logs.append(f"[{now}] [DatabaseAgent] Done. Result set in assigned_agents.")
for a in state.get("assigned_agents", []):
if a["agent"] == "coffee_db_agent":
a["result"] = result
return {"logs": logs}
|