feat: solution for 'Агент с RAG-памятью'

2026-06-30 15:24:26 +03:00
parent 88f8072c55
commit e95da4c295
6 changed files with 171 additions and 175 deletions
@@ -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 <key> <value>` | Add or update a key/value pair. | `kb add greeting "Hello, world!"` |
-| `kb query <key>` | Retrieve the value for a key. | `kb query greeting` |
-| `kb delete <key>` | 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**
+Happy coding!
@@ -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"
  }
 }
@@ -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();
+module.exports = Agent;
@@ -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;
@@ -1 +1,3 @@
-module.exports = require('./agent');
+const Agent = require('./agent');
+
+module.exports = { Agent };
@@ -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,
+  async init() {
+    this.collection = await chroma.getCollection({
+      name: this.collectionName,
+      metadata: { type: 'vector' },
    });
-  const vectorStore = await FAISS.load(embeddings, VECTORSTORE_DIR);
-  await vectorStore.addDocuments(texts);
-  await vectorStore.save();
-}
+  }

-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,
-};
+module.exports = VectorStore;