Update app.py

app.py

@@ -57,9 +57,14 @@ def create_agent_app(db_path: str):
     except Exception as e:
         flash(f"[ERROR]: Failed to connect to DB: {e}", "error")
         raise
-
+    
     @tool
     def db_query_tool(query: str) -> str:
+        """
+        Execute a SQL query against the database and return the result.
+        If the query is invalid or returns no result, an error message will be returned.
+        In case of an error, the user is advised to rewrite the query and try again.
+        """
         try:
             result = db_instance.run_no_throw(query)
             return result or "Error: Query failed. Please rewrite your query and try again."
@@ -74,13 +79,54 @@ def create_agent_app(db_path: str):
         messages: Annotated[list[AnyMessage], add_messages]
 
     try:
+        query_check_system = """You are a SQL expert with a strong attention to detail.
+                                Double check the SQLite query for common mistakes, including:
+                                - Using NOT IN with NULL values
+                                - Using UNION when UNION ALL should have been used
+                                - Using BETWEEN for exclusive ranges
+                                - Data type mismatch in predicates
+                                - Properly quoting identifiers
+                                - Using the correct number of arguments for functions
+                                - Casting to the correct data type
+                                - Using the proper columns for joins
+                                
+                                If there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.
+                                
+                                You will call the appropriate tool to execute the query after running this check.
+                             """
+        
         query_check = ChatPromptTemplate.from_messages([
-            ("system", "You are a SQL expert. Fix common issues in SQLite queries."),
+            ("system", query_check_system),
             ("placeholder", "{messages}")
         ]) | llm.bind_tools([db_query_tool])
-    
+
+        query_gen_system = """You are a SQL expert with a strong attention to detail.
+
+                            Given an input question, output a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
+                            
+                            DO NOT call any tool besides SubmitFinalAnswer to submit the final answer.
+                            
+                            When generating the query:
+                            
+                            Output the SQL query that answers the input question without a tool call.
+                            
+                            Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.
+                            You can order the results by a relevant column to return the most interesting examples in the database.
+                            Never query for all the columns from a specific table, only ask for the relevant columns given the question.
+                            
+                            If you get an error while executing a query, rewrite the query and try again.
+                            
+                            If you get an empty result set, you should try to rewrite the query to get a non-empty result set.
+                            NEVER make stuff up if you don't have enough information to answer the query... just say you don't have enough information.
+                            
+                            If you have enough information to answer the input question, simply invoke the appropriate tool to submit the final answer to the user.
+                            
+                            DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database. Do not return any sql query except answer.
+                           """
+
+        
         query_gen = ChatPromptTemplate.from_messages([
-            ("system", "You are a SQL expert. Generate SQLite query and return answer using SubmitFinalAnswer tool."),
+            ("system", query_gen_system),
             ("placeholder", "{messages}")
         ]) | llm.bind_tools([SubmitFinalAnswer])
     except Exception as e: