瀏覽代碼

feat: add completion mode and context size options for LLM configuration (#13325)

Signed-off-by: -LAN- <laipz8200@outlook.com>
-LAN- 2 月之前
父節點
當前提交
413dfd5628

+ 2 - 1
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 @@ __all__ = [
     "AudioPromptMessageContent",
     "DocumentPromptMessageContent",
     "ImagePromptMessageContent",
+    "LLMMode",
     "LLMResult",
     "LLMResultChunk",
     "LLMResultChunkDelta",

+ 34 - 0
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:

+ 6 - 2
api/core/model_runtime/model_providers/azure_ai_studio/llm/llm.py

@@ -20,7 +20,7 @@ from azure.core.exceptions import (
 )
 
 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 @@ from core.model_runtime.entities.model_entities import (
     AIModelEntity,
     FetchFrom,
     I18nObject,
+    ModelPropertyKey,
     ModelType,
     ParameterRule,
     ParameterType,
@@ -334,7 +335,10 @@ class AzureAIStudioLargeLanguageModel(LargeLanguageModel):
             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,
         )
 

+ 2 - 3
api/core/workflow/nodes/llm/entities.py

@@ -3,8 +3,7 @@ from typing import Any, Optional
 
 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 @@ from core.workflow.nodes.base import BaseNodeData
 class ModelConfig(BaseModel):
     provider: str
     name: str
-    mode: LLMMode = LLMMode.COMPLETION
+    mode: LLMMode
     completion_params: dict[str, Any] = {}