UNPKG

jina-ai-provider

Version:

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

154 lines (150 loc) 4.32 kB
// src/jina-provider.ts import { loadApiKey, withoutTrailingSlash } from "@ai-sdk/provider-utils"; // src/jina-embedding-model.ts import { TooManyEmbeddingValuesForCallError } from "@ai-sdk/provider"; import { combineHeaders, createJsonResponseHandler, postJsonToApi } from "@ai-sdk/provider-utils"; import { z as z2 } from "zod"; // src/jina-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/jina-embedding-model.ts var JinaEmbeddingModel = class { specificationVersion = "v1"; modelId; config; settings; get provider() { return this.config.provider; } get maxEmbeddingsPerCall() { return 2048; } get supportsParallelCalls() { return false; } constructor(modelId, settings, config) { this.modelId = modelId; this.settings = settings; this.config = config; } async doEmbed({ values, headers, abortSignal }) { 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: { model: this.modelId, input: values, task: this.settings.inputType, embedding_type: this.settings.embeddingType, dimensions: this.settings.outputDimension, normalized: this.settings.normalized ?? true, late_chunking: this.settings.lateChunking, truncate: this.settings.truncate ?? false }, failedResponseHandler: voyageFailedResponseHandler, fetch: this.config.fetch, headers: combineHeaders(this.config.headers(), headers), successfulResponseHandler: createJsonResponseHandler( jinaEmbeddingResponseSchema ), 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 jinaEmbeddingResponseSchema = z2.object({ data: z2.array( z2.object({ object: z2.literal("embedding"), embedding: z2.array(z2.number()), index: z2.number().optional() }) ), usage: z2.object({ total_tokens: z2.number(), prompt_tokens: z2.number().optional() }).nullish(), model: z2.string().optional() }); // src/jina-provider.ts function createJina(options = {}) { const baseURL = withoutTrailingSlash(options.baseURL) ?? "https://api.jina.ai/v1"; const getHeaders = () => ({ Authorization: `Bearer ${loadApiKey({ apiKey: options.apiKey, environmentVariableName: "JINA_API_KEY", description: "Jina" })}`, ...options.headers }); const createTextEmbeddingModel = (modelId, settings = {}) => new JinaEmbeddingModel(modelId, settings, { provider: "jina.text.embedding", baseURL, headers: getHeaders, fetch: options.fetch }); const createMultiModalEmbeddingModel = (modelId, settings = {}) => new JinaEmbeddingModel(modelId, settings, { provider: "jina.multimodal.embedding", baseURL, headers: getHeaders, fetch: options.fetch }); const provider = function(modelId, settings) { if (new.target) { throw new Error( "The Jina model function cannot be called with the new keyword." ); } return createTextEmbeddingModel(modelId, settings); }; provider.textEmbeddingModel = createTextEmbeddingModel; provider.multiModalEmbeddingModel = createMultiModalEmbeddingModel; provider.languageModel = (modelId) => { throw new Error( `Language model '${modelId}' is not supported by Jina provider.` ); }; return provider; } var jina = createJina(); export { createJina, jina }; //# sourceMappingURL=index.js.map