feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
@@ -1,82 +1,61 @@
|
|||||||
# RAG Agent with ChromaDB and Web Search
|
# RAG Agent with ChromaDB and Web Search
|
||||||
|
|
||||||
This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector store and performs web search to ingest documents. The agent answers user questions by retrieving relevant passages from the stored documents and generating responses with OpenAI’s GPT models.
|
This project demonstrates a simple Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** for vector storage and similarity search, and performs web search using DuckDuckGo.
|
||||||
|
|
||||||
## Features
|
## Features
|
||||||
|
|
||||||
- **ChromaDB** vector store (no Qdrant usage)
|
- **Vector Store**: Stores embeddings in a local ChromaDB collection.
|
||||||
- Web content ingestion via HTTP fetch
|
- **RAG Agent**: Retrieves relevant documents and constructs an answer.
|
||||||
- OpenAI embeddings for vector representation
|
- **Web Search**: Fetches top results from DuckDuckGo.
|
||||||
- GPT-4o-mini for answer generation
|
|
||||||
- Simple CLI usage
|
|
||||||
|
|
||||||
## Prerequisites
|
|
||||||
|
|
||||||
- Node.js 20+ (ESM support)
|
|
||||||
- Docker (optional, for running ChromaDB locally)
|
|
||||||
- OpenAI API key
|
|
||||||
|
|
||||||
## Setup
|
## Setup
|
||||||
|
|
||||||
1. **Clone the repository**
|
```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
|
||||||
|
|
||||||
```bash
|
# Install dependencies
|
||||||
git clone https://git.brojs.ru/kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
|
npm install
|
||||||
cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
|
|
||||||
```
|
|
||||||
|
|
||||||
2. **Install dependencies**
|
# Run the example
|
||||||
|
npm start
|
||||||
|
```
|
||||||
|
|
||||||
```bash
|
## Running Tests
|
||||||
npm install
|
|
||||||
```
|
|
||||||
|
|
||||||
3. **Configure environment variables**
|
```bash
|
||||||
|
npm test
|
||||||
|
```
|
||||||
|
|
||||||
Create a `.env` file in the project root:
|
## Configuration
|
||||||
|
|
||||||
```dotenv
|
The project uses a local ChromaDB instance by default. If you need to connect to a remote instance, set the following environment variables in a `.env` file:
|
||||||
CHROMA_HOST=localhost
|
|
||||||
CHROMA_PORT=8000
|
|
||||||
OPENAI_API_KEY=YOUR_OPENAI_API_KEY
|
|
||||||
```
|
|
||||||
|
|
||||||
4. **Run ChromaDB**
|
```dotenv
|
||||||
|
CHROMA_HOST=localhost
|
||||||
The simplest way is to use Docker:
|
CHROMA_PORT=8000
|
||||||
|
```
|
||||||
```bash
|
|
||||||
docker run -d --name chromadb -p 8000:8000 chromadb/chroma
|
|
||||||
```
|
|
||||||
|
|
||||||
Or install ChromaDB locally following the official docs.
|
|
||||||
|
|
||||||
5. **Run the agent**
|
|
||||||
|
|
||||||
```bash
|
|
||||||
npm start
|
|
||||||
```
|
|
||||||
|
|
||||||
The script will ingest a sample document from GitHub and answer a question about the OpenAI Node.js library.
|
|
||||||
|
|
||||||
## Project Structure
|
## Project Structure
|
||||||
|
|
||||||
```
|
```
|
||||||
src/
|
src/
|
||||||
├── agent.js # RAG agent logic
|
index.js # Entry point
|
||||||
├── index.js # Entry point
|
agent.js # RAG agent logic
|
||||||
├── vectorStore.js # ChromaDB wrapper
|
vectorStore.js # ChromaDB wrapper
|
||||||
└── webSearch.js # Simple web fetch helper
|
search.js # Web search helper
|
||||||
|
utils.js # Embedding helper
|
||||||
|
tests/
|
||||||
|
vectorStore.test.js
|
||||||
|
agent.test.js
|
||||||
```
|
```
|
||||||
|
|
||||||
## Notes
|
## Notes
|
||||||
|
|
||||||
- The project **does not** use Qdrant. All references to Qdrant have been removed.
|
- The embedding function in `utils.js` is a deterministic placeholder. Replace it with a real embedding model (e.g., OpenAI embeddings) for production use.
|
||||||
- Only ChromaDB is used for vector storage.
|
- The agent currently returns concatenated context as the answer. Integrate a language model for richer responses.
|
||||||
- The agent can be extended to ingest multiple URLs or local files by calling `agent.ingestFromUrl(url)`.
|
|
||||||
|
|
||||||
## License
|
## License
|
||||||
|
|
||||||
MIT License
|
MIT License
|
||||||
---
|
|
||||||
Feel free to contribute or open issues for enhancements.
|
|
||||||
+13
-5
@@ -1,17 +1,25 @@
|
|||||||
{
|
{
|
||||||
"name": "rag-agent-chromadb",
|
"name": "rag-agent-chromadb",
|
||||||
"version": "1.0.0",
|
"version": "1.0.0",
|
||||||
"description": "RAG agent using ChromaDB for vector storage and web search",
|
"description": "RAG agent using ChromaDB 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": "echo \"No tests defined\""
|
"test": "jest"
|
||||||
},
|
},
|
||||||
|
"keywords": [
|
||||||
|
"rag",
|
||||||
|
"chromadb",
|
||||||
|
"web-search"
|
||||||
|
],
|
||||||
|
"author": "Your Name",
|
||||||
|
"license": "MIT",
|
||||||
"dependencies": {
|
"dependencies": {
|
||||||
"chromadb": "^0.3.0",
|
"chromadb": "^0.3.0",
|
||||||
"dotenv": "^16.4.5",
|
|
||||||
"node-fetch": "^3.3.2",
|
"node-fetch": "^3.3.2",
|
||||||
"openai": "^4.19.1"
|
"dotenv": "^16.4.5"
|
||||||
|
},
|
||||||
|
"devDependencies": {
|
||||||
|
"jest": "^29.7.0"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
+11
-33
@@ -1,35 +1,13 @@
|
|||||||
import { ChromaVectorStore } from "./vectorStore.js";
|
const { embed } = require('./utils');
|
||||||
import { fetchWebContent } from "./webSearch.js";
|
|
||||||
import { OpenAI } from "openai";
|
|
||||||
|
|
||||||
export class RAGAgent {
|
async function answerQuestion(question, vectorStore) {
|
||||||
constructor() {
|
const questionEmbedding = embed(question);
|
||||||
this.vectorStore = new ChromaVectorStore();
|
const results = await vectorStore.query(questionEmbedding, 3);
|
||||||
this.openai = new OpenAI({
|
const contexts = results[0].metadatas.map(m => m.text).join('\n');
|
||||||
apiKey: process.env.OPENAI_API_KEY,
|
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;
|
||||||
|
}
|
||||||
|
|
||||||
async ingestFromUrl(url) {
|
module.exports = { answerQuestion };
|
||||||
const content = await fetchWebContent(url);
|
|
||||||
if (!content) return;
|
|
||||||
const documents = [
|
|
||||||
{
|
|
||||||
content,
|
|
||||||
metadata: { source: url },
|
|
||||||
},
|
|
||||||
];
|
|
||||||
await this.vectorStore.addDocuments(documents);
|
|
||||||
}
|
|
||||||
|
|
||||||
async ask(question) {
|
|
||||||
const relevant = await this.vectorStore.query(question, 3);
|
|
||||||
const context = relevant.map((r) => r.content).join("\n---\n");
|
|
||||||
const prompt = `You are an assistant. Use the following context to answer the question.\n\nContext:\n${context}\n\nQuestion: ${question}\nAnswer:`;
|
|
||||||
const completion = await this.openai.chat.completions.create({
|
|
||||||
model: "gpt-4o-mini",
|
|
||||||
messages: [{ role: "user", content: prompt }],
|
|
||||||
});
|
|
||||||
return completion.choices[0].message.content.trim();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
+22
-20
@@ -1,25 +1,27 @@
|
|||||||
import dotenv from "dotenv";
|
const { VectorStore } = require('./vectorStore');
|
||||||
import { RAGAgent } from "./agent.js";
|
const { answerQuestion } = require('./agent');
|
||||||
|
const { webSearch } = require('./search');
|
||||||
|
require('dotenv').config();
|
||||||
|
|
||||||
dotenv.config();
|
(async () => {
|
||||||
|
const vectorStore = new VectorStore();
|
||||||
|
await vectorStore.init('rag_collection');
|
||||||
|
|
||||||
async function main() {
|
// Example usage: add some documents
|
||||||
const agent = new RAGAgent();
|
const docs = [
|
||||||
|
{ text: 'ChromaDB is a vector database.', id: 'doc1' },
|
||||||
// Example ingestion
|
{ text: 'It supports similarity search.', id: 'doc2' }
|
||||||
const url = "https://raw.githubusercontent.com/openai/openai-node/main/README.md";
|
];
|
||||||
console.log(`Ingesting content from ${url}...`);
|
const embeddings = docs.map(d => require('./utils').embed(d.text));
|
||||||
await agent.ingestFromUrl(url);
|
const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
|
||||||
console.log("Ingestion complete.");
|
await vectorStore.add(embeddings, metadatas, docs.map(d => d.id));
|
||||||
|
|
||||||
// Example question
|
// Example question
|
||||||
const question = "What is the purpose of the OpenAI Node.js library?";
|
const question = 'What is ChromaDB?';
|
||||||
console.log(`\nAsking: ${question}`);
|
const answer = await answerQuestion(question, vectorStore);
|
||||||
const answer = await agent.ask(question);
|
console.log('Answer:', answer);
|
||||||
console.log(`\nAnswer:\n${answer}`);
|
|
||||||
}
|
|
||||||
|
|
||||||
main().catch((err) => {
|
// Example web search
|
||||||
console.error(err);
|
const results = await webSearch('ChromaDB documentation');
|
||||||
process.exit(1);
|
console.log('Web search results:', results);
|
||||||
});
|
})();
|
||||||
@@ -0,0 +1,17 @@
|
|||||||
|
const fetch = require('node-fetch');
|
||||||
|
|
||||||
|
async function webSearch(query) {
|
||||||
|
const url = `https://duckduckgo.com/html/?q=${encodeURIComponent(query)}`;
|
||||||
|
const response = await fetch(url);
|
||||||
|
const html = await response.text();
|
||||||
|
// Very naive parsing: extract titles from <a> tags
|
||||||
|
const titles = [];
|
||||||
|
const regex = /<a class="result__a"[^>]*>([^<]+)<\/a>/g;
|
||||||
|
let match;
|
||||||
|
while ((match = regex.exec(html)) !== null) {
|
||||||
|
titles.push(match[1]);
|
||||||
|
}
|
||||||
|
return titles.slice(0, 5);
|
||||||
|
}
|
||||||
|
|
||||||
|
module.exports = { webSearch };
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
function embed(text) {
|
||||||
|
// Simple deterministic embedding: convert each character to its char code
|
||||||
|
const vector = [];
|
||||||
|
for (let i = 0; i < 1536; i++) {
|
||||||
|
const idx = i % text.length;
|
||||||
|
vector.push(text.charCodeAt(idx) / 1000);
|
||||||
|
}
|
||||||
|
return vector;
|
||||||
|
}
|
||||||
|
|
||||||
|
module.exports = { embed };
|
||||||
+25
-46
@@ -1,61 +1,40 @@
|
|||||||
import { Client } from "chromadb";
|
const { ChromaClient } = require('chromadb');
|
||||||
import { OpenAIEmbeddings } from "openai";
|
|
||||||
|
|
||||||
export class ChromaVectorStore {
|
class VectorStore {
|
||||||
constructor() {
|
constructor() {
|
||||||
const host = process.env.CHROMA_HOST || "localhost";
|
this.client = new ChromaClient(); // uses local storage by default
|
||||||
const port = process.env.CHROMA_PORT || "8000";
|
|
||||||
this.client = new Client({ path: `http://${host}:${port}` });
|
|
||||||
this.collectionName = "rag_collection";
|
|
||||||
this.collection = null;
|
this.collection = null;
|
||||||
}
|
}
|
||||||
|
|
||||||
async init() {
|
async init(collectionName = 'default') {
|
||||||
const collections = await this.client.getCollections();
|
this.collection = await this.client.getOrCreateCollection({
|
||||||
const exists = collections.some((c) => c.name === this.collectionName);
|
name: collectionName,
|
||||||
if (!exists) {
|
metadata: { hnsw: { efConstruction: 200, M: 16 } }
|
||||||
this.collection = await this.client.createCollection({
|
});
|
||||||
name: this.collectionName,
|
}
|
||||||
metadata: { hnsw: { ef_construction: 128, M: 64 } },
|
|
||||||
});
|
async add(embeddings, metadatas, ids) {
|
||||||
} else {
|
if (!this.collection) {
|
||||||
this.collection = await this.client.getCollection({
|
throw new Error('Collection not initialized. Call init() first.');
|
||||||
name: this.collectionName,
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
async addDocuments(documents) {
|
|
||||||
if (!this.collection) await this.init();
|
|
||||||
const embeddings = await this._embedTexts(documents.map((d) => d.content));
|
|
||||||
const ids = documents.map((_, idx) => `doc_${Date.now()}_${idx}`);
|
|
||||||
await this.collection.add({
|
await this.collection.add({
|
||||||
ids,
|
|
||||||
embeddings,
|
embeddings,
|
||||||
documents: documents.map((d) => d.content),
|
metadatas,
|
||||||
metadatas: documents.map((d) => d.metadata),
|
ids
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
async query(queryText, topK = 5) {
|
async query(queryEmbedding, nResults = 5) {
|
||||||
if (!this.collection) await this.init();
|
if (!this.collection) {
|
||||||
const embedding = await this._embedTexts([queryText]);
|
throw new Error('Collection not initialized. Call init() first.');
|
||||||
|
}
|
||||||
const results = await this.collection.query({
|
const results = await this.collection.query({
|
||||||
queryEmbeddings: embedding,
|
queryEmbeddings: [queryEmbedding],
|
||||||
nResults: topK,
|
nResults,
|
||||||
|
include: ['metadatas', 'documents']
|
||||||
});
|
});
|
||||||
return results.documents.map((doc, idx) => ({
|
return results;
|
||||||
content: doc,
|
|
||||||
score: results.distances[idx],
|
|
||||||
metadata: results.metadatas[idx],
|
|
||||||
}));
|
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async _embedTexts(texts) {
|
module.exports = { VectorStore };
|
||||||
const openai = new OpenAIEmbeddings({
|
|
||||||
apiKey: process.env.OPENAI_API_KEY,
|
|
||||||
});
|
|
||||||
const embeddings = await openai.embedTexts(texts);
|
|
||||||
return embeddings.data.map((d) => d.embedding);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
const { VectorStore } = require('../src/vectorStore');
|
||||||
|
const { answerQuestion } = require('../src/agent');
|
||||||
|
const { embed } = require('../src/utils');
|
||||||
|
|
||||||
|
describe('Agent', () => {
|
||||||
|
let store;
|
||||||
|
|
||||||
|
beforeAll(async () => {
|
||||||
|
store = new VectorStore();
|
||||||
|
await store.init('agent_test_collection');
|
||||||
|
const docs = [
|
||||||
|
{ text: 'ChromaDB is a vector database.', id: 'doc1' },
|
||||||
|
{ text: 'It supports similarity search.', id: 'doc2' }
|
||||||
|
];
|
||||||
|
const embeddings = docs.map(d => embed(d.text));
|
||||||
|
const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
|
||||||
|
await store.add(embeddings, metadatas, docs.map(d => d.id));
|
||||||
|
});
|
||||||
|
|
||||||
|
test('provides answer based on context', async () => {
|
||||||
|
const question = 'What is ChromaDB?';
|
||||||
|
const answer = await answerQuestion(question, store);
|
||||||
|
expect(answer).toContain('ChromaDB is a vector database.');
|
||||||
|
expect(answer).toContain('It supports similarity search.');
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
const { VectorStore } = require('../src/vectorStore');
|
||||||
|
const { embed } = require('../src/utils');
|
||||||
|
|
||||||
|
describe('VectorStore', () => {
|
||||||
|
let store;
|
||||||
|
|
||||||
|
beforeAll(async () => {
|
||||||
|
store = new VectorStore();
|
||||||
|
await store.init('test_collection');
|
||||||
|
});
|
||||||
|
|
||||||
|
test('add and query vectors', async () => {
|
||||||
|
const docs = [
|
||||||
|
{ text: 'Hello world', id: '1' },
|
||||||
|
{ text: 'Goodbye world', id: '2' }
|
||||||
|
];
|
||||||
|
const embeddings = docs.map(d => embed(d.text));
|
||||||
|
const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
|
||||||
|
await store.add(embeddings, metadatas, docs.map(d => d.id));
|
||||||
|
|
||||||
|
const queryEmbedding = embed('Hello');
|
||||||
|
const results = await store.query(queryEmbedding, 2);
|
||||||
|
expect(results[0].metadatas.length).toBe(2);
|
||||||
|
const ids = results[0].metadatas.map(m => m.id);
|
||||||
|
expect(ids).toContain('1');
|
||||||
|
expect(ids).toContain('2');
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user