Coverage for haystack/tools/parameters_schema_utils.py: 97%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import collections 

6import types 

7from collections.abc import Callable, Sequence 

8from collections.abc import Callable as ABCCallable 

9from dataclasses import MISSING, fields, is_dataclass 

10from inspect import getdoc 

11from types import NoneType 

12from typing import Annotated, Any, Union, get_args, get_origin 

13 

14from docstring_parser import parse 

15from pydantic import BaseModel, Field, create_model 

16 

17from haystack import logging 

18from haystack.components.agents.state.state import State 

19from haystack.dataclasses import ChatMessage 

20from haystack.utils.type_serialization import _is_union_type 

21 

22logger = logging.getLogger(__name__) 

23 

24 

25def _unwrap_optional(type_hint: Any) -> Any: 

26 """ 

27 Unwrap Optional types (i.e. ``X | None`` or ``Optional[X]``) to get the inner type. 

28 

29 :param type_hint: The type hint to unwrap. 

30 :returns: The inner type if ``type_hint`` is ``Optional[X]``, otherwise ``type_hint`` unchanged. 

31 """ 

32 origin = get_origin(type_hint) 

33 if origin is Union or origin is types.UnionType: 

34 non_none = [a for a in get_args(type_hint) if a is not NoneType] 

35 if len(non_none) == 1: 

36 return non_none[0] 

37 return type_hint 

38 

39 

40def _is_state_annotation(annotation: Any) -> bool: 

41 """Recognize State through Annotated and optional wrappers.""" 

42 if get_origin(annotation) is Annotated: 

43 return _is_state_annotation(get_args(annotation)[0]) 

44 

45 unwrapped = _unwrap_optional(annotation) 

46 if unwrapped is not annotation: 

47 return _is_state_annotation(unwrapped) 

48 

49 return annotation is State 

50 

51 

52def _contains_callable_type(type_hint: Any) -> bool: 

53 """ 

54 Check if a type hint contains a Callable type, including within Union types. 

55 

56 The purpose of this function is to help identify Callable types so they can 

57 be skipped during schema generation. 

58 

59 :param type_hint: The type hint to check. 

60 :returns: True if the type contains a Callable, False otherwise. 

61 """ 

62 origin = get_origin(type_hint) 

63 

64 # Check if it's a Callable type (direct or parameterized) 

65 if type_hint in (Callable, ABCCallable) or origin in (Callable, ABCCallable): 

66 return True 

67 

68 # Recursively check Union types (both typing.Union and types.UnionType for `X | Y` syntax) 

69 if origin in (Union, types.UnionType): 

70 return any(_contains_callable_type(arg) for arg in get_args(type_hint)) 

71 

72 return False 

73 

74 

75# Schema placeholder models for Tool and Toolset 

76# These are used during JSON schema generation to represent non-serializable types 

77class _ToolSchemaPlaceholder(BaseModel): 

78 """Placeholder model representing a Tool for JSON schema generation.""" 

79 

80 name: str = Field(description="Name of the tool") 

81 description: str = Field(description="Description of the tool") 

82 parameters: dict[str, Any] = Field(description="JSON schema of the tool parameters") 

83 

84 

85class _ToolsetSchemaPlaceholder(BaseModel): 

86 """Placeholder model representing a Toolset for JSON schema generation.""" 

87 

88 tools: list[_ToolSchemaPlaceholder] = Field(description="List of tools in the toolset") 

89 

90 

91def _get_param_descriptions(method: Callable) -> tuple[str, dict[str, str]]: 

92 """ 

93 Extracts parameter descriptions from the method's docstring using docstring_parser. 

94 

95 :param method: The method to extract parameter descriptions from. 

96 :returns: 

97 A tuple including the short description of the method and a dictionary mapping parameter names to their 

98 descriptions. 

99 """ 

100 docstring = getdoc(method) 

101 if not docstring: 

102 return "", {} 

103 

104 parsed_doc = parse(docstring) 

105 param_descriptions = {} 

106 for param in parsed_doc.params: 

107 if not param.description: 

108 logger.warning( 

109 "Missing description for parameter '{arg_name}'. Please add a description in the component's " 

110 "run() method docstring using the format ':param {arg_name}: <description>'. " 

111 "This description helps the LLM understand how to use this parameter.", 

112 arg_name=param.arg_name, 

113 ) 

114 param_descriptions[param.arg_name] = param.description.strip() if param.description else "" 

115 return parsed_doc.short_description or "", param_descriptions 

116 

117 

118def _get_component_param_descriptions(component: Any) -> dict[str, str]: 

119 """ 

120 Get parameter descriptions from a component, handling both regular Components and SuperComponents. 

121 

122 For regular components, this extracts descriptions from the run method's docstring. 

123 For SuperComponents, this extracts descriptions from the underlying pipeline components. 

124 

125 :param component: The component to extract parameter descriptions from 

126 :returns: A dictionary mapping parameter names to their descriptions 

127 """ 

