test_agent.py

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)
        
        # Basic number questions
        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())


if __name__ == "__main__":
    unittest.main()