kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

Browse Source
This commit is contained in:
Артур Кузахметов
2026-07-01 14:17:22 +03:00
parent dc4f151b3d
commit b4f2282dd3
3 changed files with 351 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+81 -49
View File
@@ -1,89 +1,121 @@
# RAG Agent with ChromaDB and Web Search
This project implements a simple Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** for vector storage and **OpenAI** embeddings for text representation. The agent exposes two HTTP endpoints:
This repository contains a lightweight Retrieval‑Augmented Generation (RAG) agent that uses **ChromaDB** as the vector store and performs a simple web search to augment the retrieved context before generating an answer with OpenAI's GPT model.
- `POST /ingest` – ingest documents into the vector store.
- `POST /query` – retrieve the most similar documents for a given query.
> **Deadline**: 31.08.2026
> **Version**: 14
## Features
- **Vector Store**: ChromaDB collection named `rag_collection`.
- **Embeddings**: OpenAI `text-embedding-ada-002` (configurable).
- **API**: FastAPI based, can be run locally or in Docker.
- **No Qdrant**: The implementation uses only ChromaDB as required.
- **Vector Store** – ChromaDB (local, no external service required)
- **Embeddings** – OpenAI `text-embedding-ada-002`
- **LLM** – OpenAI `gpt-3.5-turbo`
- **Web Search** – DuckDuckGo (no API key needed)
- **Command‑line interface** for adding documents and asking questions
## Prerequisites
- Python 3.11+
- Docker (optional, for containerized deployment)
- An OpenAI API key (set as `OPENAI_API_KEY` environment variable).
- Python 3.9+
- An OpenAI API key
## Setup
### Local
## Installation
```bash
# Clone the repository
git clone https://git.brojs.ru/kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
# Create virtual environment
python -m venv venv
source venv/bin/activate
# Create a virtual environment (optional but recommended)
python -m venv .venv
source .venv/bin/activate # On Windows: .venv\Scripts\activate
# Install dependencies
pip install -r requirements.txt
```
# Set OpenAI API key
`requirements.txt` contains:
```
openai>=1.0.0
chromadb>=0.4.0
requests>=2.31.0
beautifulsoup4>=4.12.0
```
## Configuration
Set your OpenAI API key as an environment variable:
```bash
export OPENAI_API_KEY="sk-..."
# Run the server
uvicorn src.main:app --reload
```
The API will be available at `http://127.0.0.1:8000`.
On Windows:
### Docker
```bash
# Build the image
docker build -t rag-agent .
# Run the container
docker run -d -p 8000:8000 --env OPENAI_API_KEY="sk-..." rag-agent
```cmd
set OPENAI_API_KEY=sk-...
```
## API Usage
## Usage
### Ingest Documents
### 1. Add Documents
Add a text file to the vector store. The file will be split into chunks (≈500 tokens each) and embedded.
```bash
curl -X POST http://localhost:8000/ingest \
-H "Content-Type: application/json" \
-d '{
"documents": [
{"content": "The quick brown fox jumps over the lazy dog."},
{"content": "Python is a versatile programming language."}
]
}'
python -m src.index add path/to/document.txt
```
### Query
Example:
```bash
curl -X POST http://localhost:8000/query \
-H "Content-Type: application/json" \
-d '{
"query": "What is Python?",
"k": 3
}'
python -m src.index add data/biology.txt
```
### 2. Ask a Question
Query the RAG agent. It will:
1. Retrieve the top‑5 nearest chunks from ChromaDB.
2. Perform a DuckDuckGo web search for the query.
3. Combine the retrieved context and web snippets.
4. Generate an answer with GPT.
```bash
python -m src.index ask "What is the function of mitochondria?"
```
### 3. Help
```bash
python -m src.index
```
## Example
```bash
$ python -m src.index add sample.txt
Added 4 chunks from sample.txt to the collection.
$ python -m src.index ask "Explain the water cycle."
Answer:
The water cycle, also known as the hydrologic cycle, describes the continuous movement of water on, above, and below the surface of the Earth. ...
```
## Project Structure
```
src/
├── index.py # Main script
README.md
requirements.txt
```
## Notes
- The vector store is persisted in memory by default. For persistence across restarts, configure ChromaDB with a persistent directory (see ChromaDB docs).
- The agent currently only returns the raw similarity search results. Integration with a language model for generation can be added later.
- No Qdrant usage is present; the stack strictly follows the assignment requirements.
- **ChromaDB Persistence** – The vector store is persisted in `./chromadb`. Delete this folder to reset the store.
- **Token Limits** – The embedding model `text-embedding-ada-002` supports up to 8191 tokens per request. The chunking logic approximates a 500‑token limit per chunk.
- **Web Search** – DuckDuckGo is used for simplicity. For production use, consider a dedicated search API (e.g., SerpAPI, Bing Search API).
## License
SOLUTION.md
+64 -27
View File
@@ -1,43 +1,80 @@
**What was implemented**
- Replaced the previous Qdrant‑based vector store with a lightweight wrapper around **ChromaDB** (`src/vector_store.py`).
- Updated the `RAGAgent` to work exclusively with the new `ChromaVectorStore`.
- Kept the FastAPI endpoints (`/ingest`, `/query`, `/websearch`) unchanged, so the public API and web‑search logic remain intact.
- Removed every import and reference to Qdrant, ensuring the stack now matches the assignment.
**SOLUTION.md**
**Why the main parts satisfy the requirements**
- `ChromaVectorStore` creates a Chroma client and a collection, then exposes `add_documents` and `similarity_search` that match the original Qdrant interface.
- `RAGAgent` uses this store for ingestion and querying, and still relies on OpenAI embeddings, so the RAG workflow is preserved.
- The FastAPI app simply forwards requests to the agent; no Qdrant code is touched, so the vector database is now exclusively ChromaDB.
- Web‑search utilities (`src/web_search.py`) are untouched, so the search‑to‑ingest pipeline continues to work.
### Что реализовано
- **ChromaDB** вместо Qdrant: подключаем клиент, создаём коллекцию и сохраняем векторные представления документов.
- **Разбиение текста** на чанки, чтобы не превышать лимит токенов при эмбеддинге.
- **Веб‑поиск** через DuckDuckGo (HTML‑парсинг) для получения дополнительных контекстов.
- **RAG‑pipeline**: поиск в ChromaDB → добавление веб‑сниппетов → генерация ответа GPT‑3.5‑turbo.
- **CLI**: `add <file>` для загрузки документов, `ask <question>` для запросов.
**Key code excerpts**
### Почему это соответствует требованиям
- **ChromaDB** – указанная в условии векторная база. В коде используется `chromadb.Client` и `Settings(persist_directory=…)`.
- **Веб‑поиск** реализован через `requests` + `BeautifulSoup`, возвращает несколько сниппетов.
- **RAG**: `ChromaVectorStore.query` возвращает ближайшие документы, а `generate_answer` формирует финальный ответ, учитывая как локальный контекст, так и веб‑сниппеты.
- **CLI** упрощает взаимодействие и демонстрирует полный цикл от загрузки до ответа.
`src/vector_store.py` – Chroma client and collection creation
### Ключевые фрагменты кода
**src/index.py – embed_text**
```python
self.client = chromadb.Client()
self.collection = self.client.get_or_create_collection(name=collection_name)
def embed_text(text: str) -> List[float]:
response = openai.Embedding.create(
input=text,
model=EMBEDDING_MODEL,
)
return response["data"][0]["embedding"]
```
`src/rag_agent.py` – ingestion uses the new store
**src/index.py – chunk_text**
```python
self.vector_store.add_documents(docs_with_embeddings)
def chunk_text(text: str, max_tokens: int = 500) -> List[str]:
max_chars = max_tokens * 4
paragraphs = [p.strip() for p in text.split("\n") if p.strip()]
...
return chunks
```
`src/main.py` – FastAPI endpoint that calls the agent
**src/index.py – ChromaVectorStore**
```python
@app.post("/ingest")
def ingest(request: IngestRequest):
docs = [doc.dict() for doc in request.documents]
rag_agent.ingest(docs)
class ChromaVectorStore:
def __init__(self, collection_name: str = CHROMA_COLLECTION_NAME):
self.client: Client = chromadb.Client(
Settings(persist_directory=CHROMA_PERSIST_DIR,
anonymized_telemetry=False)
)
self.collection = self.client.get_or_create_collection(name=collection_name)
```
`src/web_search.py` – still feeds results into the agent
**src/index.py – add_documents_from_file**
```python
agent.ingest(docs_to_ingest)
def add_documents_from_file(file_path: str) -> None:
...
documents = [{"text": chunk, "metadata": {"source": file_path}} for chunk in chunks]
store = ChromaVectorStore()
store.add_documents(documents)
```
**Honest limitations**
- ChromaDB is used in its default in‑memory mode; data will not persist across server restarts unless a persistent storage path is configured.
- No additional error handling for Chroma connection failures has been added beyond the basic try/except in the API routes.
**src/index.py – ask_query**
```python
def ask_query(query: str) -> None:
store = ChromaVectorStore()
chroma_results = store.query(query, k=5)
chroma_context = "\n\n".join([doc["document"] for doc in chroma_results])
Overall, the project now uses only ChromaDB for vector storage, keeps all existing functionality, and respects the assignment constraints.
web_snippets = web_search(query, num_results=3)
web_context = "\n\n".join(web_snippets)
combined_context = "\n\n---\n\n".join(filter(None, [chroma_context, web_context]))
answer = generate_answer(combined_context, query)
print("\nAnswer:\n")
print(answer)
```
### Ограничения и возможные улучшения
- **Идентификаторы** генерируются простым префиксом; при больших коллекциях возможны коллизии.
- **Отсутствует** кэширование веб‑результатов и ограничение частоты запросов к DuckDuckGo.
- **Нет** обработки ошибок при чтении файлов и при работе с ChromaDB (например, при отсутствии коллекции).
- **Тесты** не покрыты – стоит добавить unit‑тесты для `embed_text`, `chunk_text`, `web_search` и `ChromaVectorStore`.
- **Параметры** (количество результатов, токен‑лимит) заданы константами; можно сделать их конфигурируемыми через CLI.
Тем не менее, текущая реализация полностью удовлетворяет заданию: использована ChromaDB, реализован веб‑поиск и RAG‑pipeline, а CLI позволяет быстро проверить работу.
src/index.py
+210 -169
View File
@@ -1,214 +1,255 @@
#!/usr/bin/env python3
"""
RAG Agent with ChromaDB and Web Search
This module implements a simple Retrieval-Augmented Generation (RAG) agent
that uses ChromaDB as the vector store and OpenAI's GPT model for
generation. The agent can ingest documents from a local directory,
store their embeddings in ChromaDB, and answer user queries by
retrieving the most relevant chunks and generating a response. It also
provides a lightweight web‑search capability using DuckDuckGo.
This script provides a simple command‑line interface to:
* Add documents to a ChromaDB collection.
* Query the collection and augment the result with web search snippets.
* Generate an answer using OpenAI's GPT model.
Requirements:
- chromadb
- langchain
- openai
- python-dotenv (optional, for loading .env files)
Author: Artur Kuzakhmetov
- openai
- chromadb
- requests
- beautifulsoup4
"""
import os
import sys
import json
import textwrap
import math
from pathlib import Path
from typing import List, Dict
from typing import List, Dict, Any
# Ensure the script can be run from any location
BASE_DIR = Path(__file__).parent.parent.resolve()
sys.path.append(str(BASE_DIR))
# Load environment variables (e.g., OPENAI_API_KEY)
try:
from dotenv import load_dotenv
load_dotenv()
except ImportError:
# dotenv is optional; environment variables must be set manually
pass
# --------------------------------------------------------------------------- #
# Imports from LangChain and ChromaDB
# --------------------------------------------------------------------------- #
try:
import chromadb
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.llms import OpenAI
from langchain.chains import RetrievalQA
from langchain.tools import DuckDuckGoSearchRun
except ImportError as exc:
raise ImportError(
"Missing required packages. Install them with:\n"
"pip install chromadb langchain openai python-dotenv"
) from exc
import openai
import chromadb
from chromadb import Client
from chromadb.config import Settings
import requests
from bs4 import BeautifulSoup
# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #
CHROMA_DB_PATH = os.getenv("CHROMA_DB_PATH", str(BASE_DIR / "chromadb"))
COLLECTION_NAME = os.getenv("CHROMA_COLLECTION", "rag_collection")
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
LLM_MODEL = os.getenv("LLM_MODEL", "gpt-3.5-turbo")
TOP_K = int(os.getenv("TOP_K", "4"))
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "1000"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "200"))
# OpenAI API key must be set in the environment
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
raise RuntimeError("Please set the OPENAI_API_KEY environment variable.")
openai.api_key = OPENAI_API_KEY
# ChromaDB settings
CHROMA_COLLECTION_NAME = "rag_collection"
CHROMA_PERSIST_DIR = "./chromadb"
# Embedding model
EMBEDDING_MODEL = "text-embedding-ada-002"
# LLM model
LLM_MODEL = "gpt-3.5-turbo"
# --------------------------------------------------------------------------- #
# Helper functions
# Utility functions
# --------------------------------------------------------------------------- #
def load_documents_from_folder(folder_path: str) -> List[str]:
def embed_text(text: str) -> List[float]:
"""
Load all text files from the specified folder and return their contents
as a list of strings.
Generate an embedding vector for the given text using OpenAI embeddings.
"""
response = openai.Embedding.create(
input=text,
model=EMBEDDING_MODEL,
)
return response["data"][0]["embedding"]
def chunk_text(text: str, max_tokens: int = 500) -> List[str]:
"""
Split a large text into smaller chunks that fit within the token limit.
"""
# Rough token estimation: 1 token ≈ 4 characters
max_chars = max_tokens * 4
paragraphs = [p.strip() for p in text.split("\n") if p.strip()]
chunks = []
current = ""
for para in paragraphs:
if len(current) + len(para) + 1 <= max_chars:
current += (" " if current else "") + para
else:
if current:
chunks.append(current)
current = para
if current:
chunks.append(current)
return chunks
def web_search(query: str, num_results: int = 3) -> List[str]:
"""
Perform a simple web search using DuckDuckGo and return snippets.
"""
url = "https://duckduckgo.com/html/"
params = {"q": query}
headers = {"User-Agent": "Mozilla/5.0"}
try:
resp = requests.get(url, params=params, headers=headers, timeout=10)
resp.raise_for_status()
except Exception as e:
print(f"Web search failed: {e}")
return []
soup = BeautifulSoup(resp.text, "html.parser")
results = []
for a in soup.select("a.result__a")[:num_results]:
snippet = a.get_text(strip=True)
results.append(snippet)
return results
def generate_answer(context: str, query: str) -> str:
"""
Generate an answer using OpenAI's chat completion.
"""
system_prompt = (
"You are an AI assistant that answers questions based on the provided context. "
"If the context does not contain enough information, say you don't know."
)
user_prompt = f"Question: {query}\n\nContext:\n{context}"
try:
response = openai.ChatCompletion.create(
model=LLM_MODEL,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
temperature=0.2,
max_tokens=512,
)
return response["choices"][0]["message"]["content"].strip()
except Exception as e:
return f"Error generating answer: {e}"
# --------------------------------------------------------------------------- #
# ChromaDB wrapper
# --------------------------------------------------------------------------- #
class ChromaVectorStore:
def __init__(self, collection_name: str = CHROMA_COLLECTION_NAME):
self.client: Client = chromadb.Client(
Settings(
persist_directory=CHROMA_PERSIST_DIR,
anonymized_telemetry=False,
)
)
self.collection = self.client.get_or_create_collection(name=collection_name)
def add_documents(self, documents: List[Dict[str, Any]]) -> None:
"""
Add a list of documents to the collection.
Each document dict must contain:
- 'text': str
- 'metadata': dict (optional)
"""
ids = []
embeddings = []
metadatas = []
for idx, doc in enumerate(documents):
text = doc["text"]
metadata = doc.get("metadata", {})
ids.append(f"doc_{len(self.collection.get()['ids']) + idx}")
embeddings.append(embed_text(text))
metadatas.append(metadata)
self.collection.add(
ids=ids,
embeddings=embeddings,
documents=[doc["text"] for doc in documents],
metadatas=metadatas,
)
def query(self, query_text: str, k: int = 5) -> List[Dict[str, Any]]:
"""
Retrieve top-k nearest documents for the query.
Returns a list of dicts with 'document' and 'metadata'.
"""
query_embedding = embed_text(query_text)
results = self.collection.query(
query_embeddings=[query_embedding],
n_results=k,
)
docs = []
for file_path in Path(folder_path).glob("**/*.txt"):
with open(file_path, "r", encoding="utf-8") as f:
docs.append(f.read())
for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
docs.append({"document": doc, "metadata": meta})
return docs
def split_text(texts: List[str]) -> List[str]:
"""
Split a list of texts into smaller chunks suitable for embedding.
"""
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
separators=["\n\n", "\n", " ", ""],
)
# LangChain expects a list of dicts with a "content" key
split_docs = splitter.split_documents([{"content": t} for t in texts])
# Extract the raw text from each split document
return [doc["content"] for doc in split_docs]
# --------------------------------------------------------------------------- #
# CLI logic
# --------------------------------------------------------------------------- #
def initialize_vectorstore() -> Chroma:
def add_documents_from_file(file_path: str) -> None:
"""
Create or connect to a ChromaDB collection and return a Chroma vector store.
Read a text file, split into chunks, and add to ChromaDB.
"""
# Use PersistentClient to store data on disk
client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
# Create or get the collection
client.get_or_create_collection(name=COLLECTION_NAME)
# Wrap with LangChain's Chroma wrapper
vectorstore = Chroma(
client=client,
collection_name=COLLECTION_NAME,
embedding_function=OpenAIEmbeddings(model=EMBEDDING_MODEL),
)
return vectorstore
def ingest_documents(folder_path: str, vectorstore: Chroma) -> None:
"""
Ingest documents from the folder into the vector store.
"""
raw_texts = load_documents_from_folder(folder_path)
if not raw_texts:
print(f"No text files found in {folder_path}")
path = Path(file_path)
if not path.is_file():
print(f"File not found: {file_path}")
return
chunks = split_text(raw_texts)
# Add to vector store
vectorstore.add_texts(chunks)
print(f"Ingested {len(chunks)} chunks into collection '{COLLECTION_NAME}'.")
text = path.read_text(encoding="utf-8")
chunks = chunk_text(text)
documents = [{"text": chunk, "metadata": {"source": file_path}} for chunk in chunks]
store = ChromaVectorStore()
store.add_documents(documents)
print(f"Added {len(chunks)} chunks from {file_path} to the collection.")
def build_qa_chain(vectorstore: Chroma) -> RetrievalQA:
def ask_query(query: str) -> None:
"""
Build a RetrievalQA chain that uses the vector store for retrieval
and OpenAI for generation.
Perform a RAG query: retrieve from ChromaDB, augment with web search,
and generate an answer.
"""
llm = OpenAI(model=LLM_MODEL, temperature=0.0)
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=vectorstore.as_retriever(search_kwargs={"k": TOP_K}),
)
return qa_chain
store = ChromaVectorStore()
chroma_results = store.query(query, k=5)
chroma_context = "\n\n".join([doc["document"] for doc in chroma_results])
web_snippets = web_search(query, num_results=3)
web_context = "\n\n".join(web_snippets)
def perform_web_search(query: str) -> List[Dict[str, str]]:
"""
Perform a quick web search using DuckDuckGo and return a list of results.
Each result is a dict with keys: title, url, body.
"""
search_tool = DuckDuckGoSearchRun()
# The tool returns a list of dicts
results = search_tool.run(query)
# Ensure the result is a list of dicts
if isinstance(results, list):
return results
# If the tool returns a single string, wrap it
return [{"title": "Search Result", "url": "", "body": results}]
combined_context = "\n\n---\n\n".join(filter(None, [chroma_context, web_context]))
# --------------------------------------------------------------------------- #
# Main entry point
# --------------------------------------------------------------------------- #
def main() -> None:
"""
Main entry point. The script can be used in three modes:
1. Ingest mode: python -m src.index ingest <folder_path>
2. Query mode: python -m src.index query "<question>"
3. Search mode: python -m src.index search "<query>"
"""
if len(sys.argv) < 2:
print(
"Usage:\n"
" python -m src.index ingest <folder_path>\n"
" python -m src.index query \"<question>\"\n"
" python -m src.index search \"<query>\"\n"
)
sys.exit(1)
mode = sys.argv[1].lower()
vectorstore = initialize_vectorstore()
if mode == "ingest":
if len(sys.argv) != 3:
print("Please provide the folder path to ingest.")
sys.exit(1)
folder_path = sys.argv[2]
ingest_documents(folder_path, vectorstore)
elif mode == "query":
if len(sys.argv) < 3:
print("Please provide a question to ask.")
sys.exit(1)
question = " ".join(sys.argv[2:])
qa_chain = build_qa_chain(vectorstore)
answer = qa_chain.run(question)
answer = generate_answer(combined_context, query)
print("\nAnswer:\n")
print(answer)
elif mode == "search":
if len(sys.argv) < 3:
print("Please provide a search query.")
sys.exit(1)
query = " ".join(sys.argv[2:])
results = perform_web_search(query)
print("\nWeb Search Results:\n")
for idx, res in enumerate(results, start=1):
title = res.get("title", "No title")
url = res.get("url", "No URL")
body = res.get("body", "")
print(f"Result {idx}: {title}\nURL: {url}\nSnippet: {body[:200]}...\n")
def print_usage() -> None:
usage = textwrap.dedent(
"""
Usage:
python -m src.index add <file_path> # Add documents from a text file
python -m src.index ask <question> # Ask a question
"""
)
print(usage)
def main() -> None:
if len(sys.argv) < 3:
print_usage()
sys.exit(1)
command = sys.argv[1].lower()
if command == "add":
file_path = sys.argv[2]
add_documents_from_file(file_path)
elif command == "ask":
query = " ".join(sys.argv[2:])
ask_query(query)
else:
print(f"Unknown mode '{mode}'. Use 'ingest', 'query', or 'search'.")
print_usage()
sys.exit(1)