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_ollama import OllamaLLM
-
-from langchain_openai import ChatOpenAI
-from langchain.agents import create_agent
 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 via Ollama
-llm = ChatOpenAI(
+llm = OllamaLLM(model="llama3")
-    model="ollama/llama3",
-    base_url="http://localhost:11434/v1",
-    api_key=None,
-    temperature=0.2,
-)
 
+# 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.
+"""
+
+def create_rag_agent():
     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=SYSTEM_PROMPT,
     )
+    return agent
 
-memory = MemorySaver()
+if __name__ == "__main__":
-
+    ag = create_rag_agent()
-async def run_agent(messages: List[HumanMessage]):
+    # Simple demo loop
-    result = await agent.ainvoke({"messages": messages}, {"configurable": {"thread_id": "session-1"}})
+    while True:
-    return result["messages"][-1].content
+        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)