UNPKG

voyage-ai-provider

Version:

Voyage AI Provider for running Voyage AI models with Vercel AI SDK

133 lines (129 loc) 3.79 kB
// src/voyage-provider.ts import { loadApiKey, withoutTrailingSlash } from "@ai-sdk/provider-utils"; // src/voyage-embedding-model.ts import { TooManyEmbeddingValuesForCallError } from "@ai-sdk/provider"; import { createJsonResponseHandler, postJsonToApi } from "@ai-sdk/provider-utils"; import { z as z2 } from "zod"; // src/voyage-error.ts import { createJsonErrorResponseHandler } from "@ai-sdk/provider-utils"; import { z } from "zod"; var voyageErrorDataSchema = z.object({ error: z.object({ code: z.string().nullable(), message: z.string(), param: z.any().nullable(), type: z.string() }) }); var voyageFailedResponseHandler = createJsonErrorResponseHandler({ errorSchema: voyageErrorDataSchema, errorToMessage: (data) => data.error.message }); // src/voyage-embedding-model.ts var VoyageEmbeddingModel = class { specificationVersion = "v1"; modelId; config; settings; get provider() { return this.config.provider; } get maxEmbeddingsPerCall() { return 128; } get supportsParallelCalls() { return false; } constructor(modelId, settings, config) { this.modelId = modelId; this.settings = settings; this.config = config; } async doEmbed({ abortSignal, values }) { if (values.length > this.maxEmbeddingsPerCall) { throw new TooManyEmbeddingValuesForCallError({ maxEmbeddingsPerCall: this.maxEmbeddingsPerCall, modelId: this.modelId, provider: this.provider, values }); } const { responseHeaders, value: response } = await postJsonToApi({ abortSignal, body: { input: values, model: this.modelId, input_type: this.settings.inputType, truncation: this.settings.truncation, output_dimension: this.settings.outputDimension, output_dtype: this.settings.outputDtype }, failedResponseHandler: voyageFailedResponseHandler, fetch: this.config.fetch, headers: this.config.headers(), successfulResponseHandler: createJsonResponseHandler( voyageTextEmbeddingResponseSchema ), url: `${this.config.baseURL}/embeddings` }); return { embeddings: response.data.map((item) => item.embedding), usage: response.usage ? { tokens: response.usage.total_tokens } : void 0, rawResponse: { headers: responseHeaders } }; } }; var voyageTextEmbeddingResponseSchema = z2.object({ data: z2.array(z2.object({ embedding: z2.array(z2.number()) })), usage: z2.object({ total_tokens: z2.number() }).nullish() }); // src/voyage-provider.ts function createVoyage(options = {}) { const baseURL = withoutTrailingSlash(options.baseURL) ?? "https://api.voyageai.com/v1"; const getHeaders = () => ({ Authorization: `Bearer ${loadApiKey({ apiKey: options.apiKey, environmentVariableName: "VOYAGE_API_KEY", description: "Voyage" })}`, ...options.headers }); const createEmbeddingModel = (modelId, settings = {}) => new VoyageEmbeddingModel(modelId, settings, { provider: "voyage.embedding", baseURL, headers: getHeaders, fetch: options.fetch }); const provider = function(modelId, settings) { if (new.target) { throw new Error( "The Voyage model function cannot be called with the new keyword." ); } return createEmbeddingModel(modelId, settings); }; provider.embedding = createEmbeddingModel; provider.textEmbedding = createEmbeddingModel; provider.textEmbeddingModel = createEmbeddingModel; provider.chat = provider.languageModel = () => { throw new Error("languageModel method is not implemented."); }; return provider; } var voyage = createVoyage(); export { createVoyage, voyage }; //# sourceMappingURL=index.js.map