@aitytech/agentkits-memory 1.0.1 → 2.0.1

package/src/embeddings/local-embeddings.ts

@@ -0,0 +1,419 @@
+/**
+ * Local Offline Embeddings Service
+ *
+ * Provides 100% offline text embeddings using Transformers.js (WASM-based).
+ * No external API calls. Model is downloaded once and cached locally.
+ *
+ * @module @aitytech/agentkits-memory/embeddings
+ */
+
+import type { EmbeddingGenerator } from '../types.js';
+
+/**
+ * Embedding provider type
+ */
+export type EmbeddingProvider = 'transformers' | 'mock';
+
+/**
+ * Local embeddings configuration
+ */
+export interface LocalEmbeddingsConfig {
+  /** Provider to use (default: 'transformers') */
+  provider?: EmbeddingProvider;
+
+  /**
+   * Model ID for Transformers.js
+   * Default: 'Xenova/multilingual-e5-small' (100+ languages, optimized for retrieval)
+   * Alternative: 'Xenova/paraphrase-multilingual-MiniLM-L12-v2' (50+ languages)
+   * Alternative: 'Xenova/all-MiniLM-L6-v2' (English only, faster)
+   */
+  modelId?: string;
+
+  /** Vector dimensions (default: 384) */
+  dimensions?: number;
+
+  /** Enable in-memory cache for repeated texts */
+  cacheEnabled?: boolean;
+
+  /** Maximum cache size (default: 1000) */
+  maxCacheSize?: number;
+
+  /** Show progress during model download */
+  showProgress?: boolean;
+
+  /** Custom cache directory for models */
+  cacheDir?: string;
+}
+
+/**
+ * Embedding result with metadata
+ */
+export interface EmbeddingResult {
+  /** The embedding vector */
+  embedding: Float32Array;
+
+  /** Time taken in milliseconds */
+  timeMs: number;
+
+  /** Whether result was from cache */
+  cached: boolean;
+
+  /** Token count (approximate) */
+  tokenCount?: number;
+}
+
+/**
+ * Local embeddings service statistics
+ */
+export interface EmbeddingsStats {
+  /** Total embeddings generated */
+  totalEmbeddings: number;
+
+  /** Cache hits */
+  cacheHits: number;
+
+  /** Cache misses */
+  cacheMisses: number;
+
+  /** Average time per embedding (ms) */
+  avgTimeMs: number;
+
+  /** Total time spent (ms) */
+  totalTimeMs: number;
+
+  /** Model loaded */
+  modelLoaded: boolean;
+
+  /** Provider being used */
+  provider: EmbeddingProvider;
+}
+
+/**
+ * In-memory LRU Cache for embeddings
+ */
+class InMemoryEmbeddingCache {
+  private cache = new Map<string, Float32Array>();
+  private accessOrder: string[] = [];
+  private maxSize: number;
+
+  constructor(maxSize: number) {
+    this.maxSize = maxSize;
+  }
+
+  get(key: string): Float32Array | undefined {
+    const value = this.cache.get(key);
+    if (value) {
+      // Move to end (most recently used)
+      const index = this.accessOrder.indexOf(key);
+      if (index > -1) {
+        this.accessOrder.splice(index, 1);
+      }
+      this.accessOrder.push(key);
+    }
+    return value;
+  }
+
+  set(key: string, value: Float32Array): void {
+    if (this.cache.has(key)) {
+      this.cache.set(key, value);
+      return;
+    }
+
+    // Evict if at capacity
+    while (this.cache.size >= this.maxSize && this.accessOrder.length > 0) {
+      const oldest = this.accessOrder.shift();
+      if (oldest) {
+        this.cache.delete(oldest);
+      }
+    }
+
+    this.cache.set(key, value);
+    this.accessOrder.push(key);
+  }
+
+  clear(): void {
+    this.cache.clear();
+    this.accessOrder = [];
+  }
+
+  get size(): number {
+    return this.cache.size;
+  }
+}
+
+/**
+ * Mock embedding provider for testing
+ */
+function createMockEmbedding(text: string, dimensions: number): Float32Array {
+  const embedding = new Float32Array(dimensions);
+  // Generate deterministic pseudo-random values based on text hash
+  let hash = 0;
+  for (let i = 0; i < text.length; i++) {
+    hash = ((hash << 5) - hash) + text.charCodeAt(i);
+    hash = hash & hash;
+  }
+
+  for (let i = 0; i < dimensions; i++) {
+    // Use hash to seed pseudo-random generation
+    hash = ((hash << 5) - hash) + i;
+    hash = hash & hash;
+    embedding[i] = (hash % 1000) / 1000 - 0.5;
+  }
+
+  // Normalize
+  let norm = 0;
+  for (let i = 0; i < dimensions; i++) {
+    norm += embedding[i] * embedding[i];
+  }
+  norm = Math.sqrt(norm);
+  if (norm > 0) {
+    for (let i = 0; i < dimensions; i++) {
+      embedding[i] /= norm;
+    }
+  }
+
+  return embedding;
+}
+
+/**
+ * Local Embeddings Service
+ *
+ * Provides offline text embeddings using Transformers.js.
+ * Models are downloaded once and cached locally in ~/.cache/huggingface.
+ */
+export class LocalEmbeddingsService {
+  private config: Required<LocalEmbeddingsConfig>;
+  private cache: InMemoryEmbeddingCache | null = null;
+  private pipeline: any = null;
+  private modelLoading: Promise<void> | null = null;
+  private stats = {
+    totalEmbeddings: 0,
+    cacheHits: 0,
+    cacheMisses: 0,
+    totalTimeMs: 0,
+  };
+
+  constructor(config: LocalEmbeddingsConfig = {}) {
+    this.config = {
+      provider: config.provider || 'transformers',
+      // Use multilingual-e5-small for best CJK support and retrieval accuracy
+      modelId: config.modelId || 'Xenova/multilingual-e5-small',
+      dimensions: config.dimensions || 384,
+      cacheEnabled: config.cacheEnabled ?? true,
+      maxCacheSize: config.maxCacheSize || 1000,
+      showProgress: config.showProgress ?? false,
+      cacheDir: config.cacheDir || '',
+    };
+
+    if (this.config.cacheEnabled) {
+      this.cache = new InMemoryEmbeddingCache(this.config.maxCacheSize);
+    }
+  }
+
+  /**
+   * Initialize the embeddings service (loads model)
+   */
+  async initialize(): Promise<void> {
+    if (this.config.provider === 'mock') {
+      return;
+    }
+
+    if (this.pipeline) {
+      return;
+    }
+
+    if (this.modelLoading) {
+      await this.modelLoading;
+      return;
+    }
+
+    this.modelLoading = this.loadModel();
+    await this.modelLoading;
+  }
+
+  private async loadModel(): Promise<void> {
+    try {
+      // Dynamic import for Transformers.js
+      const { pipeline } = await import('@xenova/transformers');
+
+      const progressCallback = this.config.showProgress
+        ? (progress: { status: string; progress?: number }) => {
+            if (progress.status === 'progress' && progress.progress !== undefined) {
+              process.stderr.write(
+                `\rLoading model: ${Math.round(progress.progress)}%`
+              );
+            } else if (progress.status === 'done') {
+              process.stderr.write('\rModel loaded successfully.          \n');
+            }
+          }
+        : undefined;
+
+      this.pipeline = await pipeline('feature-extraction', this.config.modelId, {
+        progress_callback: progressCallback,
+      });
+    } catch (error) {
+      // If Transformers.js is not available, fall back to mock
+      console.warn(
+        'Transformers.js not available, falling back to mock embeddings.',
+        'Install with: npm install @xenova/transformers'
+      );
+      this.config.provider = 'mock';
+    }
+  }
+
+  /**
+   * Generate embedding for text
+   */
+  async embed(text: string): Promise<EmbeddingResult> {
+    const startTime = performance.now();
+
+    // Check cache
+    if (this.cache) {
+      const cached = this.cache.get(text);
+      if (cached) {
+        this.stats.cacheHits++;
+        return {
+          embedding: cached,
+          timeMs: performance.now() - startTime,
+          cached: true,
+        };
+      }
+      this.stats.cacheMisses++;
+    }
+
+    // Generate embedding
+    let embedding: Float32Array;
+
+    if (this.config.provider === 'mock') {
+      embedding = createMockEmbedding(text, this.config.dimensions);
+    } else {
+      await this.initialize();
+
+      if (!this.pipeline) {
+        // Fall back to mock if model failed to load
+        embedding = createMockEmbedding(text, this.config.dimensions);
+      } else {
+        const output = await this.pipeline(text, {
+          pooling: 'mean',
+          normalize: true,
+        });
+        embedding = new Float32Array(output.data);
+      }
+    }
+
+    const timeMs = performance.now() - startTime;
+
+    // Update stats
+    this.stats.totalEmbeddings++;
+    this.stats.totalTimeMs += timeMs;
+
+    // Cache result
+    if (this.cache) {
+      this.cache.set(text, embedding);
+    }
+
+    return {
+      embedding,
+      timeMs,
+      cached: false,
+    };
+  }
+
+  /**
+   * Generate embeddings for multiple texts (batch)
+   */
+  async embedBatch(texts: string[]): Promise<EmbeddingResult[]> {
+    // Process in parallel for better performance
+    return Promise.all(texts.map((text) => this.embed(text)));
+  }
+
+  /**
+   * Get embedding generator function compatible with ProjectMemoryService
+   */
+  getGenerator(): EmbeddingGenerator {
+    return async (content: string): Promise<Float32Array> => {
+      const result = await this.embed(content);
+      return result.embedding;
+    };
+  }
+
+  /**
+   * Get service statistics
+   */
+  getStats(): EmbeddingsStats {
+    return {
+      totalEmbeddings: this.stats.totalEmbeddings,
+      cacheHits: this.stats.cacheHits,
+      cacheMisses: this.stats.cacheMisses,
+      avgTimeMs:
+        this.stats.totalEmbeddings > 0
+          ? this.stats.totalTimeMs / this.stats.totalEmbeddings
+          : 0,
+      totalTimeMs: this.stats.totalTimeMs,
+      modelLoaded: this.pipeline !== null || this.config.provider === 'mock',
+      provider: this.config.provider,
+    };
+  }
+
+  /**
+   * Clear the embedding cache
+   */
+  clearCache(): void {
+    if (this.cache) {
+      this.cache.clear();
+    }
+  }
+
+  /**
+   * Get vector dimensions
+   */
+  getDimensions(): number {
+    return this.config.dimensions;
+  }
+
+  /**
+   * Shutdown and cleanup
+   */
+  async shutdown(): Promise<void> {
+    this.clearCache();
+    this.pipeline = null;
+    this.modelLoading = null;
+  }
+}
+
+/**
+ * Create a local embeddings service with default configuration
+ */
+export function createLocalEmbeddings(
+  config?: LocalEmbeddingsConfig
+): LocalEmbeddingsService {
+  return new LocalEmbeddingsService(config);
+}
+
+/**
+ * Create an embedding generator function for use with ProjectMemoryService
+ *
+ * @example
+ * ```typescript
+ * import { createEmbeddingGenerator } from '@aitytech/agentkits-memory/embeddings';
+ * import { ProjectMemoryService } from '@aitytech/agentkits-memory';
+ *
+ * const embeddingGenerator = await createEmbeddingGenerator();
+ *
+ * const memory = new ProjectMemoryService({
+ *   projectPath: '/path/to/project',
+ *   enableVectorIndex: true,
+ *   embeddingGenerator,
+ * });
+ * ```
+ */
+export async function createEmbeddingGenerator(
+  config?: LocalEmbeddingsConfig
+): Promise<EmbeddingGenerator> {
+  const service = new LocalEmbeddingsService(config);
+  await service.initialize();
+  return service.getGenerator();
+}
+
+// Default export
+export default LocalEmbeddingsService;

