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

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();
+module.exports = Agent;

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;

src/index.js

@@ -1 +1,3 @@
-module.exports = require('./agent');
+const Agent = require('./agent');
+
+module.exports = { Agent };

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