From af004ec8d71f928cc7fd53179c5f898d2b7f128a 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=A0=D0=BE=D0=BC?=
 =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= <kirilinka320@gmail.com>
Date: Tue, 2 Jun 2026 07:06:57 +0000
Subject: [PATCH] Add agent.py

---
 agent.py | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)
 create mode 100644 agent.py

diff --git a/agent.py b/agent.py
new file mode 100644
index 0000000..73ba556
--- /dev/null
+++ b/agent.py
@@ -0,0 +1,77 @@
+"""
+Main entry point for the RAG agent.
+"""
+
+import asyncio
+import os
+from pathlib import Path
+
+from langchain_ollama import ChatOllama
+from langchain.agents import create_agent
+from langchain.agents.agent_toolkits import BaseToolkit
+from langchain.agents.agent_types import AgentType
+from langchain.tools import BaseTool
+
+from rag_tools import search_knowledge_base, add_to_knowledge_base
+from qdrant_store import load_directory
+
+# Load environment variables if any
+from dotenv import load_dotenv
+load_dotenv()
+
+# Configuration
+LLM_MODEL = "llama3"
+KNOWLEDGE_DIR = os.getenv("KNOWLEDGE_DIR", "./knowledge")
+
+# Ensure knowledge directory exists and load documents
+Path(KNOWLEDGE_DIR).mkdir(parents=True, exist_ok=True)
+load_directory(KNOWLEDGE_DIR)
+
+# Define tools
+class SearchTool(BaseTool):
+    name = "search_knowledge_base"
+    description = "Perform semantic search in the knowledge base."
+    func = search_knowledge_base
+
+class AddTool(BaseTool):
+    name = "add_to_knowledge_base"
+    description = "Add a new document to the knowledge base."
+    func = add_to_knowledge_base
+
+# Simple toolkit
+class RAGToolkit(BaseToolkit):
+    def get_tools(self):
+        return [SearchTool(), AddTool()]
+
+    def get_base_prompt(self):
+        return None
+
+# Create LLM
+llm = ChatOllama(model=LLM_MODEL)
+
+# System prompt instructing the agent to use the knowledge base
+SYSTEM_PROMPT = """
+You are an assistant that uses a knowledge base. When answering user queries, first search the knowledge base with the search_knowledge_base tool. If the information is not sufficient, ask the user for clarification. You can also add new documents to the knowledge base using add_to_knowledge_base.
+"""
+
+# Create agent
+agent = create_agent(
+    llm=llm,
+    toolkit=RAGToolkit(),
+    system_prompt=SYSTEM_PROMPT,
+    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+    verbose=True,
+)
+
+async def main():
+    print("RAG Agent ready. Type your query (or 'quit' to exit).")
+    while True:
+        user_input = input("\n> ")
+        if user_input.lower() in {"quit", "exit", "q"}:
+            print("Goodbye!")
+            break
+        response = await agent.ainvoke(user_input)
+        print("\nAssistant:", response)
+
+if __name__ == "__main__":
+    asyncio.run(main())