Files
dify-plugin-sdks/python/dify_plugin/core/runtime.py
T

337 lines
11 KiB
Python

from enum import Enum
from typing import Optional
from pydantic import BaseModel
from ..config.config import InstallMethod
from concurrent.futures import ThreadPoolExecutor
from abc import ABC
from collections.abc import Generator
import json
from typing import Generic, Type, TypeVar, Union
import uuid
import httpx
from yarl import URL
from ..core.server.__base.request_reader import RequestReader
from ..core.server.__base.response_writer import ResponseWriter
from ..core.server.tcp.request_reader import TCPReaderWriter
from ..core.entities.invocation import InvokeType
from ..core.entities.plugin.io import (
PluginInStream,
PluginInStreamBase,
PluginInStreamEvent,
)
#################################################
# Session
#################################################
class ModelInvocations:
def __init__(self, session: "Session") -> None:
from ..invocations.model.llm import LLMInvocation, SummaryInvocation
from ..invocations.model.moderation import ModerationInvocation
from ..invocations.model.rerank import RerankInvocation
from ..invocations.model.speech2text import Speech2TextInvocation
from ..invocations.model.text_embedding import TextEmbeddingInvocation
from ..invocations.model.tts import TTSInvocation
self.llm = LLMInvocation(session)
self.text_embedding = TextEmbeddingInvocation(session)
self.rerank = RerankInvocation(session)
self.speech2text = Speech2TextInvocation(session)
self.tts = TTSInvocation(session)
self.moderation = ModerationInvocation(session)
self.summary = SummaryInvocation(session)
class AppInvocations:
def __init__(self, session: "Session"):
from ..invocations.app.chat import ChatAppInvocation
from ..invocations.app.completion import CompletionAppInvocation
from ..invocations.app.workflow import WorkflowAppInvocation
self.chat = ChatAppInvocation(session)
self.completion = CompletionAppInvocation(session)
self.workflow = WorkflowAppInvocation(session)
class WorkflowNodeInvocations:
def __init__(self, session: "Session"):
from ..invocations.workflow_node.parameter_extractor import (
ParameterExtractorNodeInvocation,
)
from ..invocations.workflow_node.question_classifier import (
QuestionClassifierNodeInvocation,
)
self.question_classifier = QuestionClassifierNodeInvocation(session)
self.parameter_extractor = ParameterExtractorNodeInvocation(session)
class Session:
# class variable to store all sessions
_session_pool: set["Session"] = set()
def __init__(
self,
session_id: str,
executor: ThreadPoolExecutor,
reader: RequestReader,
writer: ResponseWriter,
install_method: Optional[InstallMethod] = None,
dify_plugin_daemon_url: Optional[str] = None,
conversation_id: Optional[str] = None,
message_id: Optional[str] = None,
app_id: Optional[str] = None,
endpoint_id: Optional[str] = None,
) -> None:
# current session id
self.session_id: str = session_id
# add current session to session pool
self._session_pool.add(self)
# thread pool executor
self._executor: ThreadPoolExecutor = executor
# reader and writer
self.reader: RequestReader = reader
self.writer: ResponseWriter = writer
# conversation id
self.conversation_id: Optional[str] = conversation_id
# message id
self.message_id: Optional[str] = message_id
# app id
self.app_id: Optional[str] = app_id
# endpoint id
self.endpoint_id: Optional[str] = endpoint_id
# install method
self.install_method: Optional[InstallMethod] = install_method
# dify plugin daemon url
self.dify_plugin_daemon_url: Optional[str] = dify_plugin_daemon_url
# register invocations
self._register_invocations()
def _register_invocations(self) -> None:
from ..invocations.tool import ToolInvocation
from ..invocations.storage import StorageInvocation
from ..invocations.file import File
self.model = ModelInvocations(self)
self.tool = ToolInvocation(self)
self.app = AppInvocations(self)
self.workflow_node = WorkflowNodeInvocations(self)
self.storage = StorageInvocation(self)
self.file = File(self)
@classmethod
def empty_session(cls) -> "Session":
return cls(
session_id="",
executor=ThreadPoolExecutor(),
reader=TCPReaderWriter(host="", port=0, key=""),
writer=TCPReaderWriter(host="", port=0, key=""),
install_method=None,
dify_plugin_daemon_url=None,
)
def __del__(self) -> None:
self._session_pool.remove(self)
def close(self) -> None:
self._session_pool.remove(self)
@classmethod
def get_session(cls, session_id: str) -> "Session | None":
for session in cls._session_pool:
if session.session_id == session_id:
return session
return None
#################################################
# Backwards Invocation Request
#################################################
class BackwardsInvocationResponseEvent(BaseModel):
class Event(Enum):
response = "response"
Error = "error"
End = "end"
backwards_request_id: str
event: Event
message: str
data: Optional[dict]
T = TypeVar("T", bound=Union[BaseModel, dict, str])
class BackwardsInvocation(Generic[T], ABC):
def __init__(
self,
session: Optional[Session] = None,
) -> None:
self.session = session
def _generate_backwards_request_id(self):
return uuid.uuid4().hex
def _backwards_invoke(
self,
type: InvokeType,
data_type: Type[T],
data: dict,
) -> Generator[T, None, None]:
"""
backwards invoke dify depends on current runtime type
"""
backwards_request_id = self._generate_backwards_request_id()
if not self.session:
raise Exception("current tool runtime does not support backwards invoke")
if self.session.install_method in [InstallMethod.Local, InstallMethod.Remote]:
return self._full_duplex_backwards_invoke(
backwards_request_id, type, data_type, data
)
return self._http_backwards_invoke(backwards_request_id, type, data_type, data)
def _line_converter_wrapper(
self,
generator: Generator[PluginInStreamBase | None, None, None],
data_type: Type[T],
) -> Generator[T, None, None]:
"""
convert string into type T
"""
empty_response_count = 0
for chunk in generator:
"""
accept response from input stream and wait for at most 60 seconds
"""
if chunk is None:
empty_response_count += 1
# if no response for 250 seconds, break
if empty_response_count >= 250:
raise Exception("invocation exited without response")
continue
event = BackwardsInvocationResponseEvent(**chunk.data)
if event.event == BackwardsInvocationResponseEvent.Event.End:
break
if event.event == BackwardsInvocationResponseEvent.Event.Error:
raise Exception(event.message)
if event.data is None:
break
empty_response_count = 0
try:
yield data_type(**event.data)
except Exception as e:
raise Exception(f"Failed to parse response: {str(e)}")
def _http_backwards_invoke(
self,
backwards_request_id: str,
type: InvokeType,
data_type: Type[T],
data: dict,
) -> Generator[T, None, None]:
"""
http backwards invoke
"""
if not self.session or not self.session.dify_plugin_daemon_url:
raise Exception("current tool runtime does not support backwards invoke")
url = (
URL(self.session.dify_plugin_daemon_url)
/ "backwards-invocation"
/ "transaction"
)
headers = {
"Dify-Plugin-Session-ID": self.session.session_id,
}
payload = self.session.writer.session_message_text(
session_id=self.session.session_id,
data=self.session.writer.stream_invoke_object(
data={
"type": type.value,
"backwards_request_id": backwards_request_id,
"request": data,
}
),
)
with httpx.Client() as client:
with client.stream(
method="POST",
url=str(url),
headers=headers,
content=payload,
timeout=(300, 300, 300, 300), # 300 seconds for connection, read, write, and pool
) as response:
def generator():
for line in response.iter_lines():
if not line:
continue
data = json.loads(line)
yield PluginInStreamBase(
session_id=data["session_id"],
event=PluginInStreamEvent.value_of(data["event"]),
data=data["data"],
)
yield from self._line_converter_wrapper(generator(), data_type)
def _full_duplex_backwards_invoke(
self,
backwards_request_id: str,
type: InvokeType,
data_type: Type[T],
data: dict,
) -> Generator[T, None, None]:
if not self.session:
raise Exception("current tool runtime does not support backwards invoke")
self.session.writer.session_message(
session_id=self.session.session_id,
data=self.session.writer.stream_invoke_object(
data={
"type": type.value,
"backwards_request_id": backwards_request_id,
"request": data,
}
),
)
def filter(data: PluginInStream) -> bool:
return (
data.event == PluginInStreamEvent.BackwardInvocationResponse
and data.data.get("backwards_request_id") == backwards_request_id
)
with self.session.reader.read(filter) as reader:
yield from self._line_converter_wrapper(
reader.read(timeout_for_round=1), data_type
)