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

pyproject.toml

@@ -1,17 +1,11 @@
 [project]
 name = "rag-agent"
 version = "0.1.0"
-description = "AI agent with local RAG memory using Qdrant and Ollama"
+description = "RAG agent with Qdrant and Ollama"
-authors = [{name = "Your Name", email = "you@example.com"}]
 requires-python = ">=3.10"
 dependencies = [
-    "langchain>=0.1.0",
+    "langchain>=0.2.0",
-    "langchain-qdrant>=0.1.0",
+    "langchain-qdrant",
-    "langchain-ollama>=0.1.0",
+    "langchain-ollama",
-    "qdrant-client>=1.0.0",
+    "qdrant-client"
-    "python-dotenv>=1.0.0"
+]
-]
-
-[build-system]
-requires = ["setuptools>=61.0"]
-build-backend = "setuptools.build_meta"

requirements.txt

@@ -1,5 +1,4 @@
-langchain>=0.1.0
+langchain>=0.2.0
-langchain-qdrant>=0.1.0
+langchain-qdrant
-langchain-ollama>=0.1.0
+langchain-ollama
-qdrant-client>=1.0.0
+qdrant-client
-python-dotenv>=1.0.0

src/agent.py

@@ -1,59 +1,35 @@
-from typing import List
+from langchain_ollama import ChatOllama
+from langchain.agents import initialize_agent, AgentType
+from langchain.tools import Tool
+from src.tools import search_knowledge_base, add_to_knowledge_base
 
-from langchain import LLMChain
+def create_agent():
-from langchain.chat_models import ChatOllama
-from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
-from langchain.agents import AgentExecutor, Tool
-
-from .tools import search_knowledge_base, add_to_knowledge_base
-
-
-def create_agent(
-    llm_model: str = "llama3",
-    tools: List[Tool] = None,
-    verbose: bool = True,
-) -> AgentExecutor:
     """
-    Create an AgentExecutor that uses the provided tools and a system prompt
+    Create a LangChain agent configured to use the knowledge base tools.
-    instructing the agent to use the knowledge base.
-
-    Parameters
-    ----------
-    llm_model : str
-        The Ollama model to use.
-    tools : List[Tool]
-        List of LangChain tools to expose to the agent.
-    verbose : bool
-        Whether to enable verbose output.
-
-    Returns
-    -------
-    AgentExecutor
-        Configured agent executor.
     """
