flynn/src/memory/embeddings.test.ts

import { describe, it, expect } from 'vitest';
import {
  createEmbeddingProvider,
  OpenAIEmbeddingProvider,
  GeminiEmbeddingProvider,
  OllamaEmbeddingProvider,
  LlamaCppEmbeddingProvider,
  VoyageAIEmbeddingProvider,
} from './embeddings.js';
import type { EmbeddingConfig } from '../config/schema.js';

describe('createEmbeddingProvider', () => {
  const baseConfig: EmbeddingConfig = {
    enabled: true,
    provider: 'openai',
    model: 'text-embedding-3-small',
    chunk_size: 512,
    chunk_overlap: 50,
    top_k: 5,
    hybrid_weight: 0.7,
  };

  it('creates OpenAI provider', () => {
    const provider = createEmbeddingProvider({ ...baseConfig, provider: 'openai' });
    expect(provider).toBeInstanceOf(OpenAIEmbeddingProvider);
  });

  it('creates Gemini provider', () => {
    const provider = createEmbeddingProvider({ ...baseConfig, provider: 'gemini' });
    expect(provider).toBeInstanceOf(GeminiEmbeddingProvider);
  });

  it('creates Ollama provider', () => {
    const provider = createEmbeddingProvider({ ...baseConfig, provider: 'ollama' });
    expect(provider).toBeInstanceOf(OllamaEmbeddingProvider);
  });

  it('creates LlamaCpp provider', () => {
    const provider = createEmbeddingProvider({ ...baseConfig, provider: 'llamacpp' });
    expect(provider).toBeInstanceOf(LlamaCppEmbeddingProvider);
  });

  it('creates Voyage provider', () => {
    const provider = createEmbeddingProvider({ ...baseConfig, provider: 'voyage' });
    expect(provider).toBeInstanceOf(VoyageAIEmbeddingProvider);
  });

  it('throws on unknown provider', () => {
    expect(() => createEmbeddingProvider({ ...baseConfig, provider: 'unknown' as never })).toThrow('Unknown embedding provider');
  });
});

describe('OpenAIEmbeddingProvider', () => {
  it('reports configured dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'openai',
      model: 'text-embedding-3-small',
      dimensions: 512,
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new OpenAIEmbeddingProvider(config);
    expect(provider.dimensions).toBe(512);
  });

  it('defaults to 1536 dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'openai',
      model: 'text-embedding-3-small',
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new OpenAIEmbeddingProvider(config);
    expect(provider.dimensions).toBe(1536);
  });
});

describe('GeminiEmbeddingProvider', () => {
  it('reports configured dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'gemini',
      model: 'text-embedding-004',
      dimensions: 256,
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new GeminiEmbeddingProvider(config);
    expect(provider.dimensions).toBe(256);
  });

  it('defaults to 768 dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'gemini',
      model: 'text-embedding-004',
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new GeminiEmbeddingProvider(config);
    expect(provider.dimensions).toBe(768);
  });
});

describe('OllamaEmbeddingProvider', () => {
  it('reports configured dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'ollama',
      model: 'nomic-embed-text',
      dimensions: 384,
      endpoint: 'http://localhost:11434',
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new OllamaEmbeddingProvider(config);
    expect(provider.dimensions).toBe(384);
  });
});

describe('LlamaCppEmbeddingProvider', () => {
  it('reports configured dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'llamacpp',
      model: 'unused',
      dimensions: 768,
      endpoint: 'http://localhost:8080',
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new LlamaCppEmbeddingProvider(config);
    expect(provider.dimensions).toBe(768);
  });

  it('defaults endpoint to localhost:8080', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'llamacpp',
      model: 'unused',
      dimensions: 768,
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    // Provider should be constructable without endpoint
    const provider = new LlamaCppEmbeddingProvider(config);
    expect(provider.dimensions).toBe(768);
  });
});

describe('VoyageAIEmbeddingProvider', () => {
  it('defaults to 1024 dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'voyage',
      model: 'voyage-3',
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new VoyageAIEmbeddingProvider(config);
    expect(provider.dimensions).toBe(1024);
  });

  it('reports configured dimensions', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'voyage',
      model: 'voyage-3-lite',
      dimensions: 512,
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    const provider = new VoyageAIEmbeddingProvider(config);
    expect(provider.dimensions).toBe(512);
  });

  it('uses custom endpoint if provided', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'voyage',
      model: 'voyage-3',
      endpoint: 'https://custom.proxy.example.com/v1',
      api_key: 'test-key',
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    // Should not throw when constructing with custom endpoint
    const provider = new VoyageAIEmbeddingProvider(config);
    expect(provider.dimensions).toBe(1024);
  });

  it('uses api_key from config', () => {
    const config: EmbeddingConfig = {
      enabled: true,
      provider: 'voyage',
      model: 'voyage-3',
      api_key: 'voy-test-key-123',
      chunk_size: 512,
      chunk_overlap: 50,
      top_k: 5,
      hybrid_weight: 0.7,
    };
    // Should construct without error when api_key is provided
    const provider = new VoyageAIEmbeddingProvider(config);
    expect(provider.dimensions).toBe(1024);
  });
});