feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
@@ -1,109 +1,51 @@
|
||||
# RAG Agent with ChromaDB & Tavily
|
||||
# RAG Agent with ChromaDB and Web Search
|
||||
|
||||
This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** for vector storage and **Tavily** for web search.
|
||||
The agent can ingest arbitrary text or web pages, store embeddings in a local Chroma collection, and answer questions by retrieving relevant documents and passing them to an OpenAI LLM.
|
||||
This project demonstrates a simple Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector store and performs web search as a fallback. The agent is written in Node.js and uses only the required dependencies.
|
||||
|
||||
> **Important**
|
||||
> The original repository used Qdrant. All references to Qdrant have been removed.
|
||||
> Only ChromaDB and Tavily are used.
|
||||
## Features
|
||||
|
||||
## Prerequisites
|
||||
|
||||
| Component | Version | Notes |
|
||||
|-----------|---------|-------|
|
||||
| Python | 3.9+ | Tested on 3.10 |
|
||||
| OpenAI API | Any key | Required for embeddings and LLM |
|
||||
| Tavily API | Any key | Required for web search |
|
||||
|
||||
Set the following environment variables before running:
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY="your-openai-key"
|
||||
export TAVILY_API_KEY="your-tavily-key"
|
||||
```
|
||||
- **Vector storage** with ChromaDB (in-memory by default).
|
||||
- **Simple embedding** function (placeholder) – replace with a real model for production.
|
||||
- **Web search** using DuckDuckGo’s HTML interface.
|
||||
- **RAG agent** that retrieves relevant documents or falls back to web search.
|
||||
|
||||
## 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 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
|
||||
npm install
|
||||
```
|
||||
|
||||
`requirements.txt` contains:
|
||||
|
||||
```
|
||||
chromadb>=0.4
|
||||
tavily>=0.1
|
||||
langchain>=0.0.350
|
||||
openai>=1.0
|
||||
```
|
||||
|
||||
> **Note**: The exact versions may vary; the above are the minimal compatible versions.
|
||||
|
||||
## Usage
|
||||
|
||||
The agent is a single script `src/index.py`. It supports two commands:
|
||||
|
||||
### 1. Ingest
|
||||
|
||||
```bash
|
||||
python src/index.py ingest <url_or_text>
|
||||
node src/index.js "Your query here"
|
||||
```
|
||||
|
||||
- If `<url_or_text>` starts with `http://` or `https://`, the script treats it as a URL, fetches the content via Tavily, and stores it.
|
||||
- Otherwise, it treats the argument as raw text and stores it directly.
|
||||
If no query is provided, it defaults to `"What is ChromaDB?"`.
|
||||
|
||||
Example:
|
||||
## Running Tests
|
||||
|
||||
```bash
|
||||
python src/index.py ingest https://en.wikipedia.org/wiki/OpenAI
|
||||
```
|
||||
|
||||
### 2. Query
|
||||
|
||||
```bash
|
||||
python src/index.py query "<your question>"
|
||||
```
|
||||
|
||||
The script retrieves relevant documents from the Chroma collection and asks OpenAI to generate an answer.
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
python src/index.py query "What is OpenAI?"
|
||||
npm test
|
||||
```
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
.
|
||||
├── src
|
||||
│ └── index.py # Main script
|
||||
├── README.md
|
||||
└── requirements.txt
|
||||
src/
|
||||
index.js # Entry point
|
||||
agent.js # RAG agent logic
|
||||
vectorStore.js # ChromaDB wrapper
|
||||
webSearch.js # Simple web search helper
|
||||
test.js # Basic test for vector store
|
||||
```
|
||||
|
||||
## How It Works
|
||||
## Extending
|
||||
|
||||
1. **Embedding** – The script uses `OpenAIEmbeddings` from LangChain to convert text into vectors.
|
||||
2. **Vector Store** – `Chroma` stores these vectors locally in `~/.rag_agent/chromadb`.
|
||||
3. **Retrieval** – When a query is made, the nearest vectors are fetched.
|
||||
4. **Generation** – The retrieved documents are fed into an OpenAI LLM to produce a final answer.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **No results from Tavily** – Ensure your Tavily API key is valid and that the URL is reachable.
|
||||
- **OpenAI errors** – Check that your OpenAI key has the necessary permissions and quota.
|
||||
- **Chroma storage issues** – The data directory is `~/.rag_agent/chromadb`. Delete it to reset the collection.
|
||||
- Replace the `embed` function in `vectorStore.js` with a real embedding model (e.g., OpenAI, HuggingFace).
|
||||
- Persist the ChromaDB collection by configuring the client with a storage path.
|
||||
- Add a language model to generate responses from retrieved documents.
|
||||
|
||||
## License
|
||||
|
||||
This project is released under the MIT License.
|
||||
MIT
|
||||
+5
-15
@@ -1,25 +1,15 @@
|
||||
{
|
||||
"name": "rag-agent-chromadb",
|
||||
"version": "1.0.0",
|
||||
"description": "RAG agent using ChromaDB and web search",
|
||||
"description": "A simple RAG agent using ChromaDB for vector storage and web search.",
|
||||
"main": "src/index.js",
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"start": "node src/index.js",
|
||||
"test": "jest"
|
||||
"test": "node src/test.js"
|
||||
},
|
||||
"keywords": [
|
||||
"rag",
|
||||
"chromadb",
|
||||
"web-search"
|
||||
],
|
||||
"author": "Your Name",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"chromadb": "^0.3.0",
|
||||
"node-fetch": "^3.3.2",
|
||||
"dotenv": "^16.4.5"
|
||||
},
|
||||
"devDependencies": {
|
||||
"jest": "^29.7.0"
|
||||
"@chromadb/chromadb": "^0.1.0",
|
||||
"node-fetch": "^3.3.2"
|
||||
}
|
||||
}
|
||||
+25
-10
@@ -1,13 +1,28 @@
|
||||
const { embed } = require('./utils');
|
||||
import { VectorStore } from "./vectorStore.js";
|
||||
import { webSearch } from "./webSearch.js";
|
||||
|
||||
async function answerQuestion(question, vectorStore) {
|
||||
const questionEmbedding = embed(question);
|
||||
const results = await vectorStore.query(questionEmbedding, 3);
|
||||
const contexts = results[0].metadatas.map(m => m.text).join('\n');
|
||||
const prompt = `Answer the question based on the following context:\n\n${contexts}\n\nQuestion: ${question}\nAnswer:`;
|
||||
// For simplicity, we just return the context as the answer.
|
||||
// In a real scenario, you would pass the prompt to a language model.
|
||||
return contexts;
|
||||
/**
|
||||
* A simple RAG agent that retrieves relevant documents from the vector store
|
||||
* and optionally performs a web search if no relevant documents are found.
|
||||
*/
|
||||
export class Agent {
|
||||
constructor(vectorStore) {
|
||||
this.vectorStore = vectorStore;
|
||||
}
|
||||
|
||||
module.exports = { answerQuestion };
|
||||
/**
|
||||
* Processes a user query and returns the best answer.
|
||||
* @param {string} query
|
||||
* @returns {Promise<string>}
|
||||
*/
|
||||
async answer(query) {
|
||||
const results = await this.vectorStore.query(query, 3);
|
||||
if (results.length > 0 && results[0].score < 0.5) {
|
||||
// Return the most relevant document text
|
||||
return results[0].text;
|
||||
}
|
||||
// Fallback to web search
|
||||
const html = await webSearch(query);
|
||||
return `No relevant local documents found. Here is the raw web search result:\n${html}`;
|
||||
}
|
||||
}
|
||||
+23
-21
@@ -1,27 +1,29 @@
|
||||
const { VectorStore } = require('./vectorStore');
|
||||
const { answerQuestion } = require('./agent');
|
||||
const { webSearch } = require('./search');
|
||||
require('dotenv').config();
|
||||
import { VectorStore } from "./vectorStore.js";
|
||||
import { Agent } from "./agent.js";
|
||||
|
||||
(async () => {
|
||||
const vectorStore = new VectorStore();
|
||||
await vectorStore.init('rag_collection');
|
||||
async function main() {
|
||||
const store = new VectorStore();
|
||||
await store.init();
|
||||
|
||||
// Example usage: add some documents
|
||||
// Sample documents to index
|
||||
const docs = [
|
||||
{ text: 'ChromaDB is a vector database.', id: 'doc1' },
|
||||
{ text: 'It supports similarity search.', id: 'doc2' }
|
||||
{ id: "1", text: "ChromaDB is a fast, lightweight vector database." },
|
||||
{ id: "2", text: "It supports in-memory and persistent storage." },
|
||||
{ id: "3", text: "You can use it with various embedding models." },
|
||||
];
|
||||
const embeddings = docs.map(d => require('./utils').embed(d.text));
|
||||
const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
|
||||
await vectorStore.add(embeddings, metadatas, docs.map(d => d.id));
|
||||
|
||||
// Example question
|
||||
const question = 'What is ChromaDB?';
|
||||
const answer = await answerQuestion(question, vectorStore);
|
||||
console.log('Answer:', answer);
|
||||
await store.addDocuments(docs);
|
||||
|
||||
// Example web search
|
||||
const results = await webSearch('ChromaDB documentation');
|
||||
console.log('Web search results:', results);
|
||||
})();
|
||||
const agent = new Agent(store);
|
||||
|
||||
const query = process.argv[2] || "What is ChromaDB?";
|
||||
console.log(`Query: ${query}`);
|
||||
const answer = await agent.answer(query);
|
||||
console.log("\nAnswer:");
|
||||
console.log(answer);
|
||||
}
|
||||
|
||||
main().catch((err) => {
|
||||
console.error(err);
|
||||
process.exit(1);
|
||||
});
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
import { VectorStore } from "./vectorStore.js";
|
||||
|
||||
async function testVectorStore() {
|
||||
const store = new VectorStore();
|
||||
await store.init();
|
||||
const docs = [
|
||||
{ id: "a", text: "Hello world" },
|
||||
{ id: "b", text: "Goodbye world" },
|
||||
];
|
||||
await store.addDocuments(docs);
|
||||
const results = await store.query("Hello", 2);
|
||||
console.log("Test results:", results);
|
||||
}
|
||||
|
||||
testVectorStore().catch((err) => console.error(err));
|
||||
+45
-16
@@ -1,40 +1,69 @@
|
||||
const { ChromaClient } = require('chromadb');
|
||||
import { Client } from "@chromadb/chromadb";
|
||||
|
||||
class VectorStore {
|
||||
/**
|
||||
* Simple embedding function that converts text into a fixed-length numeric vector.
|
||||
* This is a placeholder and should be replaced with a real embedding model for production use.
|
||||
*/
|
||||
function embed(text) {
|
||||
const vector = Array.from(text)
|
||||
.map((c) => c.charCodeAt(0))
|
||||
.slice(0, 10);
|
||||
while (vector.length < 10) {
|
||||
vector.push(0);
|
||||
}
|
||||
return vector;
|
||||
}
|
||||
|
||||
export class VectorStore {
|
||||
constructor() {
|
||||
this.client = new ChromaClient(); // uses local storage by default
|
||||
this.client = new Client();
|
||||
this.collection = null;
|
||||
}
|
||||
|
||||
async init(collectionName = 'default') {
|
||||
async init() {
|
||||
this.collection = await this.client.getOrCreateCollection({
|
||||
name: collectionName,
|
||||
metadata: { hnsw: { efConstruction: 200, M: 16 } }
|
||||
name: "rag_collection",
|
||||
});
|
||||
}
|
||||
|
||||
async add(embeddings, metadatas, ids) {
|
||||
/**
|
||||
* Adds an array of documents to the collection.
|
||||
* @param {Array<{id: string, text: string}>} docs
|
||||
*/
|
||||
async addDocuments(docs) {
|
||||
if (!this.collection) {
|
||||
throw new Error('Collection not initialized. Call init() first.');
|
||||
throw new Error("VectorStore not initialized. Call init() first.");
|
||||
}
|
||||
const ids = docs.map((d) => d.id);
|
||||
const embeddings = docs.map((d) => embed(d.text));
|
||||
const documents = docs.map((d) => d.text);
|
||||
await this.collection.add({
|
||||
ids,
|
||||
embeddings,
|
||||
metadatas,
|
||||
ids
|
||||
documents,
|
||||
});
|
||||
}
|
||||
|
||||
async query(queryEmbedding, nResults = 5) {
|
||||
/**
|
||||
* Queries the collection for the most relevant documents.
|
||||
* @param {string} queryText
|
||||
* @param {number} nResults
|
||||
* @returns {Promise<Array<{id: string, text: string, score: number}>>}
|
||||
*/
|
||||
async query(queryText, nResults = 3) {
|
||||
if (!this.collection) {
|
||||
throw new Error('Collection not initialized. Call init() first.');
|
||||
throw new Error("VectorStore not initialized. Call init() first.");
|
||||
}
|
||||
const queryEmbedding = embed(queryText);
|
||||
const results = await this.collection.query({
|
||||
queryEmbeddings: [queryEmbedding],
|
||||
nResults,
|
||||
include: ['metadatas', 'documents']
|
||||
});
|
||||
return results;
|
||||
// results is an array of objects with ids, documents, and scores
|
||||
return results[0].ids.map((id, idx) => ({
|
||||
id,
|
||||
text: results[0].documents[idx],
|
||||
score: results[0].distances[idx],
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = { VectorStore };
|
||||
+11
-9
@@ -1,15 +1,17 @@
|
||||
import fetch from "node-fetch";
|
||||
|
||||
export async function fetchWebContent(url) {
|
||||
try {
|
||||
/**
|
||||
* Performs a simple web search using DuckDuckGo's HTML interface.
|
||||
* This is a lightweight example and does not use an official API.
|
||||
* @param {string} query
|
||||
* @returns {Promise<string>} The raw HTML of the search results page.
|
||||
*/
|
||||
export async function webSearch(query) {
|
||||
const url = `https://duckduckgo.com/html/?q=${encodeURIComponent(query)}`;
|
||||
const response = await fetch(url);
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP error ${response.status}`);
|
||||
}
|
||||
const text = await response.text();
|
||||
return text;
|
||||
} catch (err) {
|
||||
console.error(`Failed to fetch ${url}: ${err.message}`);
|
||||
return "";
|
||||
throw new Error(`Web search failed with status ${response.status}`);
|
||||
}
|
||||
const html = await response.text();
|
||||
return html;
|
||||
}
|
||||
Reference in New Issue
Block a user