30 lines
1.0 KiB
Python
30 lines
1.0 KiB
Python
"""
|
|
Agent creation with routing between local KB and web search.
|
|
"""
|
|
from langchain.agents import initialize_agent, Tool
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_ollama import OllamaLLM
|
|
from tools import search_local_kb, web_search
|
|
|
|
# Define tools
|
|
local_tool = Tool(name="search_local_kb", func=search_local_kb, description="Search local knowledge base")
|
|
web_tool = Tool(name="web_search", func=web_search, description="Web search via Tavily")
|
|
|
|
llm = OllamaLLM(model="llama3")
|
|
|
|
prompt = ChatPromptTemplate.from_messages([
|
|
("system", "You are an assistant that can answer questions using either local knowledge or the web. Use the appropriate tool and indicate source in your response."),
|
|
("human", "{input}"),
|
|
])
|
|
|
|
agent_executor = initialize_agent([local_tool, web_tool], llm, agent="zero-shot-react-description", verbose=True)
|
|
|
|
# Example usage:
|
|
if __name__ == "__main__":
|
|
while True:
|
|
q = input("Query> ")
|
|
if q.lower() in {"exit", "quit"}:
|
|
break
|
|
res = agent_executor.run(q)
|
|
print(res)
|