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

This commit is contained in:
2026-06-30 12:09:29 +03:00
parent 4f5e00efdb
commit 136e69e967
6 changed files with 181 additions and 179 deletions
+59 -63
View File
@@ -1,91 +1,87 @@
# 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 database and performs live web searches to provide uptodate information. This project implements a Retrieval-Augmented Generation (RAG) agent that:
- Stores and retrieves embeddings from **ChromaDB**.
## Features - Performs web search using DuckDuckGo to fetch additional context.
- Generates answers with an **Ollama** language model.
- **Vector store** Documents are ingested, split into chunks, embedded with OpenAI embeddings, and stored in a persistent ChromaDB collection.
- **Web search** Uses DuckDuckGo scraping to fetch recent web snippets for a query.
- **RAG pipeline** Combines local document context and web results, then generates an answer with OpenAI GPT3.5Turbo.
- **CLI** Simple command line interface for ingestion and querying.
## Prerequisites ## Prerequisites
- Python 3.10+ - Node.js v20 or newer
- An OpenAI API key with access to `text-embedding-ada-002` and `gpt-3.5-turbo`. - ChromaDB server running locally (default URL: `chromadb://localhost:8000`)
- Ollama server running locally (default URL: `http://localhost:11434`)
## Setup ## Setup
```bash 1. **Clone the repository**
# 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) ```bash
python -m venv .venv git clone https://git.brojs.ru/kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
source .venv/bin/activate # On Windows use `.venv\Scripts\activate` cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
```
# Install dependencies 2. **Install dependencies**
pip install -r requirements.txt
```
## Configuration ```bash
npm install
```
Create a `.env` file in the project root (or set environment variables directly): 3. **Configure environment variables**
``` Create a `.env` file in the project root (or modify the existing one):
OPENAI_API_KEY=sk-...
CHROMA_DB_PATH=./chromadb
CHROMA_COLLECTION_NAME=rag_collection
```
> **Note**: Do not commit your `.env` file or API key to version control. ```dotenv
CHROMA_URL=chromadb://localhost:8000
CHROMA_COLLECTION=rag_collection
OLLAMA_HOST=http://localhost:11434
OLLAMA_MODEL=llama3
OLLAMA_EMBEDDING_MODEL=nomic-embed-text
ADD_SAMPLE_DOCS=true
```
## Usage - `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.
### 1. Ingest Documents 4. **Run the agent**
```bash ```bash
python src/main.py ingest path/to/doc1.txt path/to/doc2.txt npm start -- "Your question here"
``` ```
The script will read each file, split it into chunks, generate embeddings, and store them in ChromaDB. Example:
### 2. Query the Agent ```bash
npm start -- "What is LangChain?"
```
```bash The agent will:
python src/main.py query "What is the capital of France?" - Search the local ChromaDB collection.
``` - Perform a DuckDuckGo web search.
- Combine the results and generate an answer using Ollama.
The agent will:
1. Retrieve relevant chunks from the local vector store.
2. Perform a DuckDuckGo web search for the query.
3. Combine both sources of information.
4. Generate a response using OpenAI GPT3.5Turbo.
## Project Structure ## Project Structure
``` ```
src/ .
├── main.py # CLI entry point ├── src
├── vector_store.py # ChromaDB ingestion & retrieval │ ├── agent.js # Agent logic (retrieval + generation)
├── web_search.py # DuckDuckGo web search │ ├── index.js # CLI entry point
requirements.txt │ ├── vectorStore.js # ChromaDB wrapper
README.md │ └── webSearch.js # DuckDuckGo search helper
├── .env # Environment configuration
├── package.json # Dependencies and scripts
└── README.md # Documentation
``` ```
## Testing ## Notes
The project can be tested with `pytest` (tests are not included in this minimal example). - The agent uses **LangChain 1.x** APIs.
If you add tests, run: - No Qdrant references are present; only ChromaDB is used.
- The web search is performed via DuckDuckGos public JSON API (no API key required).
- The Ollama LLM is used for both embeddings and generation.
```bash Feel free to extend the agent with additional retrievers or custom prompts as needed.
pytest
```
## License
MIT License
---
Feel free to extend the agent with additional features such as custom embeddings, different LLMs, or alternative search APIs.
+8 -6
View File
@@ -1,15 +1,17 @@
{ {
"name": "rag-agent-chromadb", "name": "rag-agent-chromadb-websearch",
"version": "1.0.0", "version": "1.0.0",
"description": "A simple RAG agent using ChromaDB for vector storage and web search.", "description": "RAG agent using ChromaDB and web search with Ollama LLM",
"main": "src/index.js", "main": "src/index.js",
"type": "module", "type": "module",
"scripts": { "scripts": {
"start": "node src/index.js", "start": "node src/index.js"
"test": "node src/test.js"
}, },
"dependencies": { "dependencies": {
"@chromadb/chromadb": "^0.1.0", "langchain": "^1.0.0",
"node-fetch": "^3.3.2" "chromadb": "^1.0.0",
"node-fetch": "^3.3.0",
"dotenv": "^16.0.0",
"ollama": "^0.1.0"
} }
} }
+21 -23
View File
@@ -1,28 +1,26 @@
import { VectorStore } from "./vectorStore.js"; import { Ollama } from 'langchain/llms/ollama';
import { webSearch } from "./webSearch.js"; import { RetrievalQAChain } from 'langchain/chains/retrieval_qa';
import { BaseRetriever } from 'langchain/schema';
import { webSearch } from './webSearch.js';
/** export function createAgent(vectorStore) {
* A simple RAG agent that retrieves relevant documents from the vector store const llm = new Ollama({
* and optionally performs a web search if no relevant documents are found. model: process.env.OLLAMA_MODEL || 'llama3',
*/ baseUrl: process.env.OLLAMA_HOST || 'http://localhost:11434',
export class Agent { });
constructor(vectorStore) {
this.vectorStore = vectorStore;
}
/** class CombinedRetriever extends BaseRetriever {
* Processes a user query and returns the best answer. async getRelevantDocuments(query) {
* @param {string} query const chromaDocs = await vectorStore.similaritySearch(query, 3);
* @returns {Promise<string>} const webDocs = await webSearch(query, 3);
*/ return [...chromaDocs, ...webDocs];
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}`;
} }
const retriever = new CombinedRetriever();
const chain = RetrievalQAChain.fromLLMAndRetriever(llm, retriever, {
returnSourceDocuments: true,
});
return chain;
} }
+38 -17
View File
@@ -1,26 +1,47 @@
import { VectorStore } from "./vectorStore.js"; import dotenv from 'dotenv';
import { Agent } from "./agent.js"; dotenv.config();
import { createVectorStore } from './vectorStore.js';
import { createAgent } from './agent.js';
import { Document } from 'langchain/document';
async function main() { async function main() {
const store = new VectorStore(); const vectorStore = await createVectorStore();
await store.init();
// Sample documents to index if (process.env.ADD_SAMPLE_DOCS === 'true') {
const docs = [ const sampleDocs = [
{ id: "1", text: "ChromaDB is a fast, lightweight vector database." }, new Document({
{ id: "2", text: "It supports in-memory and persistent storage." }, pageContent:
{ id: "3", text: "You can use it with various embedding models." }, 'LangChain is a framework for building applications powered by language models.',
]; 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.');
}
await store.addDocuments(docs); const agent = createAgent(vectorStore);
const agent = new Agent(store); const query = process.argv[2];
if (!query) {
console.error('Please provide a query as a command line argument.');
process.exit(1);
}
const query = process.argv[2] || "What is ChromaDB?"; const result = await agent.invoke({ input: query });
console.log(`Query: ${query}`); console.log('Answer:', result.output);
const answer = await agent.answer(query); console.log(
console.log("\nAnswer:"); 'Sources:',
console.log(answer); result.sourceDocuments.map((d) => d.metadata.source).join(', ')
);
} }
main().catch((err) => { main().catch((err) => {
+34 -56
View File
@@ -1,69 +1,47 @@
import { Client } from "@chromadb/chromadb"; import { ChromaClient } from 'chromadb';
import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
import { Document } from 'langchain/document';
/** export async function createVectorStore() {
* Simple embedding function that converts text into a fixed-length numeric vector. const chroma = new ChromaClient({ path: process.env.CHROMA_URL });
* This is a placeholder and should be replaced with a real embedding model for production use. const collection = await chroma.getOrCreateCollection({
*/ name: process.env.CHROMA_COLLECTION,
function embed(text) { });
const vector = Array.from(text) const embeddings = new OllamaEmbeddings({
.map((c) => c.charCodeAt(0)) model: process.env.OLLAMA_EMBEDDING_MODEL || 'nomic-embed-text',
.slice(0, 10); });
while (vector.length < 10) { return new VectorStore(collection, embeddings);
vector.push(0);
}
return vector;
} }
export class VectorStore { class VectorStore {
constructor() { constructor(collection, embeddings) {
this.client = new Client(); this.collection = collection;
this.collection = null; this.embeddings = embeddings;
} }
async init() {
this.collection = await this.client.getOrCreateCollection({
name: "rag_collection",
});
}
/**
* Adds an array of documents to the collection.
* @param {Array<{id: string, text: string}>} docs
*/
async addDocuments(docs) { async addDocuments(docs) {
if (!this.collection) { const texts = docs.map((d) => d.pageContent);
throw new Error("VectorStore not initialized. Call init() first."); const embeddings = await this.embeddings.embedDocuments(texts);
} await this.collection.addDocuments({
const ids = docs.map((d) => d.id); documents: docs,
const embeddings = docs.map((d) => embed(d.text));
const documents = docs.map((d) => d.text);
await this.collection.add({
ids,
embeddings, embeddings,
documents,
}); });
} }
/** async similaritySearch(query, k = 4) {
* Queries the collection for the most relevant documents. const embedding = await this.embeddings.embedQuery(query);
* @param {string} queryText const results = await this.collection.getNearestNeighbors({
* @param {number} nResults queryEmbeddings: [embedding],
* @returns {Promise<Array<{id: string, text: string, score: number}>>} n: k,
*/
async query(queryText, nResults = 3) {
if (!this.collection) {
throw new Error("VectorStore not initialized. Call init() first.");
}
const queryEmbedding = embed(queryText);
const results = await this.collection.query({
queryEmbeddings: [queryEmbedding],
nResults,
}); });
// results is an array of objects with ids, documents, and scores const ids = results.ids[0];
return results[0].ids.map((id, idx) => ({ const docs = await this.collection.getDocuments({ ids });
id, return docs.map(
text: results[0].documents[idx], (doc) =>
score: results[0].distances[idx], new Document({
})); pageContent: doc.document,
metadata: doc.metadata,
})
);
} }
} }
+21 -14
View File
@@ -1,17 +1,24 @@
import fetch from "node-fetch"; import fetch from 'node-fetch';
import { Document } from 'langchain/document';
/** export async function webSearch(query, limit = 3) {
* Performs a simple web search using DuckDuckGo's HTML interface. const url = `https://api.duckduckgo.com/?q=${encodeURIComponent(
* This is a lightweight example and does not use an official API. query
* @param {string} query )}&format=json&pretty=1`;
* @returns {Promise<string>} The raw HTML of the search results page. const res = await fetch(url);
*/ const data = await res.json();
export async function webSearch(query) { const topics = data.RelatedTopics || [];
const url = `https://duckduckgo.com/html/?q=${encodeURIComponent(query)}`; const docs = [];
const response = await fetch(url); for (const topic of topics) {
if (!response.ok) { if (topic.Text) {
throw new Error(`Web search failed with status ${response.status}`); docs.push(
new Document({
pageContent: topic.Text,
metadata: { source: 'DuckDuckGo', url: topic.FirstURL },
})
);
}
if (docs.length >= limit) break;
} }
const html = await response.text(); return docs.slice(0, limit);
return html;
} }