tools.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  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)) and not func.startswith("__")
  19. ]
  20. specs = []
  21. for function_item in function_list:
  22. function_name = function_item["name"]
  23. function = function_item["function"]
  24. function_doc = doc_to_dict(function.__doc__ or function_name)
  25. specs.append(
  26. {
  27. "name": function_name,
  28. # TODO: multi-line desc?
  29. "description": function_doc.get("description", function_name),
  30. "parameters": {
  31. "type": "object",
  32. "properties": {
  33. param_name: {
  34. "type": param_annotation.__name__.lower(),
  35. **(
  36. {
  37. "enum": (
  38. str(param_annotation.__args__)
  39. if hasattr(param_annotation, "__args__")
  40. else None
  41. )
  42. }
  43. if hasattr(param_annotation, "__args__")
  44. else {}
  45. ),
  46. "description": function_doc.get("params", {}).get(
  47. param_name, param_name
  48. ),
  49. }
  50. for param_name, param_annotation in get_type_hints(
  51. function
  52. ).items()
  53. if param_name != "return" and param_name != "__user__"
  54. },
  55. "required": [
  56. name
  57. for name, param in inspect.signature(
  58. function
  59. ).parameters.items()
  60. if param.default is param.empty
  61. ],
  62. },
  63. }
  64. )
  65. return specs