youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-04-21 19:27:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add LLM quota deduction functionality and enhance model configuration handling in llm_utils.py; update test cases for LLM node context handling

Browse Source
This commit is contained in:
Novice
2026-03-24 08:45:11 +08:00
parent ed1bd338f1
commit dcd614ca77
4 changed files with 57 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
api/dify_graph/nodes/llm/llm_utils.py
View File

@ -22,6 +22,7 @@ from dify_graph.model_runtime.entities import (
TextPromptMessageContent,
ToolPromptMessage,
)
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessageContentUnionTypes,
@ -37,7 +38,7 @@ from dify_graph.runtime import VariablePool
from dify_graph.variables import ArrayFileSegment, FileSegment
from dify_graph.variables.segments import ArrayAnySegment, NoneSegment, StringSegment
from .entities import LLMNodeChatModelMessage, LLMNodeCompletionModelPromptTemplate
from .entities import LLMNodeChatModelMessage, LLMNodeCompletionModelPromptTemplate, ModelConfig
from .exc import (
InvalidVariableTypeError,
MemoryRolePrefixRequiredError,
@ -47,9 +48,7 @@ from .exc import (
from .protocols import TemplateRenderer
def fetch_model_config(
*, tenant_id: str, node_data_model: ModelConfig
) -> tuple[ModelInstance, Any]:
def fetch_model_config(*, tenant_id: str, node_data_model: ModelConfig) -> tuple[ModelInstance, Any]:
from core.app.llm.model_access import build_dify_model_access
from core.app.llm.model_access import fetch_model_config as _fetch
@ -61,6 +60,12 @@ def fetch_model_config(
)
def deduct_llm_quota(*, tenant_id: str, model_instance: ModelInstance, usage: LLMUsage) -> None:
from core.app.llm.quota import deduct_llm_quota as _deduct
_deduct(tenant_id=tenant_id, model_instance=model_instance, usage=usage)
def fetch_model_schema(*, model_instance: ModelInstance) -> AIModelEntity:
model_schema = cast(LargeLanguageModel, model_instance.model_type_instance).get_model_schema(
model_instance.model_name,

9
api/tests/unit_tests/core/workflow/nodes/test_llm_node_streaming.py
View File

@ -35,7 +35,14 @@ def patch_deduct_llm_quota(monkeypatch):
def _make_llm_node(reasoning_format: str) -> LLMNode:
node = LLMNode.__new__(LLMNode)
object.__setattr__(node, "_node_data", types.SimpleNamespace(reasoning_format=reasoning_format, tools=[]))
object.__setattr__(node, "tenant_id", "tenant")
object.__setattr__(
node,
"_run_context",
{"_dify": types.SimpleNamespace(
tenant_id="tenant", app_id="app", user_id="user",
user_from="account", invoke_from="debugger",
)},
)
return node