tools.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. import inspect
  2. from typing import get_type_hints, List, Dict, Any
  3. def doc_to_dict(docstring):
  4. lines = docstring.split("\n")
  5. description = lines[1].strip()
  6. param_dict = {}
  7. for line in lines:
  8. if ":param" in line:
  9. line = line.replace(":param", "").strip()
  10. param, desc = line.split(":", 1)
  11. param_dict[param.strip()] = desc.strip()
  12. ret_dict = {"description": description, "params": param_dict}
  13. return ret_dict
  14. def get_tools_specs(tools) -> List[dict]:
  15. function_list = [
  16. {"name": func, "function": getattr(tools, func)}
  17. for func in dir(tools)
  18. if callable(getattr(tools, func))
  19. and not func.startswith("__")
  20. and not inspect.isclass(getattr(tools, func))
  21. ]
  22. specs = []
  23. for function_item in function_list:
  24. function_name = function_item["name"]
  25. function = function_item["function"]
  26. function_doc = doc_to_dict(function.__doc__ or function_name)
  27. specs.append(
  28. {
  29. "name": function_name,
  30. # TODO: multi-line desc?
  31. "description": function_doc.get("description", function_name),
  32. "parameters": {
  33. "type": "object",
  34. "properties": {
  35. param_name: {
  36. "type": param_annotation.__name__.lower(),
  37. **(
  38. {
  39. "enum": (
  40. str(param_annotation.__args__)
  41. if hasattr(param_annotation, "__args__")
  42. else None
  43. )
  44. }
  45. if hasattr(param_annotation, "__args__")
  46. else {}
  47. ),
  48. "description": function_doc.get("params", {}).get(
  49. param_name, param_name
  50. ),
  51. }
  52. for param_name, param_annotation in get_type_hints(
  53. function
  54. ).items()
  55. if param_name != "return" and param_name != "__user__"
  56. },
  57. "required": [
  58. name
  59. for name, param in inspect.signature(
  60. function
  61. ).parameters.items()
  62. if param.default is param.empty
  63. and not (name.startswith("__") and name.endswith("__"))
  64. ],
  65. },
  66. }
  67. )
  68. return specs