Spaces:
Running
Running
| import re | |
| import sys | |
| from pathlib import Path | |
| from typing import Any | |
| from langchain_core.prompts import ChatPromptTemplate | |
| from ..config import Settings | |
| from ..database import run_sql | |
| from ..llm import build_chat_model | |
| from ..schemas import SQLPlan | |
| from ..usage import extract_token_usage | |
| # Central skill loader import | |
| _superagent_root = Path(__file__).resolve().parents[4] / "SuperAgent" | |
| if str(_superagent_root) not in sys.path: | |
| sys.path.insert(0, str(_superagent_root)) | |
| from app.services.skill_loader import ( | |
| SQL_PLANNER_PROMPT, | |
| ANSWER_FORMATTER_PROMPT, | |
| get_escaped_system_message, | |
| ) | |
| FORBIDDEN_SQL_PATTERNS = [ | |
| r"\binsert\b", r"\bupdate\b", r"\bdelete\b", r"\bdrop\b", | |
| r"\balter\b", r"\btruncate\b", r"\bcreate\b", r"\bgrant\b", | |
| r"\brevoke\b", r"\bcopy\b", r";.+\S", | |
| ] | |
| class SQLAgent: | |
| def __init__(self, settings: Settings) -> None: | |
| self.settings = settings | |
| self.llm = build_chat_model(settings, temperature=0) | |
| system_msg = get_escaped_system_message("executive", default_message="You are VisBot Executive.") | |
| self.plan_prompt = ChatPromptTemplate.from_messages( | |
| [ | |
| ("system", f"{system_msg}\n\n{SQL_PLANNER_PROMPT}"), | |
| ("human", "Question:\n{question}\n\nGenerate a plan now."), | |
| ] | |
| ) | |
| self.answer_prompt = ChatPromptTemplate.from_messages( | |
| [ | |
| ("system", f"{system_msg}\n\n{ANSWER_FORMATTER_PROMPT}"), | |
| ("human", "Question:\n{question}\n\nRows:\n{rows}\n\nAsset:\n{asset_name}"), | |
| ] | |
| ) | |
| def plan( | |
| self, | |
| question: str, | |
| user_id: str, | |
| asset_name: str, | |
| allowed_assets: list[str], | |
| conversation_summary: str, | |
| category_children: list[str] = None, | |
| schema_name: str | None = None, | |
| ) -> tuple[SQLPlan, dict[str, int]]: | |
| asset_context = asset_name | |
| if category_children: | |
| asset_context = f"{asset_name} (Category containing: {', '.join(category_children)})" | |
| schema = schema_name or self.settings.openremote_schema | |
| allowed_assets_clause = "IN (" + ", ".join([f"'{a}'" for a in allowed_assets if a != "... (and more)"]) + ")" if allowed_assets else "IN ('')" | |
| prompt = self.plan_prompt.partial( | |
| user_id=user_id, | |
| asset_name=asset_context, | |
| allowed_assets=allowed_assets_clause, | |
| conversation_summary=conversation_summary[-300:] if len(conversation_summary) > 300 else conversation_summary, | |
| schema=schema, | |
| extra_context="", | |
| ) | |
| raw_result = self.llm.invoke(prompt.format(question=question)) | |
| content = raw_result.content if hasattr(raw_result, "content") else str(raw_result) | |
| markdown_match = re.search(r"```sql\s*(.*?)\s*```", content, re.DOTALL | re.IGNORECASE) | |
| if markdown_match: | |
| sql = markdown_match.group(1).strip() | |
| else: | |
| sql_match = re.search(r"(WITH|SELECT)\b.*", content, re.DOTALL | re.IGNORECASE) | |
| sql = sql_match.group(0).strip() if sql_match else "" | |
| if sql: | |
| if ";" in sql: | |
| sql = sql[:sql.rindex(";")+1] | |
| return SQLPlan( | |
| needs_sql=True, | |
| scenario="dynamic_sql", | |
| sql=sql, | |
| answer_if_no_sql="" | |
| ), extract_token_usage(raw_result) | |
| return SQLPlan( | |
| needs_sql=False, | |
| scenario="direct_answer", | |
| sql="", | |
| answer_if_no_sql=content | |
| ), extract_token_usage(raw_result) | |
| def validate_sql(self, sql: str, schema_name: str | None = None) -> str: | |
| normalized = sql.strip() | |
| if not normalized: | |
| raise ValueError("The planner returned an empty SQL query.") | |
| schema = schema_name or self.settings.openremote_schema | |
| for table in ["asset", "asset_datapoint", "alarm"]: | |
| pattern = rf"(?<!{schema}\.)\b{table}\b" | |
| normalized = re.sub(pattern, f"{schema}.{table}", normalized, flags=re.IGNORECASE) | |
| lowered = normalized.lower() | |
| if not (lowered.startswith("select") or lowered.startswith("with")): | |
| raise ValueError("Only SELECT statements are allowed.") | |
| if f"{schema}." not in normalized.lower(): | |
| raise ValueError("Queries must stay inside the configured schema.") | |
| for pattern in FORBIDDEN_SQL_PATTERNS: | |
| if re.search(pattern, normalized, flags=re.IGNORECASE | re.DOTALL): | |
| raise ValueError("The generated SQL contains forbidden operations.") | |
| return normalized | |