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']),
            'style': 'traditional',
            'language': 'English',
            'student_age': 18
        })

        agent = initialize_agent(_cfg=cfg)
        self.assertIsInstance(agent, Agent)

        # Basic Course Questions
        self.assertIn('Michael Sandel', agent.chat('Who teaches the Justice Harvard course? Just give their name'))
        self.assertIn('government', agent.chat('In what department is the class taught? Just give the name.').lower())
        self.assertIn('brandeis university', agent.chat('At what college did Sandel earn his undergraduate degree?').lower())
        
        # Event Questions
        self.assertIn('1945', agent.chat('In what year did the Nuremberg Trials begin?'))
        self.assertIn('martin luther king jr.', agent.chat('Who led the March on Washington?').lower())
        self.assertIn('obergefell v. hodges', agent.chat('In what court case did gay marriage become legal in the United States?').lower())
        


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