package/src/hooks/__tests__/handlers.test.ts

@@ -16,6 +16,9 @@ import { SummarizeHook, createSummarizeHook } from '../summarize.js';
 
 const TEST_DIR = path.join(process.cwd(), '.test-hook-handlers');
 
+// Track hooks for cleanup
+let activeHooks: Array<{ shutdown: () => Promise<void> }> = [];
+
 function createTestInput(overrides: Partial<NormalizedHookInput> = {}): NormalizedHookInput {
   return {
     sessionId: 'test-session-123',
@@ -26,25 +29,53 @@ function createTestInput(overrides: Partial<NormalizedHookInput> = {}): Normaliz
   };
 }
 
+// Helper to track hooks for cleanup
+function trackHook<T extends { shutdown: () => Promise<void> }>(hook: T): T {
+  activeHooks.push(hook);
+  return hook;
+}
+
 describe('Hook Handlers', () => {
   beforeEach(() => {
+    activeHooks = [];
     // Clean up test directory
     if (existsSync(TEST_DIR)) {
-      rmSync(TEST_DIR, { recursive: true });
+      try {
+        rmSync(TEST_DIR, { recursive: true });
+      } catch {
+        // Ignore errors on Windows
+      }
     }
     mkdirSync(TEST_DIR, { recursive: true });
   });
 
-  afterEach(() => {
+  afterEach(async () => {
+    // Shutdown all hooks first (releases database locks)
+    for (const hook of activeHooks) {
+      try {
+        await hook.shutdown();
+      } catch {
+        // Ignore shutdown errors
+      }
+    }
+    activeHooks = [];
+
+    // Small delay for Windows file system
+    await new Promise((r) => setTimeout(r, 100));
+
     // Clean up test directory
     if (existsSync(TEST_DIR)) {
-      rmSync(TEST_DIR, { recursive: true });
+      try {
+        rmSync(TEST_DIR, { recursive: true });
+      } catch {
+        // Ignore errors on Windows - files may still be locked
+      }
     }
   });
 
   describe('ContextHook', () => {
     it('should return no context for new project', async () => {
-      const hook = createContextHook(TEST_DIR);
+      const hook = trackHook(createContextHook(TEST_DIR));
       const input = createTestInput();
 
       const result = await hook.execute(input);
@@ -63,7 +94,7 @@ describe('Hook Handlers', () => {
       await service.shutdown();
 
       // Run context hook
-      const hook = createContextHook(TEST_DIR);
+      const hook = trackHook(createContextHook(TEST_DIR));
       const input = createTestInput({ sessionId: 'new-session' });
 
       const result = await hook.execute(input);
@@ -92,7 +123,7 @@ describe('Hook Handlers', () => {
 
   describe('SessionInitHook', () => {
     it('should initialize a new session', async () => {
-      const hook = createSessionInitHook(TEST_DIR);
+      const hook = trackHook(createSessionInitHook(TEST_DIR));
       const input = createTestInput({ prompt: 'Hello Claude' });
 
       const result = await hook.execute(input);
@@ -100,6 +131,9 @@ describe('Hook Handlers', () => {
       expect(result.continue).toBe(true);
       expect(result.suppressOutput).toBe(true);
 
+      // Shutdown hook before verifying
+      await hook.shutdown();
+
       // Verify session was created
       const service = new MemoryHookService(TEST_DIR);
       await service.initialize();
@@ -112,12 +146,14 @@ describe('Hook Handlers', () => {
 
     it('should not overwrite existing session', async () => {
       // Create initial session
-      const hook1 = createSessionInitHook(TEST_DIR);
+      const hook1 = trackHook(createSessionInitHook(TEST_DIR));
       await hook1.execute(createTestInput({ prompt: 'First prompt' }));
+      await hook1.shutdown();
 
       // Try to re-init with different prompt
-      const hook2 = createSessionInitHook(TEST_DIR);
+      const hook2 = trackHook(createSessionInitHook(TEST_DIR));
       await hook2.execute(createTestInput({ prompt: 'Second prompt' }));
+      await hook2.shutdown();
 
       // Verify original prompt preserved
       const service = new MemoryHookService(TEST_DIR);
@@ -145,11 +181,12 @@ describe('Hook Handlers', () => {
   describe('ObservationHook', () => {
     it('should store observation', async () => {
       // Initialize session first
-      const initHook = createSessionInitHook(TEST_DIR);
+      const initHook = trackHook(createSessionInitHook(TEST_DIR));
       await initHook.execute(createTestInput());
+      await initHook.shutdown();
 
       // Store observation
-      const hook = createObservationHook(TEST_DIR);
+      const hook = trackHook(createObservationHook(TEST_DIR));
       const input = createTestInput({
         toolName: 'Read',
         toolInput: { file_path: '/path/to/file.ts' },
@@ -157,6 +194,7 @@ describe('Hook Handlers', () => {
       });
 
       const result = await hook.execute(input);
+      await hook.shutdown();
 
       expect(result.continue).toBe(true);
       expect(result.suppressOutput).toBe(true);
@@ -172,7 +210,7 @@ describe('Hook Handlers', () => {
     });
 
     it('should skip if no tool name', async () => {
-      const hook = createObservationHook(TEST_DIR);
+      const hook = trackHook(createObservationHook(TEST_DIR));
       const input = createTestInput({ toolName: undefined });
 
       const result = await hook.execute(input);
@@ -183,10 +221,11 @@ describe('Hook Handlers', () => {
 
     it('should skip internal tools', async () => {
       // Initialize session first
-      const initHook = createSessionInitHook(TEST_DIR);
+      const initHook = trackHook(createSessionInitHook(TEST_DIR));
       await initHook.execute(createTestInput());
+      await initHook.shutdown();
 
-      const hook = createObservationHook(TEST_DIR);
+      const hook = trackHook(createObservationHook(TEST_DIR));
 
       // Test skipped tools
       for (const tool of ['TodoWrite', 'TodoRead', 'AskFollowupQuestion', 'AttemptCompletion']) {
@@ -196,6 +235,7 @@ describe('Hook Handlers', () => {
         expect(result.continue).toBe(true);
         expect(result.suppressOutput).toBe(true);
       }
+      await hook.shutdown();
 
       // Verify no observations stored
       const service = new MemoryHookService(TEST_DIR);
@@ -207,7 +247,7 @@ describe('Hook Handlers', () => {
     });
 
     it('should create session if not exists', async () => {
-      const hook = createObservationHook(TEST_DIR);
+      const hook = trackHook(createObservationHook(TEST_DIR));
       const input = createTestInput({
         sessionId: 'new-session',
         toolName: 'Read',
@@ -216,6 +256,7 @@ describe('Hook Handlers', () => {
       });
 
       const result = await hook.execute(input);
+      await hook.shutdown();
 
       expect(result.continue).toBe(true);
 
@@ -251,8 +292,8 @@ describe('Hook Handlers', () => {
       await service.storeObservation('test-session-123', 'test-project', 'Write', { file_path: 'b.ts' }, {}, TEST_DIR);
       await service.shutdown();
 
-      // Run summarize hook
-      const hook = createSummarizeHook(TEST_DIR);
+      // Run summarize hook (it already calls shutdown internally)
+      const hook = trackHook(createSummarizeHook(TEST_DIR));
       const input = createTestInput();
 
       const result = await hook.execute(input);
@@ -272,7 +313,7 @@ describe('Hook Handlers', () => {
     });
 
     it('should handle non-existent session', async () => {
-      const hook = createSummarizeHook(TEST_DIR);
+      const hook = trackHook(createSummarizeHook(TEST_DIR));
       const input = createTestInput({ sessionId: 'non-existent' });
 
       const result = await hook.execute(input);