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']),
            'tavily_api_key': str(os.environ['TAVILY_API_KEY'])
        })

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

        # Questions about Posts/Stories
        self.assertIn('mockingbird', agent.chat('What is the name of the LLM created by Vectara? Give the name only').lower())
        self.assertIn('texas', agent.chat('Which state is Meta paying because they used facial recognition without user\'s permission? Give the state name only').lower())
        self.assertIn('nextos', agent.chat('Who wrote the post The Hitchhiker\'s Guide to Logical Verification? Give the user\'s name only.').lower())

        # Identifying Pages/Posts
        self.assertIn('41099028', agent.chat('What is the id for the post titled Tea-tasking, a Python package for the statistical analysis of A/B tests? Just give the id, nothing else'))
        self.assertIn('e10v_me', agent.chat('Who commented \'Experiments with 3 or more variants are quite rare in my practice\' on this page? Give the user name only.').lower())
        self.assertIn('word slicer', agent.chat('What is the title of the post with id 41078020 Give the name only').lower())

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