feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
@@ -1,24 +1,22 @@
|
|||||||
# RAG Agent with ChromaDB and Web Search
|
# RAG Agent with ChromaDB and Web Search
|
||||||
|
|
||||||
This project implements a Retrieval-Augmented Generation (RAG) agent that:
|
This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector database and the **OpenAI API** to generate responses based on retrieved documents. It also includes a simple web‑search component that fetches content from specified URLs for indexing.
|
||||||
- Stores and retrieves embeddings from **ChromaDB**.
|
|
||||||
- Performs web search using DuckDuckGo to fetch additional context.
|
## Features
|
||||||
- Generates answers with an **Ollama** language model.
|
|
||||||
|
- **Vector Store**: Uses ChromaDB to store embeddings of text chunks.
|
||||||
|
- **OpenAI Integration**: Generates answers using GPT‑3.5‑Turbo.
|
||||||
|
- **Web Search**: Fetches and parses HTML pages, splits them into manageable chunks.
|
||||||
|
- **Command Line Interface**: Ask questions interactively.
|
||||||
|
|
||||||
## Prerequisites
|
## Prerequisites
|
||||||
|
|
||||||
- Node.js v20 or newer
|
- Node.js v18+ (supports native ES modules and `node-fetch` v2).
|
||||||
- ChromaDB server running locally (default URL: `chromadb://localhost:8000`)
|
- An OpenAI API key.
|
||||||
- Ollama server running locally (default URL: `http://localhost:11434`)
|
|
||||||
|
|
||||||
## Setup
|
## Setup
|
||||||
|
|
||||||
1. **Clone the repository**
|
1. **Clone the repository** (or copy the files into a directory).
|
||||||
|
|
||||||
```bash
|
|
||||||
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
|
|
||||||
```
|
|
||||||
|
|
||||||
2. **Install dependencies**
|
2. **Install dependencies**
|
||||||
|
|
||||||
@@ -26,62 +24,39 @@ This project implements a Retrieval-Augmented Generation (RAG) agent that:
|
|||||||
npm install
|
npm install
|
||||||
```
|
```
|
||||||
|
|
||||||
3. **Configure environment variables**
|
3. **Configure environment**
|
||||||
|
|
||||||
Create a `.env` file in the project root (or modify the existing one):
|
Create a `.env` file in the project root (or edit the existing one) and add your OpenAI API key:
|
||||||
|
|
||||||
```dotenv
|
```dotenv
|
||||||
CHROMA_URL=chromadb://localhost:8000
|
OPENAI_API_KEY=your_api_key_here
|
||||||
CHROMA_COLLECTION=rag_collection
|
|
||||||
OLLAMA_HOST=http://localhost:11434
|
|
||||||
OLLAMA_MODEL=llama3
|
|
||||||
OLLAMA_EMBEDDING_MODEL=nomic-embed-text
|
|
||||||
ADD_SAMPLE_DOCS=true
|
|
||||||
```
|
```
|
||||||
|
|
||||||
- `CHROMA_URL`: URL of your ChromaDB instance.
|
|
||||||
- `CHROMA_COLLECTION`: Name of the collection to use.
|
|
||||||
- `OLLAMA_HOST`: URL of your Ollama server.
|
|
||||||
- `OLLAMA_MODEL`: Ollama model for generation.
|
|
||||||
- `OLLAMA_EMBEDDING_MODEL`: Ollama model for embeddings.
|
|
||||||
- `ADD_SAMPLE_DOCS`: Set to `true` to automatically add a few sample documents on startup.
|
|
||||||
|
|
||||||
4. **Run the agent**
|
4. **Run the agent**
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
npm start -- "Your question here"
|
npm start
|
||||||
```
|
```
|
||||||
|
|
||||||
Example:
|
The script will:
|
||||||
|
- Fetch and index the example URLs.
|
||||||
|
- Prompt you to enter questions.
|
||||||
|
- Display answers generated by the RAG agent.
|
||||||
|
|
||||||
```bash
|
## Customization
|
||||||
npm start -- "What is LangChain?"
|
|
||||||
```
|
|
||||||
|
|
||||||
The agent will:
|
- **Adding URLs**: Edit the `urls` array in `src/index.js` to index different web pages.
|
||||||
- Search the local ChromaDB collection.
|
- **Chunk Size**: Adjust the `size` parameter in `chunkText` inside `src/webSearch.js` if you need larger or smaller chunks.
|
||||||
- Perform a DuckDuckGo web search.
|
- **Model Parameters**: Modify temperature, max tokens, or model name in `src/agent.js`.
|
||||||
- Combine the results and generate an answer using Ollama.
|
|
||||||
|
|
||||||
## Project Structure
|
|
||||||
|
|
||||||
```
|
|
||||||
.
|
|
||||||
├── src
|
|
||||||
│ ├── agent.js # Agent logic (retrieval + generation)
|
|
||||||
│ ├── index.js # CLI entry point
|
|
||||||
│ ├── vectorStore.js # ChromaDB wrapper
|
|
||||||
│ └── webSearch.js # DuckDuckGo search helper
|
|
||||||
├── .env # Environment configuration
|
|
||||||
├── package.json # Dependencies and scripts
|
|
||||||
└── README.md # Documentation
|
|
||||||
```
|
|
||||||
|
|
||||||
## Notes
|
## Notes
|
||||||
|
|
||||||
- The agent uses **LangChain 1.x** APIs.
|
- The implementation strictly uses **ChromaDB** as the vector database; no other vector DBs are used.
|
||||||
- No Qdrant references are present; only ChromaDB is used.
|
- All dependencies are declared in `package.json` and can be installed via `npm install`.
|
||||||
- The web search is performed via DuckDuckGo’s public JSON API (no API key required).
|
- The OpenAI API key is loaded securely from the `.env` file using `dotenv`.
|
||||||
- The Ollama LLM is used for both embeddings and generation.
|
|
||||||
|
|
||||||
Feel free to extend the agent with additional retrievers or custom prompts as needed.
|
## License
|
||||||
|
|
||||||
|
MIT License
|
||||||
|
---
|
||||||
|
Enjoy building with RAG!
|
||||||
+5
-6
@@ -1,17 +1,16 @@
|
|||||||
{
|
{
|
||||||
"name": "rag-agent-chromadb-websearch",
|
"name": "rag-agent-chromadb-websearch",
|
||||||
"version": "1.0.0",
|
"version": "1.0.0",
|
||||||
"description": "RAG agent using ChromaDB and web search with Ollama LLM",
|
"description": "RAG agent using ChromaDB and OpenAI API with web search",
|
||||||
"main": "src/index.js",
|
"main": "src/index.js",
|
||||||
"type": "module",
|
"type": "commonjs",
|
||||||
"scripts": {
|
"scripts": {
|
||||||
"start": "node src/index.js"
|
"start": "node src/index.js"
|
||||||
},
|
},
|
||||||
"dependencies": {
|
"dependencies": {
|
||||||
"langchain": "^1.0.0",
|
|
||||||
"chromadb": "^1.0.0",
|
"chromadb": "^1.0.0",
|
||||||
"node-fetch": "^3.3.0",
|
"dotenv": "^16.4.5",
|
||||||
"dotenv": "^16.0.0",
|
"node-fetch": "^2.6.7",
|
||||||
"ollama": "^0.1.0"
|
"openai": "^4.12.0"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
+32
-21
@@ -1,26 +1,37 @@
|
|||||||
import { Ollama } from 'langchain/llms/ollama';
|
const { OpenAI } = require('openai');
|
||||||
import { RetrievalQAChain } from 'langchain/chains/retrieval_qa';
|
const dotenv = require('dotenv');
|
||||||
import { BaseRetriever } from 'langchain/schema';
|
dotenv.config();
|
||||||
import { webSearch } from './webSearch.js';
|
|
||||||
|
|
||||||
export function createAgent(vectorStore) {
|
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
|
||||||
const llm = new Ollama({
|
|
||||||
model: process.env.OLLAMA_MODEL || 'llama3',
|
|
||||||
baseUrl: process.env.OLLAMA_HOST || 'http://localhost:11434',
|
|
||||||
});
|
|
||||||
|
|
||||||
class CombinedRetriever extends BaseRetriever {
|
class Agent {
|
||||||
async getRelevantDocuments(query) {
|
constructor(vectorStore) {
|
||||||
const chromaDocs = await vectorStore.similaritySearch(query, 3);
|
this.vectorStore = vectorStore;
|
||||||
const webDocs = await webSearch(query, 3);
|
|
||||||
return [...chromaDocs, ...webDocs];
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
const retriever = new CombinedRetriever();
|
async ask(question) {
|
||||||
const chain = RetrievalQAChain.fromLLMAndRetriever(llm, retriever, {
|
const contextDocs = await this.vectorStore.query(question, 5);
|
||||||
returnSourceDocuments: true,
|
const context = contextDocs.join('\n\n');
|
||||||
});
|
const prompt = `
|
||||||
|
You are a helpful assistant. Use the following context to answer the question. If the context does not contain the answer, say you don't know.
|
||||||
|
|
||||||
return chain;
|
Context:
|
||||||
}
|
${context}
|
||||||
|
|
||||||
|
Question:
|
||||||
|
${question}
|
||||||
|
Answer:
|
||||||
|
`;
|
||||||
|
|
||||||
|
const completion = await openai.chat.completions.create({
|
||||||
|
model: 'gpt-3.5-turbo',
|
||||||
|
messages: [{ role: 'user', content: prompt }],
|
||||||
|
temperature: 0.2,
|
||||||
|
max_tokens: 300,
|
||||||
|
});
|
||||||
|
|
||||||
|
return completion.choices[0].message.content.trim();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
module.exports = Agent;
|
||||||
+42
-38
@@ -1,50 +1,54 @@
|
|||||||
import dotenv from 'dotenv';
|
const readline = require('readline');
|
||||||
|
const VectorStore = require('./vectorStore');
|
||||||
|
const Agent = require('./agent');
|
||||||
|
const { searchAndChunk } = require('./webSearch');
|
||||||
|
const dotenv = require('dotenv');
|
||||||
dotenv.config();
|
dotenv.config();
|
||||||
|
|
||||||
import { createVectorStore } from './vectorStore.js';
|
|
||||||
import { createAgent } from './agent.js';
|
|
||||||
import { Document } from 'langchain/document';
|
|
||||||
|
|
||||||
async function main() {
|
async function main() {
|
||||||
const vectorStore = await createVectorStore();
|
console.log('Initializing RAG agent...');
|
||||||
|
const vectorStore = new VectorStore();
|
||||||
|
|
||||||
if (process.env.ADD_SAMPLE_DOCS === 'true') {
|
// Example URLs to index
|
||||||
const sampleDocs = [
|
const urls = [
|
||||||
new Document({
|
'https://en.wikipedia.org/wiki/Artificial_intelligence',
|
||||||
pageContent:
|
'https://en.wikipedia.org/wiki/ChromaDB',
|
||||||
'LangChain is a framework for building applications powered by language models.',
|
'https://en.wikipedia.org/wiki/OpenAI',
|
||||||
metadata: { source: 'LangChain Docs' },
|
];
|
||||||
}),
|
|
||||||
new Document({
|
|
||||||
pageContent: 'ChromaDB is a vector database for storing embeddings.',
|
|
||||||
metadata: { source: 'ChromaDB Docs' },
|
|
||||||
}),
|
|
||||||
new Document({
|
|
||||||
pageContent: 'Ollama is a lightweight LLM server that can run locally.',
|
|
||||||
metadata: { source: 'Ollama Docs' },
|
|
||||||
}),
|
|
||||||
];
|
|
||||||
await vectorStore.addDocuments(sampleDocs);
|
|
||||||
console.log('Sample documents added to ChromaDB.');
|
|
||||||
}
|
|
||||||
|
|
||||||
const agent = createAgent(vectorStore);
|
console.log('Fetching and indexing web pages...');
|
||||||
|
const chunks = await searchAndChunk(urls);
|
||||||
|
await vectorStore.addDocuments(chunks);
|
||||||
|
console.log(`Indexed ${chunks.length} chunks.`);
|
||||||
|
|
||||||
const query = process.argv[2];
|
const agent = new Agent(vectorStore);
|
||||||
if (!query) {
|
|
||||||
console.error('Please provide a query as a command line argument.');
|
|
||||||
process.exit(1);
|
|
||||||
}
|
|
||||||
|
|
||||||
const result = await agent.invoke({ input: query });
|
const rl = readline.createInterface({
|
||||||
console.log('Answer:', result.output);
|
input: process.stdin,
|
||||||
console.log(
|
output: process.stdout,
|
||||||
'Sources:',
|
});
|
||||||
result.sourceDocuments.map((d) => d.metadata.source).join(', ')
|
|
||||||
);
|
const askQuestion = () => {
|
||||||
|
rl.question('\nEnter your question (or type "exit" to quit): ', async (answer) => {
|
||||||
|
if (answer.trim().toLowerCase() === 'exit') {
|
||||||
|
rl.close();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
console.log('\nGenerating answer...');
|
||||||
|
try {
|
||||||
|
const response = await agent.ask(answer);
|
||||||
|
console.log(`\nAnswer:\n${response}`);
|
||||||
|
} catch (err) {
|
||||||
|
console.error('Error generating answer:', err);
|
||||||
|
}
|
||||||
|
askQuestion();
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
askQuestion();
|
||||||
}
|
}
|
||||||
|
|
||||||
main().catch((err) => {
|
main().catch((err) => {
|
||||||
console.error(err);
|
console.error('Fatal error:', err);
|
||||||
process.exit(1);
|
process.exit(1);
|
||||||
});
|
});
|
||||||
+41
-36
@@ -1,47 +1,52 @@
|
|||||||
import { ChromaClient } from 'chromadb';
|
const chromadb = require('chromadb');
|
||||||
import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
|
const { OpenAI } = require('openai');
|
||||||
import { Document } from 'langchain/document';
|
const dotenv = require('dotenv');
|
||||||
|
dotenv.config();
|
||||||
|
|
||||||
export async function createVectorStore() {
|
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
|
||||||
const chroma = new ChromaClient({ path: process.env.CHROMA_URL });
|
|
||||||
const collection = await chroma.getOrCreateCollection({
|
|
||||||
name: process.env.CHROMA_COLLECTION,
|
|
||||||
});
|
|
||||||
const embeddings = new OllamaEmbeddings({
|
|
||||||
model: process.env.OLLAMA_EMBEDDING_MODEL || 'nomic-embed-text',
|
|
||||||
});
|
|
||||||
return new VectorStore(collection, embeddings);
|
|
||||||
}
|
|
||||||
|
|
||||||
class VectorStore {
|
class VectorStore {
|
||||||
constructor(collection, embeddings) {
|
constructor() {
|
||||||
this.collection = collection;
|
this.client = new chromadb.Client({ path: './chromadb' });
|
||||||
this.embeddings = embeddings;
|
this.collection = this.client.getCollection('rag_collection');
|
||||||
}
|
}
|
||||||
|
|
||||||
async addDocuments(docs) {
|
async getEmbedding(text) {
|
||||||
const texts = docs.map((d) => d.pageContent);
|
const response = await openai.embeddings.create({
|
||||||
const embeddings = await this.embeddings.embedDocuments(texts);
|
model: 'text-embedding-ada-002',
|
||||||
await this.collection.addDocuments({
|
input: text,
|
||||||
documents: docs,
|
});
|
||||||
|
return response.data[0].embedding;
|
||||||
|
}
|
||||||
|
|
||||||
|
async addDocuments(chunks) {
|
||||||
|
const documents = [];
|
||||||
|
const embeddings = [];
|
||||||
|
const ids = [];
|
||||||
|
|
||||||
|
for (const chunk of chunks) {
|
||||||
|
const embedding = await this.getEmbedding(chunk);
|
||||||
|
documents.push(chunk);
|
||||||
|
embeddings.push(embedding);
|
||||||
|
ids.push(`${Date.now()}-${Math.random()}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
await this.collection.add({
|
||||||
|
documents,
|
||||||
embeddings,
|
embeddings,
|
||||||
|
ids,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
async similaritySearch(query, k = 4) {
|
async query(queryText, k = 5) {
|
||||||
const embedding = await this.embeddings.embedQuery(query);
|
const queryEmbedding = await this.getEmbedding(queryText);
|
||||||
const results = await this.collection.getNearestNeighbors({
|
const results = await this.collection.query({
|
||||||
queryEmbeddings: [embedding],
|
queryEmbeddings: [queryEmbedding],
|
||||||
n: k,
|
nResults: k,
|
||||||
});
|
});
|
||||||
const ids = results.ids[0];
|
|
||||||
const docs = await this.collection.getDocuments({ ids });
|
return results.documents[0];
|
||||||
return docs.map(
|
|
||||||
(doc) =>
|
|
||||||
new Document({
|
|
||||||
pageContent: doc.document,
|
|
||||||
metadata: doc.metadata,
|
|
||||||
})
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
module.exports = VectorStore;
|
||||||
+39
-21
@@ -1,24 +1,42 @@
|
|||||||
import fetch from 'node-fetch';
|
const fetch = require('node-fetch');
|
||||||
import { Document } from 'langchain/document';
|
|
||||||
|
|
||||||
export async function webSearch(query, limit = 3) {
|
async function fetchPage(url) {
|
||||||
const url = `https://api.duckduckgo.com/?q=${encodeURIComponent(
|
try {
|
||||||
query
|
const res = await fetch(url);
|
||||||
)}&format=json&pretty=1`;
|
if (!res.ok) {
|
||||||
const res = await fetch(url);
|
console.warn(`Failed to fetch ${url}: ${res.statusText}`);
|
||||||
const data = await res.json();
|
return '';
|
||||||
const topics = data.RelatedTopics || [];
|
|
||||||
const docs = [];
|
|
||||||
for (const topic of topics) {
|
|
||||||
if (topic.Text) {
|
|
||||||
docs.push(
|
|
||||||
new Document({
|
|
||||||
pageContent: topic.Text,
|
|
||||||
metadata: { source: 'DuckDuckGo', url: topic.FirstURL },
|
|
||||||
})
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
if (docs.length >= limit) break;
|
const html = await res.text();
|
||||||
|
// Strip HTML tags
|
||||||
|
const text = html.replace(/<[^>]*>/g, ' ');
|
||||||
|
// Collapse whitespace
|
||||||
|
const cleaned = text.replace(/\s+/g, ' ').trim();
|
||||||
|
return cleaned;
|
||||||
|
} catch (err) {
|
||||||
|
console.error(`Error fetching ${url}:`, err);
|
||||||
|
return '';
|
||||||
}
|
}
|
||||||
return docs.slice(0, limit);
|
}
|
||||||
}
|
|
||||||
|
function chunkText(text, size = 500) {
|
||||||
|
const chunks = [];
|
||||||
|
for (let i = 0; i < text.length; i += size) {
|
||||||
|
chunks.push(text.slice(i, i + size));
|
||||||
|
}
|
||||||
|
return chunks;
|
||||||
|
}
|
||||||
|
|
||||||
|
async function searchAndChunk(urls) {
|
||||||
|
const allChunks = [];
|
||||||
|
for (const url of urls) {
|
||||||
|
const pageText = await fetchPage(url);
|
||||||
|
if (pageText) {
|
||||||
|
const chunks = chunkText(pageText);
|
||||||
|
allChunks.push(...chunks);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return allChunks;
|
||||||
|
}
|
||||||
|
|
||||||
|
module.exports = { searchAndChunk };
|
||||||
Reference in New Issue
Block a user