| import os |
| import sys |
| from pathlib import Path |
|
|
| |
| BACKEND_DIR = Path(__file__).resolve().parent.parent |
| if str(BACKEND_DIR) not in sys.path: |
| sys.path.insert(0, str(BACKEND_DIR)) |
|
|
| import config |
| from pydantic import BaseModel |
| from litellm import completion |
| from database.database import execute_query |
|
|
| |
| DB_SCHEMA = """ |
| CREATE TABLE customers (id INTEGER PRIMARY KEY, first_name TEXT, last_name TEXT, email TEXT, created_at TIMESTAMP); |
| CREATE TABLE orders (id INTEGER PRIMARY KEY, customer_id INTEGER, product_name TEXT, amount REAL, status TEXT, order_date TIMESTAMP, FOREIGN KEY (customer_id) REFERENCES customers(id)); |
| """ |
|
|
| class QueryToolParams(BaseModel): |
| natural_language_query: str |
|
|
| def run_query_tool(params: dict) -> str: |
| """Translates NL to SQL, validates it, and executes it against the database.""" |
| try: |
| tool_config = QueryToolParams(**params) |
| except Exception as e: |
| return f"Error parsing parameters: {e}" |
|
|
| prompt = f"""You are an expert SQL assistant. Convert the following natural language query into a raw SQL SELECT statement. |
| Schema: |
| {DB_SCHEMA} |
| |
| Query: {tool_config.natural_language_query} |
| |
| Return ONLY the SQL string, without markdown formatting or code blocks. |
| """ |
| try: |
| model_name = config.get_llm_model() |
| api_key = config.get_api_key() |
| |
| completion_kwargs = { |
| "model": model_name, |
| "messages": [{"role": "user", "content": prompt}], |
| "temperature": 0 |
| } |
| if api_key: |
| completion_kwargs["api_key"] = api_key |
|
|
| response = completion(**completion_kwargs) |
| message = response.choices[0].message |
| sql_query = (message.content or getattr(message, "reasoning_content", "") or "").strip() |
| except Exception as e: |
| return f"Query Tool Error: {str(e)}" |
| |
| try: |
| import re |
| |
| code_block_match = re.search(r"```(?:sql)?\s*([\s\S]*?)\s*```", sql_query, re.IGNORECASE) |
| if code_block_match: |
| sql_query = code_block_match.group(1).strip() |
| elif not sql_query.upper().startswith("SELECT"): |
| |
| select_match = re.search(r"(SELECT\s+[\s\S]+?)(?:;|\n\n|$)", sql_query, re.IGNORECASE) |
| if select_match: |
| sql_query = select_match.group(1).strip() |
| if sql_query.startswith("```sql"): |
| sql_query = sql_query.replace("```sql", "", 1) |
| if sql_query.endswith("```"): |
| sql_query = sql_query[::-1].replace("```", "", 1)[::-1] |
| elif sql_query.startswith("```"): |
| sql_query = sql_query.replace("```", "", 1) |
| if sql_query.endswith("```"): |
| sql_query = sql_query[::-1].replace("```", "", 1)[::-1] |
| |
| sql_query = sql_query.strip() |
| |
| |
| if not sql_query.upper().startswith("SELECT"): |
| return "Safety Violation: Only SELECT queries are allowed." |
| |
| if "LIMIT" not in sql_query.upper(): |
| sql_query = sql_query.rstrip(";") + " LIMIT 50;" |
| |
| rows = execute_query(sql_query) |
| if not rows: |
| return f"Executed `{sql_query}`: No results found." |
| |
| return f"Executed `{sql_query}`:\n" + str(rows) |
|
|
| except Exception as e: |
| return f"Query Tool Error: {str(e)}" |
|
|
