From 048522c4a434dcf9884c962bb12e66465e557f56 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 10:31:06 +0000
Subject: [PATCH] add agent.py

---
 agent.py | 101 ++++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 70 insertions(+), 31 deletions(-)

diff --git a/agent.py b/agent.py
index 8cc406a..758022a 100644
--- a/agent.py
+++ b/agent.py
@@ -1,40 +1,79 @@
 """
-Agent creation using LangChain create_agent.
+Agent definition for the RAG system.
+
+Provides two tools:
+* search_knowledge_base(query, max_results)
+* add_to_knowledge_base(content, title)
+
+The agent is created with create_agent from langchain.agents.
 """
+
 import os
-from langchain_ollama import OllamaLLM
-from langchain_core.messages import HumanMessage
-from langchain.agents import create_agent
+from typing import List
+from langchain_ollama import OllamaEmbeddings
+from chromadb import PersistentClient
+from chromadb.utils import embedding_functions as ef
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.tools import tool
-from tools import search_knowledge_base, add_to_knowledge_base
+from langchain.agents import create_agent
+from langchain_core.messages import HumanMessage
 
-# LLM via Ollama
-llm = OllamaLLM(model="llama3")
+# Initialize embeddings and Chroma client
+EMBEDDINGS = OllamaEmbeddings(model="nomic-embed-text")
+CHROMA_PATH = os.path.join(os.getcwd(), "chromadb_store")
+CLIENT = PersistentClient(path=CHROMA_PATH)
+COLLECTION_NAME = "knowledge"
+if COLLECTION_NAME not in CLIENT.list_collections():
+    CLIENT.create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)
+COLL = CLIENT.get_or_create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)
 
-# 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.
-"""
+# Chunker from chunker.py
+from chunker import CHUNKER
 
-def create_rag_agent():
-    agent = create_agent(
-        llm=llm,
-        tools=[search_knowledge_base, add_to_knowledge_base],
-        system_prompt=SYSTEM_PROMPT,
+@tool
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """
+    Add a document to the knowledge base.
+    The content is split into chunks and stored with metadata.
+    Returns confirmation message.
+    """
+    # Split content
+    chunks = CHUNKER.split_text(content)
+    ids = [f"{title}_{i}" for i in range(len(chunks))]
+    metadatas = [{"title": title} for _ in chunks]
+    COLL.add(ids=ids, documents=chunks, metadatas=metadatas)
+    return f"Added {len(chunks)} chunks from '{title}'."
+
+@tool
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """
+    Search the knowledge base for relevant documents.
+    Returns a formatted string of results.
+    """
+    results = COLL.query(
+        query_texts=[query],
+        n_results=max_results,
+        include=['documents', 'distances'],
     )
-    return agent
+    docs = results.get("documents", [])[0]
+    dists = results.get("distances", [])[0]
+    if not docs:
+        return "No relevant documents found."
+    output_lines = []
+    for i, (doc, dist) in enumerate(zip(docs, dists), 1):
+        output_lines.append(f"{i}. (score: {dist:.4f})\n{doc[:200]}...")
+    return "\n\n".join(output_lines)
 
-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)
+# Create agent
+SYSTEM_PROMPT = (
+    "You are an assistant that can search and add to a knowledge base. Use the provided tools."
+)
+AGENT = create_agent(
+    llm=None,  # No LLM needed for tool calls; agent will use system prompt only
+    tools=[add_to_knowledge_base, search_knowledge_base],
+    system_prompt=SYSTEM_PROMPT,
+)
+
+# Expose a simple invoke function
+async def run_agent(messages: List[HumanMessage]):
+    return await AGENT.ainvoke({"messages": messages}, {"configurable": {"thread_id": "rag-agent"}})