Record a model invocation.
(self, node_id: str, model_name: str,
input_data: Any = None, output: Any = None,
details: Dict[str, Any] = None,
stage: CallStage = CallStage.AFTER)
| 116 | ) |
| 117 | |
| 118 | def record_model_call(self, node_id: str, model_name: str, |
| 119 | input_data: Any = None, output: Any = None, |
| 120 | details: Dict[str, Any] = None, |
| 121 | stage: CallStage = CallStage.AFTER) -> None: |
| 122 | """Record a model invocation.""" |
| 123 | input_size = len(str(input_data)) if input_data is not None else 0 |
| 124 | output_size = len(str(output)) if output is not None else 0 |
| 125 | duration = self.logger.get_timer(f"model_{node_id}") |
| 126 | |
| 127 | call_details = { |
| 128 | "input_size": input_size, |
| 129 | "output_size": output_size, |
| 130 | **(details or {}) |
| 131 | } |
| 132 | |
| 133 | self.logger.record_model_call( |
| 134 | node_id, model_name, input_data, output, duration, call_details, stage |
| 135 | ) |
| 136 | |
| 137 | def record_tool_call(self, node_id: str, tool_name: str, |
| 138 | success: bool | None = True, tool_result: Any = None, |