128 from haystack.core.super_component.super_component import _SuperComponent 

129 

130 # Get descriptions from the component's run method 

131 _, param_descriptions = _get_param_descriptions(component.run) 

132 

133 # If it's a SuperComponent, enhance the parameter descriptions from the original components 

134 if isinstance(component, _SuperComponent): 

135 for super_param_name, pipeline_paths in component.input_mapping.items(): 

136 # Collect descriptions from all mapped components 

137 descriptions = [] 

138 for path in pipeline_paths: 

139 try: 

140 # Get the component and socket this input is mapped from 

141 comp_name, socket_name = component._split_component_path(path) 

142 pipeline_component = component.pipeline.get_component(comp_name) 

143 

144 # Add parameter description if available 

145 _, run_param_descriptions = _get_param_descriptions(pipeline_component.run) 

146 if input_param_mapping := run_param_descriptions.get(socket_name): 

147 descriptions.append(f"Provided to the '{comp_name}' component as: '{input_param_mapping}'") 

148 except Exception as e: 

149 logger.debug( 

150 "Error extracting description for {super_param_name} from {path}: {e}", 

151 super_param_name=super_param_name, 

152 path=path, 

153 e=str(e), 

154 ) 

155 

156 # A single SuperComponent input can map to multiple pipeline components, e.g. 

157 # input_mapping={"combined_input": ["comp_a.query", "comp_b.text"]} 

158 if descriptions: 

159 param_descriptions[super_param_name] = ", and ".join(descriptions) + "." 

160 

161 return param_descriptions 

162 

163 

164def _dataclass_to_pydantic_model(dc_type: Any) -> type[BaseModel]: 

165 """ 

166 Convert a Python dataclass to an equivalent Pydantic model. 

167 

168 :param dc_type: The dataclass type to convert. 

169 :returns: 

170 A dynamically generated Pydantic model class with fields and types derived from the dataclass definition. 

171 Field descriptions are extracted from docstrings when available. 

172 """ 

173 _, param_descriptions = _get_param_descriptions(dc_type) 

174 cls = dc_type if isinstance(dc_type, type) else dc_type.__class__ 

175 

176 field_defs: dict[str, Any] = {} 

177 for field in fields(dc_type): 

178 f_type = field.type if isinstance(field.type, str) else _resolve_type(field.type) 

179 default = field.default if field.default is not MISSING else ... 

180 default = field.default_factory() if callable(field.default_factory) else default 

181 

182 # Special handling for ChatMessage since pydantic doesn't allow for field names with leading underscores 

183 field_name = field.name 

184 if dc_type is ChatMessage and field_name.startswith("_"): 

185 # We remove the underscore since ChatMessage.from_dict does allow for field names without the underscore 

186 field_name = field_name[1:] 

187 

188 description = param_descriptions.get(field_name, f"Field '{field_name}' of '{cls.__name__}'.") 

189 field_defs[field_name] = (f_type, Field(default, description=description)) 

190 

191 return create_model(cls.__name__, **field_defs) 

192 

193 

194def _resolve_type(_type: Any) -> Any: # noqa: PLR0911 

195 """ 

196 Recursively resolve and convert complex type annotations, transforming dataclasses into Pydantic-compatible types. 

197 

198 This function walks through nested type annotations (e.g., List, Dict, Union) and converts any dataclass types 

199 it encounters into corresponding Pydantic models. 

200 

201 :param _type: The type annotation to resolve. If the type is a dataclass, it will be converted to a Pydantic model. 

202 For generic types (like list[SomeDataclass]), the inner types are also resolved recursively. 

203 

204 :returns: 

205 A fully resolved type, with all dataclass types converted to Pydantic models 

206 """ 

207 # Special handling for Tool and Toolset types - replace with schema placeholders 

208 # These types contain Callables which cannot be serialized to JSON Schema 

209 from haystack.tools.tool import Tool 

210 from haystack.tools.toolset import Toolset 

211 

212 if _type is Tool: 

213 return _ToolSchemaPlaceholder 

214 

215 if _type is Toolset: 

216 return _ToolsetSchemaPlaceholder 

217 

218 if is_dataclass(_type): 

219 return _dataclass_to_pydantic_model(_type) 

220 

221 origin = get_origin(_type) 

222 args = get_args(_type) 

223 

224 if origin is list: 

225 return list[_resolve_type(args[0]) if args else Any] # type: ignore[misc] 

226 

227 if origin is collections.abc.Sequence: 

228 return Sequence[_resolve_type(args[0]) if args else Any] # type: ignore[misc] 

229 

230 if _is_union_type(origin): 

231 return Union[tuple(_resolve_type(a) for a in args)] 

232 

233 if origin is dict: 

234 return dict[args[0] if args else Any, _resolve_type(args[1]) if args else Any] # type: ignore[misc] 

235 

236 return _type