Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-agents/src/lexigram/ai/agents/tools/schema.py: 7%

1"""Generate JSON schema from Python type hints for LLM function calling.""" 

2 

3from __future__ import annotations 

4 

5import inspect 

6import types 

7from typing import Any, Callable, Union, get_args, get_origin, get_type_hints 

8 

9 

10def _parse_param_descriptions(docstring: str | None) -> dict[str, str]: 

11 """Extract parameter descriptions from a Google-style docstring. 

12 

13 Parses the ``Args:`` section and returns a mapping of parameter name to 

14 description string. 

15 

16 Args: 

17 docstring: Raw docstring from a function. May be None. 

18 

19 Returns: 

20 Dict mapping each parameter name to its description text. Empty dict 

21 if the docstring is absent or has no ``Args:`` section. 

22 """ 

23 if not docstring: 

24 return {} 

25 

26 descriptions: dict[str, str] = {} 

27 in_args = False 

28 current_param: str | None = None 

29 current_desc_parts: list[str] = [] 

30 

31 for line in docstring.splitlines(): 

32 stripped = line.strip() 

33 

34 if stripped in ("Args:", "Arguments:", "Parameters:"): 

35 in_args = True 

36 continue 

37 

38 if in_args: 

39 # A new top-level section (e.g. "Returns:", "Raises:") ends args block 

40 if stripped and not line.startswith(" ") and not line.startswith("\t"): 

41 if stripped.endswith(":") and " " not in stripped.rstrip(":"): 

42 if current_param: 

43 descriptions[current_param] = " ".join(current_desc_parts).strip() 

44 in_args = False 

45 current_param = None 

46 current_desc_parts = [] 

47 continue 

48 

49 # Indented continuation of previous param description 

50 if current_param and stripped and ":" not in stripped.split()[0]: 

51 current_desc_parts.append(stripped) 

52 continue 

53 

54 # New param line: " param_name: description text" 

55 if ":" in stripped: 

56 if current_param: 

57 descriptions[current_param] = " ".join(current_desc_parts).strip() 

58 current_desc_parts = [] 

59 

60 param, _, desc = stripped.partition(":") 

61 param = param.strip() 

62 # Remove type hint from param if present: "param_name (type):" 

63 if "(" in param: 

64 param = param[: param.index("(")].strip() 

65 if param: 

66 current_param = param 

67 current_desc_parts = [desc.strip()] if desc.strip() else [] 

68 

69 if current_param: 

70 descriptions[current_param] = " ".join(current_desc_parts).strip() 

71 

72 return descriptions 

73 

74 

75def generate_json_schema( 

76 func: Callable[..., Any], 

77) -> dict[str, Any]: 

78 """Generate JSON schema from a callable's signature and type hints. 

79 

80 Maps Python types to JSON schema types: 

81 str -> string, int -> integer, float -> number, 

82 bool -> boolean, list -> array, dict -> object, 

83 Optional[T] -> nullable T 

84 

85 Parameter descriptions are extracted from the function's Google-style 

86 docstring ``Args:`` section and included in the schema. 

87 

88 Args: 

89 func: The callable to generate a schema for. 

90 

91 Returns: 

92 A JSON Schema ``object`` dict with ``properties`` and ``required`` keys. 

93 """ 

94 sig = inspect.signature(func) 

95 try: 

96 hints = get_type_hints(func) 

97 except (TypeError, NameError, AttributeError): 

98 hints = {} 

99 param_descriptions = _parse_param_descriptions(inspect.getdoc(func)) 

100 return _build_schema(sig, hints, param_descriptions) 

101 

102 

103def _build_schema( 

104 sig: inspect.Signature, 

105 hints: dict[str, Any], 

106 param_descriptions: dict[str, str] | None = None, 

107) -> dict[str, Any]: 

108 """Build JSON schema from an already-extracted signature and hints. 

109 

110 Args: 

111 sig: Parsed function signature. 

112 hints: Mapping of parameter name to resolved type annotation. 

113 param_descriptions: Optional mapping of parameter name to description 

114 extracted from the function docstring. 

115 

116 Returns: 

117 JSON Schema ``object`` dict. 

118 """ 

119 descriptions = param_descriptions or {} 

120 properties: dict[str, Any] = {} 

121 required: list[str] = [] 

122 

123 for param_name, param in sig.parameters.items(): 

124 if param_name in ("self", "cls"): 

125 continue 

126 

127 param_type = hints.get(param_name) 

128 if param_type is not None: 

129 prop = _python_type_to_json_schema(param_type) 

130 else: 

131 prop = {} 

132 

133 # Inject description from docstring 

134 if param_name in descriptions: 

135 prop["description"] = descriptions[param_name] 

136 

137 # Add default value 

138 if param.default is not inspect.Parameter.empty: 

139 prop["default"] = param.default 

140 

141 properties[param_name] = prop 

142 

143 # Required if no default 

144 if param.default is inspect.Parameter.empty: 

145 required.append(param_name) 

146 

