From ba9ef123a44cb4c4641d145173267e2ca23ce87f Mon Sep 17 00:00:00 2001 From: Zero Liu Date: Fri, 15 May 2026 01:06:33 -0700 Subject: [PATCH] chore(jina): adapt upstream community implementation with attribution Restores batching, dimensions/normalization defaults, and the multi-modal input type from the original @langchain/community JinaEmbeddings, with upstream copyright notice and MIT attribution preserved. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/LLMProviders/CustomJinaEmbeddings.ts | 206 +++++++++++++++++------ 1 file changed, 157 insertions(+), 49 deletions(-) diff --git a/src/LLMProviders/CustomJinaEmbeddings.ts b/src/LLMProviders/CustomJinaEmbeddings.ts index 7b6a2ecb..a55f950f 100644 --- a/src/LLMProviders/CustomJinaEmbeddings.ts +++ b/src/LLMProviders/CustomJinaEmbeddings.ts @@ -1,78 +1,186 @@ -import { Embeddings, EmbeddingsParams } from "@langchain/core/embeddings"; -import { requestUrl } from "obsidian"; +/* + * Adapted from @langchain/community JinaEmbeddings. + * Copyright (c) LangChain, Inc. Licensed under the MIT License. + * Source: https://github.com/langchain-ai/langchainjs-community/blob/886df5749a926f59e6fdf38a3465c62ec9e7ce32/libs/community/src/embeddings/jina.ts + */ -const DEFAULT_JINA_API_URL = "https://api.jina.ai/v1/embeddings"; +import { Embeddings, type EmbeddingsParams } from "@langchain/core/embeddings"; +import { chunkArray } from "@langchain/core/utils/chunk_array"; +import { getEnvironmentVariable } from "@langchain/core/utils/env"; -interface JinaEmbeddingsConfig extends EmbeddingsParams { - model?: string; +export interface JinaEmbeddingsParams extends EmbeddingsParams { + /** Model name to use. */ + model: string; + /** Compatibility alias used by this plugin's embedding manager. */ modelName?: string; - apiKey?: string; + /** Jina-compatible embeddings endpoint. */ baseUrl?: string; + /** Timeout to use when making requests to Jina. */ + timeout?: number; + /** The maximum number of documents to embed in a single request. */ + batchSize?: number; + /** Whether to strip new lines from the input text. */ + stripNewLines?: boolean; + /** The dimensions of the embedding. */ dimensions?: number; + /** Whether to L2-normalize the embedding vectors. */ + normalized?: boolean; } -interface JinaEmbeddingItem { - index: number; - embedding: number[]; +type JinaMultiModelInput = + | { + text: string; + image?: never; + } + | { + image: string; + text?: never; + }; + +export type JinaEmbeddingsInput = string | JinaMultiModelInput; + +interface EmbeddingCreateParams { + model: JinaEmbeddingsParams["model"]; + input: JinaEmbeddingsInput[]; + dimensions: number; + task: "retrieval.query" | "retrieval.passage"; + normalized?: boolean; } -interface JinaEmbeddingResponse { - data?: JinaEmbeddingItem[]; - detail?: unknown; +interface EmbeddingResponse { + model: string; + object: string; + usage: { + total_tokens: number; + prompt_tokens: number; + }; + data: { + object: string; + index: number; + embedding: number[]; + }[]; } -export class CustomJinaEmbeddings extends Embeddings { - private readonly model: string; - private readonly apiKey: string; - private readonly url: string; - private readonly dimensions?: number; +interface EmbeddingErrorResponse { + detail: string; +} - constructor(config: JinaEmbeddingsConfig = {}) { - super(config); - this.model = config.model ?? config.modelName ?? "jina-embeddings-v2-base-en"; - this.apiKey = config.apiKey ?? ""; - this.url = config.baseUrl ?? DEFAULT_JINA_API_URL; - this.dimensions = config.dimensions; +export class CustomJinaEmbeddings extends Embeddings implements JinaEmbeddingsParams { + model: JinaEmbeddingsParams["model"] = "jina-clip-v2"; + batchSize = 24; + baseUrl = "https://api.jina.ai/v1/embeddings"; + stripNewLines = true; + dimensions = 1024; + apiKey: string; + normalized = true; + + /** + * Creates a Jina embeddings client using local configuration or Jina environment variables. + */ + constructor( + fields?: Partial & { + apiKey?: string; + } + ) { + const fieldsWithDefaults = { maxConcurrency: 2, ...fields }; + super(fieldsWithDefaults); + + const apiKey = + fieldsWithDefaults?.apiKey || + getEnvironmentVariable("JINA_API_KEY") || + getEnvironmentVariable("JINA_AUTH_TOKEN"); + + if (!apiKey) throw new Error("Jina API key not found"); + + this.apiKey = apiKey; + this.model = fieldsWithDefaults?.model ?? fieldsWithDefaults?.modelName ?? this.model; + this.baseUrl = fieldsWithDefaults?.baseUrl ?? this.baseUrl; + this.dimensions = fieldsWithDefaults?.dimensions ?? this.dimensions; + this.batchSize = fieldsWithDefaults?.batchSize ?? this.batchSize; + this.stripNewLines = fieldsWithDefaults?.stripNewLines ?? this.stripNewLines; + this.normalized = fieldsWithDefaults?.normalized ?? this.normalized; } - async embedQuery(text: string): Promise { - const embeddings = await this._embed([text]); + /** + * Embeds passage documents with Jina retrieval-passage task parameters. + */ + async embedDocuments(input: JinaEmbeddingsInput[]): Promise { + const batches = chunkArray(this.doStripNewLines(input), this.batchSize); + const batchRequests = batches.map((batch) => { + const params = this.getParams(batch); + return this.embeddingWithRetry(params); + }); + + const batchResponses = await Promise.all(batchRequests); + const embeddings: number[][] = []; + + for (let i = 0; i < batchResponses.length; i += 1) { + const batch = batches[i]; + const batchResponse = batchResponses[i] || []; + for (let j = 0; j < batch.length; j += 1) { + embeddings.push(batchResponse[j]); + } + } + + return embeddings; + } + + /** + * Embeds a query with Jina retrieval-query task parameters. + */ + async embedQuery(input: JinaEmbeddingsInput): Promise { + const params = this.getParams(this.doStripNewLines([input]), true); + const embeddings = (await this.embeddingWithRetry(params)) || [[]]; return embeddings[0]; } - async embedDocuments(texts: string[]): Promise { - return this._embed(texts); + /** + * Removes newlines from string inputs when configured to match upstream Jina behavior. + */ + private doStripNewLines(input: JinaEmbeddingsInput[]): JinaEmbeddingsInput[] { + if (this.stripNewLines) { + return input.map((item) => { + if (typeof item === "string") { + return item.replace(/\n/g, " "); + } + if (item.text) { + return { text: item.text.replace(/\n/g, " ") }; + } + return item; + }); + } + return input; } - private async _embed(input: string[]): Promise { - const body: Record = { input, model: this.model }; - if (this.dimensions !== undefined) { - body.dimensions = this.dimensions; - } + /** + * Builds the request body for Jina's retrieval embedding API. + */ + private getParams(input: JinaEmbeddingsInput[], query?: boolean): EmbeddingCreateParams { + return { + model: this.model, + input, + dimensions: this.dimensions, + task: query ? "retrieval.query" : "retrieval.passage", + normalized: this.normalized, + }; + } - const response = await requestUrl({ - url: this.url, + /** + * Sends a single embeddings request and returns vectors in response order. + */ + private async embeddingWithRetry(body: EmbeddingCreateParams): Promise { + const response = await fetch(this.baseUrl, { method: "POST", - contentType: "application/json", headers: { + "Content-Type": "application/json", Authorization: `Bearer ${this.apiKey}`, - "Accept-Encoding": "identity", }, body: JSON.stringify(body), - throw: false, }); - - if (response.status >= 400) { - throw new Error(`Jina embedding request failed: ${response.status} - ${response.text ?? ""}`); + const embeddingData: EmbeddingResponse | EmbeddingErrorResponse = await response.json(); + if ("detail" in embeddingData && embeddingData.detail) { + throw new Error(`${embeddingData.detail}`); } - - const resp = response.json as JinaEmbeddingResponse; - if (!resp?.data) { - throw new Error( - typeof resp?.detail === "string" ? resp.detail : JSON.stringify(resp?.detail ?? resp) - ); - } - - return [...resp.data].sort((a, b) => a.index - b.index).map((item) => item.embedding); + return (embeddingData as EmbeddingResponse).data.map(({ embedding }) => embedding); } }