learning_ai_common_plat/packages/ollama-client/src/embed.test.ts
2026-03-29 12:43:01 -07:00

67 lines
1.9 KiB
TypeScript

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { getEmbedding, getEmbeddingVector } from './embed.js';
const BASE_URL = 'http://localhost:11434';
describe('getEmbedding', () => {
beforeEach(() => {
vi.restoreAllMocks();
});
it('returns embedding response', async () => {
const response = { model: 'nomic-embed-text', embeddings: [[0.1, 0.2, 0.3]] };
globalThis.fetch = vi.fn().mockResolvedValue({
ok: true,
json: () => Promise.resolve(response),
});
const result = await getEmbedding(BASE_URL, { model: 'nomic-embed-text', input: 'hello' });
expect(result).toEqual(response);
expect(globalThis.fetch).toHaveBeenCalledWith(
`${BASE_URL}/api/embed`,
expect.objectContaining({
method: 'POST',
body: JSON.stringify({ model: 'nomic-embed-text', input: 'hello' }),
})
);
});
it('throws on error response', async () => {
globalThis.fetch = vi.fn().mockResolvedValue({
ok: false,
status: 500,
text: () => Promise.resolve('internal error'),
});
await expect(
getEmbedding(BASE_URL, { model: 'nomic-embed-text', input: 'hello' })
).rejects.toThrow('Ollama embed failed (500)');
});
});
describe('getEmbeddingVector', () => {
beforeEach(() => {
vi.restoreAllMocks();
});
it('returns first embedding vector', async () => {
globalThis.fetch = vi.fn().mockResolvedValue({
ok: true,
json: () => Promise.resolve({ model: 'nomic-embed-text', embeddings: [[0.1, 0.2]] }),
});
const result = await getEmbeddingVector(BASE_URL, 'hello');
expect(result).toEqual([0.1, 0.2]);
});
it('returns empty array when no embeddings', async () => {
globalThis.fetch = vi.fn().mockResolvedValue({
ok: true,
json: () => Promise.resolve({ model: 'nomic-embed-text', embeddings: [] }),
});
const result = await getEmbeddingVector(BASE_URL, 'hello');
expect(result).toEqual([]);
});
});