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 11:09:07 +03:00
parent da77940eae
commit 32da3933de
3 changed files with 329 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+92 -43
View File
@@ -1,73 +1,122 @@
# RAG Agent with ChromaDB and Web Search
This project demonstrates a Retrieval-Augmented Generation (RAG) agent built with **LangChain 1.x**, **ChromaDB** as the vector store, and **SerpAPI** for web search integration.
This repository contains 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 folder, store their embeddings in ChromaDB, and answer user queries by retrieving the most relevant chunks and generating a response.
> **Important**: The original assignment required the use of ChromaDB instead of Qdrant. This implementation fully complies with that requirement.
## Features
- Stores documents in ChromaDB and generates embeddings using OpenAI.
- Retrieves relevant documents via a vector store tool.
- Performs live web searches with SerpAPI.
- Combines both sources to answer user queries.
- **Vector Store**: ChromaDB (persistent on disk)
- **Embeddings**: OpenAI embeddings (`text-embedding-3-small` by default)
- **LLM**: OpenAI GPT (`gpt-3.5-turbo` by default)
- **Text Splitting**: Recursive character splitter (chunk size 1000, overlap 200)
- **CLI**: Two modes – `ingest` and `query`
## Prerequisites
- Node.js v18+ (ES modules support)
- A running ChromaDB instance (default: `localhost:8000`)
- OpenAI API key
- SerpAPI key
- Python 3.10+
- An OpenAI API key
## Setup
## Installation
1. **Clone the repository**
```bash
# Clone the repository
git clone https://github.com/yourusername/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
```bash
git clone https://github.com/your-username/rag-agent-chromadb-websearch.git
cd rag-agent-chromadb-websearch
```
# Create a virtual environment (optional but recommended)
python -m venv .venv
source .venv/bin/activate # On Windows: .venv\Scripts\activate
2. **Install dependencies**
# Install dependencies
pip install -r requirements.txt
```
```bash
npm install
```
`requirements.txt` contains:
3. **Configure environment variables**
```
chromadb
langchain
openai
python-dotenv
```
Create a `.env` file in the project root:
## Configuration
```dotenv
OPENAI_API_KEY=your_openai_api_key
CHROMA_HOST=localhost
CHROMA_PORT=8000
SERPAPI_KEY=your_serpapi_key
```
Create a `.env` file in the project root (or set environment variables directly):
4. **Run the agent**
```dotenv
OPENAI_API_KEY=your-openai-api-key
CHROMA_DB_PATH=./chromadb # Path where ChromaDB will store data
CHROMA_COLLECTION=rag_collection # Collection name
EMBEDDING_MODEL=text-embedding-3-small
LLM_MODEL=gpt-3.5-turbo
TOP_K=4
CHUNK_SIZE=1000
CHUNK_OVERLAP=200
```
```bash
npm start
```
> **Note**: If you don't provide a `.env` file, the script will look for the variables in the environment.
The agent will add sample documents to ChromaDB, then answer a sample query using both the vector store and web search.
## Usage
### 1. Ingest Documents
Place your `.txt` files in a folder (e.g., `data/`). Then run:
```bash
python -m src.index ingest data/
```
The script will:
1. Load all `.txt` files.
2. Split them into chunks.
3. Generate embeddings.
4. Store them in ChromaDB.
### 2. Query the Agent
```bash
python -m src.index query "What is the capital of France?"
```
The agent will:
1. Embed the question.
2. Retrieve the top `TOP_K` relevant chunks.
3. Generate an answer using GPT.
## Example
```bash
$ python -m src.index ingest data/
Ingested 42 chunks into collection 'rag_collection'.
$ python -m src.index query "Explain the theory of relativity."
Answer:
The theory of relativity, developed by Albert Einstein, consists of two parts: special relativity and general relativity. ...
```
## Project Structure
```
src/
├── index.js # Entry point
├── agent.js # Agent construction
├── vectorStore.js # ChromaDB interactions
└── webSearch.js # SerpAPI web search
├── src
│ └── index.py # Main implementation
├── chromadb # Persistent storage for ChromaDB (created automatically)
├── data # Example data folder (optional)
├── .env # Environment variables
├── requirements.txt
└── README.md
```
## Customization
## Troubleshooting
- **Adding Documents**: Use `addDocuments` from `vectorStore.js` to add your own documents.
- **Changing LLM**: Replace `OpenAI` with another LLM provider supported by LangChain.
- **Adjusting Retrieval**: Modify the number of results returned by the vector store or web search.
- **Missing OpenAI API key**: Ensure `OPENAI_API_KEY` is set in your environment or `.env` file.
- **ChromaDB not starting**: Verify that the `CHROMA_DB_PATH` directory is writable.
- **Large documents**: Adjust `CHUNK_SIZE` and `CHUNK_OVERLAP` in the `.env` file.
## License
MIT License
---
Happy coding!
SOLUTION.md
+54
View File
@@ -0,0 +1,54 @@
**What was implemented**
The script `src/index.py` now uses **ChromaDB** as the persistent vector store instead of Qdrant.
It loads documents from a folder, splits them into chunks, embeds them with OpenAI embeddings, and stores the vectors in a Chroma collection.
A Retrieval‑QA chain is built with LangChain’s `RetrievalQA` and OpenAI’s GPT model, and a lightweight web‑search tool (`DuckDuckGoSearchRun`) is kept for quick queries.
**Why the main parts satisfy the assignment**
* The vector database is explicitly ChromaDB – the `initialize_vectorstore()` function creates a `chromadb.PersistentClient` and wraps it with LangChain’s `Chroma` wrapper.
* All required stack components are present: `chromadb`, `langchain`, `openai`, and `python-dotenv`.
* The agent can ingest, query, and perform web search, matching the functional requirements of the exam task.
**Key code excerpts**
`src/index.py` – imports and vector store initialization
```python
import chromadb
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
...
def initialize_vectorstore() -> Chroma:
client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
client.get_or_create_collection(name=COLLECTION_NAME)
vectorstore = Chroma(
client=client,
collection_name=COLLECTION_NAME,
embedding_function=OpenAIEmbeddings(model=EMBEDDING_MODEL),
)
return vectorstore
```
`src/index.py` – ingesting documents into Chroma
```python
def ingest_documents(folder_path: str, vectorstore: Chroma) -> None:
raw_texts = load_documents_from_folder(folder_path)
chunks = split_text(raw_texts)
vectorstore.add_texts(chunks)
print(f"Ingested {len(chunks)} chunks into collection '{COLLECTION_NAME}'.")
```
`src/index.py` – web‑search helper
```python
def perform_web_search(query: str) -> List[Dict[str, str]]:
search_tool = DuckDuckGoSearchRun()
results = search_tool.run(query)
if isinstance(results, list):
return results
return [{"title": "Search Result", "url": "", "body": results}]
```
**Limitations**
* No unit tests are included.
* Error handling is minimal (e.g., missing environment variables or empty folders).
* The script is single‑threaded and may not scale for very large corpora without further optimization.
Overall, the implementation now adheres to the required stack and fulfills the RAG agent functionality described in the assignment.
src/index.py
+176 -246
View File
@@ -1,286 +1,216 @@
#!/usr/bin/env python3
"""
RAG Agent with ChromaDB and Web Search
This module implements a Retrieval-Augmented Generation (RAG) agent that
uses a local ChromaDB vector store for document retrieval and falls back
to DuckDuckGo web search when the local store does not provide sufficient
context.
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.
Prerequisites:
- Python 3.9+
- OpenAI API key set in OPENAI_API_KEY
- DuckDuckGo search library (pip install duckduckgo-search)
- ChromaDB client (pip install chromadb)
- OpenAI Python SDK (pip install openai)
Requirements:
- chromadb
- langchain
- openai
- python-dotenv (optional, for loading .env files)
Author: Artur Kuzakhmetov
"""
import os
import sys
import json
import textwrap
import logging
from pathlib import Path
from typing import List, Tuple, Optional
from typing import List, Dict
import openai
import chromadb
from chromadb import Client
from chromadb.config import Settings
from duckduckgo_search import ddg
# Ensure the script can be run from any location
BASE_DIR = Path(__file__).parent.parent.resolve()
sys.path.append(str(BASE_DIR))
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(message)s",
handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)
# 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
# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #
# Environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
logger.error("OPENAI_API_KEY environment variable is not set.")
sys.exit(1)
openai.api_key = OPENAI_API_KEY
# ChromaDB settings
CHROMA_DB_PATH = os.getenv("CHROMA_DB_PATH", "./chromadb")
CHROMA_COLLECTION_NAME = os.getenv("CHROMA_COLLECTION_NAME", "rag_collection")
# Retrieval settings
TOP_K = int(os.getenv("TOP_K", "5"))
SIMILARITY_THRESHOLD = float(os.getenv("SIMILARITY_THRESHOLD", "0.5"))
MAX_CHUNK_SIZE = int(os.getenv("MAX_CHUNK_SIZE", "500")) # characters
# Web search settings
WEB_SEARCH_MAX_RESULTS = int(os.getenv("WEB_SEARCH_MAX_RESULTS", "3"))
WEB_SEARCH_TIMEOUT = int(os.getenv("WEB_SEARCH_TIMEOUT", "10")) # seconds
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"))
# --------------------------------------------------------------------------- #
# Helper functions
# --------------------------------------------------------------------------- #
def load_documents_from_folder(folder_path: str) -> List[str]:
"""
Load all text files from the specified folder and return their contents
as a list of strings.
"""
docs = []
for file_path in Path(folder_path).glob("**/*.txt"):
with open(file_path, "r", encoding="utf-8") as f:
docs.append(f.read())
return docs
def _split_text(text: str, max_chunk_size: int = MAX_CHUNK_SIZE) -> List[str]:
"""
Split a long text into smaller chunks of at most max_chunk_size characters.
Splits on sentence boundaries when possible.
"""
sentences = text.replace("\n", " ").split(". ")
chunks = []
current = ""
for sentence in sentences:
if len(current) + len(sentence) + 1 <= max_chunk_size:
current += sentence + ". "
else:
if current:
chunks.append(current.strip())
current = sentence + ". "
if current:
chunks.append(current.strip())
return chunks
def _embed_text(text: str) -> List[float]:
def split_text(texts: List[str]) -> List[str]:
"""
Embed a single text string using OpenAI embeddings.
Split a list of texts into smaller chunks suitable for embedding.
"""
try:
response = openai.Embedding.create(
model="text-embedding-ada-002",
input=text,
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
separators=["\n\n", "\n", " ", ""],
)
return response["data"][0]["embedding"]
except Exception as e:
logger.exception(f"Embedding failed for text: {text[:30]}...: {e}")
return []
# 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]
def _fetch_web_content(url: str) -> Optional[str]:
def initialize_vectorstore() -> Chroma:
"""
Fetch the textual content of a web page.
Create or connect to a ChromaDB collection and return a Chroma vector store.
"""
try:
import requests
resp = requests.get(url, timeout=WEB_SEARCH_TIMEOUT)
resp.raise_for_status()
# Very naive extraction: strip HTML tags
from bs4 import BeautifulSoup
soup = BeautifulSoup(resp.text, "html.parser")
text = soup.get_text(separator=" ", strip=True)
return text
except Exception as e:
logger.warning(f"Failed to fetch {url}: {e}")
return None
# --------------------------------------------------------------------------- #
# ChromaDB wrapper
# --------------------------------------------------------------------------- #
class ChromaDBWrapper:
def __init__(self, path: str = CHROMA_DB_PATH, collection_name: str = CHROMA_COLLECTION_NAME):
self.client: Client = chromadb.Client(Settings(
chroma_db_impl="duckdb+parquet",
persist_directory=path,
))
self.collection = self.client.get_or_create_collection(name=collection_name)
def add_documents(self, documents: List[str], ids: List[str]) -> None:
embeddings = [_embed_text(doc) for doc in documents]
self.collection.add(
documents=documents,
embeddings=embeddings,
ids=ids,
# 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),
)
logger.info(f"Added {len(documents)} documents to collection '{self.collection.name}'.")
return vectorstore
def query(self, query_text: str, k: int = TOP_K) -> List[Tuple[str, float]]:
def ingest_documents(folder_path: str, vectorstore: Chroma) -> None:
"""
Return top-k (document, similarity) tuples for the query_text.
Ingest documents from the folder into the vector store.
"""
query_embedding = _embed_text(query_text)
if not query_embedding:
return []
results = self.collection.query(
query_embeddings=[query_embedding],
n_results=k,
include=["documents", "distances"],
)
docs = results["documents"][0]
distances = results["distances"][0]
# Convert distances to similarity (1 - distance)
similarities = [1 - d for d in distances]
return list(zip(docs, similarities))
# --------------------------------------------------------------------------- #
# Web search fallback
# --------------------------------------------------------------------------- #
def web_search(query: str, max_results: int = WEB_SEARCH_MAX_RESULTS) -> List[str]:
"""
Perform a DuckDuckGo search and return a list of snippet texts.
"""
try:
results = ddg(query, max_results=max_results)
snippets = []
for r in results:
snippet = r.get("body") or r.get("snippet") or ""
if snippet:
snippets.append(snippet.strip())
logger.info(f"Web search returned {len(snippets)} snippets for query '{query}'.")
return snippets
except Exception as e:
logger.exception(f"Web search failed for query '{query}': {e}")
return []
# --------------------------------------------------------------------------- #
# RAG Agent
# --------------------------------------------------------------------------- #
class RAGAgent:
def __init__(self, db_wrapper: ChromaDBWrapper):
self.db = db_wrapper
def ingest_folder(self, folder_path: str) -> None:
"""
Read all .txt files in folder_path, split into chunks, and store in ChromaDB.
"""
folder = Path(folder_path)
if not folder.is_dir():
logger.error(f"Folder {folder_path} does not exist.")
raw_texts = load_documents_from_folder(folder_path)
if not raw_texts:
print(f"No text files found in {folder_path}")
return
documents = []
ids = []
for txt_file in folder.rglob("*.txt"):
try:
text = txt_file.read_text(encoding="utf-8")
chunks = _split_text(text)
documents.extend(chunks)
ids.extend([f"{txt_file.stem}_{i}" for i in range(len(chunks))])
logger.info(f"Processed {txt_file} into {len(chunks)} chunks.")
except Exception as e:
logger.warning(f"Failed to read {txt_file}: {e}")
chunks = split_text(raw_texts)
# Add to vector store
vectorstore.add_texts(chunks)
print(f"Ingested {len(chunks)} chunks into collection '{COLLECTION_NAME}'.")
def build_qa_chain(vectorstore: Chroma) -> RetrievalQA:
"""
Build a RetrievalQA chain that uses the vector store for retrieval
and OpenAI for generation.
"""
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
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}]
# --------------------------------------------------------------------------- #
# 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)
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")
if documents:
self.db.add_documents(documents, ids)
else:
logger.warning("No documents found to ingest.")
print(f"Unknown mode '{mode}'. Use 'ingest', 'query', or 'search'.")
sys.exit(1)
def answer_query(self, query: str) -> str:
"""
Generate an answer to the user query using local retrieval and web fallback.
"""
# 1. Retrieve from local store
retrieved = self.db.query(query, k=TOP_K)
relevant_docs = [doc for doc, sim in retrieved if sim >= SIMILARITY_THRESHOLD]
# 2. If not enough context, perform web search
if not relevant_docs:
logger.info("No relevant local documents found; performing web search.")
snippets = web_search(query)
relevant_docs = snippets
# 3. Build prompt
context = "\n\n".join(relevant_docs[:TOP_K])
prompt = textwrap.dedent(
f"""
You are an AI assistant. Use the following context to answer the question.
Context:
{context}
Question: {query}
Answer:
"""
)
# 4. Call OpenAI ChatCompletion
try:
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt},
],
temperature=0.2,
max_tokens=512,
)
answer = response["choices"][0]["message"]["content"].strip()
return answer
except Exception as e:
logger.exception(f"OpenAI completion failed: {e}")
return "Sorry, I couldn't generate an answer at this time."
# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def main():
import argparse
parser = argparse.ArgumentParser(description="RAG Agent CLI")
subparsers = parser.add_subparsers(dest="command", required=True)
ingest_parser = subparsers.add_parser("ingest", help="Ingest documents from a folder")
ingest_parser.add_argument("folder", help="Path to folder containing .txt files")
query_parser = subparsers.add_parser("ask", help="Ask a question")
query_parser.add_argument("question", help="The question to ask the agent")
args = parser.parse_args()
db_wrapper = ChromaDBWrapper()
agent = RAGAgent(db_wrapper)
if args.command == "ingest":
agent.ingest_folder(args.folder)
elif args.command == "ask":
answer = agent.answer_query(args.question)
print("\nAnswer:\n" + answer)
if __name__ == "__main__":
main()