utils/logger.py

"""
utils/logger.py — Structured timing logger for MedCheck Agent

Logs every phase of a request with wall-clock durations and
a final summary table.  All output goes to both the Python
logging system (file + console) and a per-request dict that
app.py can surface in the UI.
"""

from __future__ import annotations
import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Optional

# ─────────────────────────────────────────────
# Configure root logger once at import time
# ─────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(name)s  │  %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[
        logging.StreamHandler(),                          # console
        logging.FileHandler("medcheck.log", encoding="utf-8"),  # file
    ],
)

logger = logging.getLogger("medcheck")


# ─────────────────────────────────────────────
# Phase timings container
# ─────────────────────────────────────────────
@dataclass
class PhaseTimer:
    name: str
    start: float = field(default_factory=time.perf_counter)
    end: Optional[float] = None
    cache_hit: bool = False          # semantic cache hit?
    prompt_cached_tokens: int = 0    # tokens served from Anthropic prompt cache

    @property
    def elapsed_ms(self) -> float:
        if self.end is None:
            return (time.perf_counter() - self.start) * 1000
        return (self.end - self.start) * 1000

    def stop(self) -> "PhaseTimer":
        self.end = time.perf_counter()
        return self


@dataclass
class RequestLog:
    """
    Collects all timing/cache data for one end-to-end request.

    Thread-safe: start_phase/end_phase may be called from concurrent
    threads (one per task when parallel execution is enabled).
    The internal lock only guards the phases list append; the heavy
    work (web search, LLM call) runs outside any lock.
    """
    medicine_input: str
    provider: str
    model: str
    request_start: float = field(default_factory=time.perf_counter)
    request_end: Optional[float] = None
    phases: list[PhaseTimer] = field(default_factory=list)
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    # ── helpers ──────────────────────────────
    def start_phase(self, name: str) -> PhaseTimer:
        pt = PhaseTimer(name=name)
        with self._lock:
            self.phases.append(pt)
        logger.info("▶  %-36s  started", name)
        return pt

    def end_phase(self, pt: PhaseTimer,
                  cache_hit: bool = False,
                  prompt_cached_tokens: int = 0) -> None:
        pt.stop()
        pt.cache_hit = cache_hit
        pt.prompt_cached_tokens = prompt_cached_tokens
        cache_tag = " [SEMANTIC CACHE HIT]" if cache_hit else ""
        pct_tag   = (f" [prompt-cache tokens: {prompt_cached_tokens}]"
                     if prompt_cached_tokens > 0 else "")
        logger.info(
            "◀  %-36s  %7.0f ms%s%s",
            pt.name, pt.elapsed_ms, cache_tag, pct_tag,
        )

    def finish(self) -> None:
        self.request_end = time.perf_counter()
        total_ms = (self.request_end - self.request_start) * 1000
        self._print_summary(total_ms)

    def total_ms(self) -> float:
        if self.request_end is None:
            return (time.perf_counter() - self.request_start) * 1000
        return (self.request_end - self.request_start) * 1000

    # ── summary table ────────────────────────
    def _sorted_phases(self) -> list[PhaseTimer]:
        """
        Return phases in a stable, readable order even when tasks ran in
        parallel (phases arrive in thread-completion order, not task order).

        Sort key: (task_rank, start_time)
          spell=0, interactions=1, contraindications=2, unknown=3
        """
        _rank = {"spell": 0, "interactions": 1, "contraindications": 2}

        def _key(pt: PhaseTimer) -> tuple:
            prefix = pt.name.split(":")[0].strip()
            return (_rank.get(prefix, 3), pt.start)

        with self._lock:
            return sorted(self.phases, key=_key)

    def _print_summary(self, total_ms: float) -> None:
        sep = "─" * 72
        lines = [
            "",
            sep,
            f"  REQUEST SUMMARY  │  Input: {self.medicine_input!r}"
            f"  │  Provider: {self.provider}/{self.model}",
            sep,
            f"  {'Phase':<36}  {'Time (ms)':>9}  {'Cache':>14}",
            f"  {'─'*36}  {'─'*9}  {'─'*14}",
        ]
        for pt in self._sorted_phases():
            hit_str = "SEMANTIC HIT" if pt.cache_hit else (
                f"PCT:{pt.prompt_cached_tokens}" if pt.prompt_cached_tokens else "—"
            )
            lines.append(
                f"  {pt.name:<36}  {pt.elapsed_ms:>9.0f}  {hit_str:>14}"
            )
        lines += [
            f"  {'─'*36}  {'─'*9}  {'─'*14}",
            f"  {'TOTAL':<36}  {total_ms:>9.0f}",
            sep,
        ]
        logger.info("\n".join(lines))

    def to_ui_string(self) -> str:
        """Return a compact markdown table for the Gradio Performance tab."""
        rows = []
        for pt in self._sorted_phases():
            hit_str = "✅ semantic" if pt.cache_hit else (
                f"⚡ PCT {pt.prompt_cached_tokens} tok"
                if pt.prompt_cached_tokens else "🌐 live"
            )
            rows.append(
                f"| {pt.name} | `{pt.elapsed_ms:.0f} ms` | {hit_str} |"
            )
        total = self.total_ms()
        return (
            "| Phase | Time | Source |\n"
            "|---|---|---|\n"
            + "\n".join(rows)
            + f"\n| **TOTAL** | **`{total:.0f} ms`** | |"
        )