schema.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. """
  2. Schema Generator - 从函数签名自动生成 OpenAI Tool Schema
  3. 职责:
  4. 1. 解析函数签名(参数、类型注解、默认值)
  5. 2. 解析 docstring(Google 风格)
  6. 3. 生成 OpenAI Tool Calling 格式的 JSON Schema
  7. 从 Resonote/llm/tools/schema.py 抽取
  8. """
  9. import inspect
  10. import logging
  11. import types
  12. from typing import Any, Dict, List, Literal, Optional, Union, get_args, get_origin, get_type_hints
  13. logger = logging.getLogger(__name__)
  14. # 尝试导入 docstring_parser,如果不可用则提供降级方案
  15. try:
  16. from docstring_parser import parse as parse_docstring
  17. HAS_DOCSTRING_PARSER = True
  18. except ImportError:
  19. HAS_DOCSTRING_PARSER = False
  20. logger.warning("docstring_parser not installed, using fallback docstring parsing")
  21. def _simple_parse_docstring(docstring: str) -> tuple[str, Dict[str, str]]:
  22. """简单的 docstring 解析(降级方案)"""
  23. if not docstring:
  24. return "", {}
  25. lines = docstring.strip().split("\n")
  26. description = lines[0] if lines else ""
  27. param_descriptions = {}
  28. # 简单解析 Args: 部分
  29. in_args = False
  30. for line in lines[1:]:
  31. line = line.strip()
  32. if line.lower().startswith("args:"):
  33. in_args = True
  34. continue
  35. if line.lower().startswith(("returns:", "raises:", "example:")):
  36. in_args = False
  37. continue
  38. if in_args and ":" in line:
  39. parts = line.split(":", 1)
  40. param_name = parts[0].strip()
  41. param_desc = parts[1].strip() if len(parts) > 1 else ""
  42. param_descriptions[param_name] = param_desc
  43. return description, param_descriptions
  44. class SchemaGenerator:
  45. """从函数生成 OpenAI Tool Schema"""
  46. # Python 类型到 JSON Schema 类型的映射
  47. TYPE_MAP = {
  48. str: "string",
  49. int: "integer",
  50. float: "number",
  51. bool: "boolean",
  52. list: "array",
  53. dict: "object",
  54. List: "array",
  55. Dict: "object",
  56. }
  57. @classmethod
  58. def generate(cls, func: callable, hidden_params: Optional[List[str]] = None) -> Dict[str, Any]:
  59. """
  60. 从函数生成 OpenAI Tool Schema
  61. Args:
  62. func: 要生成 Schema 的函数
  63. hidden_params: 隐藏参数列表(不生成 schema)
  64. Returns:
  65. OpenAI Tool Schema(JSON 格式)
  66. """
  67. hidden_params = hidden_params or []
  68. # 解析函数签名
  69. sig = inspect.signature(func)
  70. func_name = func.__name__
  71. try:
  72. type_hints = get_type_hints(func)
  73. except (NameError, TypeError):
  74. # Some dynamically registered callables may refer to names which
  75. # are intentionally unavailable at registration time. Preserve
  76. # the previous best-effort behavior for those callables.
  77. type_hints = {}
  78. # 解析 docstring
  79. if HAS_DOCSTRING_PARSER:
  80. doc = parse_docstring(func.__doc__ or "")
  81. func_description = doc.short_description or doc.long_description or f"Call {func_name}"
  82. param_descriptions = {p.arg_name: p.description for p in doc.params if p.description}
  83. else:
  84. func_description, param_descriptions = _simple_parse_docstring(func.__doc__ or "")
  85. if not func_description:
  86. func_description = f"Call {func_name}"
  87. # 生成参数 Schema
  88. properties = {}
  89. required = []
  90. for param_name, param in sig.parameters.items():
  91. # 跳过特殊参数
  92. if param_name in ["self", "cls", "kwargs"]:
  93. continue
  94. # 跳过隐藏参数
  95. if param_name in hidden_params:
  96. continue
  97. # 获取类型注解
  98. param_type = type_hints.get(
  99. param_name,
  100. param.annotation if param.annotation != inspect.Parameter.empty else str,
  101. )
  102. # 生成参数 Schema
  103. param_schema = cls._type_to_schema(param_type)
  104. # 添加描述
  105. if param_name in param_descriptions:
  106. param_schema["description"] = param_descriptions[param_name]
  107. # 添加默认值
  108. if param.default != inspect.Parameter.empty:
  109. param_schema["default"] = param.default
  110. else:
  111. required.append(param_name)
  112. properties[param_name] = param_schema
  113. # 构建完整的 Schema
  114. schema = {
  115. "type": "function",
  116. "function": {
  117. "name": func_name,
  118. "description": func_description,
  119. "parameters": {
  120. "type": "object",
  121. "properties": properties,
  122. "required": required
  123. }
  124. }
  125. }
  126. return schema
  127. @classmethod
  128. def _type_to_schema(cls, python_type: Any) -> Dict[str, Any]:
  129. """将 Python 类型转换为 JSON Schema"""
  130. if python_type is Any:
  131. return {}
  132. origin = get_origin(python_type)
  133. args = get_args(python_type)
  134. # 处理 Literal[...]
  135. if origin is Literal:
  136. values = list(args)
  137. if all(isinstance(v, str) for v in values):
  138. return {"type": "string", "enum": values}
  139. elif all(isinstance(v, int) for v in values):
  140. return {"type": "integer", "enum": values}
  141. return {"enum": values}
  142. # 处理 Union[T, ...] 和 Optional[T]
  143. if origin in (Union, types.UnionType):
  144. if len(args) == 2 and type(None) in args:
  145. # Optional[T] = Union[T, None]
  146. inner = args[0] if args[1] is type(None) else args[1]
  147. return cls._type_to_schema(inner)
  148. non_none = [a for a in args if a is not type(None)]
  149. return {"oneOf": [cls._type_to_schema(a) for a in non_none]}
  150. # 处理 List[T]
  151. if origin is list or origin is List:
  152. if args:
  153. item_type = args[0]
  154. return {
  155. "type": "array",
  156. "items": cls._type_to_schema(item_type)
  157. }
  158. return {"type": "array"}
  159. # 处理 Dict[K, V]
  160. if origin is dict or origin is Dict:
  161. return {"type": "object"}
  162. # 处理基础类型
  163. if python_type in cls.TYPE_MAP:
  164. return {"type": cls.TYPE_MAP[python_type]}
  165. # 检查是否是 Protocol(如 ToolContext)
  166. # Protocol 类型用于依赖注入,不应出现在 schema 中
  167. type_name = getattr(python_type, "__name__", str(python_type))
  168. if "Protocol" in str(type(python_type)) or type_name in ("ToolContext",):
  169. logger.debug(f"Skipping Protocol type {python_type} (used for dependency injection)")
  170. return {}
  171. # 默认为 string
  172. logger.debug(f"Unknown type {python_type}, defaulting to string")
  173. return {"type": "string"}