voyage-ai-provider
Version:
Voyage AI Provider for running Voyage AI models with Vercel AI SDK
133 lines (129 loc) • 3.79 kB
JavaScript
// 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