Add Test File

agent.py

@@ -78,7 +78,7 @@ def create_assistant_tools(cfg):
             [ask_vehicles, ask_policies]
     )
 
-def initialize_agent(_cfg, update_func):
+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.

test_agent.py

@@ -0,0 +1,59 @@
+import unittest
+import os
+
+from omegaconf import OmegaConf
+from vectara_agent.agent import Agent
+
+import sqlite3
+from datasets import load_dataset
+
+from app import initialize_agent, setup_db
+
+from dotenv import load_dotenv
+load_dotenv(override=True)
+
+class TestAgentResponses(unittest.TestCase):
+
+    def test_responses(self):
+
+        cfg = OmegaConf.create({
+            'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
+        'corpus_ids': str(os.environ['VECTARA_CORPUS_IDS']).split(','),
+        'api_keys': str(os.environ['VECTARA_API_KEYS']).split(','),
+        'examples': os.environ.get('QUERY_EXAMPLES', None)
+        })
+
+        setup_db()
+
+        agent = initialize_agent(_cfg=cfg)
+        self.assertIsInstance(agent, Agent)
+
+        # Knows types of electric vehicles
+        type_output = agent.chat('What are the different types of electric vehicles? Only provide the name of each type, nothing else.').lower()
+        ev_types = ['battery', 'hybrid', 'plug-in hybrid', 'fuel cell']
+
+        for ev_type in ev_types:
+            self.assertIn(ev_type, type_output)
+
+
+        # Questions about car models - ev query tool
+        self.assertIn('mach-e', agent.chat('Which EV is made by Mustang? Provide the name of the model only.').lower())
+        self.assertIn('fuel cell', agent.chat('What EV type is the Toyota Mirai? Just give the type name.').lower())
+        
+        
+        # Incentive query tool tests
+        self.assertIn('no', agent.chat('Does the U.S. Department of Defense offer incentives for purchasing electric vehicles? Only say "yes" or "no".').lower())
+        self.assertIn('2035', agent.chat('At what year must all new passenger vehicles be zero emission vehicles in California? Give the year only.').lower())
+
+
+        # Database Tool questions
+        self.assertIn('king', agent.chat('Which county in the state of Washington had the highest number of EV registrations in 2023? Provide the name only.').lower())
+        self.assertIn('seattle', agent.chat('Which city in the state of Washington had the highest number of EV registrations in 2023? Provide the name only.').lower())
+        self.assertIn('tesla model y', agent.chat('What car was the most popular in Seattle in 2023? Provide the make and model only').lower())
+
+        # Misc. questions
+        self.assertIn('tesla', agent.chat('Which company developed a standard charging port for electric cars in California? Only provide the company name, nothing else').lower())
+
+if __name__ == "__main__":
+    setup_db()
+    unittest.main()