File size: 1,371 Bytes
16d3fdf
bc723d2
16d3fdf
5206fbb
16d3fdf
 
 
 
 
bc723d2
16d3fdf
 
 
 
 
 
0c0e0a0
bc723d2
16d3fdf
 
 
 
bc723d2
 
 
16d3fdf
 
 
 
1595682
16d3fdf
 
 
 
 
 
 
 
 
 
 
 
 
1595682
16d3fdf
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
"""An Langchain Agent that uses ChromaDB as a query tool"""
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.tools import Tool
from gnosis.search import Search


class PDFExplainer:
    """An Agent that uses ChromaDB as a query tool"""

    def __init__(self, llm, chroma_db, extra_tools=False):
        """Initialize the Agent"""
        search = Search(chroma_db)

        self.tools = [
            Tool.from_function(
                func=search.run,
                name="Search on ChromaDB",
                description="Useful when you need more context for answering a question.",
                handle_parsing_errors=True,
            )
        ]

        if extra_tools:
            self.tools.extend(load_tools(["wikipedia"]))

        self.agent = initialize_agent(
            self.tools,
            llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=False,
            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=False,
            handle_parsing_errors=True,
        )