tools.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. import inspect
  2. import logging
  3. import re
  4. from typing import Any, Awaitable, Callable, get_type_hints
  5. from functools import update_wrapper, partial
  6. from langchain_core.utils.function_calling import convert_to_openai_function
  7. from open_webui.apps.webui.models.tools import Tools
  8. from open_webui.apps.webui.models.users import UserModel
  9. from open_webui.apps.webui.utils import load_tools_module_by_id
  10. from pydantic import BaseModel, Field, create_model
  11. log = logging.getLogger(__name__)
  12. def apply_extra_params_to_tool_function(
  13. function: Callable, extra_params: dict
  14. ) -> Callable[..., Awaitable]:
  15. partial_func = partial(function, **extra_params)
  16. if inspect.iscoroutinefunction(function):
  17. update_wrapper(partial_func, function)
  18. return partial_func
  19. async def new_function(*args, **kwargs):
  20. return partial_func(*args, **kwargs)
  21. update_wrapper(new_function, function)
  22. return new_function
  23. # Mutation on extra_params
  24. def get_tools(
  25. webui_app, tool_ids: list[str], user: UserModel, extra_params: dict
  26. ) -> dict[str, dict]:
  27. tools_dict = {}
  28. for tool_id in tool_ids:
  29. tools = Tools.get_tool_by_id(tool_id)
  30. if tools is None:
  31. continue
  32. module = webui_app.state.TOOLS.get(tool_id, None)
  33. if module is None:
  34. module, _ = load_tools_module_by_id(tool_id)
  35. webui_app.state.TOOLS[tool_id] = module
  36. extra_params["__id__"] = tool_id
  37. if hasattr(module, "valves") and hasattr(module, "Valves"):
  38. valves = Tools.get_tool_valves_by_id(tool_id) or {}
  39. module.valves = module.Valves(**valves)
  40. if hasattr(module, "UserValves"):
  41. extra_params["__user__"]["valves"] = module.UserValves( # type: ignore
  42. **Tools.get_user_valves_by_id_and_user_id(tool_id, user.id)
  43. )
  44. for spec in tools.specs:
  45. # Remove internal parameters
  46. spec["parameters"]["properties"] = {
  47. key: val
  48. for key, val in spec["parameters"]["properties"].items()
  49. if not key.startswith("__")
  50. }
  51. function_name = spec["name"]
  52. # convert to function that takes only model params and inserts custom params
  53. original_func = getattr(module, function_name)
  54. callable = apply_extra_params_to_tool_function(original_func, extra_params)
  55. # TODO: This needs to be a pydantic model
  56. tool_dict = {
  57. "toolkit_id": tool_id,
  58. "callable": callable,
  59. "spec": spec,
  60. "pydantic_model": function_to_pydantic_model(callable),
  61. "file_handler": hasattr(module, "file_handler") and module.file_handler,
  62. "citation": hasattr(module, "citation") and module.citation,
  63. }
  64. # TODO: if collision, prepend toolkit name
  65. if function_name in tools_dict:
  66. log.warning(f"Tool {function_name} already exists in another tools!")
  67. log.warning(f"Collision between {tools} and {tool_id}.")
  68. log.warning(f"Discarding {tools}.{function_name}")
  69. else:
  70. tools_dict[function_name] = tool_dict
  71. return tools_dict
  72. def parse_docstring(docstring):
  73. """
  74. Parse a function's docstring to extract parameter descriptions in reST format.
  75. Args:
  76. docstring (str): The docstring to parse.
  77. Returns:
  78. dict: A dictionary where keys are parameter names and values are descriptions.
  79. """
  80. if not docstring:
  81. return {}
  82. # Regex to match `:param name: description` format
  83. param_pattern = re.compile(r":param (\w+):\s*(.+)")
  84. param_descriptions = {}
  85. for line in docstring.splitlines():
  86. match = param_pattern.match(line.strip())
  87. if match:
  88. param_name, param_description = match.groups()
  89. param_descriptions[param_name] = param_description
  90. return param_descriptions
  91. def function_to_pydantic_model(func: Callable) -> type[BaseModel]:
  92. """
  93. Converts a Python function's type hints and docstring to a Pydantic model,
  94. including support for nested types, default values, and descriptions.
  95. Args:
  96. func: The function whose type hints and docstring should be converted.
  97. model_name: The name of the generated Pydantic model.
  98. Returns:
  99. A Pydantic model class.
  100. """
  101. type_hints = get_type_hints(func)
  102. signature = inspect.signature(func)
  103. parameters = signature.parameters
  104. docstring = func.__doc__
  105. descriptions = parse_docstring(docstring)
  106. field_defs = {}
  107. for name, param in parameters.items():
  108. type_hint = type_hints.get(name, Any)
  109. default_value = param.default if param.default is not param.empty else ...
  110. description = descriptions.get(name, None)
  111. if not description:
  112. field_defs[name] = type_hint, default_value
  113. continue
  114. field_defs[name] = type_hint, Field(default_value, description=description)
  115. return create_model(func.__name__, **field_defs)
  116. def get_callable_attributes(tool: object) -> list[Callable]:
  117. return [
  118. getattr(tool, func)
  119. for func in dir(tool)
  120. if callable(getattr(tool, func))
  121. and not func.startswith("__")
  122. and not inspect.isclass(getattr(tool, func))
  123. ]
  124. def get_tools_specs(tool_class: object) -> list[dict]:
  125. function_list = get_callable_attributes(tool_class)
  126. models = map(function_to_pydantic_model, function_list)
  127. return [convert_to_openai_function(tool) for tool in models]