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 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.dataclasses import ChatMessage 

19from haystack.utils.type_serialization import _is_union_type 

20 

21logger = logging.getLogger(__name__) 

22 

23 

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

25 """ 

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

27 

28 :param type_hint: The type hint to unwrap. 

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

30 """ 

31 origin = get_origin(type_hint) 

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

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

34 if len(non_none) == 1: 

35 return non_none[0] 

36 return type_hint 

37 

38 

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

40 """ 

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

42 

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

44 be skipped during schema generation. 

45 

46 :param type_hint: The type hint to check. 

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

48 """ 

49 origin = get_origin(type_hint) 

50 

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

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

53 return True 

54 

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

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

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

58 

59 return False 

60 

61 

62# Schema placeholder models for Tool and Toolset 

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

64class _ToolSchemaPlaceholder(BaseModel): 

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

66 

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

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

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

70 

71 

72class _ToolsetSchemaPlaceholder(BaseModel): 

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

74 

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

76 

77 

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

79 """ 

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

81 

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

83 :returns: 

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

85 descriptions. 

86 """ 

87 docstring = getdoc(method) 

88 if not docstring: 

89 return "", {} 

90 

91 parsed_doc = parse(docstring) 

92 param_descriptions = {} 

93 for param in parsed_doc.params: 

94 if not param.description: 

95 logger.warning( 

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

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

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

99 arg_name=param.arg_name, 

100 ) 

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

102 return parsed_doc.short_description or "", param_descriptions 

103 

104 

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

106 """ 

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

108 

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

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

111 

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

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

114 """ 

115 from haystack.core.super_component.super_component import _SuperComponent 

116 

117 # Get descriptions from the component's run method 

118 _, param_descriptions = _get_param_descriptions(component.run) 

119 

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

121 if isinstance(component, _SuperComponent): 

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

123 # Collect descriptions from all mapped components 

124 descriptions = [] 

125 for path in pipeline_paths: 

126 try: 

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

128 comp_name, socket_name = component._split_component_path(path) 

129 pipeline_component = component.pipeline.get_component(comp_name) 

130 

131 # Add parameter description if available 

132 _, run_param_descriptions = _get_param_descriptions(pipeline_component.run) 

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

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

135 except Exception as e: 

136 logger.debug( 

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

138 super_param_name=super_param_name, 

139 path=path, 

140 e=str(e), 

141 ) 

142 

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

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

145 if descriptions: 

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

147 

148 return param_descriptions 

149 

150 

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

152 """ 

153 Convert a Python dataclass to an equivalent Pydantic model. 

154 

155 :param dc_type: The dataclass type to convert. 

156 :returns: 

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

158 Field descriptions are extracted from docstrings when available. 

159 """ 

160 _, param_descriptions = _get_param_descriptions(dc_type) 

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

162 

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

164 for field in fields(dc_type): 

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

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

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

168 

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

170 field_name = field.name 

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

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

173 field_name = field_name[1:] 

174 

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

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

177 

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

179 

180 

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

182 """ 

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

184 

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

186 it encounters into corresponding Pydantic models. 

187 

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

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

190 

191 :returns: 

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

193 """ 

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

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

196 from haystack.tools.tool import Tool 

197 from haystack.tools.toolset import Toolset 

198 

199 if _type is Tool: 

200 return _ToolSchemaPlaceholder 

201 

202 if _type is Toolset: 

203 return _ToolsetSchemaPlaceholder 

204 

205 if is_dataclass(_type): 

206 return _dataclass_to_pydantic_model(_type) 

207 

208 origin = get_origin(_type) 

209 args = get_args(_type) 

210 

211 if origin is list: 

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

213 

214 if origin is collections.abc.Sequence: 

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

216 

217 if _is_union_type(origin): 

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

219 

220 if origin is dict: 

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

222 

223 return _type