commit 5d01ac1d8bea2e83b400967ce334f418e958fcfc Author: kuzakhmetovartur Date: Thu May 28 16:32:09 2026 +0300 feat: solution for 'Практическое задание: Агент с RAG-памятью' diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b16538b --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +node_modules/ +.env +dist/ +build/ +*.log diff --git a/README.md b/README.md new file mode 100644 index 0000000..83c3f88 --- /dev/null +++ b/README.md @@ -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 ` – Add a single document. +- `/search ` – 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 +``` \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..52cd0b3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +``` +langchain +langchain-qdrant +langchain-ollama +qdrant-client +python-dotenv +``` \ No newline at end of file diff --git a/src/agent.py b/src/agent.py new file mode 100644 index 0000000..e9d7259 --- /dev/null +++ b/src/agent.py @@ -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 +``` \ No newline at end of file diff --git a/src/cli.py b/src/cli.py new file mode 100644 index 0000000..4ad4e8f --- /dev/null +++ b/src/cli.py @@ -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 - Add a document to the knowledge base.") + print(" /search - 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 ") + 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 ") + 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() +``` \ No newline at end of file diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..d20c115 --- /dev/null +++ b/src/main.py @@ -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() +``` \ No newline at end of file diff --git a/src/tools.py b/src/tools.py new file mode 100644 index 0000000..49a6297 --- /dev/null +++ b/src/tools.py @@ -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." +``` \ No newline at end of file diff --git a/src/vector_store.py b/src/vector_store.py new file mode 100644 index 0000000..f3ee1bf --- /dev/null +++ b/src/vector_store.py @@ -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() +``` \ No newline at end of file