feat: solution for 'Практическое задание: Агент с RAG-памятью'

requirements.txt

@@ -1,4 +1,7 @@
-langchain>=0.2.0
+langchain==0.2.0
-langchain-qdrant
+langchain-ollama==0.2.0
-langchain-ollama
+langchain-qdrant==0.2.0
-qdrant-client
+qdrant-client==1.7.0
+python-dotenv==1.0.0
+tqdm==4.66.1
+pydantic==2.7.0

src/agent.py

@@ -1,35 +1,31 @@
-from langchain_ollama import ChatOllama
+from langchain_ollama import Ollama
-from langchain.agents import initialize_agent, AgentType
+from langchain.agents import initialize_agent, AgentExecutor, AgentType
-from langchain.tools import Tool
+from langchain.tools import BaseTool
-from src.tools import search_knowledge_base, add_to_knowledge_base
+from typing import List
 
-def create_agent():
+def create_agent(tools: List[BaseTool]) -> AgentExecutor:
     """
-    Create a LangChain agent configured to use the knowledge base tools.
+    Create a LangChain agent that can use the provided tools.
+
+    Parameters:
+        tools (List[BaseTool]): List of tools for the agent.
+
+    Returns:
+        AgentExecutor: Configured agent.
     """
-    llm = ChatOllama(model="llama3")
+    llm = Ollama(model="llama3")
-    tools = [
-        Tool(
-            name="search_knowledge_base",
-            func=search_knowledge_base,
-            description="Search the knowledge base for relevant information."
-        ),
-        Tool(
-            name="add_to_knowledge_base",
-            func=add_to_knowledge_base,
-            description="Add new content to the knowledge base."
-        )
-    ]
     system_prompt = (
-        "You are an AI assistant that helps users by searching and adding information to a knowledge base. "
+        "You are an AI assistant with access to a knowledge base. "
-        "Use the provided tools to answer queries. If you need to add new information, call add_to_knowledge_base. "
+        "Use the following tools to answer user queries:\n"
-        "If you need to retrieve information, call search_knowledge_base. Provide concise answers."
+        "- search_knowledge_base: Search the knowledge base.\n"
+        "- add_to_knowledge_base: Add a new document to the knowledge base.\n"
+        "When you need to use a tool, call it with the appropriate arguments."
     )
     agent = initialize_agent(
         tools=tools,
         llm=llm,
         agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
         verbose=True,
-        agent_kwargs={"system_message": system_prompt}
+        agent_kwargs={"system_message": system_prompt},
     )
     return agent

src/chunking.py

@@ -1,16 +1,20 @@
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.schema import Document
-from typing import List
 
