update tools

app.py

@@ -8,40 +8,6 @@ import spaces
 from database import engine, receipts
 
 @tool
-def sql_engine(query: str) -> str:
-    """
-    Executes an SQL query on the 'receipts' table and returns formatted results.
-    
-    Table Schema:
-        - receipt_id: INTEGER
-        - customer_name: VARCHAR(16)
-        - price: FLOAT
-        - tip: FLOAT
-    Args:
-        query: The SQL query to execute.
-    
-    Returns:
-        Query result as a formatted string.
-    """
-    try:
-        with engine.connect() as con:
-            rows = con.execute(text(query)).fetchall()
-
-        if not rows:
-            return "No results found."
-
-        # Convert results to a readable string
-        return "\n".join([", ".join(map(str, row)) for row in rows])
-
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# Initialize CodeAgent to generate SQL queries from natural language
-agent = CodeAgent(
-    tools=[sql_engine],  # Ensure sql_engine is properly registered
-    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
-)
-
 def query_sql(user_query: str) -> str:
     """
     Converts natural language input to an SQL query using CodeAgent
@@ -59,13 +25,51 @@ def query_sql(user_query: str) -> str:
     # Log the generated SQL for debugging
     print(f"Generated SQL: {generated_sql}")
 
+    # Ensure we only execute valid SQL queries
+    if not generated_sql.lower().startswith(("select", "show", "pragma")):
+        return "Error: Only SELECT queries are allowed."
+
     # Execute the SQL query and return the result
     result = sql_engine(generated_sql)
 
     # Log the result for debugging
     print(f"SQL Query Result: {result}")
 
-    return result 
+    # If the result is empty, return a friendly message
+    if not result.strip():
+        return "No results found."
+
+    return result  # Return clean result, avoiding passing raw SQL back
+
+# Initialize CodeAgent to generate SQL queries from natural language
+agent = CodeAgent(
+    tools=[sql_engine],  # Ensure sql_engine is properly registered
+    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
+)
+
+@tool
+def sql_engine(query: str) -> str:
+    """
+    Executes an SQL query on the 'receipts' table and returns formatted results.
+    
+    Args:
+        query: The SQL query to execute.
+
+    Returns:
+        Query result as a formatted string.
+    """
+    try:
+        with engine.connect() as con:
+            rows = con.execute(text(query)).fetchall()
+
+        if not rows:
+            return "No results found."
+
+        # Convert results into a readable string format
+        return "\n".join([", ".join(map(str, row)) for row in rows])
+
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 # Define Gradio interface
 demo = gr.Interface(