update

agent.py

@@ -10,94 +10,101 @@ from dotenv import load_dotenv
 load_dotenv(override=True)
 
 from vectara_agentic.agent import Agent
+from vectara_agentic.agent_config import AgentConfig
 from vectara_agentic.tools import ToolsFactory, VectaraToolFactory
 
-def create_assistant_tools(cfg):    
-
-    class QueryElectricCars(BaseModel):
-        query: str = Field(description="The user query.")
-
-    vec_factory_1 = VectaraToolFactory(vectara_api_key=cfg.api_keys[0],
-                                        vectara_corpus_key=cfg.corpus_keys[0])
-    
-    summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'
-
-    ask_vehicles = vec_factory_1.create_rag_tool(
-        tool_name = "ask_vehicles",
-        tool_description = """
-        Given a user query, 
-        returns a response to a user question about electric vehicles.
-        """,
-        tool_args_schema = QueryElectricCars,
-        reranker = "chain", rerank_k = 100,
-        rerank_chain = [
-            {
-                "type": "slingshot",
-                "cutoff": 0.2
-            },
-            {
-                "type": "mmr",
-                "diversity_bias": 0.1
-            }
-        ],
-        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
-        summary_num_results = 5,
-        vectara_summarizer = summarizer,
-        include_citations = False,
-    )
-
-    vec_factory_2 = VectaraToolFactory(vectara_api_key=cfg.api_keys[1],
-                                       vectara_corpus_key=cfg.corpus_keys[1])
-    
-
-    class QueryEVLaws(BaseModel):
-        query: str = Field(description="The user query")
-        state: Optional[str] = Field(default=None,
-                                     description="The two digit state code. Optional.",
-                                     examples=['CA', 'US', 'WA'])
-        policy_type: Optional[str] = Field(default=None,
-                                           description="The type of policy. Optional",
-                                           examples = ['Laws and Regulations', 'State Incentives', 'Incentives', 'Utility / Private Incentives', 'Programs'])
-        
+class AgentTools:
+    def __init__(self, _cfg, agent_config):
+        self.tools_factory = ToolsFactory()
+        self.agent_config = agent_config
+        self.cfg = _cfg
 
-    ask_policies = vec_factory_2.create_rag_tool(
-        tool_name = "ask_policies",
-        tool_description = """
-        Given a user query,
-        returns a response to a user question about electric vehicles incentives and regulations, in the United States.
-        You can ask this tool any question about laws passed by states or the federal government related to electric vehicles.
-        """,
-        tool_args_schema = QueryEVLaws,
-        reranker = "chain", rerank_k = 100, 
-        rerank_chain = [
-            {
-                "type": "slingshot",
-                "cutoff": 0.2
-            },
-            {
-                "type": "mmr",
-                "diversity_bias": 0.1
-            }
-        ],
-        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
-        summary_num_results = 10,
-        vectara_summarizer = summarizer,
-        include_citations = False,
-    )
 
-    tools_factory = ToolsFactory()
+    def get_tools(self):
+        class QueryElectricCars(BaseModel):
+            query: str = Field(description="The user query.")
 
-    db_tools = tools_factory.database_tools(
-                tool_name_prefix = "ev",
-                content_description = 'Electric Vehicles in the state of Washington and other population information',
-                sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
-            )
+        vec_factory_1 = VectaraToolFactory(vectara_api_key=self.cfg.api_keys[0],
+                                            vectara_corpus_key=self.cfg.corpus_keys[0])
+        
+        summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'
+
+        ask_vehicles = vec_factory_1.create_rag_tool(
+            tool_name = "ask_vehicles",
+            tool_description = """
+            Given a user query, 
+            returns a response to a user question about electric vehicles.
+            """,
+            tool_args_schema = QueryElectricCars,
+            reranker = "chain", rerank_k = 100,
+            rerank_chain = [
+                {
+                    "type": "slingshot",
+                    "cutoff": 0.2
+                },
+                {
+                    "type": "mmr",
+                    "diversity_bias": 0.1
+                }
+            ],
+            n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
+            summary_num_results = 5,
+            vectara_summarizer = summarizer,
+            include_citations = False,
+        )
+
+        vec_factory_2 = VectaraToolFactory(vectara_api_key=self.cfg.api_keys[1],
+                                        vectara_corpus_key=self.cfg.corpus_keys[1])
+        
 
-    return (tools_factory.standard_tools() + 
-            tools_factory.guardrail_tools() +
-            db_tools +
-            [ask_vehicles, ask_policies]
-    )
+        class QueryEVLaws(BaseModel):
+            query: str = Field(description="The user query")
+            state: Optional[str] = Field(default=None,
+                                        description="The two digit state code. Optional.",
+                                        examples=['CA', 'US', 'WA'])
+            policy_type: Optional[str] = Field(default=None,
+                                            description="The type of policy. Optional",
+                                            examples = ['Laws and Regulations', 'State Incentives', 'Incentives', 'Utility / Private Incentives', 'Programs'])
+            
+
+        ask_policies = vec_factory_2.create_rag_tool(
+            tool_name = "ask_policies",
+            tool_description = """
+            Given a user query,
+            returns a response to a user question about electric vehicles incentives and regulations, in the United States.
+            You can ask this tool any question about laws passed by states or the federal government related to electric vehicles.
+            """,
+            tool_args_schema = QueryEVLaws,
+            reranker = "chain", rerank_k = 100, 
+            rerank_chain = [
+                {
+                    "type": "slingshot",
+                    "cutoff": 0.2
+                },
+                {
+                    "type": "mmr",
+                    "diversity_bias": 0.1
+                }
+            ],
+            n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
+            summary_num_results = 10,
+            vectara_summarizer = summarizer,
+            include_citations = False,
+        )
+
+        tools_factory = ToolsFactory()
+
+        db_tools = tools_factory.database_tools(
+                    tool_name_prefix = "ev",
+                    content_description = 'Electric Vehicles in the state of Washington and other population information',
+                    sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
+                )
+
+        return (tools_factory.standard_tools() + 
+                tools_factory.guardrail_tools() +
+                db_tools +
+                [ask_vehicles, ask_policies]
+        )
 
 def initialize_agent(_cfg, agent_progress_callback=None):
     electric_vehicle_bot_instructions = """
@@ -106,7 +113,7 @@ def initialize_agent(_cfg, agent_progress_callback=None):
     """
 
     agent = Agent(
-        tools=create_assistant_tools(_cfg),
+        tools=AgentTools(_cfg, AgentConfig()).get_tools(),
         topic="Electric vehicles in the United States",
         custom_instructions=electric_vehicle_bot_instructions,
         agent_progress_callback=agent_progress_callback

requirements.txt

@@ -7,4 +7,4 @@ langdetect==1.0.9
 langcodes==3.4.0
 datasets==2.19.2
 uuid==1.30
-vectara-agentic==0.2.0
+vectara-agentic==0.2.1