feat: add MetricsContext and token tracking (Phase A)

abevalflow/llm_client.py

@@ -13,11 +13,24 @@
 
 import logging
 import os
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from openai import OpenAI
 
+
+@dataclass
+class LLMResult:
+    """Chat completion result with token usage."""
+
+    content: str
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    model: str
+
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_BASE_URL = "http://litellm.ab-eval-flow.svc:4000/v1"
@@ -45,15 +58,15 @@ def get_model() -> str:
     return os.environ.get("LLM_MODEL", DEFAULT_MODEL)
 
 
-def chat_completion(
+def chat_completion_with_usage(
     messages: list[dict[str, str]],
     *,
     model: str | None = None,
     temperature: float = 0.3,
     max_tokens: int = 4096,
     **kwargs,
-) -> str:
-    """Send a chat completion request and return the assistant message content.
+) -> LLMResult:
+    """Send a chat completion request and return content with token usage.
 
     Raises on API errors so callers can handle retries at a higher level.
     """
@@ -75,6 +88,38 @@ def chat_completion(
         **kwargs,
     )
 
-    content = response.choices[0].message.content
-    logger.info("chat_completion ← %d chars", len(content) if content else 0)
-    return content or ""
+    content = response.choices[0].message.content or ""
+    usage = response.usage
+    prompt_tokens = usage.prompt_tokens if usage else 0
+    completion_tokens = usage.completion_tokens if usage else 0
+
+    logger.info(
+        "chat_completion ← %d chars, tokens: %d prompt + %d completion",
+        len(content),
+        prompt_tokens,
+        completion_tokens,
+    )
+
+    return LLMResult(
+        content=content,
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=prompt_tokens + completion_tokens,
+        model=resolved_model,
+    )
+
+
+def chat_completion(
+    messages: list[dict[str, str]],
+    *,
+    model: str | None = None,
+    temperature: float = 0.3,
+    max_tokens: int = 4096,
+    **kwargs,
+) -> str:
+    """Send a chat completion request and return the assistant message content.
+
+    For token usage, use ``chat_completion_with_usage()`` instead.
+    """
+    result = chat_completion_with_usage(messages, model=model, temperature=temperature, max_tokens=max_tokens, **kwargs)
+    return result.content

abevalflow/observability/__init__.py

@@ -0,0 +1,9 @@
+"""Observability layer for ABEvalFlow pipeline metrics and tracing."""
+
+from abevalflow.observability.context import MetricsContext, TimingRecord, TokenUsage
+
+__all__ = [
+    "MetricsContext",
+    "TimingRecord",
+    "TokenUsage",
+]

abevalflow/observability/context.py

@@ -0,0 +1,136 @@
+"""Metrics context for accumulating token usage and timing across a pipeline run.
+
+Append-only during execution, then serialized and persisted at the end.
+Per-gate token buckets keyed by phase/gate name for future parallelization.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from datetime import UTC, datetime
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+CHECKPOINT_FILENAME = "_metrics_checkpoint.json"
+
+
+class TokenUsage(BaseModel):
+    """Token counts for a single phase or gate."""
+
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    model_name: str | None = None
+    call_count: int = 0
+
+    def accumulate(self, prompt: int, completion: int, model: str | None = None) -> None:
+        self.prompt_tokens += prompt
+        self.completion_tokens += completion
+        self.total_tokens += prompt + completion
+        self.call_count += 1
+        if model:
+            self.model_name = model
+
+
+class TimingRecord(BaseModel):
+    """Duration record for a pipeline phase."""
+
+    name: str
+    start_time: float = Field(description="Unix timestamp")
+    end_time: float | None = None
+    duration_ms: int | None = None
+
+    def stop(self) -> None:
+        self.end_time = time.time()
+        self.duration_ms = int((self.end_time - self.start_time) * 1000)
+
+
+class MetricsContext(BaseModel):
+    """Accumulates token usage and timing across a pipeline run.
+
+    Token usage is bucketed by phase/gate name to support concurrent execution.
+    """
+
+    run_id: str = ""
+    submission_name: str = ""
+    model_name: str | None = None
+    start_time: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    timings: dict[str, TimingRecord] = Field(default_factory=dict)
+    token_usage: dict[str, TokenUsage] = Field(default_factory=dict)
+
+    def record_tokens(
+        self,
+        phase_name: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        model: str | None = None,
+    ) -> None:
+        if phase_name not in self.token_usage:
+            self.token_usage[phase_name] = TokenUsage()
+        self.token_usage[phase_name].accumulate(prompt_tokens, completion_tokens, model)
+        if model and not self.model_name:
+            self.model_name = model
+
+    def start_timing(self, phase_name: str) -> None:
+        self.timings[phase_name] = TimingRecord(name=phase_name, start_time=time.time())
+
+    def stop_timing(self, phase_name: str) -> None:
+        if phase_name in self.timings:
+            self.timings[phase_name].stop()
+
+    @property
+    def total_prompt_tokens(self) -> int:
+        return sum(u.prompt_tokens for u in self.token_usage.values())
+
+    @property
+    def total_completion_tokens(self) -> int:
+        return sum(u.completion_tokens for u in self.token_usage.values())
+
+    @property
+    def total_tokens(self) -> int:
+        return sum(u.total_tokens for u in self.token_usage.values())
+
+    @property
+    def llm_calls_count(self) -> int:
+        return sum(u.call_count for u in self.token_usage.values())
+
+    def timing_ms(self, phase_name: str) -> int | None:
+        rec = self.timings.get(phase_name)
+        return rec.duration_ms if rec else None
+
+    def checkpoint(self, workspace_path: Path) -> None:
+        path = workspace_path / CHECKPOINT_FILENAME
+        path.write_text(self.model_dump_json(indent=2))
+        logger.info("Metrics checkpoint written to %s", path)
+
+    @classmethod
+    def load_checkpoint(cls, workspace_path: Path) -> MetricsContext | None:
+        path = workspace_path / CHECKPOINT_FILENAME
+        if not path.exists():
+            return None
+        try:
+            return cls.model_validate_json(path.read_bytes())
+        except Exception:
+            logger.warning("Failed to load metrics checkpoint from %s", path, exc_info=True)
+            return None
+
+    def to_observability_dict(self) -> dict:
+        """Convert to kwargs for ObservabilityMetricsRow."""
+        return {
+            "submission_name": self.submission_name,
+            "model_name": self.model_name,
+            "pipeline_duration_ms": self.timing_ms("pipeline"),
+            "prepare_duration_ms": self.timing_ms("prepare"),
+            "test_duration_ms": self.timing_ms("test"),
+            "evaluate_duration_ms": self.timing_ms("evaluate"),
+            "analyze_duration_ms": self.timing_ms("analyze"),
+            "store_duration_ms": self.timing_ms("store"),
+            "total_prompt_tokens": self.total_prompt_tokens or None,
+            "total_completion_tokens": self.total_completion_tokens or None,
+            "total_tokens": self.total_tokens or None,
+            "llm_calls_count": self.llm_calls_count or None,
+        }

abevalflow/observability/decorators.py

@@ -0,0 +1,31 @@
+"""Observability decorators for timing and tracing."""
+
+from __future__ import annotations
+
+import functools
+import logging
+import time
+from collections.abc import Callable
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def timed_gate(func: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator that logs gate execution time in milliseconds.
+
+    Attaches ``_duration_ms`` to the return value if it has that attribute,
+    otherwise logs only.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        start = time.time()
+        result = func(*args, **kwargs)
+        duration_ms = int((time.time() - start) * 1000)
+        logger.info("Gate %s executed in %dms", func.__qualname__, duration_ms)
+        if hasattr(result, "_duration_ms"):
+            result._duration_ms = duration_ms
+        return result
+
+    return wrapper

pipeline/tasks/phases/test.yaml

@@ -460,17 +460,17 @@ spec:
           echo "Quality review: passed=$PASSED recommendation=$REC"
         fi
 
-    # Step 6: Finalize results
+    # Step 6: Finalize results and write metrics checkpoint
     - name: finalize
       image: registry.access.redhat.com/ubi9/python-311:9.6
       script: |
         #!/usr/bin/env bash
         set -euo pipefail
         echo "=== TEST PHASE: Finalize ==="
-        
+
         SECURITY_PASSED=$(cat "$(results.security-passed.path)")
         QUALITY_PASSED=$(cat "$(results.quality-passed.path)")
-        
+
         # Overall pass if both passed
         if [ "$SECURITY_PASSED" = "true" ] && [ "$QUALITY_PASSED" = "true" ]; then
           echo -n "true" > "$(results.tests-passed.path)"
@@ -479,3 +479,19 @@ spec:
           echo -n "false" > "$(results.tests-passed.path)"
           echo "Tests FAILED (security=$SECURITY_PASSED, quality=$QUALITY_PASSED)"
         fi
+
+        # Write metrics checkpoint with token data from quality review
+        PIPELINE_DIR="$(workspaces.source.path)/_pipeline"
+        REPORT_DIR="$(workspaces.source.path)/reports/$(params.submission-name)"
+        REVIEW_FILE="$(workspaces.source.path)/_ai_review.json"
+        mkdir -p "$REPORT_DIR"
+
+        pip install --quiet --no-cache-dir pydantic 2>&1 | tail -1
+        export PYTHONPATH="$PIPELINE_DIR"
+
+        python3 "$PIPELINE_DIR/scripts/write_metrics_checkpoint.py" \
+          --run-id "$(params.pipeline-run-name)" \
+          --submission-name "$(params.submission-name)" \
+          --report-dir "$REPORT_DIR" \
+          --review-file "$REVIEW_FILE" \
+          2>&1 || echo "Warning: metrics checkpoint failed (non-blocking)"

scripts/store_results.py

@@ -32,12 +32,14 @@
     GateResultRow,
     MCPCheckerRun,
     MCPCheckerTask,
+    ObservabilityMetricsRow,
     ScorecardRow,
     SecurityScan,
     Trial,
 )
 from abevalflow.db.observer import discover_observers, notify_observers
 from abevalflow.mcpchecker_report import MCPCheckerResult
+from abevalflow.observability.context import MetricsContext
 from abevalflow.report import AnalysisResult
 from abevalflow.scorecard import Scorecard
 
@@ -347,6 +349,8 @@ def store(
         ).scalar_one_or_none()
 
         if existing is not None:
+            # Note: early return skips scorecard/metrics persistence for
+            # pre-existing runs. Backfill via scripts/backfill_scorecards.py.
             logger.warning(
                 "Run %s already exists (id=%s) — skipping",
                 effective_run_id,
@@ -392,6 +396,20 @@ def store(
             except Exception:
                 logger.warning("Failed to persist scorecard — continuing without", exc_info=True)
 
+        metrics_ctx = MetricsContext.load_checkpoint(report_dir)
+        if metrics_ctx and metrics_ctx.total_tokens > 0:
+            try:
+                with session.begin_nested():
+                    session.add(
+                        ObservabilityMetricsRow(
+                            pipeline_run_id=effective_run_id,
+                            **metrics_ctx.to_observability_dict(),
+                        )
+                    )
+                logger.info("Observability metrics queued: tokens=%s", metrics_ctx.total_tokens)
+            except Exception:
+                logger.warning("Failed to persist observability metrics — continuing without", exc_info=True)
+
         try:
             session.commit()
         except IntegrityError as e:

scripts/test_quality_review.py

@@ -390,14 +390,15 @@ def review_submission(submission_dir: Path) -> dict:
             )
         engine = None
 
-    response_text = llm_client.chat_completion(
+    llm_result = llm_client.chat_completion_with_usage(
         messages=[
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": user_prompt},
         ],
         temperature=0.2,
         max_tokens=4096,
     )
+    response_text = llm_result.content
 
     try:
         assessment = json.loads(response_text)
@@ -410,7 +411,15 @@ def review_submission(submission_dir: Path) -> dict:
         else:
             raise ValueError("LLM review response is not valid JSON")
 
-    return _normalize_assessment(assessment, engine=engine)
+    assessment = _normalize_assessment(assessment, engine=engine)
+    assessment["token_usage"] = {
+        "prompt_tokens": llm_result.prompt_tokens,
+        "completion_tokens": llm_result.completion_tokens,
+        "total_tokens": llm_result.total_tokens,
+        "model": llm_result.model,
+    }
+
+    return assessment
 
 
 def main(argv: list[str] | None = None) -> int: