added get_sample_data

agent.py

@@ -4,7 +4,7 @@ from pydantic import Field, BaseModel
 from omegaconf import OmegaConf
 
 from llama_index.core.utilities.sql_wrapper import SQLDatabase
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, text
 
 from dotenv import load_dotenv
 load_dotenv(override=True)
@@ -69,27 +69,39 @@ def create_assistant_tools(cfg):
 
     tools_factory = ToolsFactory()
 
-    return (tools_factory.standard_tools() + 
-            tools_factory.guardrail_tools() +
-            tools_factory.database_tools(
+    db_tools = tools_factory.database_tools(
                 tool_name_prefix = "ev",
                 content_description = 'Electric Vehicles in the state of Washington',
                 sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
-            ) +
+            )
+
+    def ev_load_sample_data(table_name: str):
+        """
+        Given a database table name, returns the first 25 rows of the table.
+        """
+        ev_load_data = [db_tools[i] for i in range(len(db_tools)) if db_tools[i]._metadata.name == 'ev_load_data'][0]
+        return ev_load_data(f"SELECT * FROM {table_name} LIMIT 25")
+
+
+    return ([tools_factory.create_tool(ev_load_sample_data)] + 
+            tools_factory.standard_tools() + 
+            tools_factory.guardrail_tools() +
+            db_tools +
             [ask_vehicles, ask_policies]
     )
 
 def initialize_agent(_cfg, update_func=None):
     electric_vehicle_bot_instructions = """
     - You are a helpful research assistant, with expertise in electric vehicles, in conversation with a user.
-    - Before answering any user query, get sample data from each table in the database, so that you can understand NULL and unique values for each column.
+    - Before answering any user query, use ev_describe_tables to understand schema of each table, and use get_sample_data 
+      to get sample data from each table in the database, so that you can understand NULL and unique values for each column.
     - For a query with multiple sub-questions, break down the query into the sub-questions, 
       and make separate calls to the ask_vehicles or ask_policies tool to answer each sub-question,
       then combine the answers to provide a complete response.
     - Use the database tools (ev_load_data, ev_describe_tables and ev_list_tables) to answer analytical queries.
-    - IMPORTANT: When using database tools, always "query SELECT * FROM (table_name) LIMIT 25;" first to figure out the format of the columns and
-      then call the tool again to try to answer the user's query.
-    - Avoid "SELECT *" queries on tables, as they can be slow, instead craft the correct query to get the required information.
+      If you cannot find the information in one of the tables, try using the other tables in the database.
+    - IMPORTANT: When using database_tools, always call the ev_load_sample_data tool with the table you want to query
+      to understand the table structure, column naming, and values in the table. Never call the ev_load_data tool for a query until you have called ev_load_sample_data.
     - When providing links, try to put the name of the website or source of information for the displayed text. Don't just say 'Source'.
     - Never discuss politics, and always respond politely.
     """
@@ -115,4 +127,3 @@ def get_agent_config() -> OmegaConf:
         'demo_description': "This assistant can help you learn about electric vehicles in the United States, including how they work, the advantages of purchasing them, and recent trends based on data in the state of Washington.",
     })
     return cfg
-

requirements.txt

@@ -7,4 +7,4 @@ langdetect==1.0.9
 langcodes==3.4.0
 datasets==2.14.7
 uuid==1.30
-vectara-agentic==0.1.2
+vectara-agentic==0.1.4

utils.py

@@ -53,9 +53,9 @@ def send_amplitude_data(user_query, bot_response, demo_name, feedback=None):
         print(f"Amplitude request failed with status code {response.status_code}. Response Text: {response.text}")
 
 def escape_dollars_outside_latex(text):
-    # Define a regex pattern to find LaTeX equations (either single $ or double $$)
-    pattern = re.compile(r'(\$\$.*?\$\$|\$.*?\$)')
-    latex_matches = pattern.findall(text)
+    # Define a regex pattern to find LaTeX equations (double $$ only)
+    pattern = r'\$\$.*?\$\$'
+    latex_matches = re.findall(pattern, text, re.DOTALL)
     
     # Placeholder to temporarily store LaTeX equations
     placeholders = {}