Justice-Harvard / test_agent.py
david-oplatka's picture
Add test file
d1788ed
raw
history blame
1.51 kB
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()