diff --git a/README.md b/README.md index 0e6c852..b854c04 100644 --- a/README.md +++ b/README.md @@ -1,110 +1,73 @@ # RAG Agent with ChromaDB and Web Search -This project implements a Retrieval-Augmented Generation (RAG) agent that uses a local ChromaDB vector store for document retrieval and falls back to DuckDuckGo web search when the local store does not provide sufficient context. +This project demonstrates a Retrieval-Augmented Generation (RAG) agent built with **LangChain 1.x**, **ChromaDB** as the vector store, and **SerpAPI** for web search integration. ## Features -- **Local Retrieval** – Store and query embeddings in a persistent ChromaDB collection. -- **Web Search Fallback** – If local retrieval fails to find relevant context, the agent performs a DuckDuckGo search and uses the snippets. -- **OpenAI Integration** – Uses OpenAI embeddings (`text-embedding-ada-002`) and the `gpt-3.5-turbo` model for generation. -- **CLI** – Simple command line interface for ingesting documents and asking questions. +- Stores documents in ChromaDB and generates embeddings using OpenAI. +- Retrieves relevant documents via a vector store tool. +- Performs live web searches with SerpAPI. +- Combines both sources to answer user queries. ## Prerequisites -- Python 3.9+ -- An OpenAI API key -- (Optional) Internet access for web search +- Node.js v18+ (ES modules support) +- A running ChromaDB instance (default: `localhost:8000`) +- OpenAI API key +- SerpAPI key -## Installation +## Setup -```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 +1. **Clone the repository** -# Create a virtual environment (recommended) -python -m venv .venv -source .venv/bin/activate # On Windows use `.venv\Scripts\activate` + ```bash + git clone https://github.com/your-username/rag-agent-chromadb-websearch.git + cd rag-agent-chromadb-websearch + ``` -# Install dependencies -pip install -r requirements.txt -``` +2. **Install dependencies** -`requirements.txt` contains: + ```bash + npm install + ``` + +3. **Configure environment variables** + + Create a `.env` file in the project root: + + ```dotenv + OPENAI_API_KEY=your_openai_api_key + CHROMA_HOST=localhost + CHROMA_PORT=8000 + SERPAPI_KEY=your_serpapi_key + ``` + +4. **Run the agent** + + ```bash + npm start + ``` + + The agent will add sample documents to ChromaDB, then answer a sample query using both the vector store and web search. + +## Project Structure ``` -openai -chromadb -duckduckgo-search -beautifulsoup4 -requests +src/ +├── index.js # Entry point +├── agent.js # Agent construction +├── vectorStore.js # ChromaDB interactions +└── webSearch.js # SerpAPI web search ``` -## Environment Variables +## Customization -| Variable | Description | Example | -|----------|-------------|---------| -| `OPENAI_API_KEY` | Your OpenAI API key | `sk-...` | -| `CHROMA_DB_PATH` | Directory where ChromaDB stores data | `./chromadb` | -| `CHROMA_COLLECTION_NAME` | Name of the collection | `rag_collection` | -| `TOP_K` | Number of top documents to retrieve | `5` | -| `SIMILARITY_THRESHOLD` | Minimum similarity to consider a document relevant | `0.5` | -| `WEB_SEARCH_MAX_RESULTS` | Max number of web snippets to fetch | `3` | - -Set them in your shell or create a `.env` file and load with `dotenv` (optional). - -## Usage - -### Ingest Documents - -Place your plain text files (`.txt`) in a folder, then run: - -```bash -python src/index.py ingest /path/to/text/files -``` - -The script will read all `.txt` files, split them into chunks, embed them, and store them in ChromaDB. - -### Ask a Question - -```bash -python src/index.py ask "What is the capital of France?" -``` - -The agent will: - -1. Query the local vector store for relevant passages. -2. If none are found above the similarity threshold, perform a DuckDuckGo search. -3. Combine the retrieved context into a prompt. -4. Call OpenAI’s `gpt-3.5-turbo` to generate an answer. - -## Example - -```bash -$ python src/index.py ingest ./data -INFO:root:Added 12 documents to collection 'rag_collection'. - -$ python src/index.py ask "Explain the theory of relativity." -Answer: -The theory of relativity, developed by Albert Einstein, consists of two parts: special relativity and general relativity. ... -``` - -## Testing - -Unit tests are provided in the `tests/` directory. To run them: - -```bash -pytest tests/ -``` - -(If you don't have `pytest` installed, run `pip install pytest`.) - -## Troubleshooting - -- **No documents ingested** – Ensure the folder path is correct and contains `.txt` files. -- **OpenAI errors** – Verify that `OPENAI_API_KEY` is set and that you have sufficient quota. -- **Web search fails** – Check your internet connection and that DuckDuckGo is reachable. +- **Adding Documents**: Use `addDocuments` from `vectorStore.js` to add your own documents. +- **Changing LLM**: Replace `OpenAI` with another LLM provider supported by LangChain. +- **Adjusting Retrieval**: Modify the number of results returned by the vector store or web search. ## License -MIT License \ No newline at end of file +MIT License +--- +Happy coding! \ No newline at end of file diff --git a/package.json b/package.json index 1822234..0c2f5a0 100644 --- a/package.json +++ b/package.json @@ -1,16 +1,17 @@ { "name": "rag-agent-chromadb-websearch", "version": "1.0.0", - "description": "RAG agent using ChromaDB and OpenAI API with web search", + "description": "RAG agent using ChromaDB and web search with LangChain 1.x", "main": "src/index.js", - "type": "commonjs", + "type": "module", "scripts": { "start": "node src/index.js" }, "dependencies": { "chromadb": "^1.0.0", - "dotenv": "^16.4.5", - "node-fetch": "^2.6.7", - "openai": "^4.12.0" + "dotenv": "^16.0.3", + "langchain": "^1.0.0", + "openai": "^3.3.0", + "serpapi": "^2.0.0" } } \ No newline at end of file diff --git a/src/agent.js b/src/agent.js index 785501c..211aec4 100644 --- a/src/agent.js +++ b/src/agent.js @@ -1,37 +1,54 @@ -const { OpenAI } = require('openai'); -const dotenv = require('dotenv'); -dotenv.config(); +import { OpenAI } from 'langchain/llms/openai'; +import { Tool } from 'langchain/tools'; +import { AgentExecutor } from 'langchain/agents'; +import { query } from './vectorStore.js'; +import { webSearch } from './webSearch.js'; -const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY }); +const llm = new OpenAI({ + temperature: 0, + openAIApiKey: process.env.OPENAI_API_KEY, +}); -class Agent { - constructor(vectorStore) { - this.vectorStore = vectorStore; - } +/** + * Tool that retrieves relevant documents from ChromaDB. + */ +const vectorStoreTool = new Tool({ + name: 'VectorStore', + description: 'Retrieve relevant documents from the vector store.', + func: async (input) => { + const result = await query(input); + if (!result.documents || result.documents.length === 0) { + return 'No relevant documents found.'; + } + return result.documents.map((doc, idx) => `(${idx + 1}) ${doc}`).join('\n'); + }, +}); - async ask(question) { - const contextDocs = await this.vectorStore.query(question, 5); - 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. +/** + * Tool that performs a web search. + */ +const webSearchTool = new Tool({ + name: 'WebSearch', + description: 'Search the web for up-to-date information.', + func: async (input) => { + const results = await webSearch(input); + if (!results || results.length === 0) { + return 'No web results found.'; + } + return results + .map((r, idx) => `(${idx + 1}) ${r.title}: ${r.link}`) + .join('\n'); + }, +}); -Context: -${context} +const tools = [vectorStoreTool, webSearchTool]; -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; \ No newline at end of file +/** + * Create and return a LangChain 1.x AgentExecutor. + */ +export async function createAgent() { + const agent = await AgentExecutor.fromLLMAndTools(llm, tools, { + verbose: true, + }); + return agent; +} \ No newline at end of file diff --git a/src/index.js b/src/index.js index 1f6945a..a1d6108 100644 --- a/src/index.js +++ b/src/index.js @@ -1,54 +1,28 @@ -const readline = require('readline'); -const VectorStore = require('./vectorStore'); -const Agent = require('./agent'); -const { searchAndChunk } = require('./webSearch'); -const dotenv = require('dotenv'); +import dotenv from 'dotenv'; dotenv.config(); +import { createAgent } from './agent.js'; +import { addDocuments } from './vectorStore.js'; + async function main() { - console.log('Initializing RAG agent...'); - const vectorStore = new VectorStore(); - - // Example URLs to index - const urls = [ - 'https://en.wikipedia.org/wiki/Artificial_intelligence', - 'https://en.wikipedia.org/wiki/ChromaDB', - 'https://en.wikipedia.org/wiki/OpenAI', + // Add sample documents to the vector store + const sampleDocs = [ + 'LangChain is a framework for building applications powered by language models.', + 'ChromaDB is an open-source vector database that stores embeddings.', + 'Web search can provide up-to-date information that may not be in the vector store.', ]; + await addDocuments(sampleDocs); - console.log('Fetching and indexing web pages...'); - const chunks = await searchAndChunk(urls); - await vectorStore.addDocuments(chunks); - console.log(`Indexed ${chunks.length} chunks.`); + const agent = await createAgent(); - const agent = new Agent(vectorStore); + const query = 'Explain how LangChain can use ChromaDB and web search together.'; + const result = await agent.invoke({ input: query }); - const rl = readline.createInterface({ - input: process.stdin, - output: process.stdout, - }); - - 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(); + console.log('\n=== Agent Response ===\n'); + console.log(result.output); } main().catch((err) => { - console.error('Fatal error:', err); + console.error('Error running the agent:', err); process.exit(1); }); \ No newline at end of file diff --git a/src/vectorStore.js b/src/vectorStore.js index 9ef262d..e9c144f 100644 --- a/src/vectorStore.js +++ b/src/vectorStore.js @@ -1,52 +1,58 @@ -const chromadb = require('chromadb'); -const { OpenAI } = require('openai'); -const dotenv = require('dotenv'); -dotenv.config(); +import { ChromaClient } from 'chromadb'; +import { OpenAIEmbeddings } from 'langchain/embeddings/openai'; -const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY }); +const chroma = new ChromaClient({ + host: process.env.CHROMA_HOST || 'localhost', + port: parseInt(process.env.CHROMA_PORT, 10) || 8000, +}); -class VectorStore { - constructor() { - this.client = new chromadb.Client({ path: './chromadb' }); - this.collection = this.client.getCollection('rag_collection'); - } +const embeddings = new OpenAIEmbeddings({ + openAIApiKey: process.env.OPENAI_API_KEY, +}); - async getEmbedding(text) { - const response = await openai.embeddings.create({ - model: 'text-embedding-ada-002', - input: text, - }); - return response.data[0].embedding; - } +const COLLECTION_NAME = 'rag_collection'; - 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, - ids, - }); - } - - async query(queryText, k = 5) { - const queryEmbedding = await this.getEmbedding(queryText); - const results = await this.collection.query({ - queryEmbeddings: [queryEmbedding], - nResults: k, - }); - - return results.documents[0]; +/** + * Ensure the collection exists in ChromaDB. + */ +async function ensureCollection() { + const collections = await chroma.listCollections(); + if (!collections.includes(COLLECTION_NAME)) { + await chroma.createCollection({ name: COLLECTION_NAME }); } } -module.exports = VectorStore; \ No newline at end of file +/** + * Add an array of documents to the vector store. + * @param {string[]} docs + */ +export async function addDocuments(docs) { + await ensureCollection(); + const ids = docs.map((_, idx) => `doc-${Date.now()}-${idx}`); + const embeddingsResult = await embeddings.embedDocuments(docs); + await chroma.add({ + collection_name: COLLECTION_NAME, + ids, + documents: docs, + embeddings: embeddingsResult, + }); +} + +/** + * Query the vector store for the most relevant documents. + * @param {string} queryText + * @param {number} nResults + * @returns {Promise<{documents: string[]}>} + */ +export async function query(queryText, nResults = 3) { + await ensureCollection(); + const embedding = await embeddings.embedQuery(queryText); + const results = await chroma.query({ + collection_name: COLLECTION_NAME, + query_embeddings: [embedding], + n_results: nResults, + }); + // ChromaDB returns an array of objects; extract documents + const docs = results.documents || []; + return { documents: docs }; +} \ No newline at end of file diff --git a/src/webSearch.js b/src/webSearch.js index 7595eb2..add1666 100644 --- a/src/webSearch.js +++ b/src/webSearch.js @@ -1,42 +1,19 @@ -const fetch = require('node-fetch'); +import { GoogleSearchResults } from 'serpapi'; -async function fetchPage(url) { - try { - const res = await fetch(url); - if (!res.ok) { - console.warn(`Failed to fetch ${url}: ${res.statusText}`); - return ''; - } - 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 ''; - } -} - -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 }; \ No newline at end of file +/** + * Perform a web search using SerpAPI. + * @param {string} query + * @returns {Promise>} + */ +export async function webSearch(query) { + const search = new GoogleSearchResults({ + api_key: process.env.SERPAPI_KEY, + engine: 'google', + }); + const results = await search.get({ q: query, num: 5 }); + // Return only the organic results with title and link + return (results.organic_results || []).map((r) => ({ + title: r.title, + link: r.link, + })); +} \ No newline at end of file