Skip to content

Implement counters to track token usage across all LLM backends with model and backend labels #462

Description

@psschwei

Description:
Implement counters to track token usage across all LLM backends with model and backend labels.

Detailed Requirements:

  1. Create counters in mellea/telemetry/metrics.py:
    • mellea.llm.tokens.input - Input/prompt tokens
    • mellea.llm.tokens.output - Output/completion tokens
    • mellea.llm.tokens.total - Total tokens (optional, can be computed)
  2. Add attributes/labels:
    • gen_ai.system - Backend system (openai, ollama, etc.)
    • gen_ai.request.model - Model identifier
    • mellea.backend - Backend class name
  3. Create helper function record_token_usage(input_tokens, output_tokens, model, backend)
  4. Integrate with all backends that report token usage:
    • OpenAI: usage.prompt_tokens, usage.completion_tokens
    • Ollama: response.prompt_eval_count, response.eval_count
    • WatsonX: input_token_count, generated_token_count
    • LiteLLM: Via litellm's usage tracking
    • HuggingFace: Limited support (local models)

Files to Modify:

  • mellea/telemetry/metrics.py - Add token counters
  • mellea/telemetry/backend_instrumentation.py - Add metric recording alongside trace recording
  • mellea/backends/openai.py - Record token metrics
  • mellea/backends/ollama.py - Record token metrics
  • mellea/backends/watsonx.py - Record token metrics
  • mellea/backends/litellm.py - Record token metrics

Metric Definitions:

Metric Type Unit Attributes
mellea.llm.tokens.input Counter {token} model, backend, system
mellea.llm.tokens.output Counter {token} model, backend, system

Acceptance Criteria:

  • Token counters increment on each backend call
  • Metrics have correct model/backend labels
  • All backends with usage reporting integrated
  • Zero overhead when metrics disabled

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions