Update app.py

app.py

@@ -49,79 +49,78 @@ os.environ["MISTRAL_API_KEY"] = MISTRAL_API_KEY
 from langchain_groq import ChatGroq
 from langchain_mistralai.chat_models import ChatMistralAI
 
-#llm = ChatMistralAI(model="mistral-large-latest")
-llm = ChatGroq(model="llama3-70b-8192")
-
-@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.
-    """
-    result = db_instance.run_no_throw(query)
-    if not result:
-        return "Error: Query failed. Please rewrite your query and try again."
-    return result
-
-# Define a Pydantic model for submitting the final answer
-class SubmitFinalAnswer(BaseModel):
-    """Submit the final answer to the user based on the query results."""
-    final_answer: str = Field(..., description="The final answer to the user")
-
-# Define the state type
-class State(TypedDict):
-    messages: Annotated[list[AnyMessage], add_messages]
-
-# Define prompt templates for query checking and query generation
-from langchain_core.prompts import ChatPromptTemplate
-
-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_prompt = ChatPromptTemplate.from_messages([("system", query_check_system), ("placeholder", "{messages}")])
-query_check = query_check_prompt | 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_prompt = ChatPromptTemplate.from_messages([("system", query_gen_system), ("placeholder", "{messages}")])
-query_gen = query_gen_prompt | llm.bind_tools([SubmitFinalAnswer])
-
-
 def create_agent_app(db_path: str):
     # Construct the SQLite URI from the given file path.
     # Ensure the db_path is absolute so that SQLAlchemy can locate the file.
+    #llm = ChatMistralAI(model="mistral-large-latest")
+    llm = ChatGroq(model="llama3-70b-8192")
+    
+    @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.
+        """
+        result = db_instance.run_no_throw(query)
+        if not result:
+            return "Error: Query failed. Please rewrite your query and try again."
+        return result
+    
+    # Define a Pydantic model for submitting the final answer
+    class SubmitFinalAnswer(BaseModel):
+        """Submit the final answer to the user based on the query results."""
+        final_answer: str = Field(..., description="The final answer to the user")
+    
+    # Define the state type
+    class State(TypedDict):
+        messages: Annotated[list[AnyMessage], add_messages]
+    
+    # Define prompt templates for query checking and query generation
+    from langchain_core.prompts import ChatPromptTemplate
+    
+    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_prompt = ChatPromptTemplate.from_messages([("system", query_check_system), ("placeholder", "{messages}")])
+    query_check = query_check_prompt | 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_prompt = ChatPromptTemplate.from_messages([("system", query_gen_system), ("placeholder", "{messages}")])
+    query_gen = query_gen_prompt | llm.bind_tools([SubmitFinalAnswer])
+    
     abs_db_path = os.path.abspath(db_path)
     global DATABASE_URI
     DATABASE_URI = abs_db_path