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 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 <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
```
# Python пакеты
pi
requirements.txt
+2 -5
View File
@@ -1,7 +1,4 @@
```
langchain
langchain-qdrant
langchain-ollama
qdrant-client
python-dotenv
```
langchain-text-splitters
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
"""
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
```
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
"""
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 <file_path> - Add a document to the knowledge base.")
print("Welcome to the RAG Agent CLI.")
print("Commands:")
print(" /add <file_path> - Add a text file to the knowledge base.")
print(" /search <query> - 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 <file_path>")
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 <query>")
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()
```
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
"""
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()
```
main()
src/tools.py
+46 -36
View File
@@ -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."
```
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."
src/vector_store.py
+27 -59
View File
@@ -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()
```
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