Merge pull request #59 from bucolucas/add-metrics-tool

Add MetricsTool class

telegram_inference_bot.py

@@ -8,6 +8,7 @@ from telegram import error as TelegramErrors, Update, __version__ as telegram_ve
 from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
 from dotenv import load_dotenv
 from tools.base_tool import BaseTool
+from tools.metrics_tool import MetricsTool
 from anthropic import Anthropic
 
 # Load environment variables
@@ -38,10 +39,10 @@ conversation_history = {}
 processing_status = {}
 
 # Load tools
-tools = []
+tools = [MetricsTool()]  # Add MetricsTool instance
 tools_dir = os.path.join(os.path.dirname(__file__), 'tools')
 for filename in os.listdir(tools_dir):
-    if filename.endswith('.py') and filename != '__init__.py' and filename != 'base_tool.py':
+    if filename.endswith('.py') and filename not in ['__init__.py', 'base_tool.py', 'metrics_tool.py']:
         module_name = f'tools.{filename[:-3]}'
         module = importlib.import_module(module_name)
         for name, obj in inspect.getmembers(module):

tools/metrics_tool.py

@@ -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])