Spaces:
Running
Running
| import unittest | |
| import os | |
| from omegaconf import OmegaConf | |
| from vectara_agent.agent import Agent | |
| import sqlite3 | |
| from datasets import load_dataset | |
| from app import setup_db | |
| from agent import initialize_agent | |
| 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() |