feat(telemetry): wire metrics plugins to generation_batch_* hooks

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

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.

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}"