src/analyst.py

"""AI safety analyst: Claude chat over the compliance event log (plan 4.2).

The LLM never sits in the per-frame pixel loop — YOLO is the always-on sensor,
this module is the on-demand analyst. Claude answers natural-language questions
("how many no-vest violations this week? worst source?") by calling tools that
query the SQLite event log from src/events.py, and can pull up a violation's
snapshot frame to answer open-ended visual questions the detector can't.

Headless like the rest of src/ — the Streamlit chat tab is just a thin UI over
Analyst.respond(). Auth comes from the ANTHROPIC_API_KEY env var (locally via
shell, on the HF Space via a Space secret).
"""

from __future__ import annotations

import base64
import json
import logging
import os
import time
from pathlib import Path

import anthropic

from src import events

logger = logging.getLogger(__name__)

DEFAULT_MODEL = "claude-opus-5"
MAX_TOOL_ROUNDS = 8  # backstop against tool-call loops

SYSTEM = """\
You are the safety analyst for a PPE (personal protective equipment) compliance system on \
construction sites. A YOLOv8 detector with ByteTrack processes video and logs discrete \
*events* to a database: one row per (tracked person/object, class) sighting span, with \
unix timestamps, frame counts, peak confidence, and — for violations — a snapshot image \
of the annotated frame.

Classes: Hardhat, Mask, Safety Vest, Person, Safety Cone, machinery, vehicle, and the \
violation classes NO-Hardhat, NO-Mask, NO-Safety Vest (a worker missing that equipment).

Ground every claim in tool results — query before answering, and say so plainly when the \
log has no data for a question. Durations are ended_at - started_at in seconds. The \
detector misses 28-35% of true violations (per its model card), so zero logged violations \
never proves zero occurred — mention this when it matters. Keep answers concise and \
concrete; use short tables for per-class breakdowns. When asked something visual about a \
specific violation, view its snapshot."""

TOOLS = [
    {
        "name": "get_summary",
        "description": (
            "Headline statistics for the compliance event log: total events, violation counts and "
            "total violation seconds, per-class event counts, time range covered, number of sources, "
            "and the current unix time. Call this first for any broad or vague question."
        ),
        "input_schema": {"type": "object", "properties": {}},
    },
    {
        "name": "query_events",
        "description": (
            "Filtered list of events, newest first. Each row: id, source, track_id, cls, started_at, "
            "ended_at, duration (s), frames, max_conf, and whether a snapshot exists."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "cls": {"type": "string", "description": "Exact class name, e.g. 'NO-Hardhat'"},
                "source": {"type": "string", "description": "Video/stream name, e.g. '2.mp4'"},
                "since_hours": {"type": "number", "description": "Only events from the last N hours"},
                "violations_only": {"type": "boolean", "description": "Restrict to NO-* violation classes"},
                "limit": {"type": "integer", "description": "Max rows (default 50, cap 200)"},
            },
        },
    },
    {
        "name": "view_snapshot",
        "description": (
            "Returns the annotated snapshot frame for a violation event (by event id) so you can "
            "visually inspect the scene — use it for open-ended questions the detector's classes "
            "can't answer (context, posture, what else is happening in frame)."
        ),
        "input_schema": {
            "type": "object",
            "properties": {"event_id": {"type": "integer", "description": "The event's id"}},
            "required": ["event_id"],
        },
    },
]


class Analyst:
    def __init__(self, db_path: str | Path = events.DEFAULT_DB, model: str | None = None):
        self.db_path = Path(db_path)
        self.model = model or os.environ.get("ANTHROPIC_MODEL", DEFAULT_MODEL)
        self._client = None

    @staticmethod
    def available() -> bool:
        return bool(os.environ.get("ANTHROPIC_API_KEY"))

    @property
    def client(self) -> anthropic.Anthropic:
        if self._client is None:
            self._client = anthropic.Anthropic()
        return self._client

    def respond(self, messages: list[dict]) -> tuple[str, list[dict], list[str]]:
        """Run one assistant turn over `messages` (Anthropic format, ending on a user turn).

        Returns (answer_text, updated_messages, tool_calls_made). `messages` is not
        mutated; the returned list includes all intermediate tool traffic so the
        caller can keep it as the conversation history.
        """
        messages = list(messages)
        tool_calls: list[str] = []
        response = self._create(messages)
        for _ in range(MAX_TOOL_ROUNDS):
            if response.stop_reason != "tool_use":
                break
            messages.append({"role": "assistant", "content": response.content})
            results = []
            for block in response.content:
                if block.type != "tool_use":
                    continue
                tool_calls.append(block.name)
                try:
                    content = self.run_tool(block.name, block.input)
                    results.append({"type": "tool_result", "tool_use_id": block.id, "content": content})
                except Exception as e:  # tool bugs become model-visible errors, not crashes
                    logger.exception("Tool %s failed", block.name)
                    results.append(
                        {"type": "tool_result", "tool_use_id": block.id, "content": f"Error: {e}", "is_error": True}
                    )
            messages.append({"role": "user", "content": results})
            response = self._create(messages)

        if response.stop_reason == "refusal":
            text = "The model declined to answer this request."
        else:
            text = "\n".join(b.text for b in response.content if b.type == "text")
        messages.append({"role": "assistant", "content": response.content})
        return text, messages, tool_calls

    def _create(self, messages: list[dict]):
        return self.client.messages.create(
            model=self.model,
            max_tokens=4096,
            system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}],
            tools=TOOLS,
            messages=messages,
        )

    # --- Tool implementations (public so tests can hit them without an API key) ---

    def run_tool(self, name: str, tool_input: dict):
        conn = events.connect(self.db_path)
        try:
            if name == "get_summary":
                return json.dumps(
                    {
                        "now": time.time(),
                        **events.summary(conn),
                        "events_by_class": events.class_counts(conn),
                    }
                )
            if name == "query_events":
                since = None
                if tool_input.get("since_hours") is not None:
                    since = time.time() - float(tool_input["since_hours"]) * 3600
                rows = events.query_events(
                    conn,
                    cls=tool_input.get("cls"),
                    source=tool_input.get("source"),
                    since=since,
                    violations_only=bool(tool_input.get("violations_only")),
                    limit=int(tool_input.get("limit") or 50),
                )
                for r in rows:
                    r["has_snapshot"] = bool(r.pop("snapshot", None))
                return json.dumps({"count": len(rows), "events": rows})
            if name == "view_snapshot":
                row = conn.execute(
                    "SELECT * FROM events WHERE id = ?", (int(tool_input["event_id"]),)
                ).fetchone()
                if row is None:
                    return "No event with that id."
                if not row["snapshot"] or not Path(row["snapshot"]).exists():
                    return "That event has no snapshot (only violation events get one)."
                data = base64.standard_b64encode(Path(row["snapshot"]).read_bytes()).decode()
                meta = {k: row[k] for k in ("id", "source", "track_id", "cls", "started_at", "ended_at", "max_conf")}
                return [
                    {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": data}},
                    {"type": "text", "text": f"Snapshot for event: {json.dumps(meta)}"},
                ]
            raise ValueError(f"Unknown tool: {name}")
        finally:
            conn.close()