From d4b1b4ce8b7de85956fb091de5c863abf48dbc30 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Istv=C3=A1n=20Zolt=C3=A1n=20Szab=C3=B3?=
 <szabosteve@gmail.com>
Date: Wed, 21 May 2025 12:56:07 +0200
Subject: [PATCH] Adds custom InferenceEndpointInfo class for JinaAI (#4391)

* Adds custom InferenceEndpointInfo class for JinaAI.

* Format fix.

(cherry picked from commit 6e9264533e8ea359ea68f6af8efdfa4da916d323)
---
 specification/inference/_types/Services.ts          | 13 ++++++++++++-
 specification/inference/_types/TaskType.ts          |  5 +++++
 .../inference/put_jinaai/PutJinaAiResponse.ts       |  4 ++--
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/specification/inference/_types/Services.ts b/specification/inference/_types/Services.ts
index 7c3ea7fc9d..110ad4a5e6 100644
--- a/specification/inference/_types/Services.ts
+++ b/specification/inference/_types/Services.ts
@@ -19,7 +19,7 @@
 
 import { UserDefinedValue } from '@spec_utils/UserDefinedValue'
 import { integer } from '@_types/Numeric'
-import { TaskType } from '../_types/TaskType'
+import { TaskType, TaskTypeJinaAi } from '../_types/TaskType'
 
 /**
  * Configuration options when storing the inference endpoint
@@ -57,6 +57,17 @@ export class InferenceEndpointInfo extends InferenceEndpoint {
   task_type: TaskType
 }
 
+export class InferenceEndpointInfoJinaAi extends InferenceEndpoint {
+  /**
+   * The inference Id
+   */
+  inference_id: string
+  /**
+   * The task type
+   */
+  task_type: TaskTypeJinaAi
+}
+
 /**
  * Chunking configuration object
  */
diff --git a/specification/inference/_types/TaskType.ts b/specification/inference/_types/TaskType.ts
index 7299480cbe..9eba8b1fc1 100644
--- a/specification/inference/_types/TaskType.ts
+++ b/specification/inference/_types/TaskType.ts
@@ -27,3 +27,8 @@ export enum TaskType {
   completion,
   chat_completion
 }
+
+export enum TaskTypeJinaAi {
+  text_embedding,
+  rerank
+}
diff --git a/specification/inference/put_jinaai/PutJinaAiResponse.ts b/specification/inference/put_jinaai/PutJinaAiResponse.ts
index 581909e650..a2123296af 100644
--- a/specification/inference/put_jinaai/PutJinaAiResponse.ts
+++ b/specification/inference/put_jinaai/PutJinaAiResponse.ts
@@ -17,9 +17,9 @@
  * under the License.
  */
 
-import { InferenceEndpointInfo } from '@inference/_types/Services'
+import { InferenceEndpointInfoJinaAi } from '@inference/_types/Services'
 
 export class Response {
   /** @codegen_name endpoint_info */
-  body: InferenceEndpointInfo
+  body: InferenceEndpointInfoJinaAi
 }