mirror of
https://github.com/langgenius/dify-plugin-sdks.git
synced 2026-07-22 10:25:23 -04:00
Yeuoly
117 lines
4.3 KiB
Python
117 lines
4.3 KiB
Python
import logging
|
|
import sys
|
|
from typing import Type
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from dify_plugin.config.config import DifyPluginConfig
|
|
from dify_plugin.core.runtime.entities.request import PluginInvokeType, ToolInvokeRequest
|
|
from dify_plugin.core.runtime.session import Session
|
|
from dify_plugin.core.server.io_server import IOServer
|
|
from dify_plugin.logger_format import DifyPluginLoggerFormatter
|
|
from dify_plugin.model.model import Model, ModelProvider
|
|
from dify_plugin.tool.entities import ToolRuntime
|
|
from dify_plugin.tool.tool import Tool, ToolProvider
|
|
from dify_plugin.utils.io_writer import PluginOutputStream
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
handler = logging.StreamHandler(sys.stdout)
|
|
handler.setLevel(logging.INFO)
|
|
handler.setFormatter(DifyPluginLoggerFormatter())
|
|
|
|
logger.addHandler(handler)
|
|
|
|
class Plugin(IOServer):
|
|
class ToolRegistration(BaseModel):
|
|
cls: Type[ToolProvider]
|
|
tools: list[Type[Tool]]
|
|
|
|
tools: list[ToolRegistration]
|
|
models: list[Type[ModelProvider]]
|
|
|
|
def __init__(self, config: DifyPluginConfig) -> None:
|
|
self.tools = []
|
|
self.models = []
|
|
|
|
super().__init__(config)
|
|
|
|
def register_model_provider(self, provider: Type[ModelProvider]):
|
|
pass
|
|
|
|
def register_model(self, tool: Type[Model]):
|
|
pass
|
|
|
|
def register_tool_provider(self, tool: Type[ToolProvider]):
|
|
if not issubclass(tool, ToolProvider):
|
|
raise ValueError('Tool must be a subclass of ToolProvider')
|
|
logger.info(f'Registering tool provider {tool.__name__}')
|
|
self.tools.append(self.ToolRegistration(cls=tool, tools=[]))
|
|
|
|
return tool
|
|
|
|
def register_tool(self, provider: Type[ToolProvider]):
|
|
def decorator(cls: Type[Tool]):
|
|
if not issubclass(cls, Tool):
|
|
raise ValueError('Tool must be a subclass of Tool')
|
|
logger.info(f'Registering tool {cls.__name__} for provider {provider.__name__}')
|
|
for tool_registration in self.tools:
|
|
if tool_registration.cls == provider:
|
|
tool_registration.tools.append(cls)
|
|
break
|
|
else:
|
|
raise ValueError(f'Provider {provider.__name__} not found')
|
|
|
|
return cls
|
|
|
|
return decorator
|
|
|
|
def get_tool_provider_cls(self, provider: str):
|
|
for provider_registration in self.tools:
|
|
if provider_registration.cls.configuration().name == provider:
|
|
return provider_registration.cls
|
|
return None
|
|
|
|
def get_tool_cls(self, provider: Type[ToolProvider], tool: str):
|
|
for tool_registration in self.tools:
|
|
if tool_registration.cls == provider:
|
|
for tool_cls in tool_registration.tools:
|
|
if tool_cls.configuration().name == tool:
|
|
return tool_cls
|
|
return None
|
|
|
|
async def _execute_request(self, session_id: str, data: dict):
|
|
session = Session(session_id=session_id, executor=self.executer)
|
|
if data.get('type') == PluginInvokeType.Tool.value:
|
|
request = ToolInvokeRequest(**data)
|
|
|
|
provider_cls = self.get_tool_provider_cls(request.provider)
|
|
if provider_cls is None:
|
|
raise ValueError(f'Provider {request.provider} not found')
|
|
|
|
tool_cls = self.get_tool_cls(provider_cls, request.tool)
|
|
if tool_cls is None:
|
|
raise ValueError(f'Tool {request.tool} not found for provider {request.provider}')
|
|
|
|
provider = provider_cls()
|
|
tool = tool_cls(runtime=ToolRuntime(
|
|
credentials=request.credentials,
|
|
user_id=request.user_id
|
|
))
|
|
|
|
response = session.run_tool(action=request.action,
|
|
provider=provider,
|
|
tool=tool,
|
|
parameters=request.parameters)
|
|
|
|
async for message in response:
|
|
PluginOutputStream.session_message(
|
|
session_id=session_id,
|
|
data=PluginOutputStream.stream_object(
|
|
data=message
|
|
)
|
|
)
|
|
|
|
elif data.get('type') == PluginInvokeType.Model.value:
|
|
request = ToolInvokeRequest(**data) |