From c8357da13ba867730a85671c032403ac1da75c10 Mon Sep 17 00:00:00 2001
From: ybalbert001 <120714773+ybalbert001@users.noreply.github.com>
Date: Mon, 10 Feb 2025 12:25:04 +0800
Subject: [PATCH] =?UTF-8?q?[Fix]=20Sagemaker=20LLM=20Provider=20can't=20ad?=
 =?UTF-8?q?just=20context=20size,=20it'a=20always=202=E2=80=A6=20(#13462)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Yuanbo Li <ybalbert@amazon.com>
---
 .../model_providers/sagemaker/llm/llm.py            |  4 ++--
 .../model_providers/sagemaker/sagemaker.yaml        | 13 +++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/api/core/model_runtime/model_providers/sagemaker/llm/llm.py b/api/core/model_runtime/model_providers/sagemaker/llm/llm.py
index b8c979b1f5..3f12a68201 100644
--- a/api/core/model_runtime/model_providers/sagemaker/llm/llm.py
+++ b/api/core/model_runtime/model_providers/sagemaker/llm/llm.py
@@ -430,7 +430,7 @@ class SageMakerLargeLanguageModel(LargeLanguageModel):
                 type=ParameterType.INT,
                 use_template="max_tokens",
                 min=1,
-                max=credentials.get("context_length", 2048),
+                max=int(credentials.get("context_length", 2048)),
                 default=512,
                 label=I18nObject(zh_Hans="最大生成长度", en_US="Max Tokens"),
             ),
@@ -448,7 +448,7 @@ class SageMakerLargeLanguageModel(LargeLanguageModel):
         if support_vision:
             features.append(ModelFeature.VISION)
 
-        context_length = credentials.get("context_length", 2048)
+        context_length = int(credentials.get("context_length", 2048))
 
         entity = AIModelEntity(
             model=model,
diff --git a/api/core/model_runtime/model_providers/sagemaker/sagemaker.yaml b/api/core/model_runtime/model_providers/sagemaker/sagemaker.yaml
index 87cd50f50c..c4c94ea409 100644
--- a/api/core/model_runtime/model_providers/sagemaker/sagemaker.yaml
+++ b/api/core/model_runtime/model_providers/sagemaker/sagemaker.yaml
@@ -59,6 +59,19 @@ model_credential_schema:
       placeholder:
         zh_Hans: 请输出你的Sagemaker推理端点
         en_US: Enter your Sagemaker Inference endpoint
+    - variable: context_length
+      show_on:
+        - variable: __model_type
+          value: llm
+      label:
+        zh_Hans: 模型上下文长度
+        en_US: Model context size
+      type: text-input
+      default: '4096'
+      required: true
+      placeholder:
+        zh_Hans: 在此输入您的模型上下文长度
+        en_US: Enter your Model context size
     - variable: audio_s3_cache_bucket
       show_on:
         - variable: __model_type