main.py 934 B

123456789101112131415161718192021222324252627
  1. from typing import Generator
  2. from dify_plugin import ToolProvider, Plugin
  3. from dify_plugin.tool.entities import ToolInvokeTextMessage, ToolProviderConfiguration
  4. from dify_plugin.tool.tool import Tool
  5. from dify_plugin.tool.entities import ToolConfiguration, ToolInvokeMessage
  6. plugin = Plugin()
  7. @plugin.register_tool_provider
  8. class BasicMath(ToolProvider):
  9. @classmethod
  10. def configuration(cls) -> ToolProviderConfiguration:
  11. return ToolProviderConfiguration(name='basic_math')
  12. @plugin.register_tool(BasicMath)
  13. class Add(Tool):
  14. @classmethod
  15. def configuration(cls) -> ToolConfiguration:
  16. return ToolConfiguration(name='add')
  17. def _invoke(self, user_id: str, tool_parameter: dict) -> Generator[ToolInvokeMessage, None, None]:
  18. result = tool_parameter['a'] + tool_parameter['b']
  19. yield ToolInvokeTextMessage(message={'result': str(result)})
  20. if __name__ == '__main__':
  21. plugin.run()