From 560d375e0f3f1404cde54a012ad7602e9b8cca53 Mon Sep 17 00:00:00 2001
From: sino <sino2322@gmail.com>
Date: Thu, 19 Dec 2024 17:49:31 +0800
Subject: [PATCH] feat(ark): add doubao-pro-256k and doubao-embedding-large
 (#11831)

---
 .../model_providers/volcengine_maas/llm/models.py    |  4 ++++
 .../volcengine_maas/text_embedding/models.py         |  3 ++-
 .../volcengine_maas/volcengine_maas.yaml             | 12 ++++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/api/core/model_runtime/model_providers/volcengine_maas/llm/models.py b/api/core/model_runtime/model_providers/volcengine_maas/llm/models.py
index f7698f944..cf3cf23cf 100644
--- a/api/core/model_runtime/model_providers/volcengine_maas/llm/models.py
+++ b/api/core/model_runtime/model_providers/volcengine_maas/llm/models.py
@@ -40,6 +40,10 @@ configs: dict[str, ModelConfig] = {
         properties=ModelProperties(context_size=32768, max_tokens=4096, mode=LLMMode.CHAT),
         features=[ModelFeature.TOOL_CALL],
     ),
+    "Doubao-pro-256k": ModelConfig(
+        properties=ModelProperties(context_size=262144, max_tokens=4096, mode=LLMMode.CHAT),
+        features=[],
+    ),
     "Doubao-pro-128k": ModelConfig(
         properties=ModelProperties(context_size=131072, max_tokens=4096, mode=LLMMode.CHAT),
         features=[ModelFeature.TOOL_CALL],
diff --git a/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/models.py b/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/models.py
index 4a6f5b6f7..be9bba5f2 100644
--- a/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/models.py
+++ b/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/models.py
@@ -12,6 +12,7 @@ class ModelConfig(BaseModel):
 
 ModelConfigs = {
     "Doubao-embedding": ModelConfig(properties=ModelProperties(context_size=4096, max_chunks=32)),
+    "Doubao-embedding-large": ModelConfig(properties=ModelProperties(context_size=4096, max_chunks=32)),
 }
 
 
@@ -21,7 +22,7 @@ def get_model_config(credentials: dict) -> ModelConfig:
     if not model_configs:
         return ModelConfig(
             properties=ModelProperties(
-                context_size=int(credentials.get("context_size", 0)),
+                context_size=int(credentials.get("context_size", 4096)),
                 max_chunks=int(credentials.get("max_chunks", 1)),
             )
         )
diff --git a/api/core/model_runtime/model_providers/volcengine_maas/volcengine_maas.yaml b/api/core/model_runtime/model_providers/volcengine_maas/volcengine_maas.yaml
index 57492b6d9..2ddb61254 100644
--- a/api/core/model_runtime/model_providers/volcengine_maas/volcengine_maas.yaml
+++ b/api/core/model_runtime/model_providers/volcengine_maas/volcengine_maas.yaml
@@ -166,6 +166,12 @@ model_credential_schema:
           show_on:
             - variable: __model_type
               value: llm
+        - label:
+            en_US: Doubao-pro-256k
+          value: Doubao-pro-256k
+          show_on:
+            - variable: __model_type
+              value: llm
         - label:
             en_US: Llama3-8B
           value: Llama3-8B
@@ -220,6 +226,12 @@ model_credential_schema:
           show_on:
             - variable: __model_type
               value: text-embedding
+        - label:
+            en_US: Doubao-embedding-large
+          value: Doubao-embedding-large
+          show_on:
+            - variable: __model_type
+              value: text-embedding
         - label:
             en_US: Custom
             zh_Hans: 自定义