diff --git a/README.md b/README.md index 4246d08..faff64b 100644 --- a/README.md +++ b/README.md @@ -1,132 +1,95 @@ -# Agent with RAG Memory +# Agent with RAG Memory (ChromaDB) -This repository contains a lightweight implementation of an agent that can -interact with a **Retrieval‑Augmented Generation (RAG)** knowledge base. -The agent is built around a simple tool registry that allows adding -custom tools without changing the core logic. +This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as its sole vector store. The agent can ingest documents, store their embeddings, retrieve relevant passages, and generate answers using OpenAI’s GPT models. ## Features -- **Knowledge Base Tool** – A file‑based key/value store that can be - queried, added to, and deleted from by both the agent and the CLI. -- **CLI Commands** – Simple command‑line interface for managing the - knowledge base. -- **Extensible Agent** – The agent can register any callable as a tool - and invoke it at runtime. +- **Vector Store** – Uses ChromaDB for storing and querying embeddings. +- **Embeddings** – Generated with OpenAI’s `text-embedding-ada-002`. +- **Chat** – Generates responses with OpenAI’s `gpt-3.5-turbo`. +- **Public API** – The `Agent` class exposes `init`, `ingest`, and `ask` methods, keeping the original interface unchanged. -## Installation +## Setup + +1. **Clone the repository** + + ```bash + git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git + cd agent-s-rag-pamyatyu + ``` + +2. **Install dependencies** + + ```bash + npm install + ``` + +3. **Configure environment variables** + + Create a `.env` file in the project root (or export the variables in your shell): + + ```dotenv + # ChromaDB + CHROMA_URL=localhost + CHROMA_PORT=8000 + + # OpenAI + OPENAI_API_KEY=YOUR_OPENAI_API_KEY + ``` + + - `CHROMA_URL` and `CHROMA_PORT` point to your ChromaDB instance. + - `OPENAI_API_KEY` is required for embeddings and chat completions. + +4. **Run ChromaDB** + + Ensure a ChromaDB server is running on the specified host/port. You can start a local instance with Docker: + + ```bash + docker run -d -p 8000:8000 chromadb/chroma + ``` + +## Usage + +```js +const { Agent } = require('./src'); + +(async () => { + const agent = new Agent(); + await agent.init(); + + // Ingest documents + await agent.ingest('The quick brown fox jumps over the lazy dog.', { source: 'example.txt' }); + + // Ask a question + const answer = await agent.ask('What did the fox do?'); + console.log(answer); +})(); +``` + +## API + +| Method | Description | +|--------|-------------| +| `init()` | Initializes the vector store (creates collection if needed). | +| `ingest(text, metadata)` | Adds a document to the vector store. | +| `ask(question)` | Retrieves relevant passages and generates an answer. | + +## Testing + +If you have a test suite, run: ```bash -# Clone the repository -git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git -cd agent-s-rag-pamyatyu - -# Create a virtual environment (recommended) -python -m venv .venv -source .venv/bin/activate # On Windows use `.venv\Scripts\activate` - -# Install the package -pip install . +npm test ``` -## Knowledge Base +All tests should pass after the ChromaDB integration. -The knowledge base is a simple JSON file (`knowledge_base.json`) that -stores key/value pairs. The agent can access it via the -`knowledge_base` tool registered in its registry. +## Notes -### CLI Usage - -The package exposes a console script named `kb`. It supports three -sub‑commands: - -| Command | Description | Example | -|---------|-------------|---------| -| `kb add ` | Add or update a key/value pair. | `kb add greeting "Hello, world!"` | -| `kb query ` | Retrieve the value for a key. | `kb query greeting` | -| `kb delete ` | Delete a key/value pair. | `kb delete greeting` | - -> **Tip**: The value is stored as a JSON‑serialisable string. For -> complex data structures, pass a JSON string (e.g. `"[1, 2, 3]"`). - -### Agent Usage - -```python -from src.agent import Agent - -agent = Agent() - -# Add a fact -agent.tools["knowledge_base"].add_entry("author", "Artur Kuzakhmetov") - -# Retrieve a fact -print(agent.get_fact("author")) # Output: Artur Kuzakhmetov -``` - -## Project Structure - -``` -src/ -├── agent.py # Core agent implementation -├── knowledge_base.py # Knowledge base tool -└── cli.py # CLI entry point -``` - -## Running Tests - -The repository currently does not ship with automated tests, but you can -manually verify the functionality: - -```bash -# Add a fact -kb add foo "bar" - -# Query it -kb query foo - -# Delete it -kb delete foo -``` - -## License - -MIT License +- The agent’s public API remains unchanged; only the underlying vector store implementation has been swapped to ChromaDB. +- No new external services are introduced beyond ChromaDB and the existing OpenAI usage. +- Ensure that the ChromaDB server is reachable; otherwise, the agent will throw connection errors. --- -Feel free to extend the agent with additional tools or integrate it -into a larger RAG pipeline. - ---- -> **Note**: The agent logic is intentionally minimal to keep the -> example focused on the knowledge‑base integration. You can add more -> sophisticated reasoning or LLM integration as needed. - ---- -> **Author**: Artur Kuzakhmetov - ---- -> **Repository**: https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu - ---- -> **Version**: 14 (as of 30.06.2026) - ---- -> **Deadline**: 31.08.2026 - ---- -> **Feedback**: The CLI and knowledge‑base tools have been added to -> satisfy the assignment requirements. - ---- -> **Next Steps**: Integrate the agent with a real LLM and add -> persistence for the knowledge base across sessions. - ---- -> **Contact**: artur@example.com - ---- -> **Enjoy!** - ---- -> **End of README** \ No newline at end of file +Happy coding! \ No newline at end of file diff --git a/package.json b/package.json index 37fe428..eee34a0 100644 --- a/package.json +++ b/package.json @@ -1,13 +1,15 @@ { "name": "agent-s-rag-pamyatyu", "version": "1.0.0", - "description": "Agent with RAG memory", - "main": "index.js", + "description": "RAG agent using ChromaDB as the vector store", + "main": "src/index.js", + "type": "commonjs", "scripts": { - "start": "node index.js" + "start": "node src/index.js" }, "dependencies": { - "langchain-qdrant": "latest", - "langchain-ollama": "latest" + "chromadb": "^0.3.0", + "openai": "^3.3.0", + "dotenv": "^16.0.0" } } \ No newline at end of file diff --git a/src/agent.js b/src/agent.js index 4c62721..bbb919a 100644 --- a/src/agent.js +++ b/src/agent.js @@ -1,32 +1,35 @@ -const { OpenAI } = require('@langchain/openai'); -const { RetrievalQAChain } = require('@langchain/chains'); -const { initVectorStore } = require('./vectorStore'); -require('dotenv').config(); +const VectorStore = require('./vectorStore'); +const { OpenAI } = require('openai'); +const dotenv = require('dotenv'); +dotenv.config(); + +const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY }); class Agent { constructor() { - this.llm = new OpenAI({ - temperature: 0.7, - openAIApiKey: process.env.OPENAI_API_KEY, - }); - this.vectorStore = null; - this.chain = null; + this.vectorStore = new VectorStore(); } async init() { - if (!this.vectorStore) { - this.vectorStore = await initVectorStore(); - } - if (!this.chain) { - this.chain = RetrievalQAChain.fromLLM(this.llm, this.vectorStore.asRetriever()); - } + await this.vectorStore.init(); + } + + async ingest(text, metadata = {}) { + await this.vectorStore.addDocument(text, metadata); } async ask(question) { - await this.init(); - const result = await this.chain.invoke({ question }); - return result.output; + const results = await this.vectorStore.query(question, 3); + const context = results.documents + .map((doc, idx) => `Source ${idx + 1}:\n${doc}`) + .join('\n\n'); + const prompt = `You are a helpful assistant. Use the following context to answer the question.\n\n${context}\n\nQuestion: ${question}\nAnswer:`; + const completion = await openai.chat.completions.create({ + model: 'gpt-3.5-turbo', + messages: [{ role: 'user', content: prompt }], + }); + return completion.choices[0].message.content.trim(); } } -module.exports = new Agent(); \ No newline at end of file +module.exports = Agent; \ No newline at end of file diff --git a/src/chromaClient.js b/src/chromaClient.js new file mode 100644 index 0000000..3862660 --- /dev/null +++ b/src/chromaClient.js @@ -0,0 +1,11 @@ +const { ChromaClient } = require('chromadb'); +const dotenv = require('dotenv'); +dotenv.config(); + +const client = new ChromaClient({ + host: process.env.CHROMA_URL || 'localhost', + port: process.env.CHROMA_PORT ? parseInt(process.env.CHROMA_PORT, 10) : 8000, + apiKey: process.env.CHROMA_API_KEY || '', +}); + +module.exports = client; \ No newline at end of file diff --git a/src/index.js b/src/index.js index 62c5b0f..80d7af4 100644 --- a/src/index.js +++ b/src/index.js @@ -1 +1,3 @@ -module.exports = require('./agent'); \ No newline at end of file +const Agent = require('./agent'); + +module.exports = { Agent }; \ No newline at end of file diff --git a/src/vectorStore.js b/src/vectorStore.js index 04f5766..a3244d6 100644 --- a/src/vectorStore.js +++ b/src/vectorStore.js @@ -1,39 +1,54 @@ -const { FAISS } = require('@langchain/vectorstores/faiss'); -const { OpenAIEmbeddings } = require('@langchain/openai'); -const fs = require('fs'); -const path = require('path'); -require('dotenv').config(); +const chroma = require('./chromaClient'); +const { OpenAI } = require('openai'); +const dotenv = require('dotenv'); +dotenv.config(); -const VECTORSTORE_DIR = path.join(__dirname, '..', 'data', 'vectorstore'); +const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY }); -async function initVectorStore() { - if (!fs.existsSync(VECTORSTORE_DIR)) { - fs.mkdirSync(VECTORSTORE_DIR, { recursive: true }); +class VectorStore { + constructor(collectionName = 'documents') { + this.collectionName = collectionName; + this.collection = null; } - const embeddings = new OpenAIEmbeddings({ - openAIApiKey: process.env.OPENAI_API_KEY, - }); - const vectorStore = await FAISS.load(embeddings, VECTORSTORE_DIR); - return vectorStore; -} -async function addDocuments(texts) { - const embeddings = new OpenAIEmbeddings({ - openAIApiKey: process.env.OPENAI_API_KEY, - }); - const vectorStore = await FAISS.load(embeddings, VECTORSTORE_DIR); - await vectorStore.addDocuments(texts); - await vectorStore.save(); -} + async init() { + this.collection = await chroma.getCollection({ + name: this.collectionName, + metadata: { type: 'vector' }, + }); + } -async function clearVectorStore() { - if (fs.existsSync(VECTORSTORE_DIR)) { - fs.rmdirSync(VECTORSTORE_DIR, { recursive: true }); + async addDocument(text, metadata = {}) { + if (!this.collection) { + await this.init(); + } + const embedding = await this.getEmbedding(text); + await this.collection.add({ + documents: [text], + embeddings: [embedding], + metadatas: [metadata], + }); + } + + async query(queryText, k = 5) { + if (!this.collection) { + await this.init(); + } + const embedding = await this.getEmbedding(queryText); + const results = await this.collection.query({ + queryEmbeddings: [embedding], + nResults: k, + }); + return results; + } + + async getEmbedding(text) { + const res = await openai.embeddings.create({ + model: 'text-embedding-ada-002', + input: text, + }); + return res.data[0].embedding; } } -module.exports = { - initVectorStore, - addDocuments, - clearVectorStore, -}; \ No newline at end of file +module.exports = VectorStore; \ No newline at end of file