david-oplatka commited on
Commit
d1788ed
1 Parent(s): 7735a45

Add test file

Browse files
Files changed (1) hide show
  1. test_agent.py +41 -0
test_agent.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import unittest
2
+ import os
3
+
4
+ from omegaconf import OmegaConf
5
+ from vectara_agent.agent import Agent
6
+
7
+ from agent import initialize_agent
8
+
9
+ from dotenv import load_dotenv
10
+ load_dotenv(override=True)
11
+
12
+ class TestAgentResponses(unittest.TestCase):
13
+
14
+ def test_responses(self):
15
+
16
+ cfg = OmegaConf.create({
17
+ 'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
18
+ 'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
19
+ 'api_key': str(os.environ['VECTARA_API_KEY']),
20
+ 'style': 'traditional',
21
+ 'language': 'English',
22
+ 'student_age': 18
23
+ })
24
+
25
+ agent = initialize_agent(_cfg=cfg)
26
+ self.assertIsInstance(agent, Agent)
27
+
28
+ # Basic Course Questions
29
+ self.assertIn('Michael Sandel', agent.chat('Who teaches the Justice Harvard course? Just give their name'))
30
+ self.assertIn('government', agent.chat('In what department is the class taught? Just give the name.').lower())
31
+ self.assertIn('brandeis university', agent.chat('At what college did Sandel earn his undergraduate degree?').lower())
32
+
33
+ # Event Questions
34
+ self.assertIn('1945', agent.chat('In what year did the Nuremberg Trials begin?'))
35
+ self.assertIn('martin luther king jr.', agent.chat('Who led the March on Washington?').lower())
36
+ self.assertIn('obergefell v. hodges', agent.chat('In what court case did gay marriage become legal in the United States?').lower())
37
+
38
+
39
+
40
+ if __name__ == "__main__":
41
+ unittest.main()