From 966c86598089dd707baeb9a1a76f96d8e8b0e6f2 Mon Sep 17 00:00:00 2001
From: Juan Perez Tejada <jpereztejada@apexsystems.com>
Date: Sat, 9 Dec 2023 13:41:33 -0600
Subject: [PATCH] Solve pylint complains

---
 app.py       |  2 +-
 src/agent.py | 20 +++++++++++++++++---
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/app.py b/app.py
index 1ef338a..f05be37 100644
--- a/app.py
+++ b/app.py
@@ -79,7 +79,7 @@ def set_api_key():
 else:
     st.write("Please upload a file of type: pdf")
 
-st.subheader("Search your knowledge base")
+st.subheader("Search on your knowledge base")
 # if st.button("Chroma data collection"):
 #     st.write(collection)
 
diff --git a/src/agent.py b/src/agent.py
index c4b11ef..fec7376 100644
--- a/src/agent.py
+++ b/src/agent.py
@@ -11,19 +11,33 @@ def __init__(self, llm, chroma_db):
         """Initialize the Agent"""
         search = Search(chroma_db)
 
-        tools = [
+        self.tools = [
             Tool.from_function(
                 func=search.run,
                 name="Search DB",
-                description="Useful when you need to answer questions about a PDF in the collection.",
+                description="Useful when you need more context about a specific topic.",
                 handle_parsing_errors=True,
             )
         ]
 
         self.agent = initialize_agent(
-            tools,
+            self.tools,
             llm,
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
             verbose=True,
             handle_parsing_errors=True,
         )
+
+    def add_tools(self, tools: list[Tool]):
+        """Add tools to the Agent"""
+        self.tools.extend(tools)
+
+    def replace_agent(self, agent: AgentType, llm):
+        """Replace the Agent"""
+        self.agent = initialize_agent(
+            self.tools,
+            llm,
+            agent=agent,
+            verbose=True,
+            handle_parsing_errors=True,
+        )