add agent.py

agent.py

@@ -1,31 +1,40 @@
 """
-Agent creation with RAG integration.
+Agent creation using LangChain create_agent.
 """
 import os
-from typing import List
-
-from langchain_openai import ChatOpenAI
-from langchain.agents import create_agent
+from langchain_ollama import OllamaLLM
 from langchain_core.messages import HumanMessage
-from langgraph.checkpoint.memory import MemorySaver
+from langchain.agents import create_agent
+from langchain.tools import tool
 from tools import search_knowledge_base, add_to_knowledge_base
 
-# LLM via Ollama (llama3)
-llm = ChatOpenAI(
-    model="ollama/llama3",
-    base_url="http://localhost:11434/v1",
-    api_key=None,
-    temperature=0.2,
-)
+# LLM via Ollama
+llm = OllamaLLM(model="llama3")
 
-agent = create_agent(
-    llm=llm,
-    tools=[search_knowledge_base, add_to_knowledge_base],
-    system_prompt="You are a helpful assistant that can search and add documents to the knowledge base.",
-)
+# System prompt instructing to use knowledge base tools
+SYSTEM_PROMPT = """
+You are a helpful assistant that can store and retrieve information.
+Use the provided tools search_knowledge_base and add_to_knowledge_base.
+When answering, prefer to call the tools if needed.
+"""
 
-memory = MemorySaver()
+def create_rag_agent():
+    agent = create_agent(
+        llm=llm,
+        tools=[search_knowledge_base, add_to_knowledge_base],
+        system_prompt=SYSTEM_PROMPT,
+    )
+    return agent
 
-async def run_agent(messages: List[HumanMessage]):
-    result = await agent.ainvoke({"messages": messages}, {"configurable": {"thread_id": "session-1"}})
-    return result["messages"][-1].content
+if __name__ == "__main__":
+    ag = create_rag_agent()
+    # Simple demo loop
+    while True:
+        user_input = input("User: ")
+        if user_input.lower() in ("quit", "exit"):
+            break
+        result = ag.ainvoke(
+            {"messages": [HumanMessage(content=user_input)]},
+            {"configurable": {"thread_id": "demo"}},
+        )
+        print(result["messages"][-1].content)