@aitytech/agentkits-memory 1.0.1 → 2.0.1

package/src/__tests__/embedding-integration.test.ts

@@ -0,0 +1,481 @@
+/**
+ * Embedding Integration Tests
+ *
+ * Tests for embedding support across ProjectMemoryService,
+ * HybridSearchEngine, and MCP Server integration.
+ *
+ * @module @aitytech/agentkits-memory/__tests__/embedding-integration.test
+ */
+
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import * as os from 'node:os';
+import Database from 'better-sqlite3';
+import {
+  ProjectMemoryService,
+  LocalEmbeddingsService,
+  HybridSearchEngine,
+  type MemoryEntry,
+} from '../index.js';
+
+describe('Embedding Integration', () => {
+  let tempDir: string;
+
+  beforeEach(() => {
+    tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'embed-integration-'));
+  });
+
+  afterEach(() => {
+    fs.rmSync(tempDir, { recursive: true, force: true });
+  });
+
+  describe('ProjectMemoryService with LocalEmbeddingsService', () => {
+    it('should generate embeddings when storing entries', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const embeddingGenerator = async (text: string): Promise<Float32Array> => {
+        const result = await embeddingsService.embed(text);
+        return result.embedding;
+      };
+
+      const service = new ProjectMemoryService({
+        baseDir: tempDir,
+        dbFilename: 'test.db',
+        embeddingGenerator,
+      });
+      await service.initialize();
+
+      const entry = await service.storeEntry({
+        key: 'test-entry',
+        content: 'This is a test content for embedding generation',
+        namespace: 'test',
+      });
+
+      expect(entry.embedding).toBeDefined();
+      expect(entry.embedding).toBeInstanceOf(Float32Array);
+      expect(entry.embedding!.length).toBe(384); // Default dimension
+
+      await service.shutdown();
+    });
+
+    it('should perform semantic search with embeddings', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const embeddingGenerator = async (text: string): Promise<Float32Array> => {
+        const result = await embeddingsService.embed(text);
+        return result.embedding;
+      };
+
+      const service = new ProjectMemoryService({
+        baseDir: tempDir,
+        dbFilename: 'test.db',
+        embeddingGenerator,
+      });
+      await service.initialize();
+
+      // Store entries with different content
+      await service.storeEntry({
+        key: 'auth-pattern',
+        content: 'Authentication using JWT tokens with refresh mechanism',
+        namespace: 'patterns',
+      });
+
+      await service.storeEntry({
+        key: 'db-pattern',
+        content: 'Database connection pooling for PostgreSQL',
+        namespace: 'patterns',
+      });
+
+      await service.storeEntry({
+        key: 'error-handling',
+        content: 'Global error handler for API exceptions',
+        namespace: 'errors',
+      });
+
+      // Semantic search should find related content
+      const results = await service.semanticSearch('JWT authentication', 5);
+
+      expect(results.length).toBeGreaterThan(0);
+      expect(results[0].entry.key).toBe('auth-pattern');
+      expect(results[0].score).toBeGreaterThan(0);
+
+      await service.shutdown();
+    });
+
+    it('should update embeddings when content changes', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const embeddingGenerator = async (text: string): Promise<Float32Array> => {
+        const result = await embeddingsService.embed(text);
+        return result.embedding;
+      };
+
+      const service = new ProjectMemoryService({
+        baseDir: tempDir,
+        dbFilename: 'test.db',
+        embeddingGenerator,
+      });
+      await service.initialize();
+
+      const entry = await service.storeEntry({
+        key: 'updateable',
+        content: 'Original content about cats',
+        namespace: 'test',
+      });
+
+      const originalEmbedding = entry.embedding!.slice();
+
+      // Update content
+      const updated = await service.update(entry.id, {
+        content: 'Updated content about dogs',
+      });
+
+      expect(updated).not.toBeNull();
+      expect(updated!.embedding).toBeDefined();
+
+      // Embedding should be different after content change
+      const isDifferent = !originalEmbedding.every(
+        (val, i) => val === updated!.embedding![i]
+      );
+      expect(isDifferent).toBe(true);
+
+      await service.shutdown();
+    });
+  });
+
+  describe('HybridSearchEngine with embeddings', () => {
+    it('should perform semantic search and return scored results', async () => {
+      const dbPath = path.join(tempDir, 'hybrid.db');
+      const db = new Database(dbPath);
+      db.pragma('journal_mode = WAL');
+
+      // Create table with rowid for FTS sync
+      db.exec(`
+        CREATE TABLE memory_entries (
+          id TEXT PRIMARY KEY,
+          key TEXT NOT NULL,
+          content TEXT NOT NULL,
+          type TEXT DEFAULT 'semantic',
+          namespace TEXT DEFAULT 'general',
+          tags TEXT DEFAULT '[]',
+          embedding BLOB,
+          created_at INTEGER NOT NULL,
+          updated_at INTEGER NOT NULL
+        )
+      `);
+
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const embeddingGenerator = async (text: string): Promise<Float32Array> => {
+        const result = await embeddingsService.embed(text);
+        return result.embedding;
+      };
+
+      const engine = new HybridSearchEngine(db, {}, embeddingGenerator);
+      await engine.initialize();
+
+      // Insert test data with embeddings
+      const now = Date.now();
+      const entries = [
+        { id: '1', key: 'react-hooks', content: 'React hooks for state management with useState and useEffect' },
+        { id: '2', key: 'vue-composition', content: 'Vue 3 composition API for reactive state' },
+        { id: '3', key: 'angular-services', content: 'Angular dependency injection and services' },
+      ];
+
+      for (const entry of entries) {
+        const embedding = await embeddingGenerator(entry.content);
+        const embeddingBuffer = Buffer.from(embedding.buffer);
+
+        db.prepare(`
+          INSERT INTO memory_entries (id, key, content, embedding, created_at, updated_at)
+          VALUES (?, ?, ?, ?, ?, ?)
+        `).run(entry.id, entry.key, entry.content, embeddingBuffer, now, now);
+      }
+
+      await engine.rebuildFtsIndex();
+
+      // Search using semantic only (more reliable in tests)
+      const results = await engine.searchCompact('React state hooks', {
+        limit: 10,
+        includeKeyword: false,
+        includeSemantic: true,
+      });
+
+      expect(results.length).toBeGreaterThan(0);
+      expect(results[0].id).toBe('1'); // React entry should be first
+
+      // Semantic score should be present
+      expect(results[0].semanticScore).toBeGreaterThan(0);
+      expect(results[0].score).toBeGreaterThan(0);
+
+      db.close();
+    });
+
+    it('should work with text-only search when no embeddings', async () => {
+      const dbPath = path.join(tempDir, 'text-only.db');
+      const db = new Database(dbPath);
+      db.pragma('journal_mode = WAL');
+
+      db.exec(`
+        CREATE TABLE memory_entries (
+          id TEXT PRIMARY KEY,
+          key TEXT NOT NULL,
+          content TEXT NOT NULL,
+          type TEXT DEFAULT 'semantic',
+          namespace TEXT DEFAULT 'general',
+          tags TEXT DEFAULT '[]',
+          embedding BLOB,
+          created_at INTEGER NOT NULL,
+          updated_at INTEGER NOT NULL
+        )
+      `);
+
+      // Engine without embedding generator
+      const engine = new HybridSearchEngine(db);
+      await engine.initialize();
+
+      const now = Date.now();
+      db.prepare(`
+        INSERT INTO memory_entries (id, key, content, created_at, updated_at)
+        VALUES (?, ?, ?, ?, ?)
+      `).run('1', 'test', 'Test content with keywords', now, now);
+
+      await engine.rebuildFtsIndex();
+
+      const results = await engine.searchCompact('keywords', {
+        limit: 10,
+        includeKeyword: true,
+        includeSemantic: true, // Should not fail even without embedding generator
+      });
+
+      expect(results.length).toBeGreaterThan(0);
+      expect(results[0].keywordScore).toBeGreaterThan(0);
+      expect(results[0].semanticScore).toBe(0); // No semantic without embeddings
+
+      db.close();
+    });
+
+    it('should support vector-only search mode', async () => {
+      const dbPath = path.join(tempDir, 'vector-only.db');
+      const db = new Database(dbPath);
+      db.pragma('journal_mode = WAL');
+
+      db.exec(`
+        CREATE TABLE memory_entries (
+          id TEXT PRIMARY KEY,
+          key TEXT NOT NULL,
+          content TEXT NOT NULL,
+          type TEXT DEFAULT 'semantic',
+          namespace TEXT DEFAULT 'general',
+          tags TEXT DEFAULT '[]',
+          embedding BLOB,
+          created_at INTEGER NOT NULL,
+          updated_at INTEGER NOT NULL
+        )
+      `);
+
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const embeddingGenerator = async (text: string): Promise<Float32Array> => {
+        const result = await embeddingsService.embed(text);
+        return result.embedding;
+      };
+
+      const engine = new HybridSearchEngine(db, {}, embeddingGenerator);
+      await engine.initialize();
+
+      const now = Date.now();
+      const embedding = await embeddingGenerator('Machine learning algorithms');
+      const embeddingBuffer = Buffer.from(embedding.buffer);
+
+      db.prepare(`
+        INSERT INTO memory_entries (id, key, content, embedding, created_at, updated_at)
+        VALUES (?, ?, ?, ?, ?, ?)
+      `).run('1', 'ml', 'Machine learning algorithms', embeddingBuffer, now, now);
+
+      // Vector-only search (no keyword)
+      const results = await engine.searchCompact('AI neural networks', {
+        limit: 10,
+        includeKeyword: false,
+        includeSemantic: true,
+      });
+
+      expect(results.length).toBeGreaterThan(0);
+      expect(results[0].keywordScore).toBe(0); // Keyword disabled
+      expect(results[0].semanticScore).toBeGreaterThan(0);
+
+      db.close();
+    });
+  });
+
+  describe('CJK language support with embeddings', () => {
+    it('should generate embeddings for Japanese text', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const result = await embeddingsService.embed('これは日本語のテストです');
+
+      expect(result.embedding).toBeDefined();
+      expect(result.embedding.length).toBe(384);
+    });
+
+    it('should generate embeddings for Chinese text', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const result = await embeddingsService.embed('这是中文测试');
+
+      expect(result.embedding).toBeDefined();
+      expect(result.embedding.length).toBe(384);
+    });
+
+    it('should find semantically similar CJK content', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const embeddingGenerator = async (text: string): Promise<Float32Array> => {
+        const result = await embeddingsService.embed(text);
+        return result.embedding;
+      };
+
+      const service = new ProjectMemoryService({
+        baseDir: tempDir,
+        dbFilename: 'cjk.db',
+        embeddingGenerator,
+      });
+      await service.initialize();
+
+      await service.storeEntry({
+        key: 'japanese-greeting',
+        content: 'おはようございます。今日はいい天気ですね。',
+        namespace: 'test',
+      });
+
+      await service.storeEntry({
+        key: 'japanese-farewell',
+        content: 'さようなら。また会いましょう。',
+        namespace: 'test',
+      });
+
+      // Search for morning greeting
+      const results = await service.semanticSearch('朝の挨拶', 5);
+
+      expect(results.length).toBeGreaterThan(0);
+
+      await service.shutdown();
+    });
+  });
+
+  describe('Embedding caching', () => {
+    it('should cache embeddings and return faster on second call', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      const text = 'This is a test sentence for caching';
+
+      // First call - should compute
+      const result1 = await embeddingsService.embed(text);
+      expect(result1.cached).toBe(false);
+
+      // Second call - should be cached
+      const result2 = await embeddingsService.embed(text);
+      expect(result2.cached).toBe(true);
+
+      // Embeddings should be identical
+      expect(result1.embedding.length).toBe(result2.embedding.length);
+      const areSame = result1.embedding.every(
+        (val, i) => val === result2.embedding[i]
+      );
+      expect(areSame).toBe(true);
+    });
+
+    it('should use in-memory cache within same service instance', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+        maxCacheSize: 100,
+      });
+      await embeddingsService.initialize();
+
+      const text = 'Cache consistency test';
+
+      // First call
+      const result1 = await embeddingsService.embed(text);
+      expect(result1.cached).toBe(false);
+
+      // Second call with same text - should hit cache
+      const result2 = await embeddingsService.embed(text);
+      expect(result2.cached).toBe(true);
+
+      // Third call with different text - cache miss
+      const result3 = await embeddingsService.embed('Different text');
+      expect(result3.cached).toBe(false);
+
+      // Same different text again - cache hit
+      const result4 = await embeddingsService.embed('Different text');
+      expect(result4.cached).toBe(true);
+    });
+  });
+
+  describe('Error handling', () => {
+    it('should handle embedding generation failure gracefully', async () => {
+      const failingGenerator = async (): Promise<Float32Array> => {
+        throw new Error('Embedding service unavailable');
+      };
+
+      const service = new ProjectMemoryService({
+        baseDir: tempDir,
+        dbFilename: 'fail.db',
+        embeddingGenerator: failingGenerator,
+      });
+      await service.initialize();
+
+      // Should still store entry even if embedding fails
+      const entry = await service.storeEntry({
+        key: 'no-embedding',
+        content: 'Content without embedding due to failure',
+        namespace: 'test',
+      });
+
+      expect(entry.id).toBeDefined();
+      expect(entry.embedding).toBeUndefined();
+
+      await service.shutdown();
+    });
+
+    it('should handle empty content gracefully', async () => {
+      const embeddingsService = new LocalEmbeddingsService({
+        cacheDir: path.join(tempDir, 'cache'),
+      });
+      await embeddingsService.initialize();
+
+      // Empty string should still work
+      const result = await embeddingsService.embed('');
+      expect(result.embedding).toBeDefined();
+      expect(result.embedding.length).toBe(384);
+    });
+  });
+});