Spaces:
Running
Running
| import unittest | |
| import os | |
| from omegaconf import OmegaConf | |
| from vectara_agent.agent import Agent | |
| 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_id': str(os.environ['VECTARA_CORPUS_ID']), | |
| 'api_key': str(os.environ['VECTARA_API_KEY']), | |
| 'examples': os.environ.get('QUERY_EXAMPLES', None) | |
| }) | |
| agent = initialize_agent(_cfg=cfg) | |
| self.assertIsInstance(agent, Agent) | |
| # Test RAG tool | |
| self.assertIn('elon musk', agent.chat('Who is the CEO of Tesla? Just give their name.').lower()) | |
| self.assertIn('david zinsner', agent.chat('Who is the CFO of Intel? Just give their name.').lower()) | |
| # Test Yahoo Finance tool | |
| self.assertIn('274.52B', agent.chat('Was was the revenue for Apple in 2020? Just provide the number.')) | |
| self.assertIn('amazon', agent.chat('Which company had the highest revenue in 2023? Just provide the name.').lower()) | |
| self.assertIn('amazon', agent.chat('Which company had the lowest profits in 2022?').lower()) | |
| self.assertIn('10.46B', agent.chat('What was AMD\'s gross profit in 2023?. Just provide the number.')) | |
| # Test other custom tools | |
| self.assertIn('2020', agent.chat('What is the earliest year you can provide information for? Just give the year.').lower()) | |
| self.assertIn('2024', agent.chat('What is the most recent year you can provide information for? Just give the year.').lower()) | |
| self.assertIn('AAPL', agent.chat('What is the stock ticker symbol for Apple? Just give the abbreiated ticker.')) | |
| if __name__ == "__main__": | |
| unittest.main() | |