From 413dfd56283febdedf5ebb25ce27ff07fd86363b Mon Sep 17 00:00:00 2001
From: -LAN- <laipz8200@outlook.com>
Date: Fri, 7 Feb 2025 15:08:53 +0800
Subject: [PATCH] feat: add completion mode and context size options for LLM
 configuration (#13325)

Signed-off-by: -LAN- <laipz8200@outlook.com>
---
 api/core/model_runtime/entities/__init__.py   |  3 +-
 .../azure_ai_studio/azure_ai_studio.yaml      | 34 +++++++++++++++++++
 .../azure_ai_studio/llm/llm.py                |  8 +++--
 api/core/workflow/nodes/llm/entities.py       |  5 ++-
 4 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/api/core/model_runtime/entities/__init__.py b/api/core/model_runtime/entities/__init__.py
index c3e1351e3b6d26..4746ddedcf3b8b 100644
--- a/api/core/model_runtime/entities/__init__.py
+++ b/api/core/model_runtime/entities/__init__.py
@@ -1,4 +1,4 @@
-from .llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
+from .llm_entities import LLMMode, LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
 from .message_entities import (
     AssistantPromptMessage,
     AudioPromptMessageContent,
@@ -23,6 +23,7 @@
     "AudioPromptMessageContent",
     "DocumentPromptMessageContent",
     "ImagePromptMessageContent",
+    "LLMMode",
     "LLMResult",
     "LLMResultChunk",
     "LLMResultChunkDelta",
diff --git a/api/core/model_runtime/model_providers/azure_ai_studio/azure_ai_studio.yaml b/api/core/model_runtime/model_providers/azure_ai_studio/azure_ai_studio.yaml
index 9e17ba088480db..6d9516b6711831 100644
--- a/api/core/model_runtime/model_providers/azure_ai_studio/azure_ai_studio.yaml
+++ b/api/core/model_runtime/model_providers/azure_ai_studio/azure_ai_studio.yaml
@@ -51,6 +51,40 @@ model_credential_schema:
       show_on:
         - variable: __model_type
           value: llm
+    - variable: mode
+      show_on:
+        - variable: __model_type
+          value: llm
+      label:
+        en_US: Completion mode
+      type: select
+      required: false
+      default: chat
+      placeholder:
+        zh_Hans: 选择对话类型
+        en_US: Select completion mode
+      options:
+        - value: completion
+          label:
+            en_US: Completion
+            zh_Hans: 补全
+        - value: chat
+          label:
+            en_US: Chat
+            zh_Hans: 对话
+    - variable: context_size
+      label:
+        zh_Hans: 模型上下文长度
+        en_US: Model context size
+      required: true
+      show_on:
+        - variable: __model_type
+          value: llm
+      type: text-input
+      default: "4096"
+      placeholder:
+        zh_Hans: 在此输入您的模型上下文长度
+        en_US: Enter your Model context size
     - variable: jwt_token
       required: true
       label:
diff --git a/api/core/model_runtime/model_providers/azure_ai_studio/llm/llm.py b/api/core/model_runtime/model_providers/azure_ai_studio/llm/llm.py
index 393f8494dc6868..97c88d0756842d 100644
--- a/api/core/model_runtime/model_providers/azure_ai_studio/llm/llm.py
+++ b/api/core/model_runtime/model_providers/azure_ai_studio/llm/llm.py
@@ -20,7 +20,7 @@
 )
 
 from core.model_runtime.callbacks.base_callback import Callback
-from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
+from core.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
 from core.model_runtime.entities.message_entities import (
     AssistantPromptMessage,
     PromptMessage,
@@ -30,6 +30,7 @@
     AIModelEntity,
     FetchFrom,
     I18nObject,
+    ModelPropertyKey,
     ModelType,
     ParameterRule,
     ParameterType,
@@ -334,7 +335,10 @@ def get_customizable_model_schema(self, model: str, credentials: dict) -> Option
             fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
             model_type=ModelType.LLM,
             features=[],
-            model_properties={},
+            model_properties={
+                ModelPropertyKey.CONTEXT_SIZE: int(credentials.get("context_size", "4096")),
+                ModelPropertyKey.MODE: credentials.get("mode", LLMMode.CHAT),
+            },
             parameter_rules=rules,
         )
 
diff --git a/api/core/workflow/nodes/llm/entities.py b/api/core/workflow/nodes/llm/entities.py
index d6fff2a793de04..bf54fdb80c630f 100644
--- a/api/core/workflow/nodes/llm/entities.py
+++ b/api/core/workflow/nodes/llm/entities.py
@@ -3,8 +3,7 @@
 
 from pydantic import BaseModel, Field, field_validator
 
-from core.model_runtime.entities import ImagePromptMessageContent
-from core.model_runtime.entities.llm_entities import LLMMode
+from core.model_runtime.entities import ImagePromptMessageContent, LLMMode
 from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
 from core.workflow.entities.variable_entities import VariableSelector
 from core.workflow.nodes.base import BaseNodeData
@@ -13,7 +12,7 @@
 class ModelConfig(BaseModel):
     provider: str
     name: str
-    mode: LLMMode = LLMMode.COMPLETION
+    mode: LLMMode
     completion_params: dict[str, Any] = {}