diff --git a/docs/docs/observability/metrics.md b/docs/docs/observability/metrics.md index 1f3363ca9a..705a793c52 100644 --- a/docs/docs/observability/metrics.md +++ b/docs/docs/observability/metrics.md @@ -56,9 +56,6 @@ All token metrics include these attributes following Gen-AI semantic conventions | LiteLLM | Yes | Yes | `usage.prompt_tokens` and `usage.completion_tokens` | | HuggingFace | Yes | Yes | Calculated from input_ids and output sequences | -> **Note:** Token usage metrics are only tracked for `generate_from_context` -> requests. `generate_from_raw` calls do not record token metrics. - ### Token recording timing Token metrics are recorded **after the full response is received**, not diff --git a/mellea/telemetry/metrics_plugins.py b/mellea/telemetry/metrics_plugins.py index 0608082290..f91557d2d0 100644 --- a/mellea/telemetry/metrics_plugins.py +++ b/mellea/telemetry/metrics_plugins.py @@ -3,7 +3,7 @@ This module contains plugins that hook into the generation pipeline to automatically record metrics when enabled. Currently includes: -- TokenMetricsPlugin: Records token usage statistics from ModelOutputThunk.generation.usage +- TokenMetricsPlugin: Records token usage statistics from generation usage data - LatencyMetricsPlugin: Records request duration and TTFB latency histograms - ErrorMetricsPlugin: Records LLM error counts categorized by semantic error type - CostMetricsPlugin: Records estimated request cost in USD from pricing registry @@ -23,6 +23,8 @@ if TYPE_CHECKING: from mellea.core.base import GenerationMetadata from mellea.plugins.hooks.generation import ( + GenerationBatchErrorPayload, + GenerationBatchPostCallPayload, GenerationErrorPayload, GenerationPostCallPayload, ) @@ -37,9 +39,9 @@ class TokenMetricsPlugin(Plugin, name="token_metrics", priority=50): """Records token usage metrics from generation outputs. - This plugin hooks into the generation_post_call event to automatically - record token usage metrics when the usage field is populated on - ModelOutputThunk instances. + This plugin hooks into the generation_post_call and + generation_batch_post_call events to automatically record token usage + metrics when usage data is present. The plugin reads the standardized usage field (OpenAI-compatible format) and records metrics following OpenTelemetry Gen-AI semantic conventions. @@ -69,13 +71,36 @@ async def record_token_metrics( provider=gen.provider or "unknown", ) + @hook("generation_batch_post_call", mode=PluginMode.FIRE_AND_FORGET) + async def record_batch_token_metrics( + self, payload: GenerationBatchPostCallPayload, context: dict[str, Any] + ) -> None: + """Record token metrics after a batch generation completes. + + Args: + payload: Contains the batch-level usage dict, model, and provider. + context: Plugin context (unused). + """ + from mellea.telemetry.metrics import record_token_usage_metrics + + if payload.usage is None: + return + + record_token_usage_metrics( + input_tokens=payload.usage.get("prompt_tokens"), + output_tokens=payload.usage.get("completion_tokens"), + model=payload.model or "unknown", + provider=payload.provider or "unknown", + ) + class LatencyMetricsPlugin(Plugin, name="latency_metrics", priority=51): """Records request duration and TTFB latency metrics from generation outputs. - This plugin hooks into the generation_post_call event to automatically - record latency metrics. It records total request duration for every request - and time-to-first-token (TTFB) for streaming requests. + This plugin hooks into the generation_post_call and + generation_batch_post_call events to automatically record latency + metrics. It records total request duration for every request and + time-to-first-token (TTFB) for streaming requests. """ @hook("generation_post_call", mode=PluginMode.FIRE_AND_FORGET) @@ -106,12 +131,35 @@ async def record_latency_metrics( if gen.streaming and gen.ttfb_ms is not None: record_ttfb(ttfb_s=gen.ttfb_ms / 1000.0, model=model, provider=provider) + @hook("generation_batch_post_call", mode=PluginMode.FIRE_AND_FORGET) + async def record_batch_latency_metrics( + self, payload: GenerationBatchPostCallPayload, context: dict[str, Any] + ) -> None: + """Record request duration after a batch generation completes. + + Batch generations (`generate_from_raw`) are non-streaming, so only the + total request duration is recorded; TTFB does not apply. + + Args: + payload: Contains latency_ms, model, and provider for the batch. + context: Plugin context (unused). + """ + from mellea.telemetry.metrics import record_request_duration + + record_request_duration( + duration_s=payload.latency_ms / 1000.0, + model=payload.model or "unknown", + provider=payload.provider or "unknown", + streaming=False, + ) + class ErrorMetricsPlugin(Plugin, name="error_metrics", priority=52): """Records LLM error counts from generation errors. - This plugin hooks into the generation_error event to classify exceptions - by semantic error type and increment the `mellea.llm.errors` counter. + This plugin hooks into the generation_error and generation_batch_error + events to classify exceptions by semantic error type and increment the + `mellea.llm.errors` counter. """ @hook("generation_error", mode=PluginMode.FIRE_AND_FORGET) @@ -140,13 +188,34 @@ async def record_error_metrics( exception_class=type(payload.exception).__name__, ) + @hook("generation_batch_error", mode=PluginMode.FIRE_AND_FORGET) + async def record_batch_error_metrics( + self, payload: GenerationBatchErrorPayload, context: dict[str, Any] + ) -> None: + """Record error metrics when a batch generation fails. + + Args: + payload: Contains the exception, model, and provider for the batch. + context: Plugin context (unused). + """ + from mellea.telemetry.metrics import classify_error, record_error + + error_type = classify_error(payload.exception) + record_error( + error_type=error_type, + model=payload.model or "unknown", + provider=payload.provider or "unknown", + exception_class=type(payload.exception).__name__, + ) + class CostMetricsPlugin(Plugin, name="cost_metrics", priority=53): """Records estimated request cost metrics from generation outputs. - This plugin hooks into the generation_post_call event to automatically - record cost metrics when token usage and model pricing data are available. - Cost is skipped and a warning is logged for models not in the pricing registry. + This plugin hooks into the generation_post_call and + generation_batch_post_call events to automatically record cost metrics + when token usage and model pricing data are available. Cost is skipped + and a warning is logged for models not in the pricing registry. """ @hook("generation_post_call", mode=PluginMode.FIRE_AND_FORGET) @@ -185,6 +254,41 @@ async def record_cost_metrics( if cost is not None: record_cost(cost=cost, model=model, provider=provider) + @hook("generation_batch_post_call", mode=PluginMode.FIRE_AND_FORGET) + async def record_batch_cost_metrics( + self, payload: GenerationBatchPostCallPayload, context: dict[str, Any] + ) -> None: + """Record cost metrics after a batch generation completes. + + Args: + payload: Contains the batch-level usage dict, model, and provider. + context: Plugin context (unused). + """ + from mellea.telemetry.metrics import record_cost + from mellea.telemetry.pricing import compute_cost + + if payload.usage is None: + return + + model = payload.model or "unknown" + provider = payload.provider or "unknown" + details = payload.usage.get("prompt_tokens_details") + cached_tokens = ( + details.get("cached_tokens") if isinstance(details, dict) else 0 + ) or 0 + cache_creation = payload.usage.get("cache_creation_input_tokens") or 0 + prompt_tokens = payload.usage.get("prompt_tokens") or 0 + cost = compute_cost( + model=model, + provider=payload.provider, + prompt_tokens=prompt_tokens, + completion_tokens=payload.usage.get("completion_tokens"), + cached_tokens=cached_tokens, + cache_creation_tokens=cache_creation, + ) + if cost is not None: + record_cost(cost=cost, model=model, provider=provider) + class SamplingMetricsPlugin(Plugin, name="sampling_metrics", priority=54): """Records sampling loop attempt and outcome metrics. diff --git a/test/telemetry/test_metrics_backend.py b/test/telemetry/test_metrics_backend.py index 2119501cf5..b18f482450 100644 --- a/test/telemetry/test_metrics_backend.py +++ b/test/telemetry/test_metrics_backend.py @@ -463,7 +463,7 @@ async def test_error_metrics_on_backend_failure(enable_metrics, metric_reader): backend = OpenAIBackend( model_id="nonexistent-model-xyz", - base_url="http://localhost:11434/v1", + base_url=f"http://{os.environ.get('OLLAMA_HOST', 'localhost:11434')}/v1", api_key="dummy", ) ctx = SimpleContext() @@ -535,3 +535,86 @@ async def test_ollama_sampling_metrics_integration(enable_metrics, metric_reader ) assert result.success assert successes == 1, f"Expected 1 success, got {successes}" + + +@pytest.mark.asyncio +@pytest.mark.ollama +async def test_ollama_generate_from_raw_metrics_integration( + enable_metrics, metric_reader +): + """Token and latency metrics are recorded for `generate_from_raw` calls.""" + from mellea.backends.ollama import OllamaModelBackend + from mellea.core import CBlock + from mellea.stdlib.context import SimpleContext + from mellea.telemetry import metrics as metrics_module + + provider = _setup_metrics_provider(metrics_module, metric_reader) + + backend = OllamaModelBackend(model_id=IBM_GRANITE_4_1_3B.ollama_name) # type: ignore + ctx = SimpleContext() + actions = [CBlock("Say 'hi' and nothing else."), CBlock("Say 'hello'.")] + + results = await backend.generate_from_raw(actions, ctx=ctx) + assert len(results) == len(actions) + + await drain_background_tasks() + provider.force_flush() + metrics_data = metric_reader.get_metrics_data() + + input_tokens = get_metric_value( + metrics_data, "mellea.llm.tokens.input", {"gen_ai.provider.name": "ollama"} + ) + output_tokens = get_metric_value( + metrics_data, "mellea.llm.tokens.output", {"gen_ai.provider.name": "ollama"} + ) + assert input_tokens is not None and input_tokens > 0 + assert output_tokens is not None and output_tokens > 0 + + duration_dp = _find_histogram_data_point( + metrics_data, + "mellea.llm.request.duration", + {"gen_ai.provider.name": "ollama", "streaming": False}, + ) + assert duration_dp is not None, "Request duration should be recorded for batch" + assert duration_dp.sum > 0 + + +@pytest.mark.asyncio +@pytest.mark.openai +@pytest.mark.ollama +async def test_generate_from_raw_error_metrics_integration( + enable_metrics, metric_reader +): + """Test that error metrics are recorded when `generate_from_raw` fails.""" + from mellea.backends.openai import OpenAIBackend + from mellea.core import CBlock + from mellea.stdlib.context import SimpleContext + from mellea.telemetry import metrics as metrics_module + + provider = _setup_metrics_provider(metrics_module, metric_reader) + + backend = OpenAIBackend( + model_id="nonexistent-model-xyz", + base_url=f"http://{os.environ.get('OLLAMA_HOST', 'localhost:11434')}/v1", + api_key="dummy", + ) + ctx = SimpleContext() + actions = [CBlock("Say 'hi'.")] + + with pytest.raises(Exception): + await backend.generate_from_raw(actions, ctx=ctx) + + await drain_background_tasks() + provider.force_flush() + metrics_data = metric_reader.get_metrics_data() + + error_count = get_metric_value( + metrics_data, + "mellea.llm.errors", + { + "gen_ai.provider.name": "openai", + "gen_ai.request.model": "nonexistent-model-xyz", + }, + ) + assert error_count is not None, "Error counter should have been recorded" + assert error_count == 1, f"Expected 1 error, got {error_count}" diff --git a/test/telemetry/test_metrics_plugins.py b/test/telemetry/test_metrics_plugins.py index d55597446f..b2ab128cc9 100644 --- a/test/telemetry/test_metrics_plugins.py +++ b/test/telemetry/test_metrics_plugins.py @@ -8,6 +8,8 @@ from mellea.core.base import GenerationMetadata, ModelOutputThunk from mellea.plugins.hooks.generation import ( + GenerationBatchErrorPayload, + GenerationBatchPostCallPayload, GenerationErrorPayload, GenerationPostCallPayload, ) @@ -99,6 +101,69 @@ async def test_record_token_metrics_missing_model_provider(token_plugin): ) +# TokenMetricsPlugin batch tests + + +def _make_batch_post_payload( + usage=None, model="batch-model", provider="batch-provider", latency_ms=300.0 +): + """Create a GenerationBatchPostCallPayload.""" + return GenerationBatchPostCallPayload( + generation_id="batch-1", + model_outputs=[], + usage=usage, + model=model, + provider=provider, + latency_ms=latency_ms, + ) + + +@pytest.mark.asyncio +async def test_record_batch_token_metrics_with_usage(token_plugin): + """Batch token metrics are recorded from payload.usage.""" + payload = _make_batch_post_payload( + usage={"prompt_tokens": 50, "completion_tokens": 25, "total_tokens": 75} + ) + + with patch("mellea.telemetry.metrics.record_token_usage_metrics") as mock_record: + await token_plugin.record_batch_token_metrics(payload, {}) + + mock_record.assert_called_once_with( + input_tokens=50, + output_tokens=25, + model="batch-model", + provider="batch-provider", + ) + + +@pytest.mark.asyncio +async def test_record_batch_token_metrics_no_usage(token_plugin): + """Batch token metrics are not recorded when usage is None.""" + payload = _make_batch_post_payload(usage=None) + + with patch("mellea.telemetry.metrics.record_token_usage_metrics") as mock_record: + await token_plugin.record_batch_token_metrics(payload, {}) + + mock_record.assert_not_called() + + +@pytest.mark.asyncio +async def test_record_batch_token_metrics_missing_model_provider(token_plugin): + """Batch fallback to 'unknown' when model/provider are None.""" + payload = _make_batch_post_payload( + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + model=None, + provider=None, + ) + + with patch("mellea.telemetry.metrics.record_token_usage_metrics") as mock_record: + await token_plugin.record_batch_token_metrics(payload, {}) + + mock_record.assert_called_once_with( + input_tokens=10, output_tokens=5, model="unknown", provider="unknown" + ) + + # LatencyMetricsPlugin tests @@ -194,6 +259,42 @@ async def test_latency_missing_model_provider(latency_plugin): ) +# LatencyMetricsPlugin batch tests + + +@pytest.mark.asyncio +async def test_batch_latency_records_duration_non_streaming(latency_plugin): + """Batch generations record duration with streaming=False.""" + payload = _make_batch_post_payload(latency_ms=1500.0) + + with ( + patch("mellea.telemetry.metrics.record_request_duration") as mock_dur, + patch("mellea.telemetry.metrics.record_ttfb") as mock_ttfb, + ): + await latency_plugin.record_batch_latency_metrics(payload, {}) + + mock_dur.assert_called_once_with( + duration_s=1.5, + model="batch-model", + provider="batch-provider", + streaming=False, + ) + mock_ttfb.assert_not_called() + + +@pytest.mark.asyncio +async def test_batch_latency_missing_model_provider(latency_plugin): + """Batch model/provider default to 'unknown' when None.""" + payload = _make_batch_post_payload(latency_ms=400.0, model=None, provider=None) + + with patch("mellea.telemetry.metrics.record_request_duration") as mock_dur: + await latency_plugin.record_batch_latency_metrics(payload, {}) + + mock_dur.assert_called_once_with( + duration_s=0.4, model="unknown", provider="unknown", streaming=False + ) + + # ErrorMetricsPlugin tests @@ -275,6 +376,70 @@ async def test_error_plugin_handles_none_model_output(error_plugin): ) +# ErrorMetricsPlugin batch tests + + +def _make_batch_error_payload(exc, model="batch-model", provider="batch-provider"): + """Create a GenerationBatchErrorPayload wrapping the given exception.""" + return GenerationBatchErrorPayload( + generation_id="batch-err", + exception=exc, + model=model, + provider=provider, + latency_ms=10.0, + ) + + +@pytest.mark.asyncio +async def test_batch_error_plugin_records_correct_type(error_plugin): + """Batch plugin classifies the exception and calls record_error correctly.""" + payload = _make_batch_error_payload(TimeoutError("timed out")) + + with patch("mellea.telemetry.metrics.record_error") as mock_record: + await error_plugin.record_batch_error_metrics(payload, {}) + + mock_record.assert_called_once_with( + error_type=ERROR_TYPE_TIMEOUT, + model="batch-model", + provider="batch-provider", + exception_class="TimeoutError", + ) + + +@pytest.mark.asyncio +async def test_batch_error_plugin_unknown_exception(error_plugin): + """Batch plugin classifies unknown exceptions correctly.""" + payload = _make_batch_error_payload(ValueError("nope")) + + with patch("mellea.telemetry.metrics.record_error") as mock_record: + await error_plugin.record_batch_error_metrics(payload, {}) + + mock_record.assert_called_once_with( + error_type=ERROR_TYPE_UNKNOWN, + model="batch-model", + provider="batch-provider", + exception_class="ValueError", + ) + + +@pytest.mark.asyncio +async def test_batch_error_plugin_unknown_model_provider_fallback(error_plugin): + """Batch model/provider fall back to 'unknown' when None.""" + payload = _make_batch_error_payload( + ConnectionError("refused"), model=None, provider=None + ) + + with patch("mellea.telemetry.metrics.record_error") as mock_record: + await error_plugin.record_batch_error_metrics(payload, {}) + + mock_record.assert_called_once_with( + error_type=ERROR_TYPE_TRANSPORT_ERROR, + model="unknown", + provider="unknown", + exception_class="ConnectionError", + ) + + # CostMetricsPlugin tests @@ -410,6 +575,125 @@ async def test_cost_plugin_unknown_model_provider_fallback(cost_plugin): ) +# CostMetricsPlugin batch tests + + +@pytest.mark.asyncio +async def test_batch_cost_plugin_records_cost_for_known_model(cost_plugin): + """Batch plugin calls record_cost when compute_cost returns a value.""" + payload = _make_batch_post_payload( + usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150} + ) + + with ( + patch( + "mellea.telemetry.pricing.compute_cost", return_value=0.0042 + ) as mock_cost, + patch("mellea.telemetry.metrics.record_cost") as mock_record, + ): + await cost_plugin.record_batch_cost_metrics(payload, {}) + + mock_cost.assert_called_once_with( + model="batch-model", + provider="batch-provider", + prompt_tokens=100, + completion_tokens=50, + cached_tokens=0, + cache_creation_tokens=0, + ) + mock_record.assert_called_once_with( + cost=0.0042, model="batch-model", provider="batch-provider" + ) + + +@pytest.mark.asyncio +async def test_batch_cost_plugin_cache_tokens_forwarded(cost_plugin): + """Batch plugin forwards cache token fields to compute_cost correctly.""" + payload = _make_batch_post_payload( + usage={ + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "prompt_tokens_details": {"cached_tokens": 50}, + "cache_creation_input_tokens": 10, + } + ) + + with ( + patch("mellea.telemetry.pricing.compute_cost", return_value=0.005) as mock_cost, + patch("mellea.telemetry.metrics.record_cost"), + ): + await cost_plugin.record_batch_cost_metrics(payload, {}) + + mock_cost.assert_called_once_with( + model="batch-model", + provider="batch-provider", + prompt_tokens=100, + completion_tokens=20, + cached_tokens=50, + cache_creation_tokens=10, + ) + + +@pytest.mark.asyncio +async def test_batch_cost_plugin_skips_unknown_model(cost_plugin): + """Batch plugin does not call record_cost when compute_cost returns None.""" + payload = _make_batch_post_payload( + usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150} + ) + + with ( + patch("mellea.telemetry.pricing.compute_cost", return_value=None), + patch("mellea.telemetry.metrics.record_cost") as mock_record, + ): + await cost_plugin.record_batch_cost_metrics(payload, {}) + + mock_record.assert_not_called() + + +@pytest.mark.asyncio +async def test_batch_cost_plugin_skips_none_usage(cost_plugin): + """Batch plugin does not call record_cost when payload.usage is None.""" + payload = _make_batch_post_payload(usage=None) + + with ( + patch("mellea.telemetry.pricing.compute_cost") as mock_cost, + patch("mellea.telemetry.metrics.record_cost") as mock_record, + ): + await cost_plugin.record_batch_cost_metrics(payload, {}) + + mock_cost.assert_not_called() + mock_record.assert_not_called() + + +@pytest.mark.asyncio +async def test_batch_cost_plugin_unknown_model_provider_fallback(cost_plugin): + """Batch model/provider fall back to 'unknown' when None on the payload.""" + payload = _make_batch_post_payload( + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + model=None, + provider=None, + ) + + with ( + patch("mellea.telemetry.pricing.compute_cost", return_value=0.001) as mock_cost, + patch("mellea.telemetry.metrics.record_cost") as mock_record, + ): + await cost_plugin.record_batch_cost_metrics(payload, {}) + + mock_cost.assert_called_once_with( + model="unknown", + provider=None, + prompt_tokens=10, + completion_tokens=5, + cached_tokens=0, + cache_creation_tokens=0, + ) + mock_record.assert_called_once_with( + cost=0.001, model="unknown", provider="unknown" + ) + + # SamplingMetricsPlugin tests