From ef8bb1bd2c090958b88411f021fd0d5a780a59be Mon Sep 17 00:00:00 2001 From: kuzakhmetovartur Date: Tue, 30 Jun 2026 11:57:34 +0300 Subject: [PATCH] =?UTF-8?q?feat:=20solution=20for=20'=D0=AD=D0=BA=D0=B7?= =?UTF-8?q?=D0=B0=D0=BC=D0=B5=D0=BD:=20RAG-=D0=B0=D0=B3=D0=B5=D0=BD=D1=82?= =?UTF-8?q?=20=D1=81=20ChromaDB=20=D0=B8=20=D0=B2=D0=B5=D0=B1-=D0=BF=D0=BE?= =?UTF-8?q?=D0=B8=D1=81=D0=BA=D0=BE=D0=BC'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 104 ++++++++++----------------------------------- package.json | 20 +++------ src/agent.js | 37 +++++++++++----- src/index.js | 44 ++++++++++--------- src/test.js | 15 +++++++ src/vectorStore.js | 63 +++++++++++++++++++-------- src/webSearch.js | 24 ++++++----- 7 files changed, 151 insertions(+), 156 deletions(-) create mode 100644 src/test.js diff --git a/README.md b/README.md index 5d5b2f9..94936b1 100644 --- a/README.md +++ b/README.md @@ -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 +node src/index.js "Your query here" ``` -- If `` 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 "" -``` - -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. \ No newline at end of file +MIT \ No newline at end of file diff --git a/package.json b/package.json index e4e9442..7dfa5d2 100644 --- a/package.json +++ b/package.json @@ -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" } } \ No newline at end of file diff --git a/src/agent.js b/src/agent.js index 9d48910..0e01f42 100644 --- a/src/agent.js +++ b/src/agent.js @@ -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 }; \ No newline at end of file + /** + * Processes a user query and returns the best answer. + * @param {string} query + * @returns {Promise} + */ + 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}`; + } +} \ No newline at end of file diff --git a/src/index.js b/src/index.js index e04dc46..73e5ef8 100644 --- a/src/index.js +++ b/src/index.js @@ -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); -})(); \ No newline at end of file + 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); +}); \ No newline at end of file diff --git a/src/test.js b/src/test.js new file mode 100644 index 0000000..b66ff71 --- /dev/null +++ b/src/test.js @@ -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)); \ No newline at end of file diff --git a/src/vectorStore.js b/src/vectorStore.js index 8f29924..a2e4c6a 100644 --- a/src/vectorStore.js +++ b/src/vectorStore.js @@ -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>} + */ + 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 }; \ No newline at end of file +} \ No newline at end of file diff --git a/src/webSearch.js b/src/webSearch.js index 90f5b1d..e90bf6a 100644 --- a/src/webSearch.js +++ b/src/webSearch.js @@ -1,15 +1,17 @@ import fetch from "node-fetch"; -export async function fetchWebContent(url) { - try { - 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 ""; +/** + * 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} 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(`Web search failed with status ${response.status}`); } + const html = await response.text(); + return html; } \ No newline at end of file