models/observability.py

"""
observability.py
~~~~~~~~~~~~~~~~
Observability, tracking, and analytics logger for AI assistant interactions.
Records details like response time, token estimation, safety violations,
refusals, and memory utilization to both CSV and JSON log files.
Provides convenient analytical helper functions.

Usage
-----
    from models.observability import ObservabilityLogger

    obs = ObservabilityLogger()
    obs.log_turn(
        assistant_name="OSS Qwen (Local)",
        prompt="Explain quantum physics.",
        response="Quantum physics is...",
        latency_ms=1240.5,
        token_estimate=45,
        is_harmful=False,
        is_refusal=False,
        memory_usage_bytes=1024 * 1024 * 150  # 150 MB
    )
"""

from __future__ import annotations

import os
import csv
import json
import time
from datetime import datetime
from typing import Dict, List, Any
import pandas as pd

from models.logger_config import logger

LOG_DIR = "logs"
CSV_PATH = os.path.join(LOG_DIR, "observability_turns.csv")
JSON_PATH = os.path.join(LOG_DIR, "observability_turns.json")


class ObservabilityLogger:
    """
    Manages operational logs for tracking request-level metrics.
    Ensures thread-safe writing of transaction history to JSON and CSV files.
    """

    def __init__(self) -> None:
        os.makedirs(LOG_DIR, exist_ok=True)
        self._init_csv()

    def _init_csv(self) -> None:
        """Create CSV file and write header if it does not exist."""
        if not os.path.exists(CSV_PATH):
            headers = [
                "timestamp", "assistant_name", "prompt_len", "response_len",
                "latency_ms", "token_estimate", "is_harmful", "is_refusal",
                "memory_usage_bytes"
            ]
            try:
                with open(CSV_PATH, "w", newline="", encoding="utf-8") as f:
                    writer = csv.writer(f)
                    writer.writerow(headers)
                logger.info(f"Initialized CSV log file at {CSV_PATH}")
            except Exception as e:
                logger.error(f"Failed to initialize CSV log: {e}")

    def log_turn(
        self,
        assistant_name: str,
        prompt: str,
        response: str,
        latency_ms: float,
        token_estimate: int,
        is_harmful: bool,
        is_refusal: bool,
        memory_usage_bytes: int
    ) -> None:
        """
        Record a single conversation turn to both JSON and CSV files.
        """
        timestamp = datetime.utcnow().isoformat() + "Z"
        
        # 1. Log to CSV
        row = [
            timestamp,
            assistant_name,
            len(prompt),
            len(response),
            round(latency_ms, 2),
            token_estimate,
            int(is_harmful),
            int(is_refusal),
            memory_usage_bytes
        ]
        
        try:
            with open(CSV_PATH, "a", newline="", encoding="utf-8") as f:
                writer = csv.writer(f)
                writer.writerow(row)
        except Exception as e:
            logger.error(f"Failed to write turn to CSV log: {e}")

        # 2. Log to JSON
        record = {
            "timestamp": timestamp,
            "assistant_name": assistant_name,
            "prompt_snippet": prompt[:100],
            "response_snippet": response[:100],
            "latency_ms": round(latency_ms, 2),
            "token_estimate": token_estimate,
            "is_harmful": is_harmful,
            "is_refusal": is_refusal,
            "memory_usage_bytes": memory_usage_bytes
        }
        
        try:
            records = []
            if os.path.exists(JSON_PATH):
                with open(JSON_PATH, "r", encoding="utf-8") as f:
                    try:
                        records = json.load(f)
                        if not isinstance(records, list):
                            records = []
                    except json.JSONDecodeError:
                        records = []

            records.append(record)
            with open(JSON_PATH, "w", encoding="utf-8") as f:
                json.dump(records, f, indent=2)
        except Exception as e:
            logger.error(f"Failed to write turn to JSON log: {e}")


# ── Analytics Helper Functions ───────────────────────────────────────────────

def get_logs_df() -> pd.DataFrame:
    """Load the CSV logs as a pandas DataFrame."""
    if not os.path.exists(CSV_PATH):
        # Return empty df with schema if file does not exist
        return pd.DataFrame(columns=[
            "timestamp", "assistant_name", "prompt_len", "response_len",
            "latency_ms", "token_estimate", "is_harmful", "is_refusal",
            "memory_usage_bytes"
        ])
    try:
        return pd.read_csv(CSV_PATH)
    except Exception as e:
        logger.error(f"Failed to read logs: {e}")
        return pd.DataFrame()


def get_average_latency(assistant_name: str) -> float:
    """Return average latency in ms for the specified assistant."""
    df = get_logs_df()
    if df.empty:
        return 0.0
    sub = df[df["assistant_name"] == assistant_name]
    if sub.empty:
        return 0.0
    return float(sub["latency_ms"].mean())


def get_total_tokens(assistant_name: str) -> int:
    """Return total accumulated token estimate for the specified assistant."""
    df = get_logs_df()
    if df.empty:
        return 0
    sub = df[df["assistant_name"] == assistant_name]
    if sub.empty:
        return 0
    return int(sub["token_estimate"].sum())


def get_harmful_count(assistant_name: str) -> int:
    """Return the total number of flagged harmful prompts for the assistant."""
    df = get_logs_df()
    if df.empty:
        return 0
    sub = df[df["assistant_name"] == assistant_name]
    if sub.empty:
        return 0
    return int(sub["is_harmful"].sum())


def get_refusal_rate(assistant_name: str) -> float:
    """Return the refusal rate (ratio of refusals to total turns) for the assistant."""
    df = get_logs_df()
    if df.empty:
        return 0.0
    sub = df[df["assistant_name"] == assistant_name]
    if sub.empty:
        return 0.0
    return float(sub["is_refusal"].mean())


def get_summary_stats() -> Dict[str, Any]:
    """
    Compile general summary metrics across all logged turns.
    Returns:
    --------
    Dict containing overall latency, total tokens, refusal counts, and memory peaks.
    """
    df = get_logs_df()
    if df.empty:
        return {"total_turns": 0, "assistants": {}}

    stats: Dict[str, Any] = {
        "total_turns": len(df),
        "assistants": {}
    }

    for name in df["assistant_name"].unique():
        sub = df[df["assistant_name"] == name]
        stats["assistants"][name] = {
            "turns_count": len(sub),
            "avg_latency_ms": round(sub["latency_ms"].mean(), 1) if not sub["latency_ms"].empty else 0.0,
            "total_tokens_est": int(sub["token_estimate"].sum()),
            "harmful_prompts_detected": int(sub["is_harmful"].sum()),
            "refusals_triggered": int(sub["is_refusal"].sum()),
            "max_memory_mb": round(sub["memory_usage_bytes"].max() / (1024 * 1024), 2) if not sub["memory_usage_bytes"].empty else 0.0
        }

    return stats