|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import asyncio |
| 4 | +import copy |
| 5 | +import logging |
4 | 6 | import re |
5 | 7 | from typing import TYPE_CHECKING, Literal |
6 | 8 |
|
|
25 | 27 | from lfx.schema.content_block import ContentBlock |
26 | 28 | from lfx.schema.dotdict import dotdict |
27 | 29 |
|
| 30 | +logger = logging.getLogger(__name__) |
| 31 | + |
28 | 32 | TOOL_TYPES_SET = {"Tool", "BaseTool", "StructuredTool"} |
29 | 33 |
|
30 | 34 |
|
@@ -168,6 +172,54 @@ def _should_skip_output(self, output: Output) -> bool: |
168 | 172 | output.name == TOOL_OUTPUT_NAME or any(tool_type in output.types for tool_type in TOOL_TYPES_SET) |
169 | 173 | ) |
170 | 174 |
|
| 175 | + def _attach_runtime_metadata( |
| 176 | + self, |
| 177 | + tool: BaseTool | StructuredTool, |
| 178 | + output: Output, |
| 179 | + *, |
| 180 | + is_async: bool, |
| 181 | + args_schema: type | None, |
| 182 | + ) -> None: |
| 183 | + """Annotate tool metadata with component/runtime information.""" |
| 184 | + metadata = tool.metadata or {} |
| 185 | + metadata["_component_method"] = output.method |
| 186 | + metadata["_component_is_async"] = is_async |
| 187 | + metadata["component_class"] = self.component.__class__.__name__ |
| 188 | + metadata["component_module"] = self.component.__class__.__module__ |
| 189 | + metadata["component_id"] = self.component.get_id() |
| 190 | + metadata["stream_topic"] = self.component._build_stream_topic() # noqa: SLF001 |
| 191 | + |
| 192 | + try: |
| 193 | + component_state = self.component._serialize_for_executor() # noqa: SLF001 |
| 194 | + metadata["_component_state"] = component_state |
| 195 | + except Exception as exc: # noqa: BLE001 |
| 196 | + logger.warning( |
| 197 | + "Failed to serialize component state for tool '%s': %s", |
| 198 | + tool.name, |
| 199 | + exc, |
| 200 | + ) |
| 201 | + |
| 202 | + metadata["_tool_input_keys"] = self._extract_input_keys(args_schema, output) |
| 203 | + |
| 204 | + executor_meta = self.component.get_executor_node_metadata() |
| 205 | + if executor_meta: |
| 206 | + metadata["executor_node"] = executor_meta |
| 207 | + tool.metadata = metadata |
| 208 | + |
| 209 | + @staticmethod |
| 210 | + def _extract_input_keys(args_schema: type | None, output: Output) -> list[str]: |
| 211 | + """Return list of input field names for this tool.""" |
| 212 | + keys: list[str] = [] |
| 213 | + if args_schema is not None: |
| 214 | + for attr_name in ("model_fields", "__fields__"): |
| 215 | + schema_fields = getattr(args_schema, attr_name, None) |
| 216 | + if schema_fields: |
| 217 | + keys = list(schema_fields.keys()) |
| 218 | + break |
| 219 | + if not keys and getattr(output, "required_inputs", None): |
| 220 | + keys = list(output.required_inputs) |
| 221 | + return keys |
| 222 | + |
171 | 223 | def get_tools( |
172 | 224 | self, |
173 | 225 | tool_name: str | None = None, |
@@ -228,37 +280,37 @@ def get_tools( |
228 | 280 | formatted_name = _format_tool_name(name) |
229 | 281 | event_manager = self.component.get_event_manager() |
230 | 282 | if asyncio.iscoroutinefunction(output_method): |
231 | | - tools.append( |
232 | | - StructuredTool( |
233 | | - name=formatted_name, |
234 | | - description=build_description(self.component), |
235 | | - coroutine=_build_output_async_function(self.component, output_method, event_manager), |
236 | | - args_schema=args_schema, |
237 | | - handle_tool_error=True, |
238 | | - callbacks=callbacks, |
239 | | - tags=[formatted_name], |
240 | | - metadata={ |
241 | | - "display_name": formatted_name, |
242 | | - "display_description": build_description(self.component), |
243 | | - }, |
244 | | - ) |
| 283 | + structured_tool = StructuredTool( |
| 284 | + name=formatted_name, |
| 285 | + description=build_description(self.component), |
| 286 | + coroutine=_build_output_async_function(self.component, output_method, event_manager), |
| 287 | + args_schema=args_schema, |
| 288 | + handle_tool_error=True, |
| 289 | + callbacks=callbacks, |
| 290 | + tags=[formatted_name], |
| 291 | + metadata={ |
| 292 | + "display_name": formatted_name, |
| 293 | + "display_description": build_description(self.component), |
| 294 | + }, |
245 | 295 | ) |
| 296 | + tools.append(structured_tool) |
| 297 | + self._attach_runtime_metadata(structured_tool, output, is_async=True, args_schema=args_schema) |
246 | 298 | else: |
247 | | - tools.append( |
248 | | - StructuredTool( |
249 | | - name=formatted_name, |
250 | | - description=build_description(self.component), |
251 | | - func=_build_output_function(self.component, output_method, event_manager), |
252 | | - args_schema=args_schema, |
253 | | - handle_tool_error=True, |
254 | | - callbacks=callbacks, |
255 | | - tags=[formatted_name], |
256 | | - metadata={ |
257 | | - "display_name": formatted_name, |
258 | | - "display_description": build_description(self.component), |
259 | | - }, |
260 | | - ) |
| 299 | + structured_tool = StructuredTool( |
| 300 | + name=formatted_name, |
| 301 | + description=build_description(self.component), |
| 302 | + func=_build_output_function(self.component, output_method, event_manager), |
| 303 | + args_schema=args_schema, |
| 304 | + handle_tool_error=True, |
| 305 | + callbacks=callbacks, |
| 306 | + tags=[formatted_name], |
| 307 | + metadata={ |
| 308 | + "display_name": formatted_name, |
| 309 | + "display_description": build_description(self.component), |
| 310 | + }, |
261 | 311 | ) |
| 312 | + tools.append(structured_tool) |
| 313 | + self._attach_runtime_metadata(structured_tool, output, is_async=False, args_schema=args_schema) |
262 | 314 | if len(tools) == 1 and (tool_name or tool_description): |
263 | 315 | tool = tools[0] |
264 | 316 | tool.name = _format_tool_name(str(tool_name)) or tool.name |
|
0 commit comments