AI System Observability and Monitoring

import hashlib
import json
import time
import uuid
from datetime import datetime, timezone
from typing import Any, Callable

# ── Log record structure ───────────────────────────────────────────────────────
def make_log_record(
  model: str,
  rendered_prompt: str,
  response_text: str,
  input_tokens: int,
  output_tokens: int,
  latency_ms: float,
  finish_reason: str,
  temperature: float,
  eval_fn: Callable[[str, str], dict] | None = None,
  session_id: str | None = None,
  error: str | None = None,
) -> dict[str, Any]:
  prompt_hash = hashlib.sha256(rendered_prompt.encode()).hexdigest()[:16]

  eval_scores = {}
  if eval_fn and response_text:
      try:
          eval_scores = eval_fn(rendered_prompt, response_text)
      except Exception as e:
          eval_scores = {"eval_error": str(e)}

  return {
      "request_id": str(uuid.uuid4()),
      "session_id": session_id or str(uuid.uuid4()),
      "timestamp": datetime.now(timezone.utc).isoformat(),
      "model": model,
      "prompt_hash": prompt_hash,
      "rendered_prompt": rendered_prompt,  # FULL rendered prompt  -  not template
      "response": response_text,
      "input_tokens": input_tokens,
      "output_tokens": output_tokens,
      "total_tokens": input_tokens + output_tokens,
      "latency_ms": round(latency_ms, 2),
      "temperature": temperature,
      "finish_reason": finish_reason,
      "eval_scores": eval_scores,
      "error": error,
  }


# ── LLM call wrapper ───────────────────────────────────────────────────────────
class ObservableLLM:
  def __init__(
      self,
      model: str = "gpt-4o-2024-11-20",
      temperature: float = 0.0,
      eval_fn: Callable | None = None,
      log_sink: Callable[[dict], None] | None = None,
  ):
      self.model = model
      self.temperature = temperature
      self.eval_fn = eval_fn
      self.log_sink = log_sink or self._default_log_sink

  def _default_log_sink(self, record: dict) -> None:
      # In production: send to your log aggregator
      # e.g., structlog, CloudWatch, Datadog, etc.
      print(json.dumps({
          "ts": record["timestamp"],
          "model": record["model"],
          "prompt_hash": record["prompt_hash"],
          "latency_ms": record["latency_ms"],
          "tokens": record["total_tokens"],
          "finish": record["finish_reason"],
          "evals": record["eval_scores"],
      }, indent=2))

  def call(
      self,
      rendered_prompt: str,
      session_id: str | None = None,
  ) -> str:
      """
      Wraps an LLM call with structured logging.
      In production: replace the simulation block with your API call.
      """
      start = time.perf_counter()
      error = None
      response_text = ""
      input_tokens = 0
      output_tokens = 0
      finish_reason = "stop"

      try:
          # ── Production: replace this block ───────────────────────────────
          # from openai import OpenAI
          # client = OpenAI()
          # completion = client.chat.completions.create(
          #     model=self.model,
          #     messages=[{"role": "user", "content": rendered_prompt}],
          #     temperature=self.temperature,
          # )
          # response_text = completion.choices[0].message.content
          # input_tokens = completion.usage.prompt_tokens
          # output_tokens = completion.usage.completion_tokens
          # finish_reason = completion.choices[0].finish_reason
          # ─────────────────────────────────────────────────────────────────

          # Simulation
          time.sleep(0.05)
          response_text = f"Simulated response to: '{rendered_prompt[:40]}...'"
          input_tokens = len(rendered_prompt.split()) * 4 // 3
          output_tokens = 48
          finish_reason = "stop"

      except Exception as e:
          error = f"{type(e).__name__}: {e}"
          finish_reason = "error"

      latency_ms = (time.perf_counter() - start) * 1000

      record = make_log_record(
          model=self.model,
          rendered_prompt=rendered_prompt,
          response_text=response_text,
          input_tokens=input_tokens,
          output_tokens=output_tokens,
          latency_ms=latency_ms,
          finish_reason=finish_reason,
          temperature=self.temperature,
          eval_fn=self.eval_fn,
          session_id=session_id,
          error=error,
      )
      self.log_sink(record)

      if error:
          raise RuntimeError(error)

      return response_text


# ── Example eval function ──────────────────────────────────────────────────────
def simple_eval(prompt: str, response: str) -> dict[str, float]:
  """
  In production: use your actual eval suite.
  This demo checks for minimal quality signals.
  """
  checks = {
      "non_empty": 1.0 if len(response.strip()) > 0 else 0.0,
      "no_apology": 0.0 if "i apologize" in response.lower() else 1.0,
      "min_length": 1.0 if len(response.split()) >= 5 else 0.0,
      "no_refusal": 0.0 if "i cannot" in response.lower() else 1.0,
  }
  checks["overall"] = sum(checks.values()) / len(checks)
  return checks


# ── Demo ───────────────────────────────────────────────────────────────────────
llm = ObservableLLM(
  model="gpt-4o-2024-11-20",
  temperature=0.0,
  eval_fn=simple_eval,
)

session = str(uuid.uuid4())[:8]
response = llm.call(
  rendered_prompt="Explain what BM25 is in one paragraph.",
  session_id=session,
)
print(f"\nResponse: {response}")