(self)
| 109 | return self.node_call_counts[node_id] |
| 110 | |
| 111 | def get_token_usage(self) -> Dict[str, Any]: |
| 112 | data = { |
| 113 | "workflow_id": self.workflow_id, |
| 114 | "total_usage": { |
| 115 | "input_tokens": self.total_usage.input_tokens, |
| 116 | "output_tokens": self.total_usage.output_tokens, |
| 117 | "total_tokens": self.total_usage.total_tokens, |
| 118 | }, |
| 119 | "node_usages": { |
| 120 | node_id: { |
| 121 | "input_tokens": usage.input_tokens, |
| 122 | "output_tokens": usage.output_tokens, |
| 123 | "total_tokens": usage.total_tokens, |
| 124 | } |
| 125 | for node_id, usage in self.node_usages.items() |
| 126 | }, |
| 127 | "model_usages": { |
| 128 | model_name: { |
| 129 | "input_tokens": usage.input_tokens, |
| 130 | "output_tokens": usage.output_tokens, |
| 131 | "total_tokens": usage.total_tokens, |
| 132 | } |
| 133 | for model_name, usage in self.model_usages.items() |
| 134 | }, |
| 135 | "node_execution_counts": dict(self.node_call_counts), |
| 136 | "call_history": self.call_history, |
| 137 | } |
| 138 | return data |
| 139 | |
| 140 | def export_to_file(self, filepath: str): |
| 141 | """Export token usage data to a JSON file.""" |
no test coverage detected