From acd0c52d0131a9b5eb0c6010daadfde8160b5241 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?=
 =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?=
 <kirillkutlahmetov10@gmail.com>
Date: Thu, 28 May 2026 09:51:46 +0000
Subject: [PATCH] add agent.py

---
 agent.py | 53 +++++++++++++++++++++++++++++++----------------------
 1 file changed, 31 insertions(+), 22 deletions(-)

diff --git a/agent.py b/agent.py
index 7233f37..8cc406a 100644
--- a/agent.py
+++ b/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)