147 return { 

148 "type": "object", 

149 "properties": properties, 

150 "required": required, 

151 } 

152 

153 

154def _python_type_to_json_schema(py_type: Any) -> dict[str, Any]: 

155 """Convert a Python type annotation to its JSON Schema equivalent. 

156 

157 Handles primitives, ``list[T]``, ``dict[K, V]``, ``X | Y`` (union), 

158 and ``X | None`` (optional). Pydantic models are converted via 

159 ``model_json_schema()``. 

160 

161 Args: 

162 py_type: A Python type annotation (from ``typing.get_type_hints``). 

163 

164 Returns: 

165 A JSON Schema dict for the given type. 

166 """ 

167 origin = get_origin(py_type) 

168 

169 # X | None or Optional[X] → schema with nullable=True 

170 # Handles both typing.Union (Optional[X]) and 3.10+ X | None (types.UnionType) 

171 if origin is Union or isinstance(py_type, types.UnionType): 

172 args = get_args(py_type) 

173 non_none = [a for a in args if a is not type(None)] 

174 if len(non_none) == 1: 

175 optional_schema = _python_type_to_json_schema(non_none[0]) 

176 optional_schema["nullable"] = True 

177 return optional_schema 

178 return {"anyOf": [_python_type_to_json_schema(a) for a in non_none]} 

179 

180 # list[T] → {"type": "array", "items": {...}} 

181 if origin is list: 

182 item_types = get_args(py_type) 

183 schema: dict[str, Any] = {"type": "array"} 

184 if item_types: 

185 schema["items"] = _python_type_to_json_schema(item_types[0]) 

186 return schema 

187 

188 # set[T] → {"type": "array", "items": {...}} 

189 if origin is set: 

190 item_types = get_args(py_type) 

191 schema = {"type": "array"} 

192 if item_types: 

193 schema["items"] = _python_type_to_json_schema(item_types[0]) 

194 return schema 

195 

196 # dict[K, V] → {"type": "object"} 

197 if origin is dict: 

198 return {"type": "object"} 

199 

200 # Pydantic models 

201 if hasattr(py_type, "model_json_schema"): 

202 return py_type.model_json_schema() 

203 

204 # Plain type(None) 

205 if py_type is type(None): 

206 return {"type": "null"} 

207 

208 # Annotated[T, ...] → extract inner type T 

209 import typing 

210 if origin is getattr(typing, "Annotated", type(None)) or str(origin).startswith("typing.Annotated"): 

211 # get_args(py_type)[0] is the base type 

212 args = get_args(py_type) 

213 if args: 

214 return _python_type_to_json_schema(args[0]) 

215 

216 # Literal[...] → {"type": "...", "enum": [...]} 

217 if origin is getattr(typing, "Literal", type(None)) or str(origin).startswith("typing.Literal"): 

218 args = get_args(py_type) 

219 if not args: 

220 return {} 

221 # Try to infer the type from the first value 

222 first_val = args[0] 

223 base_type = "string" 

224 if isinstance(first_val, int): 

225 base_type = "integer" 

226 elif isinstance(first_val, float): 

227 base_type = "number" 

228 elif isinstance(first_val, bool): 

229 base_type = "boolean" 

230 

231 return {"type": base_type, "enum": list(args)} 

232 

233 # enum.Enum 

234 import enum 

235 if isinstance(py_type, type) and issubclass(py_type, enum.Enum): 

236 values = [e.value for e in py_type] 

237 if not values: 

238 return {"type": "string"} 

239 

240 first_val = values[0] 

241 base_type = "string" 

242 if isinstance(first_val, int): 

243 base_type = "integer" 

244 elif isinstance(first_val, float): 

245 base_type = "number" 

246 elif isinstance(first_val, bool): 

247 base_type = "boolean" 

248 

249 return {"type": base_type, "enum": values} 

250 

251 _PYTHON_TO_JSON: dict[type, str] = { 

252 str: "string", 

253 int: "integer", 

254 float: "number", 

255 bool: "boolean", 

256 list: "array", 

257 dict: "object", 

258 bytes: "string", 

259 } 

260 if py_type in _PYTHON_TO_JSON: 

261 return {"type": _PYTHON_TO_JSON[py_type]} 

262 

263 # Unknown type — leave untyped 

264 return {} 

265 

266 

267# --------------------------------------------------------------------------- 

268# Backward-compat alias — internal callers may use _python_type_to_json 

269# --------------------------------------------------------------------------- 

270 

271def _python_type_to_json(py_type: Any) -> str: 

272 """Return the JSON schema type string for a Python type. 

273 

274 Legacy helper retained for internal use. Prefer 

275 :func:`_python_type_to_json_schema` for new code. 

276 

277 Args: 

278 py_type: A Python type annotation. 

279 

280 Returns: 

281 A JSON schema type string (``"string"``, ``"integer"``, etc.). 

282 """ 

283 schema_dict = _python_type_to_json_schema(py_type) 

284 schema_type = schema_dict.get("type") 

285 return schema_type if isinstance(schema_type, str) else "string"