diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..05814c1 --- /dev/null +++ b/agent.py @@ -0,0 +1,28 @@ +from langchain_ollama import ChatOllama +from langchain.agents import create_tool_calling_agent +from langchain.agents import AgentExecutor +from tools import search_local_kb, web_search + + +def create_agent(): + """Create an agent with tools for local and web search.""" + llm = ChatOllama(model="llama3", temperature=0) + + tools = [search_local_kb, web_search] + + system_prompt = """You are a helpful AI assistant with access to two search tools: + +1. search_local_kb - Search in the local knowledge base (ChromaDB) for information from stored documents +2. web_search - Search the web via Tavily for current information and recent news + +Choose the appropriate tool based on the user's question: +- For questions about local documents, notes, or stored knowledge: use search_local_kb +- For questions about current events, recent news, or facts that may not be in local documents: use web_search + +Always indicate the source of your answer in the response: 'Source: chromadb' or 'Source: tavily' +""" + + agent = create_tool_calling_agent(llm, tools, system_prompt) + agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True) + + return agent_executor \ No newline at end of file