Add metrics_tool.py

2024-08-18 18:54:40 -05:00
parent 01d7c508d9
commit 1bd7f6de89
1 changed files with 75 additions and 0 deletions
@@ -0,0 +1,75 @@
+# tools/metrics_tool.py
+
+from .base_tool import BaseTool
+from .metrics import metrics
+
+class MetricsTool(BaseTool):
+    def __init__(self):
+        self.metrics = metrics
+
+    def clear(self):
+        pass
+
+    def get_functions(self):
+        return [
+            {
+                "name": "get_function_metrics",
+                "description": "Get metrics for all measured functions.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {},
+                    "required": []
+                }
+            },
+            {
+                "name": "get_specific_function_metrics",
+                "description": "Get metrics for a specific function.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "function_name": {
+                            "type": "string",
+                            "description": "Name of the function to get metrics for"
+                        }
+                    },
+                    "required": ["function_name"]
+                }
+            },
+            {
+                "name": "get_top_n_functions",
+                "description": "Get the top N functions by total execution time.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "n": {
+                            "type": "integer",
+                            "description": "Number of top functions to retrieve"
+                        }
+                    },
+                    "required": ["n"]
+                }
+            }
+        ]
+
+    @metrics.measure
+    def execute(self, function_name, **kwargs):
+        if function_name == "get_function_metrics":
+            return self._get_function_metrics()
+        elif function_name == "get_specific_function_metrics":
+            return self._get_specific_function_metrics(kwargs.get("function_name"))
+        elif function_name == "get_top_n_functions":
+            return self._get_top_n_functions(kwargs.get("n"))
+        else:
+            return f"Unknown function: {function_name}"
+
+    def _get_function_metrics(self):
+        return self.metrics.get_metrics()
+
+    def _get_specific_function_metrics(self, function_name):
+        all_metrics = self.metrics.get_metrics()
+        return all_metrics.get(function_name, f"No metrics found for function: {function_name}")
+
+    def _get_top_n_functions(self, n):
+        all_metrics = self.metrics.get_metrics()
+        sorted_metrics = sorted(all_metrics.items(), key=lambda x: x[1]['total_time'], reverse=True)
+        return dict(sorted_metrics[:n])