From 61d5771c9bf5d4cd751af6bd102d6848c0896bb8 Mon Sep 17 00:00:00 2001 From: kuzakhmetovartur Date: Thu, 28 May 2026 17:55:09 +0300 Subject: [PATCH] =?UTF-8?q?feat:=20solution=20for=20'=D0=9F=D1=80=D0=B0?= =?UTF-8?q?=D0=BA=D1=82=D0=B8=D1=87=D0=B5=D1=81=D0=BA=D0=BE=D0=B5=20=D0=B7?= =?UTF-8?q?=D0=B0=D0=B4=D0=B0=D0=BD=D0=B8=D0=B5:=20=D0=90=D0=B3=D0=B5?= =?UTF-8?q?=D0=BD=D1=82=20=D1=81=20RAG-=D0=BF=D0=B0=D0=BC=D1=8F=D1=82?= =?UTF-8?q?=D1=8C=D1=8E'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 109 +++++++++----------------------------------- requirements.txt | 7 +-- src/__init__.py | 1 + src/agent.py | 57 ++++++++++++----------- src/chunking.py | 6 +++ src/cli.py | 63 +++++++++++++------------ src/config.py | 3 ++ src/loader.py | 35 ++++++++++++++ src/main.py | 42 +---------------- src/tools.py | 82 ++++++++++++++++++--------------- src/vector_store.py | 86 +++++++++++----------------------- 11 files changed, 208 insertions(+), 283 deletions(-) create mode 100644 src/__init__.py create mode 100644 src/chunking.py create mode 100644 src/config.py create mode 100644 src/loader.py diff --git a/README.md b/README.md index 83c3f88..3a808dc 100644 --- a/README.md +++ b/README.md @@ -1,93 +1,26 @@ -```markdown -# RAG Agent with Qdrant and Ollama +# Практическое задание: Агент с RAG-памятью -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: +Главная +Мои задания +Агент с RAG-памятью +5Д +EN +Агент с RAG-памятью -- `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. +Практическое задание: Агент с RAG-памятью +Цель -## Features +Построить AI-агента с локальным RAG-хранилищем знаний на базе Qdrant и Ollama. Агент должен уметь искать и сохранять информацию в векторной базе. -- **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. +Стек +Python 3.10+ +Qdrant — векторная база данных +Ollama — локальные LLM и эмбеддинги (llama3, nomic-embed-text) +LangChain — фреймворк для агентов и RAG +Установка +# Ollama +ollama pull llama3 +ollama pull nomic-embed-text -## 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 +# Python пакеты +pi \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 52cd0b3..fc17789 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,4 @@ -``` langchain -langchain-qdrant langchain-ollama -qdrant-client -python-dotenv -``` \ No newline at end of file +langchain-text-splitters +chromadb \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..c6d0da1 --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +# Empty init file to make src a package \ No newline at end of file diff --git a/src/agent.py b/src/agent.py index e9d7259..f0493b9 100644 --- a/src/agent.py +++ b/src/agent.py @@ -1,40 +1,45 @@ -```python -""" -Agent creation with RAG integration. -""" - -from langchain.llms import Ollama +from langchain_ollama import Ollama from langchain.agents import initialize_agent, AgentType from langchain.tools import Tool -from tools import search_knowledge_base, add_to_knowledge_base +from .tools import search_knowledge_base, add_to_knowledge_base +from .config import LLM_MODEL def create_agent(): """ - Create and configure the LangChain agent. + Create an RAG-enabled agent that can search and add to a knowledge base. - Returns: - AgentExecutor instance ready to run queries. + Returns + ------- + AgentExecutor + The configured agent. """ - llm = Ollama(model="llama3") + llm = Ollama(model=LLM_MODEL) tools = [ - Tool.from_function(search_knowledge_base), - Tool.from_function(add_to_knowledge_base), + Tool( + name="search_knowledge_base", + func=search_knowledge_base, + 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." + ), ] - 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." - ) - }, + 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." ) - return agent -``` \ No newline at end of file + + agent = initialize_agent( + tools=tools, + llm=llm, + agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION, + verbose=True, + system_message=system_prompt, + ) + return agent \ No newline at end of file diff --git a/src/chunking.py b/src/chunking.py new file mode 100644 index 0000000..9e16375 --- /dev/null +++ b/src/chunking.py @@ -0,0 +1,6 @@ +from langchain_text_splitters import RecursiveCharacterTextSplitter + + +def chunk_text(text: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> list[str]: + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + return splitter.split_text(text) \ No newline at end of file diff --git a/src/cli.py b/src/cli.py index 4ad4e8f..db936a1 100644 --- a/src/cli.py +++ b/src/cli.py @@ -1,62 +1,67 @@ -```python -""" -Interactive CLI for the RAG agent. -""" - -import argparse +import sys from pathlib import Path -from tools import add_to_knowledge_base -from agent import create_agent +from .agent import create_agent +from .tools import add_to_knowledge_base, search_knowledge_base def main(): agent = create_agent() - print("RAG Agent CLI. Commands:") - print(" /add - Add a document to the knowledge base.") + print("Welcome to the RAG Agent CLI.") + print("Commands:") + print(" /add - Add a text file to the knowledge base.") print(" /search - Search the knowledge base.") print(" /quit - Exit the program.") + print("Any other input will be sent to the agent for general processing.\n") while True: try: - inp = input("> ").strip() + user_input = input(">> ").strip() except EOFError: break - if not inp: + if not user_input: continue - if inp.startswith("/add"): - parts = inp.split(maxsplit=1) + if user_input.startswith("/add"): + parts = user_input.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.") + file_path = parts[1] + path_obj = Path(file_path) + if not path_obj.is_file(): + print(f"File not found: {file_path}") continue - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - result = add_to_knowledge_base(content, title) - print(result) + 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 inp.startswith("/search"): - parts = inp.split(maxsplit=1) + elif user_input.startswith("/search"): + parts = user_input.split(maxsplit=1) if len(parts) < 2: print("Usage: /search ") continue query = parts[1] - response = agent.run(query) - print(response) + result = search_knowledge_base(query, max_results=5) + print(result) - elif inp.startswith("/quit"): + elif user_input.startswith("/quit"): print("Goodbye!") break else: - print("Unknown command. Use /add, /search, /quit.") + # General query to the agent + try: + response = agent.run(user_input) + print(response) + except Exception as e: + print(f"Agent error: {e}") if __name__ == "__main__": - main() -``` \ No newline at end of file + main() \ No newline at end of file diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..acb38d0 --- /dev/null +++ b/src/config.py @@ -0,0 +1,3 @@ +CHROMA_DB_PATH = "./chroma_db" +EMBEDDING_MODEL = "nomic-embed-text" +LLM_MODEL = "llama3" \ No newline at end of file diff --git a/src/loader.py b/src/loader.py new file mode 100644 index 0000000..bccbdce --- /dev/null +++ b/src/loader.py @@ -0,0 +1,35 @@ +import os +import uuid +from typing import List, Dict + +from .vector_store import ChromaVectorStore +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. + + Parameters + ---------- + directory_path : str + Path to the directory containing text files. + """ + for root, _, files in os.walk(directory_path): + for file in files: + if file.lower().endswith(".txt"): + file_path = os.path.join(root, file) + try: + with open(file_path, "r", encoding="utf-8") as f: + content = f.read() + title = file + 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}") \ No newline at end of file diff --git a/src/main.py b/src/main.py index d20c115..27d14a7 100644 --- a/src/main.py +++ b/src/main.py @@ -1,42 +1,4 @@ -```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) - +from .cli import main if __name__ == "__main__": - main() -``` \ No newline at end of file + main() \ No newline at end of file diff --git a/src/tools.py b/src/tools.py index 49a6297..92d8c35 100644 --- a/src/tools.py +++ b/src/tools.py @@ -1,55 +1,65 @@ -```python -""" -Tools for the RAG agent: searching and adding to the knowledge base. -""" +import uuid +from typing import List, Dict from langchain.tools import tool -from langchain_text_splitter import RecursiveCharacterTextSplitter +from .vector_store import ChromaVectorStore +from .chunking import chunk_text +from .config import CHROMA_DB_PATH -from vector_store import vector_store - -# Chunking configuration -splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) +# Initialize a single vector store instance +store = ChromaVectorStore(CHROMA_DB_PATH) -@tool +@tool("search_knowledge_base") def search_knowledge_base(query: str, max_results: int = 5) -> str: """ - Semantic search in the knowledge base. + Search the knowledge base for the most relevant documents. - Args: - query: Search query. - max_results: Maximum number of results to return. + Parameters + ---------- + query : str + The search query. + max_results : int, optional + Number of top results to return (default is 5). - Returns: - Formatted string with search results. + Returns + ------- + str + Formatted 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')})" + results = store.search(query, limit=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(formatted) + return "\n".join(output) -@tool +@tool("add_to_knowledge_base") def add_to_knowledge_base(content: str, title: str) -> str: """ - Add a document to the knowledge base. + Add a new document to the knowledge base. - Args: - content: Full text of the document. - title: Title of the document. + Parameters + ---------- + content : str + The full text content of the document. + title : str + A title or identifier for the document. - Returns: + Returns + ------- + str 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 + chunks = chunk_text(content) + ids = [str(uuid.uuid4()) for _ in chunks] + metadatas = [{"title": title} for _ in chunks] + store.add_documents(chunks, metadatas, ids) + return f"Added {len(chunks)} chunks to the knowledge base." \ No newline at end of file diff --git a/src/vector_store.py b/src/vector_store.py index f3ee1bf..495d645 100644 --- a/src/vector_store.py +++ b/src/vector_store.py @@ -1,66 +1,34 @@ -```python -""" -Vector store implementation using Qdrant and Ollama embeddings. -""" - -from typing import List +import uuid +from typing import List, Dict +from chromadb import Client +from chromadb.config import Settings from langchain_ollama import OllamaEmbeddings -from langchain_qdrant import QdrantVectorStore + +from .config import CHROMA_DB_PATH, EMBEDDING_MODEL -class QdrantVectorStoreWrapper: - """ - Wrapper around LangChain's QdrantVectorStore. - Handles initialization, document addition, and similarity search. - """ +class ChromaVectorStore: + def __init__(self, db_path: str = CHROMA_DB_PATH): + self.client = Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=db_path)) + self.collection_name = "knowledge_base" + self.collection = self.client.get_or_create_collection(name=self.collection_name) + self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) - 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], metadatas: List[Dict], ids: List[str]): + embeddings = self.embeddings.embed_documents(documents) + self.collection.add( + documents=documents, + embeddings=embeddings, + metadatas=metadatas, + ids=ids ) - 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 + def search(self, query: str, limit: int = 5): + query_embedding = self.embeddings.embed_query(query) + results = self.collection.query( + query_embeddings=[query_embedding], + n_results=limit, + include=["documents", "distances", "metadatas"] + ) + return results \ No newline at end of file