File size: 8,875 Bytes
cabc6bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
"""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()