-def chunk_document(content: str, title: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> List[Document]:
+def chunk_document(text: str, chunk_size: int = 1000, chunk_overlap: int = 200):
     """
-    Split a document into chunks suitable for vector storage.
+    Split a document into chunks using RecursiveCharacterTextSplitter.
-    Each chunk is stored as a Document with metadata.
+
+    Args:
+        text: The full text of the document.
+        chunk_size: Maximum size of each chunk in characters.
+        chunk_overlap: Number of characters to overlap between chunks.
+
+    Returns:
+        List of chunk strings.
     """
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    splitter = RecursiveCharacterTextSplitter(
-    texts = splitter.split_text(content)
+        chunk_size=chunk_size,
-    documents = []
+        chunk_overlap=chunk_overlap,
-    for i, text in enumerate(texts):
+        separators=["\n\n", "\n", " ", ""],
-        metadata = {"source": title, "chunk_id": i}
+    )
-        documents.append(Document(page_content=text, metadata=metadata))
+    return splitter.split_text(text)
-    return documents

src/cli.py

@@ -1,56 +1,91 @@
-import argparse
+import sys
-from src.loader import load_documents_from_directory
+from typing import List
-from src.tools import add_to_knowledge_base, search_knowledge_base
 from src.agent import create_agent
+from src.tools import search_knowledge_base, add_to_knowledge_base
+from src.loader import load_documents_from_directory
+from src.vector_store import QdrantVectorStore
+from src.tools import vector_store as global_vector_store
 
 def main():
-    parser = argparse.ArgumentParser(description="RAG Agent CLI")
+    # Initialize vector store
-    parser.add_argument("--init-dir", type=str, help="Directory to load documents from")
+    vector_store = QdrantVectorStore()
-    args = parser.parse_args()
+    # Override global reference for tools
+    global global_vector_store
+    global_vector_store = vector_store
 
-    if args.init_dir:
+    # Prepare tools
-        docs = load_documents_from_directory(args.init_dir)
+    tools = [search_knowledge_base, add_to_knowledge_base]
-        for doc in docs:
+
-            add_to_knowledge_base(doc.page_content, doc.metadata.get("source", "unknown"))
+    # Create agent
-        print(f"Loaded {len(docs)} documents into the knowledge base.")
+    agent = create_agent(tools)
+
+    print("Welcome to the RAG Agent CLI.")
+    print("Commands:")
+    print("  /load <directory>   - Load all .txt and .md files from directory into the knowledge base.")
+    print("  /add <title>        - Add a new document. You will be prompted for content.")
+    print("  /search <query>     - Search the knowledge base.")
+    print("  /quit               - Exit the program.")
+    print("")
 
-    agent = create_agent()
-    print("Agent ready. Type your queries. Use /add <title> <content>, /search <query>, or /quit to exit.")
     while True:
         try:
             user_input = input(">> ").strip()
         except (EOFError, KeyboardInterrupt):
-            print("\nGoodbye!")
+            print("\nExiting.")
             break
 
         if not user_input:
             continue
 
-        if user_input.lower() == "/quit":
+        if user_input.startswith("/"):
+            parts = user_input.split(maxsplit=1)
+            command = parts[0][1:].lower()
+            arg = parts[1] if len(parts) > 1 else ""
+
+            if command == "quit":
                 print("Goodbye!")
                 break
 
-        if user_input.lower().startswith("/add"):
+            elif command == "load":
-            parts = user_input.split(" ", 2)
+                if not arg:
-            if len(parts) < 3:
+                    print("Please provide a directory path.")
-                print("Usage: /add <title> <content>")
                     continue
-            _, title, content = parts
+                docs = load_documents_from_directory(arg)
-            response = add_to_knowledge_base(content, title)
+                if not docs:
-            print(response)
+                    print("No documents found.")
                     continue
+                vector_store.add_documents(docs)
+                print(f"Loaded {len(docs)} documents into the knowledge base.")
 
-        if user_input.lower().startswith("/search"):
+            elif command == "add":
-            parts = user_input.split(" ", 1)
+                title = arg.strip()
-            if len(parts) < 2:
+                if not title:
-                print("Usage: /search <query>")
+                    print("Please provide a title for the document.")
-                continue
-            _, query = parts
-            response = search_knowledge_base(query)
-            print(response)
                     continue
+                print("Enter document content. Finish with a single line containing only 'END'.")
+                lines = []
+                while True:
+                    line = input()
+                    if line.strip() == "END":
+                        break
+                    lines.append(line)
+                content = "\n".join(lines)
+                result = add_to_knowledge_base(content=content, title=title)
+                print(result)
 
-        # Treat as normal query for the agent
+            elif command == "search":
+                query = arg.strip()
+                if not query:
+                    print("Please provide a search query.")
+                    continue
+                result = search_knowledge_base(query=query, max_results=5)
+                print(result)
+
+            else:
+                print(f"Unknown command: {command}")
+
+        else:
+            # Treat as a natural language query for the agent
             response = agent.run(user_input)
             print(response)
 

src/loader.py

@@ -1,22 +1,21 @@
 import os
 from pathlib import Path
-from typing import List
+from typing import List, Dict
-from langchain.schema import Document
+import re
 
-def load_documents_from_directory(directory: str) -> List[Document]:
+def load_documents_from_directory(directory: str) -> List[Dict[str, str]]:
     """
-    Recursively load all text-based files from a directory into Documents.
+    Load all .txt and .md files from the given directory.
-    Supported extensions: .txt, .md, .py, .json, .csv
+
+    Returns a list of dicts with keys: title, content.
     """
     docs = []
-    for root, _, files in os.walk(directory):
+    for file_path in Path(directory).rglob("*"):
-        for file in files:
+        if file_path.suffix.lower() in {".txt", ".md"}:
-            if file.lower().endswith(('.txt', '.md', '.py', '.json', '.csv')):
-                path = Path(root) / file
             try:
-                    with open(path, 'r', encoding='utf-8') as f:
+                content = file_path.read_text(encoding="utf-8")
-                        content = f.read()
+                title = file_path.stem
-                    docs.append(Document(page_content=content, metadata={"source": str(path)}))
+                docs.append({"title": title, "content": content})
             except Exception as e:
-                    print(f"Failed to read {path}: {e}")
+                print(f"Failed to read {file_path}: {e}")
     return docs

src/main.py

@@ -1,4 +1,4 @@
-from src.cli import main
+# This file is intentionally left blank.
-
+# The CLI functionality is provided in src/cli.py.
-if __name__ == "__main__":
+# To run the application, execute:
-    main()
+#   python -m src.cli

src/tools.py

@@ -1,30 +1,52 @@
 from langchain.tools import tool
-from src.vector_store import vector_store
+from typing import List, Dict, Any
-from src.chunking import chunk_document
+from src.vector_store import QdrantVectorStore
-from typing import List
 
-@tool
+# Instantiate a global vector store (will be overridden in main)
+vector_store: QdrantVectorStore = None
+
+@tool("search_knowledge_base")
 def search_knowledge_base(query: str, max_results: int = 5) -> str:
     """
-    Search the knowledge base for relevant information.
+    Search the knowledge base for the given query and return top results.
-    Returns a formatted string of results.
+
+    Parameters:
+        query (str): The search query.
+        max_results (int): Number of results to return.
+
+    Returns:
+        str: Formatted search results.
     """
-    results = vector_store.semantic_search(query, k=max_results)
+    if vector_store is None:
+        raise ValueError("Vector store not initialized.")
+    results = vector_store.search(query, max_results)
     if not results:
         return "No relevant documents found."
     formatted = []
-    for i, doc in enumerate(results, 1):
+    for idx, res in enumerate(results, 1):
-        source = doc.metadata.get("source", "unknown")
+        formatted.append(
-        snippet = doc.page_content[:500].replace("\n", " ")
+            f"{idx}. Title: {res['title']}\n   Score: {res['score']:.4f}\n   Snippet: {res['content'][:200]}..."
-        formatted.append(f"{i}. Source: {source}\n{snippet}...")
+        )
     return "\n\n".join(formatted)
 
-@tool
+@tool("add_to_knowledge_base")
 def add_to_knowledge_base(content: str, title: str) -> str:
     """
-    Add new content to the knowledge base.
+    Add a new document to the knowledge base.
-    The content is split into chunks before being stored.
+
+    Parameters:
+        content (str): Full text of the document.
+        title (str): Title of the document.
+
+    Returns:
+        str: Confirmation message.
     """
-    docs = chunk_document(content, title)
+    if vector_store is None:
-    vector_store.add_documents(docs)
+        raise ValueError("Vector store not initialized.")
-    return f"Added {len(docs)} chunks from '{title}' to the knowledge base."
+    document = {
+        "content": content,
+        "title": title,
+        "metadata": {},
+    }
+    vector_store.add_documents([document])
+    return f"Document '{title}' added to the knowledge base."

src/vector_store.py

@@ -1,51 +1,94 @@
-from typing import List
+import os
-from langchain_ollama import OllamaEmbeddings
+from typing import List, Dict, Any
-from langchain_qdrant import QdrantVectorStore
-from langchain.schema import Document
 from qdrant_client import QdrantClient
 from qdrant_client.http import models as qdrant_models
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import Qdrant
+from src.chunking import chunk_document
 
-class QdrantStore:
+class QdrantVectorStore:
-    """
+    def __init__(
-    Wrapper around QdrantVectorStore that handles collection creation,
+        self,
-    embedding generation via Ollama, and semantic search.
+        collection_name: str = "knowledge_base",
-    """
+        host: str = "localhost",
-    def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333):
+        port: int = 6333,
-        self.client = QdrantClient(host=host, port=port)
+        embedding_model: str = "nomic-embed-text",
+    ):
         self.collection_name = collection_name
