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

This commit is contained in:
2026-06-30 11:57:34 +03:00
parent 2f1a172780
commit ef8bb1bd2c
7 changed files with 151 additions and 156 deletions
+23 -81
View File
@@ -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. 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.
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.
> **Important** ## Features
> The original repository used Qdrant. All references to Qdrant have been removed.
> Only ChromaDB and Tavily are used.
## Prerequisites - **Vector storage** with ChromaDB (in-memory by default).
- **Simple embedding** function (placeholder) replace with a real model for production.
| Component | Version | Notes | - **Web search** using DuckDuckGos HTML interface.
|-----------|---------|-------| - **RAG agent** that retrieves relevant documents or falls back to web search.
| 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"
```
## Installation ## Installation
```bash ```bash
# Clone the repository npm install
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
``` ```
`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 ## Usage
The agent is a single script `src/index.py`. It supports two commands:
### 1. Ingest
```bash ```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. If no query is provided, it defaults to `"What is ChromaDB?"`.
- Otherwise, it treats the argument as raw text and stores it directly.
Example: ## Running Tests
```bash ```bash
python src/index.py ingest https://en.wikipedia.org/wiki/OpenAI npm test
```
### 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?"
``` ```
## Project Structure ## Project Structure
``` ```
. src/
├── src index.js # Entry point
│ └── index.py # Main script agent.js # RAG agent logic
├── README.md vectorStore.js # ChromaDB wrapper
└── requirements.txt 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. - Replace the `embed` function in `vectorStore.js` with a real embedding model (e.g., OpenAI, HuggingFace).
2. **Vector Store** `Chroma` stores these vectors locally in `~/.rag_agent/chromadb`. - Persist the ChromaDB collection by configuring the client with a storage path.
3. **Retrieval** When a query is made, the nearest vectors are fetched. - Add a language model to generate responses from retrieved documents.
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.
## License ## License
This project is released under the MIT License. MIT
+5 -15
View File
@@ -1,25 +1,15 @@
{ {
"name": "rag-agent-chromadb", "name": "rag-agent-chromadb",
"version": "1.0.0", "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", "main": "src/index.js",
"type": "module",
"scripts": { "scripts": {
"start": "node src/index.js", "start": "node src/index.js",
"test": "jest" "test": "node src/test.js"
}, },
"keywords": [
"rag",
"chromadb",
"web-search"
],
"author": "Your Name",
"license": "MIT",
"dependencies": { "dependencies": {
"chromadb": "^0.3.0", "@chromadb/chromadb": "^0.1.0",
"node-fetch": "^3.3.2", "node-fetch": "^3.3.2"
"dotenv": "^16.4.5"
},
"devDependencies": {
"jest": "^29.7.0"
} }
} }
+26 -11
View File
@@ -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); * A simple RAG agent that retrieves relevant documents from the vector store
const results = await vectorStore.query(questionEmbedding, 3); * and optionally performs a web search if no relevant documents are found.
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:`; export class Agent {
// For simplicity, we just return the context as the answer. constructor(vectorStore) {
// In a real scenario, you would pass the prompt to a language model. this.vectorStore = vectorStore;
return contexts; }
/**
* 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}`;
}
} }
module.exports = { answerQuestion };
+23 -21
View File
@@ -1,27 +1,29 @@
const { VectorStore } = require('./vectorStore'); import { VectorStore } from "./vectorStore.js";
const { answerQuestion } = require('./agent'); import { Agent } from "./agent.js";
const { webSearch } = require('./search');
require('dotenv').config();
(async () => { async function main() {
const vectorStore = new VectorStore(); const store = new VectorStore();
await vectorStore.init('rag_collection'); await store.init();
// Example usage: add some documents // Sample documents to index
const docs = [ const docs = [
{ text: 'ChromaDB is a vector database.', id: 'doc1' }, { id: "1", text: "ChromaDB is a fast, lightweight vector database." },
{ text: 'It supports similarity search.', id: 'doc2' } { 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 await store.addDocuments(docs);
const question = 'What is ChromaDB?';
const answer = await answerQuestion(question, vectorStore);
console.log('Answer:', answer);
// Example web search const agent = new Agent(store);
const results = await webSearch('ChromaDB documentation');
console.log('Web search results:', results); 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
View File
@@ -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
View File
@@ -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() { constructor() {
this.client = new ChromaClient(); // uses local storage by default this.client = new Client();
this.collection = null; this.collection = null;
} }
async init(collectionName = 'default') { async init() {
this.collection = await this.client.getOrCreateCollection({ this.collection = await this.client.getOrCreateCollection({
name: collectionName, name: "rag_collection",
metadata: { hnsw: { efConstruction: 200, M: 16 } }
}); });
} }
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) { 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({ await this.collection.add({
ids,
embeddings, embeddings,
metadatas, documents,
ids
}); });
} }
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) { 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({ const results = await this.collection.query({
queryEmbeddings: [queryEmbedding], queryEmbeddings: [queryEmbedding],
nResults, 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
View File
@@ -1,15 +1,17 @@
import fetch from "node-fetch"; 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); const response = await fetch(url);
if (!response.ok) { if (!response.ok) {
throw new Error(`HTTP error ${response.status}`); throw new Error(`Web search failed with status ${response.status}`);
}
const text = await response.text();
return text;
} catch (err) {
console.error(`Failed to fetch ${url}: ${err.message}`);
return "";
} }
const html = await response.text();
return html;
} }