diff --git a/telegram_inference_bot.py b/telegram_inference_bot.py index fa0a39f..bb8c5c2 100644 --- a/telegram_inference_bot.py +++ b/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): diff --git a/tools/metrics_tool.py b/tools/metrics_tool.py new file mode 100644 index 0000000..4db57be --- /dev/null +++ b/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]) \ No newline at end of file