kuzakhmetovartur/prakticheskoe-zadanie-agent-s-rag-pamyat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Артур Кузахметов
2026-05-28 17:55:09 +03:00
parent 5d01ac1d8b
commit 61d5771c9b
11 changed files with 208 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+21 -88
View File
@@ -1,93 +1,26 @@
```markdown # Практическое задание: Агент с RAG-памятью
# 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: Главная
Мои задания
Агент с RAG-памятью
5Д
EN
Агент с RAG-памятью
- `search_knowledge_base(query, max_results)` – semantic search in the knowledge base. Практическое задание: Агент с RAG-памятью
- `add_to_knowledge_base(content, title)` – add a new document to the knowledge base. Цель
## Features Построить AI-агента с локальным RAG-хранилищем знаний на базе Qdrant и Ollama. Агент должен уметь искать и сохранять информацию в векторной базе.
- **Vector store**: Qdrant with Ollama embeddings (`nomic-embed-text`). Стек
- **Chunking**: Recursive character splitter with overlap. Python 3.10+
- **Agent**: Zero-shot React agent that uses the two tools. Qdrant — векторная база данных
- **CLI**: Interactive command line interface to add documents and query the agent. Ollama — локальные LLM и эмбеддинги (llama3, nomic-embed-text)
- **Batch loading**: Script to load all text files from a directory into the knowledge base. LangChain — фреймворк для агентов и RAG
Установка
# Ollama
ollama pull llama3
ollama pull nomic-embed-text
## Prerequisites # Python пакеты
pi
- 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
```
requirements.txt
+2 -5
View File
@@ -1,7 +1,4 @@
```
langchain langchain
langchain-qdrant
langchain-ollama langchain-ollama
qdrant-client langchain-text-splitters
python-dotenv chromadb
```
src/__init__.py
+1
View File
@@ -0,0 +1 @@
# Empty init file to make src a package
src/agent.py
+31 -26
View File
@@ -1,40 +1,45 @@
```python from langchain_ollama import Ollama
"""
Agent creation with RAG integration.
"""
from langchain.llms import Ollama
from langchain.agents import initialize_agent, AgentType from langchain.agents import initialize_agent, AgentType
from langchain.tools import Tool 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(): def create_agent():
""" """
Create and configure the LangChain agent. Create an RAG-enabled agent that can search and add to a knowledge base.
Returns: Returns
AgentExecutor instance ready to run queries. -------
AgentExecutor
The configured agent.
""" """
llm = Ollama(model="llama3") llm = Ollama(model=LLM_MODEL)
tools = [ tools = [
Tool.from_function(search_knowledge_base), Tool(
Tool.from_function(add_to_knowledge_base), 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( system_prompt = (
tools, "You are an AI assistant that can search and add information to a knowledge base. "
llm, "Use the provided tools to answer user queries."
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
``` agent = initialize_agent(
tools=tools,
llm=llm,
agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
verbose=True,
system_message=system_prompt,
)
return agent
src/chunking.py
+6
View File
@@ -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)
src/cli.py
+34 -29
View File
@@ -1,62 +1,67 @@
```python import sys
"""
Interactive CLI for the RAG agent.
"""
import argparse
from pathlib import Path 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(): def main():
agent = create_agent() agent = create_agent()
print("RAG Agent CLI. Commands:") print("Welcome to the RAG Agent CLI.")
print(" /add <file_path> - Add a document to the knowledge base.") print("Commands:")
print(" /add <file_path> - Add a text file to the knowledge base.")
print(" /search <query> - Search the knowledge base.") print(" /search <query> - Search the knowledge base.")
print(" /quit - Exit the program.") print(" /quit - Exit the program.")
print("Any other input will be sent to the agent for general processing.\n")
while True: while True:
try: try:
inp = input("> ").strip() user_input = input(">> ").strip()
except EOFError: except EOFError:
break break
if not inp: if not user_input:
continue continue
if inp.startswith("/add"): if user_input.startswith("/add"):
parts = inp.split(maxsplit=1) parts = user_input.split(maxsplit=1)
if len(parts) < 2: if len(parts) < 2:
print("Usage: /add <file_path>") print("Usage: /add <file_path>")
continue continue
file_path = Path(parts[1]) file_path = parts[1]
if not file_path.is_file(): path_obj = Path(file_path)
print(f"File {file_path} does not exist.") if not path_obj.is_file():
print(f"File not found: {file_path}")
continue continue
content = file_path.read_text(encoding="utf-8") try:
title = file_path.stem content = path_obj.read_text(encoding="utf-8")
result = add_to_knowledge_base(content, title) title = path_obj.name
print(result) result = add_to_knowledge_base(content, title)
print(result)
except Exception as e:
print(f"Error reading file: {e}")
elif inp.startswith("/search"): elif user_input.startswith("/search"):
parts = inp.split(maxsplit=1) parts = user_input.split(maxsplit=1)
if len(parts) < 2: if len(parts) < 2:
print("Usage: /search <query>") print("Usage: /search <query>")
continue continue
query = parts[1] query = parts[1]
response = agent.run(query) result = search_knowledge_base(query, max_results=5)
print(response) print(result)
elif inp.startswith("/quit"): elif user_input.startswith("/quit"):
print("Goodbye!") print("Goodbye!")
break break
else: 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__": if __name__ == "__main__":
main() main()
```
src/config.py
+3
View File
@@ -0,0 +1,3 @@
CHROMA_DB_PATH = "./chroma_db"
EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3"
src/loader.py
+35
View File
@@ -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}")
src/main.py
+2 -40
View File
@@ -1,42 +1,4 @@
```python from .cli import main
"""
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__": if __name__ == "__main__":
main() main()
```
src/tools.py
+46 -36
View File
@@ -1,55 +1,65 @@
```python import uuid
""" from typing import List, Dict
Tools for the RAG agent: searching and adding to the knowledge base.
"""
from langchain.tools import tool 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 # Initialize a single vector store instance
store = ChromaVectorStore(CHROMA_DB_PATH)
# Chunking configuration
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
@tool @tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> str: 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: Parameters
query: Search query. ----------
max_results: Maximum number of results to return. query : str
The search query.
max_results : int, optional
Number of top results to return (default is 5).
Returns: Returns
Formatted string with search results. -------
str
Formatted search results.
""" """
results = vector_store.search(query, max_results) results = store.search(query, limit=max_results)
if not results: docs = results["documents"][0]
return "No relevant documents found." distances = results["distances"][0]
metadatas = results["metadatas"][0]
formatted = [] output = []
for i, doc in enumerate(results): for doc, dist, meta in zip(docs, distances, metadatas):
snippet = doc.page_content[:200].replace("\n", " ") output.append(
formatted.append( f"Title: {meta.get('title', 'N/A')}\n"
f"{i + 1}. {snippet} (Title: {doc.metadata.get('title', 'N/A')})" 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: 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: Parameters
content: Full text of the document. ----------
title: Title of the document. content : str
The full text content of the document.
title : str
A title or identifier for the document.
Returns: Returns
-------
str
Confirmation message. Confirmation message.
""" """
chunks = splitter.split_text(content) chunks = chunk_text(content)
vector_store.add_documents(chunks, [title] * len(chunks)) ids = [str(uuid.uuid4()) for _ in chunks]
return f"Document '{title}' added to knowledge base with {len(chunks)} chunks." 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
+27 -59
View File
@@ -1,66 +1,34 @@
```python import uuid
""" from typing import List, Dict
Vector store implementation using Qdrant and Ollama embeddings.
"""
from typing import List
from chromadb import Client
from chromadb.config import Settings
from langchain_ollama import OllamaEmbeddings from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from .config import CHROMA_DB_PATH, EMBEDDING_MODEL
class QdrantVectorStoreWrapper: class ChromaVectorStore:
""" def __init__(self, db_path: str = CHROMA_DB_PATH):
Wrapper around LangChain's QdrantVectorStore. self.client = Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=db_path))
Handles initialization, document addition, and similarity search. 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__( def add_documents(self, documents: List[str], metadatas: List[Dict], ids: List[str]):
self, embeddings = self.embeddings.embed_documents(documents)
collection_name: str = "knowledge_base", self.collection.add(
host: str = "localhost", documents=documents,
port: int = 6333, embeddings=embeddings,
): metadatas=metadatas,
""" ids=ids
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: def search(self, query: str, limit: int = 5):
""" query_embedding = self.embeddings.embed_query(query)
Add documents to the vector store with metadata. results = self.collection.query(
query_embeddings=[query_embedding],
Args: n_results=limit,
documents: List of document texts. include=["documents", "distances", "metadatas"]
titles: List of titles corresponding to each document. )
""" return results
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()
```