46 lines
1.3 KiB
TypeScript
46 lines
1.3 KiB
TypeScript
import { VoyageAIClient } from "voyageai";
|
|
import { pipeline } from "node:stream/promises";
|
|
import { Readable } from "node:stream";
|
|
import { batch } from "../lib/streams.js";
|
|
import type { EmbeddingMap } from "../types/domain.js";
|
|
|
|
const client = new VoyageAIClient({
|
|
apiKey: process.env.VOYAGEAI_API_KEY,
|
|
});
|
|
|
|
const EMBED_BATCH_SIZE = 128;
|
|
|
|
type EmbeddableNote = { id: string; text: string };
|
|
|
|
/**
|
|
* @returns noteId to embedding vector
|
|
*/
|
|
export const embedNotes = async (notes: EmbeddableNote[]): Promise<EmbeddingMap> => {
|
|
const embeddingMap: EmbeddingMap = new Map();
|
|
|
|
if (!notes || notes.length === 0) {
|
|
return embeddingMap;
|
|
}
|
|
|
|
await pipeline(
|
|
Readable.from(notes, { objectMode: true }),
|
|
batch<EmbeddableNote>(EMBED_BATCH_SIZE),
|
|
async (batches: AsyncIterable<EmbeddableNote[]>) => {
|
|
for await (const chunk of batches) {
|
|
const response = await client.embed({
|
|
input: chunk.map((n) => n.text),
|
|
model: "voyage-3.5",
|
|
});
|
|
|
|
response.data?.forEach((item, i) => {
|
|
if (item.embedding) {
|
|
embeddingMap.set(chunk[i].id, item.embedding);
|
|
}
|
|
});
|
|
}
|
|
}
|
|
);
|
|
|
|
return embeddingMap;
|
|
};
|