Spaces:
Running
Running
File size: 4,800 Bytes
89acd07 | 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 | 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
|