update model_provider jina to support custom url and model

langgenius · May 7, 2024 · c82b139 · c82b139
1 parent 6f19115
commit c82b139
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 1 deletion.
diff --git a/api/core/model_runtime/model_providers/jina/jina.yaml b/api/core/model_runtime/model_providers/jina/jina.yaml
@@ -19,6 +19,7 @@ supported_model_types:
   - rerank
 configurate_methods:
   - predefined-model
+  - customizable-model
 provider_credential_schema:
   credential_form_schemas:
     - variable: api_key
@@ -29,3 +30,32 @@ provider_credential_schema:
       placeholder:
         zh_Hans: 在此输入您的 API Key
         en_US: Enter your API Key
+model_credential_schema:
+  model:
+    label:
+      en_US: Model Name
+      zh_Hans: 模型名称
+    placeholder:
+      en_US: Enter your model name
+      zh_Hans: 输入模型名称
+  credential_form_schemas:
+    - variable: base_url
+      label:
+        zh_Hans: 服务器URL
+        en_US: Base URL
+      type: text-input
+      required: true
+      placeholder:
+        zh_Hans: Base URL, e.g. https://api.jina.ai/v1
+        en_US: Base URL, e.g. https://api.jina.ai/v1
+      default: 'https://api.jina.ai/v1'
+    - variable: context_size
+      label:
+        zh_Hans: 上下文大小
+        en_US: Context size
+      placeholder:
+        zh_Hans: 输入上下文大小
+        en_US: Enter context size
+      required: false
+      type: text-input
+      default: '8192'
diff --git a/api/core/model_runtime/model_providers/jina/rerank/rerank.py b/api/core/model_runtime/model_providers/jina/rerank/rerank.py
@@ -2,6 +2,8 @@
 
 import httpx
 
+from core.model_runtime.entities.common_entities import I18nObject
+from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType
 from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
 from core.model_runtime.errors.invoke import (
     InvokeAuthorizationError,
@@ -38,9 +40,13 @@ def _invoke(self, model: str, credentials: dict,
         if len(docs) == 0:
             return RerankResult(model=model, docs=[])
 
+        base_url = credentials.get('base_url', 'https://api.jina.ai/v1')
+        if base_url.endswith('/'):
+            base_url = base_url[:-1]
+
         try:
             response = httpx.post(
-                "https://api.jina.ai/v1/rerank",
+                base_url + '/rerank',
                 json={
                     "model": model,
                     "query": query,
@@ -103,3 +109,19 @@ def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]
             InvokeAuthorizationError: [httpx.HTTPStatusError],  
             InvokeBadRequestError: [httpx.RequestError]
         }
+
+    def get_customizable_model_schema(self, model: str, credentials: dict) -> AIModelEntity:
+        """
+            generate custom model entities from credentials
+        """
+        entity = AIModelEntity(
+            model=model,
+            label=I18nObject(en_US=model),
+            model_type=ModelType.RERANK,
+            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
+            model_properties={
+                ModelPropertyKey.CONTEXT_SIZE: int(credentials.get('context_size'))
+            }
+        )
+
+        return entity