diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..f282cd3 --- /dev/null +++ b/agent.py @@ -0,0 +1,29 @@ +""" +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)