-        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
+        self.client = QdrantClient(host=host, port=port)
+        self.embedding_model = embedding_model
+        self.embeddings = OllamaEmbeddings(model=self.embedding_model)
 
-        # Determine vector size from a sample embedding
+        # Ensure collection exists
-        sample_vector = self.embeddings.embed_query("")
+        if not self.client.collection_exists(self.collection_name):
-        vector_size = len(sample_vector)
+            self.client.create_collection(
-
+                collection_name=self.collection_name,
-        # Create collection if it does not exist
-        if not self.client.has_collection(collection_name):
-            self.client.recreate_collection(
-                collection_name=collection_name,
                 vectors_config=qdrant_models.VectorParams(
-                    size=vector_size,
+                    size=self.embeddings.embedding_size,
-                    distance="Cosine"
+                    distance=qdrant_models.Distance.COSINE,
-                )
+                ),
             )
 
-        self.store = QdrantVectorStore(
+        self.qdrant = Qdrant(
             client=self.client,
-            collection_name=collection_name,
+            collection_name=self.collection_name,
-            embeddings=self.embeddings
+            embeddings=self.embeddings,
         )
 
-    def add_documents(self, documents: List[Document]):
+    def add_documents(self, documents: List[Dict[str, Any]]):
         """
-        Add a list of Documents to the vector store.
+        Add documents to the vector store. Each document dict should contain:
+            - content: full text
+            - title: document title
+            - metadata: optional dict
         """
-        self.store.add_documents(documents)
+        points = []
+        for doc in documents:
+            content = doc.get("content", "")
+            title = doc.get("title", "")
+            metadata = doc.get("metadata", {})
+            chunks = chunk_document(content)
+            for idx, chunk in enumerate(chunks):
+                point = qdrant_models.PointStruct(
+                    id=None,
+                    vector=self.embeddings.embed_query(chunk),
+                    payload={
+                        "title": title,
+                        "chunk_index": idx,
+                        "content": chunk,
+                        **metadata,
+                    },
+                )
+                points.append(point)
 
-    def semantic_search(self, query: str, k: int = 5):
+        if points:
-        """
+            self.client.upsert(
-        Perform a semantic similarity search and return top-k Documents.
+                collection_name=self.collection_name,
-        """
+                points=points,
-        return self.store.similarity_search(query, k=k)
+            )
 
-# Global instance used by tools and agent
+    def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
-vector_store = QdrantStore()
+        """
+        Perform a semantic search in the vector store.
+
+        Returns a list of dicts with keys: title, content, score.
+        """
+        query_vector = self.embeddings.embed_query(query)
+        search_result = self.client.search(
+            collection_name=self.collection_name,
+            query_vector=query_vector,
+            limit=max_results,
+            with_payload=True,
+            with_vector=False,
+        )
+        results = []
+        for hit in search_result:
+            payload = hit.payload
+            results.append(
+                {
+                    "title": payload.get("title", ""),
+                    "content": payload.get("content", ""),
+                    "score": hit.score,
+                }
+            )
+        return results