feat(ai): ollama embed-text wrapper
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
24
lib/ai/ollama.js
Normal file
24
lib/ai/ollama.js
Normal file
@@ -0,0 +1,24 @@
|
||||
export class OllamaError extends Error {
|
||||
constructor(status, body) { super(`ollama ${status}: ${body}`); this.status = status; }
|
||||
}
|
||||
|
||||
export async function embedText(text, { model = 'nomic-embed-text', timeoutMs = 60_000 } = {}) {
|
||||
const url = (process.env.OLLAMA_URL || 'http://192.168.1.185:11434') + '/api/embeddings';
|
||||
const res = await fetch(url, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ model, prompt: text }),
|
||||
signal: AbortSignal.timeout(timeoutMs)
|
||||
});
|
||||
if (!res.ok) throw new OllamaError(res.status, await res.text());
|
||||
const j = await res.json();
|
||||
return j.embedding;
|
||||
}
|
||||
|
||||
export function padTo(vector, dim) {
|
||||
if (vector.length === dim) return vector;
|
||||
if (vector.length > dim) return vector.slice(0, dim);
|
||||
const out = vector.slice();
|
||||
while (out.length < dim) out.push(0);
|
||||
return out;
|
||||
}
|
||||
36
tests/ai/ollama.test.js
Normal file
36
tests/ai/ollama.test.js
Normal file
@@ -0,0 +1,36 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { embedText, padTo, OllamaError } from '../../lib/ai/ollama.js';
|
||||
|
||||
beforeEach(() => {
|
||||
global.fetch = vi.fn(async () => new Response(
|
||||
JSON.stringify({ embedding: new Array(768).fill(0.1) }),
|
||||
{ status: 200, headers: { 'content-type': 'application/json' }}
|
||||
));
|
||||
});
|
||||
afterEach(() => vi.restoreAllMocks());
|
||||
|
||||
describe('ollama.embedText', () => {
|
||||
it('returns 768-dim vector', async () => {
|
||||
const v = await embedText('hello');
|
||||
expect(v).toHaveLength(768);
|
||||
expect(v[0]).toBeCloseTo(0.1, 5);
|
||||
});
|
||||
|
||||
it('throws OllamaError on non-200', async () => {
|
||||
global.fetch = vi.fn(async () => new Response('boom', { status: 502 }));
|
||||
await expect(embedText('x')).rejects.toThrow(OllamaError);
|
||||
});
|
||||
});
|
||||
|
||||
describe('padTo', () => {
|
||||
it('pads short vectors with zeros', () => {
|
||||
const v = padTo([1, 2, 3], 5);
|
||||
expect(v).toEqual([1, 2, 3, 0, 0]);
|
||||
});
|
||||
it('truncates long vectors', () => {
|
||||
expect(padTo([1, 2, 3, 4, 5], 3)).toEqual([1, 2, 3]);
|
||||
});
|
||||
it('returns same on exact length', () => {
|
||||
expect(padTo([1, 2], 2)).toEqual([1, 2]);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user