langgenius/dify-plugin-sdks
1
0
Fork 0
mirror of https://github.com/langgenius/dify-plugin-sdks.git synced 2026-07-22 18:35:29 -04:00
Code Issues 58 Packages Projects Releases 10 Wiki Activity
Files
feat/polling-trigger
dify-plugin-sdks/python/dify_plugin/interfaces/model/large_language_model.py
T
immortal-wm 2fc9753673 fix: corrected latency accuracy (#189) 
* fix latency calc

* feat: add timing context management to AIModel for latency tracking

- Introduced a new TimingContextRaceConditionError to handle race conditions in multi-threaded environments.
- Implemented a timing_context method in AIModel to track request timing and prevent concurrent usage.
- Updated various model classes (LargeLanguageModel, ModerationModel, RerankModel, Speech2TextModel, TextEmbeddingModel, TTSModel) to utilize the new timing context for latency calculations.
- Added tests to validate timing context behavior under concurrent and sequential usage scenarios.

* feat: implement ModelFactory for dynamic model instance creation

- Added ModelFactory class to generate stateless model instances based on provider configurations and model classes.
- Updated PluginRegistration to utilize ModelFactory for managing model instances, enhancing code organization and maintainability.

* feat: add unit tests for model registry and mock model provider

- Introduced a new test file for validating model registry functionality.
- Implemented mock classes for model provider and LLM to facilitate testing.
- Added tests to ensure correct model instance retrieval from the registry.

* apply ruff

---------

Co-authored-by: Yeuoly <admin@srmxy.cn>
2025-09-16 16:15:20 +08:00

618 lines
23 KiB
Python
Raw Permalink Blame History

import logging
import re
import time
from abc import abstractmethod
from collections.abc import Generator, Mapping
from typing import Union
from pydantic import ConfigDict
from dify_plugin.entities.model import (
ModelPropertyKey,
ModelType,
ParameterRule,
ParameterType,
PriceType,
)
from dify_plugin.entities.model.llm import (
LLMMode,
LLMResult,
LLMResultChunk,
LLMResultChunkDelta,
LLMUsage,
)
from dify_plugin.entities.model.message import (
AssistantPromptMessage,
PromptMessage,
PromptMessageContentType,
PromptMessageTool,
SystemPromptMessage,
UserPromptMessage,
)
from dify_plugin.interfaces.model.ai_model import AIModel
logger = logging.getLogger(__name__)
class LargeLanguageModel(AIModel):
"""
Model class for large language model.
"""
model_type: ModelType = ModelType.LLM
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
############################################################
# Methods that can be implemented by plugin #
############################################################
@abstractmethod
def _invoke(
self,
model: str,
credentials: dict,
prompt_messages: list[PromptMessage],
model_parameters: dict,
tools: list[PromptMessageTool] | None = None,
stop: list[str] | None = None,
stream: bool = True,
user: str | None = None,
) -> Union[LLMResult, Generator[LLMResultChunk, None, None]]:
"""
Invoke large language model
:param model: model name
:param credentials: model credentials
:param prompt_messages: prompt messages
:param model_parameters: model parameters
:param tools: tools for tool calling
:param stop: stop words
:param stream: is stream response
:param user: unique user id
:return: full response or stream response chunk generator result
"""
raise NotImplementedError
@abstractmethod
def get_num_tokens(
self,
model: str,
credentials: dict,
prompt_messages: list[PromptMessage],
tools: list[PromptMessageTool] | None = None,
) -> int:
"""
Get number of tokens for given prompt messages
:param model: model name
:param credentials: model credentials
:param prompt_messages: prompt messages
:param tools: tools for tool calling
:return:
"""
raise NotImplementedError
############################################################
# For plugin implementation use only #
############################################################
def enforce_stop_tokens(self, text: str, stop: list[str]) -> str:
"""Cut off the text as soon as any stop words occur."""
return re.split("|".join(stop), text, maxsplit=1)[0]
def get_parameter_rules(self, model: str, credentials: dict) -> list[ParameterRule]:
"""
Get parameter rules
:param model: model name
:param credentials: model credentials
:return: parameter rules
"""
model_schema = self.get_model_schema(model, credentials)
if model_schema:
return model_schema.parameter_rules
return []
def get_model_mode(self, model: str, credentials: Mapping | None = None) -> LLMMode:
"""
Get model mode
:param model: model name
:param credentials: model credentials
:return: model mode
"""
model_schema = self.get_model_schema(model, credentials)
mode = LLMMode.CHAT
if model_schema and model_schema.model_properties.get(ModelPropertyKey.MODE):
mode = LLMMode.value_of(model_schema.model_properties[ModelPropertyKey.MODE])
return mode
def _calc_response_usage(
self, model: str, credentials: dict, prompt_tokens: int, completion_tokens: int
) -> LLMUsage:
"""
Calculate response usage
:param model: model name
:param credentials: model credentials
:param prompt_tokens: prompt tokens
:param completion_tokens: completion tokens
:return: usage
"""
# get prompt price info
prompt_price_info = self.get_price(
model=model,
credentials=credentials,
price_type=PriceType.INPUT,
tokens=prompt_tokens,
)
# get completion price info
completion_price_info = self.get_price(
model=model,
credentials=credentials,
price_type=PriceType.OUTPUT,
tokens=completion_tokens,
)
# Calculate latency from thread-local storage
current_time = time.perf_counter()
latency = current_time - self.started_at
# transform usage
usage = LLMUsage(
prompt_tokens=prompt_tokens,
prompt_unit_price=prompt_price_info.unit_price,
prompt_price_unit=prompt_price_info.unit,
prompt_price=prompt_price_info.total_amount,
completion_tokens=completion_tokens,
completion_unit_price=completion_price_info.unit_price,
completion_price_unit=completion_price_info.unit,
completion_price=completion_price_info.total_amount,
total_tokens=prompt_tokens + completion_tokens,
total_price=prompt_price_info.total_amount + completion_price_info.total_amount,
currency=prompt_price_info.currency,
latency=latency,
)
return usage
def _validate_and_filter_model_parameters(self, model: str, model_parameters: dict, credentials: dict) -> dict:
"""
Validate model parameters
:param model: model name
:param model_parameters: model parameters
:param credentials: model credentials
:return:
"""
parameter_rules = self.get_parameter_rules(model, credentials)
# validate model parameters
filtered_model_parameters = {}
for parameter_rule in parameter_rules:
parameter_name = parameter_rule.name
parameter_value = model_parameters.get(parameter_name)
if parameter_value is None:
if parameter_rule.use_template and parameter_rule.use_template in model_parameters:
# if parameter value is None, use template value variable name instead
parameter_value = model_parameters[parameter_rule.use_template]
else:
if parameter_rule.required:
if parameter_rule.default is not None:
filtered_model_parameters[parameter_name] = parameter_rule.default
continue
else:
raise ValueError(f"Model Parameter {parameter_name} is required.")
else:
continue
# validate parameter value type
if parameter_rule.type == ParameterType.INT:
if not isinstance(parameter_value, int):
raise ValueError(f"Model Parameter {parameter_name} should be int.")
# validate parameter value range
if parameter_rule.min is not None and parameter_value < parameter_rule.min:
raise ValueError(
f"Model Parameter {parameter_name} should be greater than or equal to {parameter_rule.min}."
)
if parameter_rule.max is not None and parameter_value > parameter_rule.max:
raise ValueError(
f"Model Parameter {parameter_name} should be less than or equal to {parameter_rule.max}."
)
elif parameter_rule.type == ParameterType.FLOAT:
if not isinstance(parameter_value, float | int):
raise ValueError(f"Model Parameter {parameter_name} should be float.")
# validate parameter value precision
if parameter_rule.precision is not None:
if parameter_rule.precision == 0:
if parameter_value != int(parameter_value):
raise ValueError(f"Model Parameter {parameter_name} should be int.")
else:
if parameter_value != round(parameter_value, parameter_rule.precision):
raise ValueError(
f"Model Parameter {parameter_name} should be round to "
f"{parameter_rule.precision} decimal places."
)
# validate parameter value range
if parameter_rule.min is not None and parameter_value < parameter_rule.min:
raise ValueError(
f"Model Parameter {parameter_name} should be greater than or equal to {parameter_rule.min}."
)
if parameter_rule.max is not None and parameter_value > parameter_rule.max:
raise ValueError(
f"Model Parameter {parameter_name} should be less than or equal to {parameter_rule.max}."
)
elif parameter_rule.type == ParameterType.BOOLEAN:
if not isinstance(parameter_value, bool):
raise ValueError(f"Model Parameter {parameter_name} should be bool.")
elif parameter_rule.type == ParameterType.STRING:
if not isinstance(parameter_value, str):
raise ValueError(f"Model Parameter {parameter_name} should be string.")
# validate options
if parameter_rule.options and parameter_value not in parameter_rule.options:
raise ValueError(f"Model Parameter {parameter_name} should be one of {parameter_rule.options}.")
elif parameter_rule.type == ParameterType.TEXT:
if not isinstance(parameter_value, str):
raise ValueError(f"Model Parameter {parameter_name} should be string.")
else:
raise ValueError(f"Model Parameter {parameter_name} type {parameter_rule.type} is not supported.")
filtered_model_parameters[parameter_name] = parameter_value
return filtered_model_parameters
def _code_block_mode_wrapper(
self,
model: str,
credentials: dict,
prompt_messages: list[PromptMessage],
model_parameters: dict,
tools: list[PromptMessageTool] | None = None,
stop: list[str] | None = None,
stream: bool = True,
user: str | None = None,
) -> Union[LLMResult, Generator[LLMResultChunk, None, None]]:
"""
Code block mode wrapper, ensure the response is a code block with output markdown quote
:param model: model name
:param credentials: model credentials
:param prompt_messages: prompt messages
:param model_parameters: model parameters
:param tools: tools for tool calling
:param stop: stop words
:param stream: is stream response
:param user: unique user id
:param callbacks: callbacks
:return: full response or stream response chunk generator result
"""
block_prompts = """You should always follow the instructions and output a valid {{block}} object.
The structure of the {{block}} object you can found in the instructions, use {"answer": "$your_answer"} as the default structure
if you are not sure about the structure.
<instructions>
{{instructions}}
</instructions>
""" # noqa: E501
code_block = model_parameters.get("response_format", "")
if not code_block:
return self._invoke(
model=model,
credentials=credentials,
prompt_messages=prompt_messages,
model_parameters=model_parameters,
tools=tools,
stop=stop,
stream=stream,
user=user,
)
model_parameters.pop("response_format")
stop = stop or []
stop.extend(["\n```", "```\n"])
block_prompts = block_prompts.replace("{{block}}", code_block)
# check if there is a system message
if len(prompt_messages) > 0 and isinstance(prompt_messages[0], SystemPromptMessage):
# override the system message
prompt_messages[0] = SystemPromptMessage(
content=block_prompts.replace("{{instructions}}", str(prompt_messages[0].content))
)
else:
# insert the system message
prompt_messages.insert(
0,
SystemPromptMessage(
content=block_prompts.replace(
"{{instructions}}",
f"Please output a valid {code_block} object.",
)
),
)
if len(prompt_messages) > 0 and isinstance(prompt_messages[-1], UserPromptMessage):
# add ```JSON\n to the last text message
if isinstance(prompt_messages[-1].content, str):
prompt_messages[-1].content += f"\n```{code_block}\n"
elif isinstance(prompt_messages[-1].content, list):
for i in range(len(prompt_messages[-1].content) - 1, -1, -1):
if prompt_messages[-1].content[i].type == PromptMessageContentType.TEXT:
prompt_messages[-1].content[i].data += f"\n```{code_block}\n"
break
else:
# append a user message
prompt_messages.append(UserPromptMessage(content=f"```{code_block}\n"))
response = self._invoke(
model=model,
credentials=credentials,
prompt_messages=prompt_messages,
model_parameters=model_parameters,
tools=tools,
stop=stop,
stream=stream,
user=user,
)
if isinstance(response, Generator):
first_chunk = next(response)
def new_generator():
yield first_chunk
yield from response
if (
first_chunk.delta.message.content
and isinstance(first_chunk.delta.message.content, str)
and first_chunk.delta.message.content.startswith("`")
):
return self._code_block_mode_stream_processor_with_backtick(
model=model,
prompt_messages=prompt_messages,
input_generator=new_generator(),
)
else:
return self._code_block_mode_stream_processor(
model=model,
prompt_messages=prompt_messages,
input_generator=new_generator(),
)
return response
def _code_block_mode_stream_processor(
self,
model: str,
prompt_messages: list[PromptMessage],
input_generator: Generator[LLMResultChunk, None, None],
) -> Generator[LLMResultChunk, None, None]:
"""
Code block mode stream processor, ensure the response is a code block with output markdown quote
:param model: model name
:param prompt_messages: prompt messages
:param input_generator: input generator
:return: output generator
"""
state = "normal"
backtick_count = 0
for piece in input_generator:
if piece.delta.message.content:
content = piece.delta.message.content
piece.delta.message.content = ""
yield piece
piece = content
else:
yield piece
continue
new_piece: str = ""
for char in piece:
char = str(char)
if state == "normal":
if char == "`":
state = "in_backticks"
backtick_count = 1
else:
new_piece += char
elif state == "in_backticks":
if char == "`":
backtick_count += 1
if backtick_count == 3:
state = "skip_content"
backtick_count = 0
else:
new_piece += "`" * backtick_count + char
state = "normal"
backtick_count = 0
elif state == "skip_content" and char.isspace():
state = "normal"
if new_piece:
yield LLMResultChunk(
model=model,
delta=LLMResultChunkDelta(
index=0,
message=AssistantPromptMessage(content=new_piece, tool_calls=[]),
),
)
def _code_block_mode_stream_processor_with_backtick(
self,
model: str,
prompt_messages: list,
input_generator: Generator[LLMResultChunk, None, None],
) -> Generator[LLMResultChunk, None, None]:
"""
Code block mode stream processor, ensure the response is a code block with output markdown quote.
This version skips the language identifier that follows the opening triple backticks.
:param model: model name
:param prompt_messages: prompt messages
:param input_generator: input generator
:return: output generator
"""
state = "search_start"
backtick_count = 0
for piece in input_generator:
if piece.delta.message.content:
content = piece.delta.message.content
# Reset content to ensure we're only processing and yielding the relevant parts
piece.delta.message.content = ""
# Yield a piece with cleared content before processing it to maintain the generator structure
yield piece
piece = content
else:
# Yield pieces without content directly
yield piece
continue
if state == "done":
continue
new_piece: str = ""
for char in piece:
if state == "search_start":
if char == "`":
backtick_count += 1
if backtick_count == 3:
state = "skip_language"
backtick_count = 0
else:
backtick_count = 0
elif state == "skip_language":
# Skip everything until the first newline, marking the end of the language identifier
if char == "\n":
state = "in_code_block"
elif state == "in_code_block":
if char == "`":
backtick_count += 1
if backtick_count == 3:
state = "done"
break
else:
if backtick_count > 0:
# If backticks were counted but we're still collecting content, it was a false start
new_piece += "`" * backtick_count
backtick_count = 0
new_piece += str(char)
elif state == "done":
break
if new_piece:
# Only yield content collected within the code block
yield LLMResultChunk(
model=model,
delta=LLMResultChunkDelta(
index=0,
message=AssistantPromptMessage(content=new_piece, tool_calls=[]),
),
)
def _wrap_thinking_by_reasoning_content(self, delta: dict, is_reasoning: bool) -> tuple[str, bool]:
"""
If the reasoning response is from delta.get("reasoning_content"), we wrap
it with HTML think tag.
:param delta: delta dictionary from LLM streaming response
:param is_reasoning: is reasoning
:return: tuple of (processed_content, is_reasoning)
"""
content = delta.get("content") or ""
reasoning_content = delta.get("reasoning_content")
if reasoning_content:
if not is_reasoning:
content = "<think>\n" + reasoning_content
is_reasoning = True
else:
content = reasoning_content
elif is_reasoning and content:
content = "\n</think>" + content
is_reasoning = False
return content, is_reasoning
############################################################
# For executor use only #
############################################################
def invoke(
self,
model: str,
credentials: dict,
prompt_messages: list[PromptMessage],
model_parameters: dict | None = None,
tools: list[PromptMessageTool] | None = None,
stop: list[str] | None = None,
stream: bool = True,
user: str | None = None,
) -> Generator[LLMResultChunk, None, None]:
"""
Invoke large language model
:param model: model name
:param credentials: model credentials
:param prompt_messages: prompt messages
:param model_parameters: model parameters
:param tools: tools for tool calling
:param stop: stop words
:param stream: is stream response
:param user: unique user id
:param callbacks: callbacks
:return: full response or stream response chunk generator result
"""
# validate and filter model parameters
if model_parameters is None:
model_parameters = {}
model_parameters = self._validate_and_filter_model_parameters(model, model_parameters, credentials)
with self.timing_context():
try:
if "response_format" in model_parameters and model_parameters["response_format"] in {"JSON", "XML"}:
result = self._code_block_mode_wrapper(
model=model,
credentials=credentials,
prompt_messages=prompt_messages,
model_parameters=model_parameters,
tools=tools,
stop=stop,
stream=stream,
user=user,
)
else:
result = self._invoke(
model,
credentials,
prompt_messages,
model_parameters,
tools,
stop,
stream,
user,
)
except Exception as e:
raise self._transform_invoke_error(e) from e
if isinstance(result, LLMResult):
yield result.to_llm_result_chunk()
else:
# NOTE: `yield from` cannot been replaced by `return` because of `timing_context`
yield from result
Reference in New Issue View Git Blame Copy Permalink