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

2026-05-28 16:32:09 +03:00
commit 5d01ac1d8b
8 changed files with 370 additions and 0 deletions
@@ -0,0 +1,5 @@
+node_modules/
+.env
+dist/
+build/
+*.log
@@ -0,0 +1,93 @@
+```markdown
+# RAG Agent with Qdrant and Ollama
+
+This project implements an AI agent that can search and add documents to a local knowledge base using **Qdrant** for vector storage and **Ollama** for embeddings and LLM inference. The agent is built with **LangChain** and exposes two tools:
+
+- `search_knowledge_base(query, max_results)` – semantic search in the knowledge base.
+- `add_to_knowledge_base(content, title)` – add a new document to the knowledge base.
+
+## Features
+
+- **Vector store**: Qdrant with Ollama embeddings (`nomic-embed-text`).
+- **Chunking**: Recursive character splitter with overlap.
+- **Agent**: Zero-shot React agent that uses the two tools.
+- **CLI**: Interactive command line interface to add documents and query the agent.
+- **Batch loading**: Script to load all text files from a directory into the knowledge base.
+
+## Prerequisites
+
+- Python 3.10+
+- Docker (for Qdrant) or a running Qdrant instance.
+- Ollama installed locally with the following models:
+  ```bash
+  ollama pull llama3
+  ollama pull nomic-embed-text
+  ```
+
+## Setup
+
+```bash
+# Clone the repository
+git clone https://github.com/your-username/rag-agent.git
+cd rag-agent
+
+# Create a virtual environment
+python -m venv venv
+source venv/bin/activate  # On Windows: venv\Scripts\activate
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Start Qdrant (Docker example)
+docker run -p 6333:6333 qdrant/qdrant
+```
+
+## Usage
+
+### 1. Load documents into the knowledge base
+
+```bash
+python src/main.py /path/to/documents
+```
+
+Supported file types: `.txt`, `.md`. (PDF support can be added with an additional parser.)
+
+### 2. Start the interactive CLI
+
+```bash
+python src/cli.py
+```
+
+Commands:
+
+- `/add <file_path>` – Add a single document.
+- `/search <query>` – Query the agent.
+- `/quit` – Exit.
+
+### 3. Example
+
+```bash
+> /add example.txt
+Document 'example' added to knowledge base with 3 chunks.
+> /search What is the capital of France?
+1. The capital of France is Paris. (Title: example)
+```
+
+## Project Structure
+
+```
+rag-agent/
+├── src/
+│   ├── agent.py
+│   ├── cli.py
+│   ├── main.py
+│   ├── tools.py
+│   └── vector_store.py
+├── requirements.txt
+└── README.md
+```
+
+## License
+
+MIT License
+```
@@ -0,0 +1,7 @@
+```
+langchain
+langchain-qdrant
+langchain-ollama
+qdrant-client
+python-dotenv
+```
@@ -0,0 +1,40 @@
+```python
+"""
+Agent creation with RAG integration.
+"""
+
+from langchain.llms import Ollama
+from langchain.agents import initialize_agent, AgentType
+from langchain.tools import Tool
+
+from tools import search_knowledge_base, add_to_knowledge_base
+
+
+def create_agent():
+    """
+    Create and configure the LangChain agent.
+
+    Returns:
+        AgentExecutor instance ready to run queries.
+    """
+    llm = Ollama(model="llama3")
+
+    tools = [
+        Tool.from_function(search_knowledge_base),
+        Tool.from_function(add_to_knowledge_base),
+    ]
+
+    agent = initialize_agent(
+        tools,
+        llm,
+        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+        verbose=True,
+        agent_kwargs={
+            "system_message": (
+                "You are an AI assistant that can search and add documents to a knowledge base. "
+                "Use the provided tools to answer user queries."
+            )
+        },
+    )
+    return agent
+```
@@ -0,0 +1,62 @@
+```python
+"""
+Interactive CLI for the RAG agent.
+"""
+
+import argparse
+from pathlib import Path
+
+from tools import add_to_knowledge_base
+from agent import create_agent
+
+
+def main():
+    agent = create_agent()
+    print("RAG Agent CLI. Commands:")
+    print("  /add <file_path>   - Add a document to the knowledge base.")
+    print("  /search <query>    - Search the knowledge base.")
+    print("  /quit              - Exit the program.")
+
+    while True:
+        try:
+            inp = input("> ").strip()
+        except EOFError:
+            break
+
+        if not inp:
+            continue
+
+        if inp.startswith("/add"):
+            parts = inp.split(maxsplit=1)
+            if len(parts) < 2:
+                print("Usage: /add <file_path>")
+                continue
+            file_path = Path(parts[1])
+            if not file_path.is_file():
+                print(f"File {file_path} does not exist.")
+                continue
+            content = file_path.read_text(encoding="utf-8")
+            title = file_path.stem
+            result = add_to_knowledge_base(content, title)
+            print(result)
+
+        elif inp.startswith("/search"):
+            parts = inp.split(maxsplit=1)
+            if len(parts) < 2:
+                print("Usage: /search <query>")
+                continue
+            query = parts[1]
+            response = agent.run(query)
+            print(response)
+
+        elif inp.startswith("/quit"):
+            print("Goodbye!")
+            break
+
+        else:
+            print("Unknown command. Use /add, /search, /quit.")
+
+
+if __name__ == "__main__":
+    main()
+```
@@ -0,0 +1,42 @@
+```python
+"""
+Script to load documents from a directory into the vector store.
+"""
+
+import argparse
+from pathlib import Path
+
+from vector_store import vector_store
+from langchain_text_splitter import RecursiveCharacterTextSplitter
+
+# Chunking configuration
+splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+
+def load_documents_from_dir(directory: str):
+    """
+    Load all supported text files from the given directory into the knowledge base.
+
+    Args:
+        directory: Path to the directory containing documents.
+    """
+    dir_path = Path(directory)
+    for file_path in dir_path.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
+            chunks = splitter.split_text(content)
+            vector_store.add_documents(chunks, [title] * len(chunks))
+            print(f"Loaded {file_path} into knowledge base.")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Load documents into the knowledge base.")
+    parser.add_argument("directory", help="Path to directory with documents.")
+    args = parser.parse_args()
+    load_documents_from_dir(args.directory)
+
+
+if __name__ == "__main__":
+    main()
+```
@@ -0,0 +1,55 @@
+```python
+"""
+Tools for the RAG agent: searching and adding to the knowledge base.
+"""
+
+from langchain.tools import tool
+from langchain_text_splitter import RecursiveCharacterTextSplitter
+
+from vector_store import vector_store
+
+# Chunking configuration
+splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+
+@tool
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """
+    Semantic search in the knowledge base.
+
+    Args:
+        query: Search query.
+        max_results: Maximum number of results to return.
+
+    Returns:
+        Formatted string with search results.
+    """
+    results = vector_store.search(query, max_results)
+    if not results:
+        return "No relevant documents found."
+
+    formatted = []
+    for i, doc in enumerate(results):
+        snippet = doc.page_content[:200].replace("\n", " ")
+        formatted.append(
+            f"{i + 1}. {snippet} (Title: {doc.metadata.get('title', 'N/A')})"
+        )
+    return "\n".join(formatted)
+
+
+@tool
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """
+    Add a document to the knowledge base.
+
+    Args:
+        content: Full text of the document.
+        title: Title of the document.
+
+    Returns:
+        Confirmation message.
+    """
+    chunks = splitter.split_text(content)
+    vector_store.add_documents(chunks, [title] * len(chunks))
+    return f"Document '{title}' added to knowledge base with {len(chunks)} chunks."
+```
@@ -0,0 +1,66 @@
+```python
+"""
+Vector store implementation using Qdrant and Ollama embeddings.
+"""
+
+from typing import List
+
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+
+
+class QdrantVectorStoreWrapper:
+    """
+    Wrapper around LangChain's QdrantVectorStore.
+    Handles initialization, document addition, and similarity search.
+    """
+
+    def __init__(
+        self,
+        collection_name: str = "knowledge_base",
+        host: str = "localhost",
+        port: int = 6333,
+    ):
+        """
+        Initialize the vector store.
+
+        Args:
+            collection_name: Name of the Qdrant collection.
+            host: Qdrant host address.
+            port: Qdrant port.
+        """
+        self.embedding = OllamaEmbeddings(model="nomic-embed-text")
+        self.store = QdrantVectorStore(
+            url=f"http://{host}:{port}",
+            collection_name=collection_name,
+            embedding=self.embedding,
+        )
+
+    def add_documents(self, documents: List[str], titles: List[str]) -> None:
+        """
+        Add documents to the vector store with metadata.
+
+        Args:
+            documents: List of document texts.
+            titles: List of titles corresponding to each document.
+        """
+        metadatas = [{"title": title} for title in titles]
+        self.store.add_texts(documents, metadatas=metadatas)
+
+    def search(self, query: str, k: int = 5):
+        """
+        Perform a similarity search.
+
+        Args:
+            query: Query string.
+            k: Number of results to return.
+
+        Returns:
+            List of Document objects sorted by relevance.
+        """
+        return self.store.similarity_search(query, k)
+
+
+# Global instance used by tools and agent
+vector_store = QdrantVectorStoreWrapper()
+```