feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

tests/agent.test.js

@@ -0,0 +1,26 @@
+const { VectorStore } = require('../src/vectorStore');
+const { answerQuestion } = require('../src/agent');
+const { embed } = require('../src/utils');
+
+describe('Agent', () => {
+  let store;
+
+  beforeAll(async () => {
+    store = new VectorStore();
+    await store.init('agent_test_collection');
+    const docs = [
+      { text: 'ChromaDB is a vector database.', id: 'doc1' },
+      { text: 'It supports similarity search.', id: 'doc2' }
+    ];
+    const embeddings = docs.map(d => embed(d.text));
+    const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
+    await store.add(embeddings, metadatas, docs.map(d => d.id));
+  });
+
+  test('provides answer based on context', async () => {
+    const question = 'What is ChromaDB?';
+    const answer = await answerQuestion(question, store);
+    expect(answer).toContain('ChromaDB is a vector database.');
+    expect(answer).toContain('It supports similarity search.');
+  });
+});

tests/vectorStore.test.js

@@ -0,0 +1,28 @@
+const { VectorStore } = require('../src/vectorStore');
+const { embed } = require('../src/utils');
+
+describe('VectorStore', () => {
+  let store;
+
+  beforeAll(async () => {
+    store = new VectorStore();
+    await store.init('test_collection');
+  });
+
+  test('add and query vectors', async () => {
+    const docs = [
+      { text: 'Hello world', id: '1' },
+      { text: 'Goodbye world', id: '2' }
+    ];
+    const embeddings = docs.map(d => embed(d.text));
+    const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
+    await store.add(embeddings, metadatas, docs.map(d => d.id));
+
+    const queryEmbedding = embed('Hello');
+    const results = await store.query(queryEmbedding, 2);
+    expect(results[0].metadatas.length).toBe(2);
+    const ids = results[0].metadatas.map(m => m.id);
+    expect(ids).toContain('1');
+    expect(ids).toContain('2');
+  });
+});