feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
@@ -1,89 +1,121 @@
|
|||||||
# RAG Agent with ChromaDB and Web Search
|
# 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.
|
> **Deadline**: 31.08.2026
|
||||||
- `POST /query` – retrieve the most similar documents for a given query.
|
> **Version**: 14
|
||||||
|
|
||||||
## Features
|
## Features
|
||||||
|
|
||||||
- **Vector Store**: ChromaDB collection named `rag_collection`.
|
- **Vector Store** – ChromaDB (local, no external service required)
|
||||||
- **Embeddings**: OpenAI `text-embedding-ada-002` (configurable).
|
- **Embeddings** – OpenAI `text-embedding-ada-002`
|
||||||
- **API**: FastAPI based, can be run locally or in Docker.
|
- **LLM** – OpenAI `gpt-3.5-turbo`
|
||||||
- **No Qdrant**: The implementation uses only ChromaDB as required.
|
- **Web Search** – DuckDuckGo (no API key needed)
|
||||||
|
- **Command‑line interface** for adding documents and asking questions
|
||||||
|
|
||||||
## Prerequisites
|
## Prerequisites
|
||||||
|
|
||||||
- Python 3.11+
|
- Python 3.9+
|
||||||
- Docker (optional, for containerized deployment)
|
- An OpenAI API key
|
||||||
- An OpenAI API key (set as `OPENAI_API_KEY` environment variable).
|
|
||||||
|
|
||||||
## Setup
|
## Installation
|
||||||
|
|
||||||
### Local
|
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# Clone the repository
|
# Clone the repository
|
||||||
git clone https://git.brojs.ru/kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
|
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
|
cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
|
||||||
|
|
||||||
# Create virtual environment
|
# Create a virtual environment (optional but recommended)
|
||||||
python -m venv venv
|
python -m venv .venv
|
||||||
source venv/bin/activate
|
source .venv/bin/activate # On Windows: .venv\Scripts\activate
|
||||||
|
|
||||||
# Install dependencies
|
# Install dependencies
|
||||||
pip install -r requirements.txt
|
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-..."
|
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
|
```cmd
|
||||||
|
set OPENAI_API_KEY=sk-...
|
||||||
```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
|
|
||||||
```
|
```
|
||||||
|
|
||||||
## 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
|
```bash
|
||||||
curl -X POST http://localhost:8000/ingest \
|
python -m src.index add path/to/document.txt
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"documents": [
|
|
||||||
{"content": "The quick brown fox jumps over the lazy dog."},
|
|
||||||
{"content": "Python is a versatile programming language."}
|
|
||||||
]
|
|
||||||
}'
|
|
||||||
```
|
```
|
||||||
|
|
||||||
### Query
|
Example:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
curl -X POST http://localhost:8000/query \
|
python -m src.index add data/biology.txt
|
||||||
-H "Content-Type: application/json" \
|
```
|
||||||
-d '{
|
|
||||||
"query": "What is Python?",
|
### 2. Ask a Question
|
||||||
"k": 3
|
|
||||||
}'
|
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
|
## Notes
|
||||||
|
|
||||||
- The vector store is persisted in memory by default. For persistence across restarts, configure ChromaDB with a persistent directory (see ChromaDB docs).
|
- **ChromaDB Persistence** – The vector store is persisted in `./chromadb`. Delete this folder to reset the store.
|
||||||
- The agent currently only returns the raw similarity search results. Integration with a language model for generation can be added later.
|
- **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.
|
||||||
- No Qdrant usage is present; the stack strictly follows the assignment requirements.
|
- **Web Search** – DuckDuckGo is used for simplicity. For production use, consider a dedicated search API (e.g., SerpAPI, Bing Search API).
|
||||||
|
|
||||||
## License
|
## License
|
||||||
|
|
||||||
|
|||||||
+64
-27
@@ -1,43 +1,80 @@
|
|||||||
**What was implemented**
|
**SOLUTION.md**
|
||||||
- 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.
|
|
||||||
|
|
||||||
**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.
|
- **ChromaDB** вместо Qdrant: подключаем клиент, создаём коллекцию и сохраняем векторные представления документов.
|
||||||
- `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.
|
- **Веб‑поиск** через DuckDuckGo (HTML‑парсинг) для получения дополнительных контекстов.
|
||||||
- Web‑search utilities (`src/web_search.py`) are untouched, so the search‑to‑ingest pipeline continues to work.
|
- **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
|
```python
|
||||||
self.client = chromadb.Client()
|
def embed_text(text: str) -> List[float]:
|
||||||
self.collection = self.client.get_or_create_collection(name=collection_name)
|
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
|
```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
|
```python
|
||||||
@app.post("/ingest")
|
class ChromaVectorStore:
|
||||||
def ingest(request: IngestRequest):
|
def __init__(self, collection_name: str = CHROMA_COLLECTION_NAME):
|
||||||
docs = [doc.dict() for doc in request.documents]
|
self.client: Client = chromadb.Client(
|
||||||
rag_agent.ingest(docs)
|
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
|
```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**
|
**src/index.py – ask_query**
|
||||||
- ChromaDB is used in its default in‑memory mode; data will not persist across server restarts unless a persistent storage path is configured.
|
```python
|
||||||
- No additional error handling for Chroma connection failures has been added beyond the basic try/except in the API routes.
|
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 позволяет быстро проверить работу.
|
||||||
+206
-165
@@ -1,214 +1,255 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
"""
|
"""
|
||||||
RAG Agent with ChromaDB and Web Search
|
RAG Agent with ChromaDB and Web Search
|
||||||
|
|
||||||
This module implements a simple Retrieval-Augmented Generation (RAG) agent
|
This script provides a simple command‑line interface to:
|
||||||
that uses ChromaDB as the vector store and OpenAI's GPT model for
|
* Add documents to a ChromaDB collection.
|
||||||
generation. The agent can ingest documents from a local directory,
|
* Query the collection and augment the result with web search snippets.
|
||||||
store their embeddings in ChromaDB, and answer user queries by
|
* Generate an answer using OpenAI's GPT model.
|
||||||
retrieving the most relevant chunks and generating a response. It also
|
|
||||||
provides a lightweight web‑search capability using DuckDuckGo.
|
|
||||||
|
|
||||||
Requirements:
|
Requirements:
|
||||||
- chromadb
|
- openai
|
||||||
- langchain
|
- chromadb
|
||||||
- openai
|
- requests
|
||||||
- python-dotenv (optional, for loading .env files)
|
- beautifulsoup4
|
||||||
|
|
||||||
Author: Artur Kuzakhmetov
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import json
|
import json
|
||||||
|
import textwrap
|
||||||
|
import math
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import List, Dict
|
from typing import List, Dict, Any
|
||||||
|
|
||||||
# Ensure the script can be run from any location
|
import openai
|
||||||
BASE_DIR = Path(__file__).parent.parent.resolve()
|
import chromadb
|
||||||
sys.path.append(str(BASE_DIR))
|
from chromadb import Client
|
||||||
|
from chromadb.config import Settings
|
||||||
# Load environment variables (e.g., OPENAI_API_KEY)
|
import requests
|
||||||
try:
|
from bs4 import BeautifulSoup
|
||||||
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
|
# Configuration
|
||||||
# --------------------------------------------------------------------------- #
|
# --------------------------------------------------------------------------- #
|
||||||
CHROMA_DB_PATH = os.getenv("CHROMA_DB_PATH", str(BASE_DIR / "chromadb"))
|
|
||||||
COLLECTION_NAME = os.getenv("CHROMA_COLLECTION", "rag_collection")
|
# OpenAI API key must be set in the environment
|
||||||
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
|
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||||
LLM_MODEL = os.getenv("LLM_MODEL", "gpt-3.5-turbo")
|
if not OPENAI_API_KEY:
|
||||||
TOP_K = int(os.getenv("TOP_K", "4"))
|
raise RuntimeError("Please set the OPENAI_API_KEY environment variable.")
|
||||||
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "1000"))
|
openai.api_key = OPENAI_API_KEY
|
||||||
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "200"))
|
|
||||||
|
# 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]:
|
|
||||||
"""
|
|
||||||
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 embed_text(text: str) -> List[float]:
|
||||||
def split_text(texts: List[str]) -> List[str]:
|
|
||||||
"""
|
"""
|
||||||
Split a list of texts into smaller chunks suitable for embedding.
|
Generate an embedding vector for the given text using OpenAI embeddings.
|
||||||
"""
|
"""
|
||||||
splitter = RecursiveCharacterTextSplitter(
|
response = openai.Embedding.create(
|
||||||
chunk_size=CHUNK_SIZE,
|
input=text,
|
||||||
chunk_overlap=CHUNK_OVERLAP,
|
model=EMBEDDING_MODEL,
|
||||||
separators=["\n\n", "\n", " ", ""],
|
|
||||||
)
|
)
|
||||||
# LangChain expects a list of dicts with a "content" key
|
return response["data"][0]["embedding"]
|
||||||
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 initialize_vectorstore() -> Chroma:
|
def chunk_text(text: str, max_tokens: int = 500) -> List[str]:
|
||||||
"""
|
"""
|
||||||
Create or connect to a ChromaDB collection and return a Chroma vector store.
|
Split a large text into smaller chunks that fit within the token limit.
|
||||||
"""
|
"""
|
||||||
# Use PersistentClient to store data on disk
|
# Rough token estimation: 1 token ≈ 4 characters
|
||||||
client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
|
max_chars = max_tokens * 4
|
||||||
# Create or get the collection
|
paragraphs = [p.strip() for p in text.split("\n") if p.strip()]
|
||||||
client.get_or_create_collection(name=COLLECTION_NAME)
|
chunks = []
|
||||||
# Wrap with LangChain's Chroma wrapper
|
current = ""
|
||||||
vectorstore = Chroma(
|
for para in paragraphs:
|
||||||
client=client,
|
if len(current) + len(para) + 1 <= max_chars:
|
||||||
collection_name=COLLECTION_NAME,
|
current += (" " if current else "") + para
|
||||||
embedding_function=OpenAIEmbeddings(model=EMBEDDING_MODEL),
|
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."
|
||||||
)
|
)
|
||||||
return vectorstore
|
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}"
|
||||||
|
|
||||||
|
|
||||||
def ingest_documents(folder_path: str, vectorstore: Chroma) -> None:
|
# --------------------------------------------------------------------------- #
|
||||||
|
# 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 doc, meta in zip(results["documents"][0], results["metadatas"][0]):
|
||||||
|
docs.append({"document": doc, "metadata": meta})
|
||||||
|
return docs
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# CLI logic
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
def add_documents_from_file(file_path: str) -> None:
|
||||||
"""
|
"""
|
||||||
Ingest documents from the folder into the vector store.
|
Read a text file, split into chunks, and add to ChromaDB.
|
||||||
"""
|
"""
|
||||||
raw_texts = load_documents_from_folder(folder_path)
|
path = Path(file_path)
|
||||||
if not raw_texts:
|
if not path.is_file():
|
||||||
print(f"No text files found in {folder_path}")
|
print(f"File not found: {file_path}")
|
||||||
return
|
return
|
||||||
|
|
||||||
chunks = split_text(raw_texts)
|
text = path.read_text(encoding="utf-8")
|
||||||
# Add to vector store
|
chunks = chunk_text(text)
|
||||||
vectorstore.add_texts(chunks)
|
documents = [{"text": chunk, "metadata": {"source": file_path}} for chunk in chunks]
|
||||||
print(f"Ingested {len(chunks)} chunks into collection '{COLLECTION_NAME}'.")
|
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
|
Perform a RAG query: retrieve from ChromaDB, augment with web search,
|
||||||
and OpenAI for generation.
|
and generate an answer.
|
||||||
"""
|
"""
|
||||||
llm = OpenAI(model=LLM_MODEL, temperature=0.0)
|
store = ChromaVectorStore()
|
||||||
qa_chain = RetrievalQA.from_chain_type(
|
chroma_results = store.query(query, k=5)
|
||||||
llm=llm,
|
chroma_context = "\n\n".join([doc["document"] for doc in chroma_results])
|
||||||
chain_type="stuff",
|
|
||||||
retriever=vectorstore.as_retriever(search_kwargs={"k": TOP_K}),
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
"""
|
||||||
)
|
)
|
||||||
return qa_chain
|
print(usage)
|
||||||
|
|
||||||
|
|
||||||
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:
|
def main() -> None:
|
||||||
"""
|
if len(sys.argv) < 3:
|
||||||
Main entry point. The script can be used in three modes:
|
print_usage()
|
||||||
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)
|
sys.exit(1)
|
||||||
|
|
||||||
mode = sys.argv[1].lower()
|
command = sys.argv[1].lower()
|
||||||
vectorstore = initialize_vectorstore()
|
if command == "add":
|
||||||
|
file_path = sys.argv[2]
|
||||||
if mode == "ingest":
|
add_documents_from_file(file_path)
|
||||||
if len(sys.argv) != 3:
|
elif command == "ask":
|
||||||
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:])
|
query = " ".join(sys.argv[2:])
|
||||||
results = perform_web_search(query)
|
ask_query(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")
|
|
||||||
|
|
||||||
else:
|
else:
|
||||||
print(f"Unknown mode '{mode}'. Use 'ingest', 'query', or 'search'.")
|
print_usage()
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user