Spaces:
Runtime error
Runtime error
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,
)
|