feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
@@ -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 up‑to‑date 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 GPT‑3.5‑Turbo.
|
|
||||||
- **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 GPT‑3.5‑Turbo.
|
|
||||||
|
|
||||||
## 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 DuckDuckGo’s 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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;
|
|
||||||
}
|
}
|
||||||
Reference in New Issue
Block a user