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

pyproject.toml

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

requirements.txt

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

src/agent.py

@@ -1,45 +1,59 @@
-from langchain_ollama import Ollama
+from typing import List
-from langchain.agents import initialize_agent, AgentType
+
-from langchain.tools import Tool
+from langchain import LLMChain
+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
-from .config import LLM_MODEL
 
 
-def create_agent():
+def create_agent(
+    llm_model: str = "llama3",
+    tools: List[Tool] = None,
+    verbose: bool = True,
+) -> AgentExecutor:
     """
-    Create an RAG-enabled agent that can search and add to a knowledge base.
+    Create an AgentExecutor that uses the provided tools and a system prompt
+    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
-        The configured agent.
+        Configured agent executor.
     """
-    llm = Ollama(model=LLM_MODEL)
+    if tools is None:
+        tools = [search_knowledge_base, add_to_knowledge_base]
 
-    tools = [
+    # System prompt instructing the agent to use the knowledge base
-        Tool(
+    system_prompt = SystemMessagePromptTemplate.from_template(
-            name="search_knowledge_base",
+        """
-            func=search_knowledge_base,
+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.
-            description="Search the knowledge base for relevant documents."
+"""
-        ),
-        Tool(
-            name="add_to_knowledge_base",
-            func=add_to_knowledge_base,
-            description="Add a new document to the knowledge base."
-        ),
-    ]
-
-    system_prompt = (
-        "You are an AI assistant that can search and add information to a knowledge base. "
-        "Use the provided tools to answer user queries."
     )
 
-    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,
-        llm=llm,
+        verbose=verbose,
-        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
+        agent="zero-shot-react-description",
-        verbose=True,
-        system_message=system_prompt,
     )
+
     return agent

src/cli.py

@@ -1,66 +1,108 @@
+import argparse
+import os
 import sys
 from pathlib import Path
 
+from .vector_store import QdrantVectorStore
 from .agent import create_agent
-from .tools import add_to_knowledge_base, search_knowledge_base
+from .tools import vector_store as global_vector_store
 
 
-def main():
+def load_documents_from_directory(
-    agent = create_agent()
+    directory: Path, vector_store: QdrantVectorStore
-    print("Welcome to the RAG Agent CLI.")
+) -> None:
-    print("Commands:")
+    """
-    print("  /add <file_path>   - Add a text file to the knowledge base.")
+    Load all text files from the specified directory into the vector store.
-    print("  /search <query>    - Search the knowledge base.")
+    """
-    print("  /quit              - Exit the program.")
+    if not directory.is_dir():
-    print("Any other input will be sent to the agent for general processing.\n")
+        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:
+        except (EOFError, KeyboardInterrupt):
+            print("\nExiting.")
             break
 
         if not user_input:
             continue
 
-        if user_input.startswith("/add"):
+        if user_input.lower() == "/quit":
-            parts = user_input.split(maxsplit=1)
-            if len(parts) < 2:
-                print("Usage: /add <file_path>")
-                continue
-            file_path = parts[1]
-            path_obj = Path(file_path)
-            if not path_obj.is_file():
-                print(f"File not found: {file_path}")
-                continue
-            try:
-                content = path_obj.read_text(encoding="utf-8")
-                title = path_obj.name
-                result = add_to_knowledge_base(content, title)
-                print(result)
-            except Exception as e:
-                print(f"Error reading file: {e}")
-
-        elif user_input.startswith("/search"):
-            parts = user_input.split(maxsplit=1)
-            if len(parts) < 2:
-                print("Usage: /search <query>")
-                continue
-            query = parts[1]
-            result = search_knowledge_base(query, max_results=5)
-            print(result)
-
-        elif user_input.startswith("/quit"):
             print("Goodbye!")
             break
 
