tools.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. import inspect
  2. from typing import get_type_hints
  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"
  56. and not (
  57. param_name.startswith("__") and param_name.endswith("__")
  58. )
  59. },
  60. "required": [
  61. name
  62. for name, param in inspect.signature(
  63. function
  64. ).parameters.items()
  65. if param.default is param.empty
  66. and not (name.startswith("__") and name.endswith("__"))
  67. ],
  68. },
  69. }
  70. )
  71. return specs