43 lines
1.1 KiB
JavaScript
43 lines
1.1 KiB
JavaScript
import { VoyageAIClient } from "voyageai";
|
|
import { pipeline } from "node:stream/promises";
|
|
import { Readable } from "node:stream";
|
|
import { batch } from "../lib/streams.js";
|
|
|
|
const client = new VoyageAIClient({
|
|
apiKey: process.env.VOYAGEAI_API_KEY,
|
|
});
|
|
|
|
// Well under Voyage's per-request input and token ceilings.
|
|
const EMBED_BATCH_SIZE = 128;
|
|
|
|
/**
|
|
* @param {Array<{id: string, text: string}>} notes
|
|
* @returns {Promise<Map<string, number[]>>} noteId → embedding vector
|
|
*/
|
|
export const embedNotes = async (notes) => {
|
|
const embeddingMap = new Map();
|
|
|
|
if (!notes || notes.length === 0) {
|
|
return embeddingMap;
|
|
}
|
|
|
|
await pipeline(
|
|
Readable.from(notes, { objectMode: true }),
|
|
batch(EMBED_BATCH_SIZE),
|
|
async (batches) => {
|
|
for await (const chunk of batches) {
|
|
const response = await client.embed({
|
|
input: chunk.map((n) => n.text),
|
|
model: "voyage-3",
|
|
});
|
|
|
|
response.data.forEach((item, i) => {
|
|
embeddingMap.set(chunk[i].id, item.embedding);
|
|
});
|
|
}
|
|
}
|
|
);
|
|
|
|
return embeddingMap;
|
|
};
|