-        else:
+        if user_input.startswith("/add"):
-            # General query to the agent
+            parts = user_input.split(maxsplit=2)
-            try:
+            if len(parts) < 3:
-                response = agent.run(user_input)
+                print("Usage: /add <title> <content>")
-                print(response)
+                continue
-            except Exception as e:
+            title, content = parts[1], parts[2]
-                print(f"Agent error: {e}")
+            response = agent.run({"input": f"/add {title} {content}"})
+            print(response)
+            continue
+
+        if user_input.startswith("/search"):
+            query = user_input[len("/search"):].strip()
+            if not query:
+                print("Usage: /search <query>")
+                continue
+            response = agent.run({"input": f"/search {query}"})
+            print(response)
+            continue
+
+        # Default: treat as normal query
+        response = agent.run({"input": 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__":

src/main.py

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

src/tools.py

@@ -1,47 +1,34 @@
-import uuid
+from typing import List, Dict, Any
-from typing import List, Dict
 
 from langchain.tools import tool
-from .vector_store import ChromaVectorStore
-from .chunking import chunk_text
-from .config import CHROMA_DB_PATH
 
-# Initialize a single vector store instance
+from .vector_store import QdrantVectorStore
-store = ChromaVectorStore(CHROMA_DB_PATH)
+
+# Instantiate a global vector store (will be overridden in main)
+vector_store = QdrantVectorStore()
 
 
-@tool("search_knowledge_base")
+@tool
-def search_knowledge_base(query: str, max_results: int = 5) -> str:
+def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
     """
-    Search the knowledge base for the most relevant documents.
+    Search the knowledge base for the given query.
 
     Parameters
     ----------
     query : str
         The search query.
     max_results : int, optional
-        Number of top results to return (default is 5).
+        Maximum number of results to return.
 
     Returns
     -------
-    str
+    List[Dict[str, Any]]
-        Formatted search results.
+        List of search results with title, content, and score.
     """
-    results = store.search(query, limit=max_results)
+    return vector_store.semantic_search(query, max_results)
-    docs = results["documents"][0]
-    distances = results["distances"][0]
-    metadatas = results["metadatas"][0]
-    output = []
-    for doc, dist, meta in zip(docs, distances, metadatas):
-        output.append(
-            f"Title: {meta.get('title', 'N/A')}\n"
-            f"Distance: {dist:.4f}\n"
-            f"Content: {doc}\n"
-        )
-    return "\n".join(output)
 
 
-@tool("add_to_knowledge_base")
+@tool
 def add_to_knowledge_base(content: str, title: str) -> str:
     """
     Add a new document to the knowledge base.
@@ -51,15 +38,12 @@ def add_to_knowledge_base(content: str, title: str) -> str:
     content : str
         The full text content of the document.
     title : str
-        A title or identifier for the document.
+        The title of the document.
 
     Returns
     -------
     str
         Confirmation message.
     """
-    chunks = chunk_text(content)
+    vector_store.add_documents([{"content": content, "title": title}])
-    ids = [str(uuid.uuid4()) for _ in chunks]
+    return f"Document '{title}' added to the knowledge base."
-    metadatas = [{"title": title} for _ in chunks]
-    store.add_documents(chunks, metadatas, ids)
-    return f"Added {len(chunks)} chunks to the knowledge base."

src/vector_store.py

@@ -1,34 +1,104 @@
-import uuid
+import os
-from typing import List, Dict
+from typing import List, Dict, Any
 
-from chromadb import Client
-from chromadb.config import Settings
 from langchain_ollama import OllamaEmbeddings
-
+from langchain_qdrant import Qdrant
-from .config import CHROMA_DB_PATH, EMBEDDING_MODEL
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from qdrant_client import QdrantClient
+from qdrant_client.http import models as qdrant_models
 
 
-class ChromaVectorStore:
+class QdrantVectorStore:
-    def __init__(self, db_path: str = CHROMA_DB_PATH):
+    """
-        self.client = Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=db_path))
+    Wrapper around Qdrant vector store with Ollama embeddings.
-        self.collection_name = "knowledge_base"
+    """
-        self.collection = self.client.get_or_create_collection(name=self.collection_name)
-        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
 
-    def add_documents(self, documents: List[str], metadatas: List[Dict], ids: List[str]):
+    def __init__(
-        embeddings = self.embeddings.embed_documents(documents)
+        self,
-        self.collection.add(
+        host: str = "localhost",
-            documents=documents,
+        port: int = 6333,
-            embeddings=embeddings,
+        collection_name: str = "rag_collection",
-            metadatas=metadatas,
+        chunk_size: int = 1000,
-            ids=ids
+        chunk_overlap: int = 200,
+    ):
+        self.client = QdrantClient(host=host, port=port)
+        self.collection_name = collection_name
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+        # Create collection if it does not exist
+        if not self.client.has_collection(collection_name):
+            self.client.create_collection(
+                collection_name=collection_name,
+                vectors_config=qdrant_models.VectorParams(
+                    size=384,  # size of nomic-embed-text embeddings
+                    distance=qdrant_models.Distance.COSINE,
+                ),
+            )
+
+        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.chunk_size,
+            chunk_overlap=self.chunk_overlap,
         )
 
-    def search(self, query: str, limit: int = 5):
+    def add_documents(self, documents: List[Dict[str, str]]) -> None:
+        """
+        Add documents to the vector store.
+
+        Each document dict must contain 'content' and 'title'.
+        """
+        for doc in documents:
+            content = doc.get("content", "")
+            title = doc.get("title", "Untitled")
+            # 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]]:
+        """
+        Perform semantic search in the vector store.
+
+        Returns a list of dicts with 'title', 'content', and 'score'.
+        """
         query_embedding = self.embeddings.embed_query(query)
-        results = self.collection.query(
+        search_result = self.client.search(
-            query_embeddings=[query_embedding],
+            collection_name=self.collection_name,
-            n_results=limit,
+            query_vector=query_embedding,
-            include=["documents", "distances", "metadatas"]
+            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