mirror of
https://github.com/langgenius/dify-plugin-sdks.git
synced 2026-07-22 02:15:22 -04:00
immortal-wm
2fc9753673
* 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>
76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
from abc import abstractmethod
|
|
|
|
from dify_plugin.entities.model import ModelType
|
|
from dify_plugin.entities.model.rerank import RerankResult
|
|
from dify_plugin.interfaces.model.ai_model import AIModel
|
|
|
|
|
|
class RerankModel(AIModel):
|
|
"""
|
|
Base Model class for rerank model.
|
|
"""
|
|
|
|
model_type: ModelType = ModelType.RERANK
|
|
|
|
############################################################
|
|
# Methods that can be implemented by plugin #
|
|
############################################################
|
|
|
|
@abstractmethod
|
|
def _invoke(
|
|
self,
|
|
model: str,
|
|
credentials: dict,
|
|
query: str,
|
|
docs: list[str],
|
|
score_threshold: float | None = None,
|
|
top_n: int | None = None,
|
|
user: str | None = None,
|
|
) -> RerankResult:
|
|
"""
|
|
Invoke rerank model
|
|
|
|
:param model: model name
|
|
:param credentials: model credentials
|
|
:param query: search query
|
|
:param docs: docs for reranking
|
|
:param score_threshold: score threshold
|
|
:param top_n: top n
|
|
:param user: unique user id
|
|
:return: rerank result
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
############################################################
|
|
# For executor use only #
|
|
############################################################
|
|
|
|
def invoke(
|
|
self,
|
|
model: str,
|
|
credentials: dict,
|
|
query: str,
|
|
docs: list[str],
|
|
score_threshold: float | None = None,
|
|
top_n: int | None = None,
|
|
user: str | None = None,
|
|
) -> RerankResult:
|
|
"""
|
|
Invoke rerank model
|
|
|
|
:param model: model name
|
|
:param credentials: model credentials
|
|
:param query: search query
|
|
:param docs: docs for reranking
|
|
:param score_threshold: score threshold
|
|
:param top_n: top n
|
|
:param user: unique user id
|
|
:return: rerank result
|
|
"""
|
|
|
|
with self.timing_context():
|
|
try:
|
|
return self._invoke(model, credentials, query, docs, score_threshold, top_n, user)
|
|
except Exception as e:
|
|
raise self._transform_invoke_error(e) from e
|