Add test cases and fix formatting

agent.py

@@ -1,16 +1,16 @@
-
 import os
 import pandas as pd
 import requests
+from pydantic import Field, BaseModel
 
 from omegaconf import OmegaConf
 
+from vectara_agent.agent import Agent
+from vectara_agent.tools import ToolsFactory, VectaraToolFactory
+
 from dotenv import load_dotenv
 load_dotenv(override=True)
 
-from pydantic import Field, BaseModel
-from vectara_agent.agent import Agent, AgentStatusType
-from vectara_agent.tools import ToolsFactory, VectaraToolFactory
 
 tickers = {
     "AAPL": "Apple Computer", 
@@ -28,7 +28,7 @@ tickers = {
 years = [2020, 2021, 2022, 2023, 2024]
 initial_prompt = "How can I help you today?"
 
-def create_assistant_tools(cfg):    
+def create_assistant_tools(cfg):
 
     def get_company_info() -> list[str]:
         """
@@ -44,7 +44,7 @@ def create_assistant_tools(cfg):
         Always check this before using any other tool.
         """
         return years
-    
+
     # Tool to get the income statement for a given company and year using the FMP API
     def get_income_statement(
         ticker=Field(description="the ticker symbol of the company."),
@@ -68,16 +68,16 @@ def create_assistant_tools(cfg):
             ]
             values_dict = income_statement_specific_year.to_dict(orient="records")[0]
             return f"Financial results: {', '.join([f'{key}: {value}' for key, value in values_dict.items() if key not in ['date', 'cik', 'link', 'finalLink']])}"
-        else:
-            return "FMP API returned error. This tool does not work."
+        
+        return "FMP API returned error. This tool does not work."
 
     class QueryTranscriptsArgs(BaseModel):
         query: str = Field(..., description="The user query.")
         year: int = Field(..., description=f"The year. An integer between {min(years)} and {max(years)}.")
         ticker: str = Field(..., description=f"The company ticker. Must be a valid ticket symbol from the list {tickers.keys()}.")
 
-    vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key, 
-                                     vectara_customer_id=cfg.customer_id, 
+    vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key,
+                                     vectara_customer_id=cfg.customer_id,
                                      vectara_corpus_id=cfg.corpus_id)
     tools_factory = ToolsFactory()
 
@@ -85,10 +85,10 @@ def create_assistant_tools(cfg):
         tool_name = "ask_transcripts",
         tool_description = """
         Given a company name and year, responds to a user question about the company, based on analyst call transcripts about the company's financial reports for that year.
-        You can ask this tool any question about the compaany including risks, opportunities, financial performance, competitors and more.
+        You can ask this tool any question about the company including risks, opportunities, financial performance, competitors and more.
         """,
         tool_args_schema = QueryTranscriptsArgs,
-        reranker = "multilingual_reranker_v1", rerank_k = 100, 
+        reranker = "multilingual_reranker_v1", rerank_k = 100,
         n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
         summary_num_results = 10,
         vectara_summarizer = 'vectara-summary-ext-24-05-med-omni',
@@ -98,13 +98,13 @@ def create_assistant_tools(cfg):
     return (
             [tools_factory.create_tool(tool) for tool in
                 [
-                    get_company_info, 
+                    get_company_info,
                     get_valid_years,
                     get_income_statement,
                 ]
             ] +
-            tools_factory.standard_tools() + 
-            tools_factory.financial_tools() + 
+            tools_factory.standard_tools() +
+            tools_factory.financial_tools() +
             tools_factory.guardrail_tools() +
             [ask_transcripts]
     )

test_agent.py

@@ -22,11 +22,21 @@ class TestAgentResponses(unittest.TestCase):
 
         agent = initialize_agent(_cfg=cfg)
         self.assertIsInstance(agent, Agent)
-        
-        # Basic number questions
+
+        # Test RAG tool
+        self.assertIn('elon musk', agent.chat('Who is the CEO of Tesla? Just give their name.').lower())
+        self.assertIn('david zinsner', agent.chat('Who is the CFO of Intel? Just give their name.').lower())
+
+        # Test Yahoo Finance tool
         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())
+        self.assertIn('10.46B', agent.chat('What was AMD\'s gross profit in 2023?. Just provide the number.'))
+
+        # Test other custom tools
+        self.assertIn('2020', agent.chat('What is the earliest year you can provide information for? Just give the year.').lower())
+        self.assertIn('2024', agent.chat('What is the most recent year you can provide information for? Just give the year.').lower())
+        self.assertIn('AAPL', agent.chat('What is the stock ticker symbol for Apple? Just give the abbreiated ticker.'))
 
 
 if __name__ == "__main__":