Spaces:
Runtime error
Runtime error
| 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} | |