Add agent.py

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)