mirror of
https://github.com/langgenius/dify-plugin-sdks.git
synced 2026-07-22 18:35:29 -04:00
Yeuoly
f902504af0
* Inline multimodal entities into existing models * apply ruff * bump to 0.7.0b1 * fix: remove tenant_id from invoke_multimodal_embedding * tests: add rerank * apply ruff * fix * fix: typing
105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
from collections.abc import Mapping
|
|
|
|
from dify_plugin.entities.model import ModelType
|
|
from dify_plugin.entities.model.rerank import MultiModalRerankResult, RerankDocument, RerankResult
|
|
from dify_plugin.entities.model.text_embedding import MultiModalContent, MultiModalContentType
|
|
from dify_plugin.errors.model import InvokeError
|
|
from dify_plugin.interfaces.model.rerank_model import RerankModel
|
|
from tests.interfaces.model.utils import prepare_model_factory
|
|
|
|
|
|
class MockRerankModel(RerankModel):
|
|
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:
|
|
return RerankResult(
|
|
model=model,
|
|
docs=[RerankDocument(index=0, text=doc, score=0) for doc in docs],
|
|
)
|
|
|
|
@property
|
|
def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:
|
|
return {}
|
|
|
|
def validate_credentials(self, model: str, credentials: Mapping) -> None:
|
|
pass
|
|
|
|
|
|
def test_rerank():
|
|
model_factory = prepare_model_factory({ModelType.RERANK: MockRerankModel})
|
|
instance = model_factory.get_instance(ModelType.RERANK)
|
|
assert isinstance(instance, MockRerankModel)
|
|
result = instance.invoke(model="test", credentials={}, query="test", docs=["test"])
|
|
assert result.model == "test"
|
|
assert len(result.docs) == 1
|
|
assert result.docs[0].index == 0
|
|
assert result.docs[0].text == "test"
|
|
assert result.docs[0].score == 0
|
|
|
|
|
|
class MockMultiModalRerankModel(RerankModel):
|
|
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:
|
|
result = self._invoke_multimodal(
|
|
model,
|
|
credentials,
|
|
MultiModalContent(content=query, content_type=MultiModalContentType.TEXT),
|
|
[MultiModalContent(content=doc, content_type=MultiModalContentType.TEXT) for doc in docs],
|
|
score_threshold,
|
|
top_n,
|
|
user,
|
|
)
|
|
return RerankResult(
|
|
model=model,
|
|
docs=[RerankDocument(index=0, text=doc.text, score=0) for doc in result.docs],
|
|
)
|
|
|
|
def _invoke_multimodal(
|
|
self,
|
|
model: str,
|
|
credentials: dict,
|
|
query: MultiModalContent,
|
|
docs: list[MultiModalContent],
|
|
score_threshold: float | None = None,
|
|
top_n: int | None = None,
|
|
user: str | None = None,
|
|
) -> MultiModalRerankResult:
|
|
return MultiModalRerankResult(
|
|
model=model,
|
|
docs=[RerankDocument(index=0, text=doc.content, score=0) for doc in docs],
|
|
)
|
|
|
|
@property
|
|
def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:
|
|
return {}
|
|
|
|
def validate_credentials(self, model: str, credentials: Mapping) -> None:
|
|
pass
|
|
|
|
|
|
def test_rerank_multimodal():
|
|
model_factory = prepare_model_factory({ModelType.RERANK: MockMultiModalRerankModel})
|
|
instance = model_factory.get_instance(ModelType.RERANK)
|
|
assert isinstance(instance, MockMultiModalRerankModel)
|
|
result = instance.invoke(model="test", credentials={}, query="test", docs=["test"])
|
|
assert result.model == "test"
|
|
assert len(result.docs) == 1
|
|
assert result.docs[0].index == 0
|
|
assert result.docs[0].text == "test"
|
|
assert result.docs[0].score == 0
|