diff --git a/src/dify_plugin/__init__.py b/src/dify_plugin/__init__.py index 1f6662f0..2ead5444 100644 --- a/src/dify_plugin/__init__.py +++ b/src/dify_plugin/__init__.py @@ -1,5 +1,6 @@ from dify_plugin import _gevent # noqa: F401 - import applies gevent patching from dify_plugin.config.config import DifyPluginEnv +from dify_plugin.core.session_context import get_current_session from dify_plugin.interfaces.agent import AgentProvider, AgentStrategy from dify_plugin.interfaces.endpoint import Endpoint from dify_plugin.interfaces.model import ModelProvider @@ -47,4 +48,5 @@ "TextEmbeddingModel", "Tool", "ToolProvider", + "get_current_session", ] diff --git a/src/dify_plugin/core/plugin_executor.py b/src/dify_plugin/core/plugin_executor.py index 9d0a9cdd..ac5b7c31 100644 --- a/src/dify_plugin/core/plugin_executor.py +++ b/src/dify_plugin/core/plugin_executor.py @@ -53,6 +53,7 @@ ) from dify_plugin.core.plugin_registration import PluginRegistration from dify_plugin.core.runtime import Session +from dify_plugin.core.session_context import use_current_session from dify_plugin.core.utils.http_parser import deserialize_request, serialize_response from dify_plugin.entities import ParameterOption from dify_plugin.entities.agent import AgentRuntime @@ -242,22 +243,30 @@ def validate_model_credentials( return {"result": True, "credentials": data.credentials} def invoke_llm(self, session: Session, data: ModelInvokeLLMRequest) -> object: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, ) if isinstance(model_instance, LargeLanguageModel): - return model_instance.invoke( - data.model, - data.credentials, - data.prompt_messages, - data.model_parameters, - data.tools, - data.stop, - data.stream, - data.user_id, - ) + with use_current_session(session): + result = model_instance.invoke( + data.model, + data.credentials, + data.prompt_messages, + data.model_parameters, + data.tools, + data.stop, + data.stream, + data.user_id, + ) + if not isinstance(result, Generator): + return result + + def generator() -> Generator[object, None, None]: + with use_current_session(session): + yield from result + + return generator() msg = f"Model `{data.model_type}` not found for provider `{data.provider}`" raise ValueError( msg, @@ -358,18 +367,18 @@ def invoke_text_embedding( session: Session, data: ModelInvokeTextEmbeddingRequest, ) -> object: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, ) if isinstance(model_instance, TextEmbeddingModel): - return model_instance.invoke( - data.model, - data.credentials, - data.texts, - data.user_id, - ) + with use_current_session(session): + return model_instance.invoke( + data.model, + data.credentials, + data.texts, + data.user_id, + ) msg = f"Model `{data.model_type}` not found for provider `{data.provider}`" raise ValueError( msg, @@ -380,19 +389,19 @@ def invoke_multimodal_embedding( session: Session, data: ModelInvokeMultimodalEmbeddingRequest, ) -> object: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, ) if isinstance(model_instance, TextEmbeddingModel): - return model_instance.invoke_multimodal( - data.model, - data.credentials, - data.documents, - user=data.user_id, - input_type=data.input_type, - ) + with use_current_session(session): + return model_instance.invoke_multimodal( + data.model, + data.credentials, + data.documents, + user=data.user_id, + input_type=data.input_type, + ) msg = f"Model `{data.model_type}` not found for provider `{data.provider}`" raise ValueError( msg, @@ -422,21 +431,21 @@ def get_text_embedding_num_tokens( ) def invoke_rerank(self, session: Session, data: ModelInvokeRerankRequest) -> object: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, ) if isinstance(model_instance, RerankModel): - return model_instance.invoke( - data.model, - data.credentials, - data.query, - data.docs, - data.score_threshold, - data.top_n, - data.user_id, - ) + with use_current_session(session): + return model_instance.invoke( + data.model, + data.credentials, + data.query, + data.docs, + data.score_threshold, + data.top_n, + data.user_id, + ) msg = f"Model `{data.model_type}` not found for provider `{data.provider}`" raise ValueError( msg, @@ -447,21 +456,21 @@ def invoke_multimodal_rerank( session: Session, data: ModelInvokeMultimodalRerankRequest, ) -> object: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, ) if isinstance(model_instance, RerankModel): - return model_instance.invoke_multimodal( - data.model, - data.credentials, - data.query, - data.docs, - score_threshold=data.score_threshold, - top_n=data.top_n, - user=data.user_id, - ) + with use_current_session(session): + return model_instance.invoke_multimodal( + data.model, + data.credentials, + data.query, + data.docs, + score_threshold=data.score_threshold, + top_n=data.top_n, + user=data.user_id, + ) msg = f"Model `{data.model_type}` not found for provider `{data.provider}`" raise ValueError( msg, @@ -472,26 +481,26 @@ def invoke_tts( session: Session, data: ModelInvokeTTSRequest, ) -> Generator[dict[str, str], None, None]: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, ) if isinstance(model_instance, TTSModel): - b = model_instance.invoke( - data.model, - data.tenant_id, - data.credentials, - data.content_text, - data.voice, - data.user_id, - ) - if isinstance(b, bytes | bytearray | memoryview): - yield {"result": binascii.hexlify(b).decode()} - return + with use_current_session(session): + b = model_instance.invoke( + data.model, + data.tenant_id, + data.credentials, + data.content_text, + data.voice, + data.user_id, + ) + if isinstance(b, bytes | bytearray | memoryview): + yield {"result": binascii.hexlify(b).decode()} + return - for chunk in b: - yield {"result": binascii.hexlify(chunk).decode()} + for chunk in b: + yield {"result": binascii.hexlify(chunk).decode()} else: msg = f"Model `{data.model_type}` not found for provider `{data.provider}`" raise TypeError( @@ -526,7 +535,6 @@ def invoke_speech_to_text( session: Session, data: ModelInvokeSpeech2TextRequest, ) -> dict[str, str]: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, @@ -538,14 +546,14 @@ def invoke_speech_to_text( with pathlib.Path(temp.name).open("rb") as f: if isinstance(model_instance, Speech2TextModel): - return { - "result": model_instance.invoke( + with use_current_session(session): + result = model_instance.invoke( data.model, data.credentials, f, data.user_id, - ), - } + ) + return {"result": result} msg = ( f"Model `{data.model_type}` not found for provider " f"`{data.provider}`" @@ -581,21 +589,20 @@ def invoke_moderation( session: Session, data: ModelInvokeModerationRequest, ) -> dict[str, bool]: - del session model_instance = self.registration.get_model_instance( data.provider, data.model_type, ) if isinstance(model_instance, ModerationModel): - return { - "result": model_instance.invoke( + with use_current_session(session): + result = model_instance.invoke( data.model, data.credentials, data.text, data.user_id, - ), - } + ) + return {"result": result} msg = f"Model `{data.model_type}` not found for provider `{data.provider}`" raise ValueError( msg, diff --git a/src/dify_plugin/core/session_context.py b/src/dify_plugin/core/session_context.py new file mode 100644 index 00000000..71ea2296 --- /dev/null +++ b/src/dify_plugin/core/session_context.py @@ -0,0 +1,37 @@ +"""Current model invocation session access.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from dify_plugin.core.runtime import Session + +_current_session: ContextVar[Session | None] = ContextVar( + "dify_plugin_current_session", + default=None, +) + + +def get_current_session() -> Session | None: + """Return the current model invocation session, if any. + + Model plugins can read ``session.app_id`` from the returned session. + It is ``None`` for model calls outside an app execution context. + + Returns: + The current session, or ``None``. + """ + return _current_session.get() + + +@contextmanager +def use_current_session(session: Session) -> Iterator[None]: + token = _current_session.set(session) + try: + yield + finally: + _current_session.reset(token) diff --git a/tests/core/__init__.py b/tests/core/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/tests/core/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/core/test_session_context.py b/tests/core/test_session_context.py new file mode 100644 index 00000000..bc5bb2d4 --- /dev/null +++ b/tests/core/test_session_context.py @@ -0,0 +1,89 @@ +from collections.abc import Generator +from unittest.mock import MagicMock + +import dify_plugin +from dify_plugin.core.plugin_executor import PluginExecutor +from dify_plugin.core.runtime import Session +from dify_plugin.core.session_context import get_current_session +from dify_plugin.interfaces.model.large_language_model import LargeLanguageModel +from dify_plugin.interfaces.model.text_embedding_model import TextEmbeddingModel + + +def _session(app_id: str = "app-1") -> Session: + session = Session.empty_session() + session.app_id = app_id + return session + + +def _executor(model: object) -> PluginExecutor: + registration = MagicMock() + registration.get_model_instance.return_value = model + return PluginExecutor(config=MagicMock(), registration=registration) + + +def _llm_data() -> MagicMock: + data = MagicMock() + data.provider = "provider" + data.model_type = "llm" + data.model = "model" + data.credentials = {} + data.prompt_messages = [] + data.model_parameters = {} + data.tools = None + data.stop = None + data.stream = True + data.user_id = "user" + return data + + +def _text_embedding_data() -> MagicMock: + data = MagicMock() + data.provider = "provider" + data.model_type = "text-embedding" + data.model = "model" + data.credentials = {} + data.texts = ["hello"] + data.user_id = "user" + return data + + +def test_get_current_session_is_public() -> None: + assert dify_plugin.get_current_session is get_current_session + assert "get_current_session" in dify_plugin.__all__ + + +def test_llm_stream_keeps_session_until_consumed() -> None: + seen: list[Session | None] = [] + + def stream() -> Generator[str, None, None]: + seen.append(get_current_session()) + yield "chunk" + + model = MagicMock(spec=LargeLanguageModel) + model.invoke.return_value = stream() + + result = _executor(model).invoke_llm(_session(), _llm_data()) + + assert get_current_session() is None + assert list(result) == ["chunk"] + assert seen[0] is not None + assert seen[0].app_id == "app-1" + assert get_current_session() is None + + +def test_text_embedding_exposes_session_during_invoke() -> None: + seen: list[Session | None] = [] + + def invoke(*_args: object) -> list[str]: + seen.append(get_current_session()) + return ["embedding"] + + model = MagicMock(spec=TextEmbeddingModel) + model.invoke.side_effect = invoke + + assert _executor(model).invoke_text_embedding( + _session(), _text_embedding_data() + ) == ["embedding"] + assert seen[0] is not None + assert seen[0].app_id == "app-1" + assert get_current_session() is None