From 223d4788e35d6db97d6c3f41e9063ed2c2452c20 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=90=D0=B4=D0=B5=D0=BB=D0=B8=D0=BD=D0=B0=20=D0=A1=D0=B0?=
 =?UTF-8?q?=D1=82=D1=82=D0=B0=D1=80=D0=BE=D0=B2=D0=B0?=
 <sattarova2509@gmail.com>
Date: Thu, 28 May 2026 13:26:19 +0000
Subject: [PATCH] Update rag_agent

---
 rag_agent.py | 42 +++++++++++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/rag_agent.py b/rag_agent.py
index 956fb95..ca38488 100644
--- a/rag_agent.py
+++ b/rag_agent.py
@@ -1,16 +1,40 @@
 from langchain.agents import create_openai_functions_agent, AgentExecutor
 from rag_tools import search_knowledge_base, add_to_knowledge_base
-from langchain.schema import HumanMessage
+from langchain_ollama import Ollama
+
+# LLM for agent
+llm = Ollama(model="llama3")
+
+# Create agent with tools and llm
+agent = create_openai_functions_agent(tools=[search_knowledge_base, add_to_knowledge_base], llm=llm)
+executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
 
 def run_agent():
-    tools=[search_knowledge_base,add_to_knowledge_base]
-    agent=create_openai_functions_agent(tools=tools, llm=None)
-    executor=AgentExecutor(agent=agent, tools=tools, verbose=True)
+    print("RAG agent ready. Commands: /add <title> <content>, /search <query>, /quit")
     while True:
-        inp=input("> ")
-        if inp.strip().lower()=="/quit":
+        inp = input("> ")
+        if inp.strip().lower() == "/quit":
             break
-        response=executor.invoke({"input":inp})
-        print(response["output"])
-if __name__=="__main__":
+        if inp.startswith("/add"):
+            parts = inp.split(maxsplit=2)
+            if len(parts) < 3:
+                print("Usage: /add title content")
+                continue
+            _, title, content = parts
+            res = executor.invoke({"input": f"Add document '{title}'"})
+            # directly call tool
+            add_to_knowledge_base(content=content, title=title)
+            print(f"Added {title}")
+        elif inp.startswith("/search"):
+            query = inp[len("/search"):].strip()
+            if not query:
+                print("Usage: /search query")
+                continue
+            res = executor.invoke({"input": f"Search for '{query}'"})
+            print(res["output"])
+        else:
+            res = executor.invoke({"input": inp})
+            print(res["output"])
+
+if __name__ == "__main__":
     run_agent()