feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'

src/agent.js

@@ -0,0 +1,63 @@
+import { initializeAgentExecutorWithOptions } from "langchain/agents";
+import { OpenAI } from "langchain/llms/openai";
+import { RetrievalQAChain } from "langchain/chains";
+import { RetrievalQA } from "langchain/chains/retrieval_qa";
+import { OpenAIEmbeddings } from "langchain/embeddings/openai";
+import { ChromaClient } from "chromadb";
+
+const llm = new OpenAI({
+  temperature: 0,
+  openAIApiKey: process.env.OPENAI_API_KEY,
+});
+
+const client = new ChromaClient({
+  path: process.env.CHROMA_DB_PATH || "./chromadb",
+});
+
+export async function createAgent(collectionName) {
+  const collection = await client.getOrCreateCollection({
+    name: collectionName,
+  });
+
+  const retriever = {
+    async getRelevantDocuments(query) {
+      const embedding = await new OpenAIEmbeddings({
+        openAIApiKey: process.env.OPENAI_API_KEY,
+      }).embedQuery(query);
+      const results = await collection.query({
+        queryEmbeddings: [embedding],
+        nResults: 5,
+      });
+      return results.documents[0].map((doc, idx) => ({
+        pageContent: doc,
+        metadata: results.metadatas[0][idx],
+      }));
+    },
+  };
+
+  const qaChain = RetrievalQAChain.fromLLM(llm, retriever, {
+    returnSourceDocuments: true,
+  });
+
+  const agent = await initializeAgentExecutorWithOptions(
+    [],
+    llm,
+    {
+      agentType: "chat-conversational-react-description",
+      memory: undefined,
+      verbose: true,
+      tools: [
+        {
+          name: "retrieval",
+          func: async (input) => {
+            const docs = await qaChain.call({ input });
+            return docs.output;
+          },
+          description: "Use this tool to retrieve answers from the knowledge base",
+        },
+      ],
+    }
+  );
+
+  return agent;
+}

src/index.js

@@ -0,0 +1,50 @@
+import dotenv from "dotenv";
+import readline from "readline";
+import { createAgent } from "./agent.js";
+import { addDocument } from "./vectorstore.js";
+
+dotenv.config();
+
+const COLLECTION = "faq_collection";
+
+async function main() {
+  // Optional: add some sample documents
+  await addDocument(
+    COLLECTION,
+    "What is the return policy?",
+    { source: "FAQ" }
+  );
+  await addDocument(
+    COLLECTION,
+    "How can I track my order?",
+    { source: "FAQ" }
+  );
+
+  const agent = await createAgent(COLLECTION);
+
+  const rl = readline.createInterface({
+    input: process.stdin,
+    output: process.stdout,
+    prompt: "You: ",
+  });
+
+  console.log("FAQ Bot is ready. Type your question and press Enter.");
+  rl.prompt();
+
+  rl.on("line", async (line) => {
+    const question = line.trim();
+    if (!question) {
+      rl.prompt();
+      return;
+    }
+    try {
+      const result = await agent.call({ input: question });
+      console.log(`Bot: ${result.output}`);
+    } catch (err) {
+      console.error("Error:", err);
+    }
+    rl.prompt();
+  });
+}
+
+main().catch((err) => console.error(err));

src/vectorstore.js

@@ -0,0 +1,40 @@
+import { ChromaClient } from "chromadb";
+import { OpenAIEmbeddings } from "langchain/embeddings/openai";
+import { OpenAI } from "langchain/llms/openai";
+
+const client = new ChromaClient({
+  path: process.env.CHROMA_DB_PATH || "./chromadb",
+});
+
+const embeddings = new OpenAIEmbeddings({
+  openAIApiKey: process.env.OPENAI_API_KEY,
+});
+
+export async function addDocument(collectionName, text, metadata = {}) {
+  const collection = await client.getOrCreateCollection({
+    name: collectionName,
+  });
+  const embedding = await embeddings.embedQuery(text);
+  await collection.add({
+    documents: [text],
+    embeddings: [embedding],
+    metadatas: [metadata],
+  });
+}
+
+export async function getSimilarDocuments(collectionName, query, k = 5) {
+  const collection = await client.getOrCreateCollection({
+    name: collectionName,
+  });
+  const embedding = await embeddings.embedQuery(query);
+  const results = await collection.query({
+    queryEmbeddings: [embedding],
+    nResults: k,
+  });
+  return results.ids[0].map((id, idx) => ({
+    id,
+    score: results.scores[0][idx],
+    document: results.documents[0][idx],
+    metadata: results.metadatas[0][idx],
+  }));
+}