/**
* Unit tests for EmbeddingService
*/
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { EmbeddingService } from '../../src/utils/embeddings.js';
// Mock OpenAI
vi.mock('openai', () => ({
default: vi.fn().mockImplementation(() => ({
embeddings: {
create: vi.fn(),
},
})),
}));
describe('EmbeddingService', () => {
let embeddingService: EmbeddingService;
let mockOpenAI: any;
beforeEach(() => {
vi.clearAllMocks();
embeddingService = new EmbeddingService('test-api-key');
mockOpenAI = (embeddingService as any).openai;
});
describe('generateEmbedding', () => {
it('should generate embedding for text', async () => {
const mockEmbedding = [0.1, 0.2, 0.3];
mockOpenAI.embeddings.create.mockResolvedValue({
data: [{ embedding: mockEmbedding }],
});
const result = await embeddingService.generateEmbedding('test text');
expect(result).toEqual(mockEmbedding);
expect(mockOpenAI.embeddings.create).toHaveBeenCalledWith({
model: 'text-embedding-3-small',
input: 'test text',
});
});
it('should handle empty embedding response', async () => {
mockOpenAI.embeddings.create.mockResolvedValue({
data: [],
});
await expect(embeddingService.generateEmbedding('test text')).rejects.toThrow(
'OpenAI API returned invalid or empty embedding'
);
});
it('should throw error on API failure', async () => {
mockOpenAI.embeddings.create.mockRejectedValue(new Error('API Error'));
await expect(embeddingService.generateEmbedding('test text')).rejects.toThrow(
'Embedding generation failed'
);
});
});
describe('generateEmbeddings', () => {
it('should generate embeddings for multiple texts', async () => {
const mockEmbeddings = [[0.1, 0.2], [0.3, 0.4]];
mockOpenAI.embeddings.create.mockResolvedValue({
data: mockEmbeddings.map(embedding => ({ embedding })),
});
const result = await embeddingService.generateEmbeddings(['text1', 'text2']);
expect(result).toEqual(mockEmbeddings);
expect(mockOpenAI.embeddings.create).toHaveBeenCalledWith({
model: 'text-embedding-3-small',
input: ['text1', 'text2'],
});
});
});
describe('cosineSimilarity', () => {
it('should calculate cosine similarity correctly', () => {
const a = [1, 0, 0];
const b = [0, 1, 0];
const result = EmbeddingService.cosineSimilarity(a, b);
expect(result).toBe(0);
});
it('should return 1 for identical vectors', () => {
const a = [1, 2, 3];
const b = [1, 2, 3];
const result = EmbeddingService.cosineSimilarity(a, b);
expect(result).toBeCloseTo(1, 5);
});
it('should return 0 for different length vectors', () => {
const a = [1, 2];
const b = [1, 2, 3];
// Changed from throwing to returning 0 for better resilience
const result = EmbeddingService.cosineSimilarity(a, b);
expect(result).toBe(0);
});
it('should handle zero vectors', () => {
const a = [0, 0, 0];
const b = [1, 2, 3];
const result = EmbeddingService.cosineSimilarity(a, b);
expect(result).toBe(0);
});
});
describe('findMostSimilar', () => {
it('should find most similar vectors', () => {
const queryVector = [1, 0, 0];
const vectors = [
{ vector: [1, 0, 0], data: 'exact match' },
{ vector: [0.9, 0.1, 0], data: 'close match' },
{ vector: [0, 1, 0], data: 'orthogonal' },
{ vector: [-1, 0, 0], data: 'opposite' },
];
const results = EmbeddingService.findMostSimilar(queryVector, vectors, 0.5, 2);
expect(results).toHaveLength(2);
expect(results[0]?.data).toBe('exact match');
expect(results[0]?.similarity).toBeCloseTo(1, 5);
expect(results[1]?.data).toBe('close match');
});
it('should filter by threshold', () => {
const queryVector = [1, 0, 0];
const vectors = [
{ vector: [0.5, 0.5, 0], data: 'low similarity' },
{ vector: [0, 1, 0], data: 'orthogonal' },
];
const results = EmbeddingService.findMostSimilar(queryVector, vectors, 0.8, 10);
expect(results).toHaveLength(0);
});
});
describe('prepareTextForEmbedding', () => {
it('should clean and normalize text', () => {
const text = ' Hello world! \n\n This is a test. ';
const result = EmbeddingService.prepareTextForEmbedding(text);
expect(result).toBe('Hello world! This is a test.');
});
it('should truncate long text', () => {
const longText = 'a'.repeat(10000);
const result = EmbeddingService.prepareTextForEmbedding(longText, 100);
expect(result.length).toBeLessThanOrEqual(100);
});
it('should end at word boundary when truncating', () => {
const text = 'word1 word2 word3 word4 word5';
const result = EmbeddingService.prepareTextForEmbedding(text, 15);
// The function only truncates at word boundary if lastSpace > maxLength * 0.8
// For maxLength=15, that's 12, but lastSpace=11, so it doesn't truncate
expect(result).toBe('word1 word2 wor');
});
});
describe('createMemoryEmbeddingText', () => {
it('should combine memory fields correctly', () => {
const memory = {
title: 'Test Memory',
content: 'This is test content',
tags: ['tag1', 'tag2'],
memoryType: 'personal',
};
const result = EmbeddingService.createMemoryEmbeddingText(memory);
expect(result).toContain('Test Memory');
expect(result).toContain('This is test content');
expect(result).toContain('Type: personal');
expect(result).toContain('Tags: tag1, tag2');
});
it('should handle missing optional fields', () => {
const memory = {
title: 'Test Memory',
content: 'This is test content',
};
const result = EmbeddingService.createMemoryEmbeddingText(memory);
expect(result).toContain('Test Memory');
expect(result).toContain('This is test content');
expect(result).not.toContain('Type:');
expect(result).not.toContain('Tags:');
});
});
describe('createEntityEmbeddingText', () => {
it('should combine entity fields correctly', () => {
const entity = {
name: 'John Doe',
description: 'Software engineer',
company: 'Tech Corp',
title: 'Senior Developer',
notes: 'Great at TypeScript',
tags: ['developer', 'typescript'],
};
const result = EmbeddingService.createEntityEmbeddingText(entity);
expect(result).toContain('John Doe');
expect(result).toContain('Software engineer');
expect(result).toContain('Company: Tech Corp');
expect(result).toContain('Title: Senior Developer');
expect(result).toContain('Great at TypeScript');
expect(result).toContain('Tags: developer, typescript');
});
});
describe('normalizeVector', () => {
it('should normalize vector to unit length', () => {
const vector = [3, 4, 0];
const normalized = EmbeddingService.normalizeVector(vector);
// Length should be 1
const length = Math.sqrt(normalized.reduce((sum, val) => sum + val * val, 0));
expect(length).toBeCloseTo(1, 5);
// Direction should be preserved
expect(normalized[0]).toBeCloseTo(0.6, 5);
expect(normalized[1]).toBeCloseTo(0.8, 5);
expect(normalized[2]).toBe(0);
});
it('should handle zero vector', () => {
const vector = [0, 0, 0];
const normalized = EmbeddingService.normalizeVector(vector);
expect(normalized).toEqual([0, 0, 0]);
});
});
describe('dotProduct', () => {
it('should calculate dot product correctly', () => {
const a = [1, 2, 3];
const b = [4, 5, 6];
const result = EmbeddingService.dotProduct(a, b);
expect(result).toBe(32); // 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32
});
it('should throw error for different length vectors', () => {
const a = [1, 2];
const b = [1, 2, 3];
expect(() => EmbeddingService.dotProduct(a, b)).toThrow(
'Vectors must have the same length'
);
});
});
describe('euclideanDistance', () => {
it('should calculate Euclidean distance correctly', () => {
const a = [0, 0, 0];
const b = [3, 4, 0];
const result = EmbeddingService.euclideanDistance(a, b);
expect(result).toBe(5); // sqrt(3^2 + 4^2) = 5
});
it('should return 0 for identical vectors', () => {
const a = [1, 2, 3];
const b = [1, 2, 3];
const result = EmbeddingService.euclideanDistance(a, b);
expect(result).toBe(0);
});
});
});
