jina-ai-provider
Version:
Jina AI Provider for running Jina AI models with Vercel AI SDK
154 lines (150 loc) • 4.32 kB
JavaScript
// 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