Skip to content

refactor: migrate token metrics to hooks/plugins system#653

Open
ajbozarth wants to merge 4 commits intogenerative-computing:mainfrom
ajbozarth:feat/metrics-hooks-plugins
Open

refactor: migrate token metrics to hooks/plugins system#653
ajbozarth wants to merge 4 commits intogenerative-computing:mainfrom
ajbozarth:feat/metrics-hooks-plugins

Conversation

@ajbozarth
Copy link
Contributor

@ajbozarth ajbozarth commented Mar 14, 2026

Misc PR

Type of PR

  • Bug Fix
  • New Feature
  • Documentation
  • Other

Description

Refactors token metrics implementation to use the hooks/plugins system.

Token Usage Standardization (#607)

  • Added a usage field to ModelOutputThunk (uses the OpenAI-compatible format)
  • Enables programmatic access when using functions that return a ModelOutputThunk: result.usage['prompt_tokens'], result.usage['completion_tokens'], result.usage['total_tokens']

Hooks/Plugins Migration (#608)

  • Created TokenMetricsPlugin that hooks into the generation_post_call
  • Auto-registers when MELLEA_METRICS_ENABLED=true
  • Plugin uses FIRE_AND_FORGET mode (async, non-blocking)

Testing

  • Tests added to the respective file if code was changed
  • New code has 100% coverage if code as added
  • Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

- Add usage field to ModelOutputThunk with OpenAI-compatible format
- Update all 5 backends to populate usage field during post_processing

Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
Completes migration from direct backend metrics calls to hook-based plugin system.
TokenMetricsPlugin now handles all token metrics recording via generation_post_call hook.

- Remove record_token_usage_metrics() calls from OpenAI, Ollama, LiteLLM, WatsonX, HuggingFace backends
- Update metrics output format for OTel compliance

Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
Add assertions to verify usage, model, and provider fields are populated
in test_async_avalue tests for all 5 supported backends.

Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
@ajbozarth ajbozarth self-assigned this Mar 14, 2026
@ajbozarth ajbozarth requested review from a team, jakelorocco and nrfulton as code owners March 14, 2026 01:00
@github-actions
Copy link
Contributor

The PR description has been updated. Please fill out the template for your PR to be reviewed.

@mergify
Copy link

mergify bot commented Mar 14, 2026

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

  • title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Refactor: update token metrics telemetry to use hooks/plugins system Add token metrics as a field on ModelOutputThunk

1 participant