telemetry_decorator: guard record_tool_usage and milestone emits (sync/async)

main
David Sarno 2025-09-10 09:38:53 -07:00
parent 9f7308b4c2
commit 9b5488dcaf
1 changed files with 26 additions and 12 deletions

View File

@ -36,18 +36,25 @@ def telemetry_tool(tool_name: str):
_decorator_log_count += 1 _decorator_log_count += 1
result = func(*args, **kwargs) result = func(*args, **kwargs)
success = True success = True
if tool_name == "manage_script" and kwargs.get("action") == "create": action_val = sub_action or kwargs.get("action")
record_milestone(MilestoneType.FIRST_SCRIPT_CREATION) try:
elif tool_name.startswith("manage_scene"): if tool_name == "manage_script" and action_val == "create":
record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION) record_milestone(MilestoneType.FIRST_SCRIPT_CREATION)
record_milestone(MilestoneType.FIRST_TOOL_USAGE) elif tool_name.startswith("manage_scene"):
record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION)
record_milestone(MilestoneType.FIRST_TOOL_USAGE)
except Exception:
_log.debug("milestone emit failed", exc_info=True)
return result return result
except Exception as e: except Exception as e:
error = str(e) error = str(e)
raise raise
finally: finally:
duration_ms = (time.time() - start_time) * 1000 duration_ms = (time.time() - start_time) * 1000
record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action) try:
record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action)
except Exception:
_log.debug("record_tool_usage failed", exc_info=True)
@functools.wraps(func) @functools.wraps(func)
async def _async_wrapper(*args, **kwargs) -> Any: async def _async_wrapper(*args, **kwargs) -> Any:
@ -70,18 +77,25 @@ def telemetry_tool(tool_name: str):
_decorator_log_count += 1 _decorator_log_count += 1
result = await func(*args, **kwargs) result = await func(*args, **kwargs)
success = True success = True
if tool_name == "manage_script" and kwargs.get("action") == "create": action_val = sub_action or kwargs.get("action")
record_milestone(MilestoneType.FIRST_SCRIPT_CREATION) try:
elif tool_name.startswith("manage_scene"): if tool_name == "manage_script" and action_val == "create":
record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION) record_milestone(MilestoneType.FIRST_SCRIPT_CREATION)
record_milestone(MilestoneType.FIRST_TOOL_USAGE) elif tool_name.startswith("manage_scene"):
record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION)
record_milestone(MilestoneType.FIRST_TOOL_USAGE)
except Exception:
_log.debug("milestone emit failed", exc_info=True)
return result return result
except Exception as e: except Exception as e:
error = str(e) error = str(e)
raise raise
finally: finally:
duration_ms = (time.time() - start_time) * 1000 duration_ms = (time.time() - start_time) * 1000
record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action) try:
record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action)
except Exception:
_log.debug("record_tool_usage failed", exc_info=True)
return _async_wrapper if inspect.iscoroutinefunction(func) else _sync_wrapper return _async_wrapper if inspect.iscoroutinefunction(func) else _sync_wrapper
return decorator return decorator