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

src/agent.js

@@ -1,14 +1,32 @@
-const { getChatCompletion } = require('./utils');
-const retriever = require('./retriever');
+const { OpenAI } = require('@langchain/openai');
+const { RetrievalQAChain } = require('@langchain/chains');
+const { initVectorStore } = require('./vectorStore');
+require('dotenv').config();
 
-async function ask(question) {
-  const passages = await retriever.getRelevantPassages(question, 3);
-  const context = passages.join('\n---\n');
-  const prompt = `You are an assistant. Use the following context to answer the question.\n\nContext:\n${context}\n\nQuestion: ${question}\nAnswer:`;
-  const answer = await getChatCompletion(prompt);
-  return answer;
+class Agent {
+  constructor() {
+    this.llm = new OpenAI({
+      temperature: 0.7,
+      openAIApiKey: process.env.OPENAI_API_KEY,
+    });
+    this.vectorStore = null;
+    this.chain = null;
+  }
+
+  async init() {
+    if (!this.vectorStore) {
+      this.vectorStore = await initVectorStore();
+    }
+    if (!this.chain) {
+      this.chain = RetrievalQAChain.fromLLM(this.llm, this.vectorStore.asRetriever());
+    }
+  }
+
+  async ask(question) {
+    await this.init();
+    const result = await this.chain.invoke({ question });
+    return result.output;
+  }
 }
 
-module.exports = {
-  ask,
-};
+module.exports = new Agent();

src/cli.js

@@ -0,0 +1,45 @@
+#!/usr/bin/env node
+require('dotenv').config();
+const { program } = require('commander');
+const { addDocumentFromFile, clearMemory, agent } = require('./tools');
+
+program
+  .name('rag-agent')
+  .description('CLI for a RAG agent with OpenAI and FAISS')
+  .version('1.0.0');
+
+program
+  .command('add <file>')
+  .description('Add a document to the vector store')
+  .action(async (file) => {
+    try {
+      await addDocumentFromFile(file);
+    } catch (err) {
+      console.error('Error adding document:', err.message);
+    }
+  });
+
+program
+  .command('query <question...>')
+  .description('Ask a question to the agent')
+  .action(async (question) => {
+    try {
+      const answer = await agent.ask(question.join(' '));
+      console.log('Answer:', answer);
+    } catch (err) {
+      console.error('Error querying agent:', err.message);
+    }
+  });
+
+program
+  .command('clear')
+  .description('Clear all stored memory')
+  .action(async () => {
+    try {
+      await clearMemory();
+    } catch (err) {
+      console.error('Error clearing memory:', err.message);
+    }
+  });
+
+program.parseAsync(process.argv);

src/index.js

@@ -1,43 +1 @@
-const readline = require('readline');
-const agent = require('./agent');
-const retriever = require('./retriever');
-
-async function init() {
-  // Load knowledge base from ./knowledge directory
-  await retriever.loadKnowledgeBase('./knowledge');
-  console.log('Knowledge base loaded.');
-}
-
-async function main() {
-  await init();
-
-  const rl = readline.createInterface({
-    input: process.stdin,
-    output: process.stdout,
-    prompt: 'You> ',
-  });
-
-  rl.prompt();
-
-  rl.on('line', async (line) => {
-    const trimmed = line.trim();
-    if (trimmed.toLowerCase() === 'exit') {
-      rl.close();
-      process.exit(0);
-    }
-    try {
-      const answer = await agent.ask(trimmed);
-      console.log(`Assistant: ${answer}\n`);
-    } catch (err) {
-      console.error(`Error: ${err.message}\n`);
-    }
-    rl.prompt();
-  });
-
-  rl.on('close', () => {
-    console.log('Goodbye!');
-    process.exit(0);
-  });
-}
-
-main();
+module.exports = require('./agent');

src/tools.js

@@ -0,0 +1,29 @@
+const fs = require('fs');
+const { RecursiveCharacterTextSplitter } = require('@langchain/textsplitter');
+const { addDocuments, clearVectorStore } = require('./vectorStore');
+const agent = require('./agent');
+
+async function addDocumentFromFile(filePath) {
+  if (!fs.existsSync(filePath)) {
+    throw new Error(`File not found: ${filePath}`);
+  }
+  const content = fs.readFileSync(filePath, 'utf-8');
+  const splitter = new RecursiveCharacterTextSplitter({
+    chunkSize: 1000,
+    chunkOverlap: 200,
+  });
+  const docs = await splitter.splitText(content);
+  await addDocuments(docs);
+  console.log(`Added ${docs.length} chunks from ${filePath}`);
+}
+
+async function clearMemory() {
+  await clearVectorStore();
+  console.log('Cleared vector store memory.');
+}
+
+module.exports = {
+  addDocumentFromFile,
+  clearMemory,
+  agent,
+};

src/vectorStore.js

@@ -1,33 +1,39 @@
-class VectorStore {
-  constructor() {
-    this.documents = [];
-  }
+const { FAISS } = require('@langchain/vectorstores/faiss');
+const { OpenAIEmbeddings } = require('@langchain/openai');
+const fs = require('fs');
+const path = require('path');
+require('dotenv').config();
 
-  addDocument(id, embedding, text) {
-    this.documents.push({ id, embedding, text });
-  }
+const VECTORSTORE_DIR = path.join(__dirname, '..', 'data', 'vectorstore');
 
-  cosineSimilarity(a, b) {
-    let dot = 0;
-    let normA = 0;
-    let normB = 0;
-    for (let i = 0; i < a.length; i++) {
-      dot += a[i] * b[i];
-      normA += a[i] * a[i];
-      normB += b[i] * b[i];
-    }
-    return dot / (Math.sqrt(normA) * Math.sqrt(normB));
+async function initVectorStore() {
+  if (!fs.existsSync(VECTORSTORE_DIR)) {
+    fs.mkdirSync(VECTORSTORE_DIR, { recursive: true });
   }
+  const embeddings = new OpenAIEmbeddings({
+    openAIApiKey: process.env.OPENAI_API_KEY,
+  });
+  const vectorStore = await FAISS.load(embeddings, VECTORSTORE_DIR);
+  return vectorStore;
+}
 
-  query(queryEmbedding, k) {
-    const sims = this.documents.map((doc) => ({
-      doc,
-      similarity: this.cosineSimilarity(queryEmbedding, doc.embedding),
-    }));
-    sims.sort((a, b) => b.similarity - a.similarity);
-    return sims.slice(0, k).map((s) => s.doc);
+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 function clearVectorStore() {
+  if (fs.existsSync(VECTORSTORE_DIR)) {
+    fs.rmdirSync(VECTORSTORE_DIR, { recursive: true });
   }
 }
 
-const store = new VectorStore();
-module.exports = store;
+module.exports = {
+  initVectorStore,
+  addDocuments,
+  clearVectorStore,
+};