feat: add completion mode and context size options for LLM configurat…

…ion (#13325) Signed-off-by: -LAN- <laipz8200@outlook.com>
langgenius · Feb 7, 2025 · 413dfd5 · 413dfd5
1 parent f951590
commit 413dfd5
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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] = {}