-    if tools is None:
+    llm = ChatOllama(model="llama3")
-        tools = [search_knowledge_base, add_to_knowledge_base]
+    tools = [
-
+        Tool(
-    # System prompt instructing the agent to use the knowledge base
+            name="search_knowledge_base",
-    system_prompt = SystemMessagePromptTemplate.from_template(
+            func=search_knowledge_base,
-        """
+            description="Search the knowledge base for relevant information."
-You are an AI assistant that has access to a knowledge base. Use the provided tools to search the knowledge base or add new documents. When answering user queries, first decide if you need to search the knowledge base. If so, use the `search_knowledge_base` tool. If you need to add new information, use the `add_to_knowledge_base` tool. Always provide a concise answer after retrieving 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. "
+        "Use the provided tools to answer queries. If you need to add new information, call add_to_knowledge_base. "
+        "If you need to retrieve information, call search_knowledge_base. Provide concise answers."
     )
-
+    agent = initialize_agent(
-    human_prompt = HumanMessagePromptTemplate.from_template("{input}")
-
-    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])
-
-    llm = ChatOllama(model=llm_model)
-
-    llm_chain = LLMChain(llm=llm, prompt=chat_prompt)
-
-    agent = AgentExecutor.from_llm_and_tools(
-        llm=llm_chain,
         tools=tools,
-        verbose=verbose,
+        llm=llm,
-        agent="zero-shot-react-description",
+        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        verbose=True,
+        agent_kwargs={"system_message": system_prompt}
     )
-
     return agent

src/chunking.py

@@ -1,6 +1,16 @@
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.text_splitter 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_text(text: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> list[str]:
+    """
+    Split a document into chunks suitable for vector storage.
+    Each chunk is stored as a Document with metadata.
+    """
     splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    return splitter.split_text(text)
+    texts = splitter.split_text(content)
+    documents = []
+    for i, text in enumerate(texts):
+        metadata = {"source": title, "chunk_id": i}
+        documents.append(Document(page_content=text, metadata=metadata))
+    return documents

src/cli.py

@@ -1,50 +1,26 @@
 import argparse
-import os
+from src.loader import load_documents_from_directory
-import sys
+from src.tools import add_to_knowledge_base, search_knowledge_base
-from pathlib import Path
+from src.agent import create_agent
 
-from .vector_store import QdrantVectorStore
+def main():
-from .agent import create_agent
+    parser = argparse.ArgumentParser(description="RAG Agent CLI")
-from .tools import vector_store as global_vector_store
+    parser.add_argument("--init-dir", type=str, help="Directory to load documents from")
+    args = parser.parse_args()
 
+    if args.init_dir:
+        docs = load_documents_from_directory(args.init_dir)
+        for doc in docs:
+            add_to_knowledge_base(doc.page_content, doc.metadata.get("source", "unknown"))
+        print(f"Loaded {len(docs)} documents into the knowledge base.")
 
-def load_documents_from_directory(
+    agent = create_agent()
-    directory: Path, vector_store: QdrantVectorStore
+    print("Agent ready. Type your queries. Use /add <title> <content>, /search <query>, or /quit to exit.")
-) -> None:
-    """
-    Load all text files from the specified directory into the vector store.
-    """
-    if not directory.is_dir():
-        print(f"Directory {directory} does not exist.")
-        return
-
-    documents = []
-    for file_path in directory.rglob("*"):
-        if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}:
-            content = file_path.read_text(encoding="utf-8")
-            title = file_path.stem
-            documents.append({"content": content, "title": title})
-    if documents:
-        vector_store.add_documents(documents)
-        print(f"Loaded {len(documents)} documents into the knowledge base.")
-    else:
-        print("No documents found in the directory.")
-
-
-def interactive_loop(agent, max_results: int = 5):
-    """
-    Simple interactive CLI loop for the agent.
-    Commands:
-        /add <title> <content>   - Add a new document
-        /search <query>          - Search the knowledge base
-        /quit                    - Exit
-    """
-    print("Welcome to the RAG Agent CLI. Type /quit to exit.")
     while True:
         try:
             user_input = input(">> ").strip()
         except (EOFError, KeyboardInterrupt):
-            print("\nExiting.")
+            print("\nGoodbye!")
             break
 
         if not user_input:
@@ -54,56 +30,29 @@ def interactive_loop(agent, max_results: int = 5):
             print("Goodbye!")
             break
 
-        if user_input.startswith("/add"):
+        if user_input.lower().startswith("/add"):
-            parts = user_input.split(maxsplit=2)
+            parts = user_input.split(" ", 2)
             if len(parts) < 3:
                 print("Usage: /add <title> <content>")
                 continue
-            title, content = parts[1], parts[2]
+            _, title, content = parts
-            response = agent.run({"input": f"/add {title} {content}"})
+            response = add_to_knowledge_base(content, title)
             print(response)
             continue
 
-        if user_input.startswith("/search"):
+        if user_input.lower().startswith("/search"):
-            query = user_input[len("/search"):].strip()
+            parts = user_input.split(" ", 1)
-            if not query:
+            if len(parts) < 2:
                 print("Usage: /search <query>")
                 continue
-            response = agent.run({"input": f"/search {query}"})
+            _, query = parts
+            response = search_knowledge_base(query)
             print(response)
             continue
 
-        # Default: treat as normal query
+        # Treat as normal query for the agent
-        response = agent.run({"input": user_input})
+        response = agent.run(user_input)
         print(response)
 
-
-def main():
-    parser = argparse.ArgumentParser(description="RAG Agent CLI")
-    parser.add_argument(
-        "--load-dir",
-        type=str,
-        help="Directory containing documents to load into the knowledge base",
-    )
-    parser.add_argument(
-        "--max-results",
-        type=int,
-        default=5,
-        help="Maximum number of search results to return",
-    )
-    args = parser.parse_args()
-
-    vector_store = QdrantVectorStore()
-    # Override global vector store used by tools
-    global_vector_store.__dict__.update(vector_store.__dict__)
-
-    if args.load_dir:
-        load_documents_from_directory(Path(args.load_dir), vector_store)
-
-    agent = create_agent(verbose=True)
-
-    interactive_loop(agent, max_results=args.max_results)
-
-
 if __name__ == "__main__":
     main()

src/loader.py

@@ -1,35 +1,22 @@
 import os
-import uuid
+from pathlib import Path
-from typing import List, Dict
+from typing import List
+from langchain.schema import Document
 
-from .vector_store import ChromaVectorStore
+def load_documents_from_directory(directory: str) -> List[Document]:
-from .chunking import chunk_text
-from .config import CHROMA_DB_PATH
-
-store = ChromaVectorStore(CHROMA_DB_PATH)
-
-
-def load_documents_from_directory(directory_path: str):
     """
-    Load all .txt files from a directory into the vector store.
+    Recursively load all text-based files from a directory into Documents.
-
+    Supported extensions: .txt, .md, .py, .json, .csv
-    Parameters
-    ----------
-    directory_path : str
-        Path to the directory containing text files.
     """
-    for root, _, files in os.walk(directory_path):
+    docs = []
+    for root, _, files in os.walk(directory):
         for file in files:
-            if file.lower().endswith(".txt"):
+            if file.lower().endswith(('.txt', '.md', '.py', '.json', '.csv')):
-                file_path = os.path.join(root, file)
+                path = Path(root) / file
                 try:
-                    with open(file_path, "r", encoding="utf-8") as f:
+                    with open(path, 'r', encoding='utf-8') as f:
                         content = f.read()
-                    title = file
+                    docs.append(Document(page_content=content, metadata={"source": str(path)}))
-                    chunks = chunk_text(content)
-                    ids = [str(uuid.uuid4()) for _ in chunks]
-                    metadatas = [{"title": title} for _ in chunks]
-                    store.add_documents(chunks, metadatas, ids)
-                    print(f"Loaded {len(chunks)} chunks from {file_path}")
                 except Exception as e:
-                    print(f"Failed to load {file_path}: {e}")
+                    print(f"Failed to read {path}: {e}")
+    return docs

src/main.py

@@ -1,3 +1,4 @@
-# This file is intentionally left empty.
+from src.cli import main
-# The application entry point is defined in src/cli.py.
+
-# Importing src.cli.main will start the CLI when executed as a script.
+if __name__ == "__main__":
+    main()

src/tools.py

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

src/vector_store.py

@@ -1,104 +1,51 @@
-import os
+from typing import List
-from typing import List, Dict, Any
-
 from langchain_ollama import OllamaEmbeddings
-from langchain_qdrant import Qdrant
+from langchain_qdrant import QdrantVectorStore
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.schema import Document
 from qdrant_client import QdrantClient
 from qdrant_client.http import models as qdrant_models
 
-
+class QdrantStore:
-class QdrantVectorStore:
     """
-    Wrapper around Qdrant vector store with Ollama embeddings.
+    Wrapper around QdrantVectorStore that handles collection creation,
+    embedding generation via Ollama, and semantic search.
     """
-
+    def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333):
-    def __init__(
-        self,
-        host: str = "localhost",
-        port: int = 6333,
-        collection_name: str = "rag_collection",
-        chunk_size: int = 1000,
-        chunk_overlap: int = 200,
-    ):
         self.client = QdrantClient(host=host, port=port)
         self.collection_name = collection_name
-        self.chunk_size = chunk_size
+        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
-        self.chunk_overlap = chunk_overlap
+
+        # Determine vector size from a sample embedding
+        sample_vector = self.embeddings.embed_query("")
+        vector_size = len(sample_vector)
 
         # Create collection if it does not exist
         if not self.client.has_collection(collection_name):
-            self.client.create_collection(
+            self.client.recreate_collection(
                 collection_name=collection_name,
                 vectors_config=qdrant_models.VectorParams(
-                    size=384,  # size of nomic-embed-text embeddings
+                    size=vector_size,
-                    distance=qdrant_models.Distance.COSINE,
+                    distance="Cosine"
-                ),
+                )
             )
 
-        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
+        self.store = QdrantVectorStore(
-        self.text_splitter = RecursiveCharacterTextSplitter(
+            client=self.client,
-            chunk_size=self.chunk_size,
+            collection_name=collection_name,
-            chunk_overlap=self.chunk_overlap,
+            embeddings=self.embeddings
         )
 
-    def add_documents(self, documents: List[Dict[str, str]]) -> None:
+    def add_documents(self, documents: List[Document]):
         """
-        Add documents to the vector store.
+        Add a list of Documents to the vector store.
+        """
+        self.store.add_documents(documents)
 
-        Each document dict must contain 'content' and 'title'.
+    def semantic_search(self, query: str, k: int = 5):
         """
-        for doc in documents:
+        Perform a semantic similarity search and return top-k Documents.
-            content = doc.get("content", "")
+        """
-            title = doc.get("title", "Untitled")
+        return self.store.similarity_search(query, k=k)
-            # Split into chunks
-            chunks = self.text_splitter.split_text(content)
-            # Embed each chunk
-            embeddings = self.embeddings.embed_documents(chunks)
-            # Prepare payloads
-            payloads = [
-                {
-                    "title": title,
-                    "chunk_index": idx,
-                    "content": chunk,
-                }
-                for idx, chunk in enumerate(chunks)
-            ]
-            # Upsert into Qdrant
-            self.client.upsert(
-                collection_name=self.collection_name,
-                points=[
-                    qdrant_models.PointStruct(
-                        id=None,
-                        vector=emb,
-                        payload=payload,
-                    )
-                    for emb, payload in zip(embeddings, payloads)
-                ],
-            )
 
-    def semantic_search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
+# Global instance used by tools and agent
-        """
+vector_store = QdrantStore()
-        Perform semantic search in the vector store.
-
-        Returns a list of dicts with 'title', 'content', and 'score'.
-        """
-        query_embedding = self.embeddings.embed_query(query)
-        search_result = self.client.search(
-            collection_name=self.collection_name,
-            query_vector=query_embedding,
-            limit=max_results,
-            with_payload=True,
-            score_threshold=0.0,
-        )
-        results = []
-        for point in search_result:
-            payload = point.payload
-            results.append(
-                {
-                    "title": payload.get("title", "Untitled"),
-                    "content": payload.get("content", ""),
-                    "score": point.score,
-                }
-            )
-        return results