render.py

"""Turn whatever someone posted into something a template can loop over.

Agent traces arrive in a handful of shapes. We recognise the common message-list
dialects (OpenAI chat completions, Anthropic messages, and the LangChain-ish
variants that mostly look like them) and fall back to pretty-printed JSON or
plain text. Nothing here trusts its input: every value ends up escaped by the
Jinja autoescaper, and unknown block types are shown as JSON rather than
guessed at.
"""

from __future__ import annotations

import json
from typing import Any

ROLE_ALIASES = {
    "human": "user",
    "ai": "assistant",
    "model": "assistant",
    "function": "tool",
}

KNOWN_ROLES = {"system", "developer", "user", "assistant", "tool"}


def _pretty(value: Any) -> str:
    if isinstance(value, str):
        return value
    return json.dumps(value, indent=2, ensure_ascii=False, default=str)


def _blocks_from_content(content: Any) -> list[dict[str, Any]]:
    """Normalise a message's ``content`` into a list of display blocks."""
    if content is None:
        return []
    if isinstance(content, str):
        return [{"type": "text", "text": content}] if content.strip() else []
    if isinstance(content, dict):
        content = [content]
    if not isinstance(content, list):
        return [{"type": "json", "text": _pretty(content)}]

    blocks: list[dict[str, Any]] = []
    for item in content:
        if isinstance(item, str):
            blocks.append({"type": "text", "text": item})
            continue
        if not isinstance(item, dict):
            blocks.append({"type": "json", "text": _pretty(item)})
            continue

        kind = item.get("type")
        if kind == "text":
            blocks.append({"type": "text", "text": item.get("text", "")})
        elif kind == "thinking":
            blocks.append(
                {"type": "thinking", "text": item.get("thinking") or item.get("text", "")}
            )
        elif kind == "tool_use":
            blocks.append(
                {
                    "type": "tool_use",
                    "name": item.get("name", "tool"),
                    "text": _pretty(item.get("input", {})),
                }
            )
        elif kind == "tool_result":
            blocks.append(
                {
                    "type": "tool_result",
                    "name": item.get("tool_use_id", ""),
                    "text": _pretty(item.get("content", "")),
                    "error": bool(item.get("is_error")),
                }
            )
        elif kind in ("image", "image_url"):
            # Never render remote images: it would leak the secret URL through
            # the Referer header and defeat the whole point.
            blocks.append({"type": "note", "text": "[image omitted]"})
        else:
            blocks.append({"type": "json", "text": _pretty(item)})
    return blocks


def _normalize_message(raw: Any) -> dict[str, Any]:
    if not isinstance(raw, dict):
        return {"role": "unknown", "blocks": [{"type": "json", "text": _pretty(raw)}]}

    role = str(raw.get("role") or raw.get("type") or "unknown").lower()
    role = ROLE_ALIASES.get(role, role)
    blocks = _blocks_from_content(raw.get("content"))

    # OpenAI-style tool calls live beside content, not inside it.
    for call in raw.get("tool_calls") or []:
        if not isinstance(call, dict):
            continue
        fn = call.get("function") or {}
        args = fn.get("arguments")
        if isinstance(args, str):
            try:
                args = json.loads(args)
            except (ValueError, TypeError):
                pass  # leave the raw string; malformed args are worth seeing
        blocks.append(
            {
                "type": "tool_use",
                "name": fn.get("name") or call.get("name") or "tool",
                "text": _pretty(args if args is not None else {}),
            }
        )

    if role == "tool" and not blocks and raw.get("name"):
        blocks.append({"type": "note", "text": str(raw["name"])})

    # Anthropic puts tool results in a *user* turn. Labelling that column "user"
    # is technically faithful and completely misleading to read, so show what
    # the turn actually contains.
    if role == "user" and blocks and all(b["type"] == "tool_result" for b in blocks):
        role = "tool"

    return {
        "role": role,
        "role_class": role if role in KNOWN_ROLES else "unknown",
        "name": raw.get("name") or "",
        "blocks": blocks,
    }


def _find_messages(payload: Any) -> list | None:
    """Locate a message list in a few common wrapper shapes."""
    if isinstance(payload, list) and payload and isinstance(payload[0], dict):
        if any(k in payload[0] for k in ("role", "type", "content")):
            return payload
        return None
    if isinstance(payload, dict):
        for key in ("messages", "trace", "conversation", "history", "turns", "steps"):
            value = payload.get(key)
            if isinstance(value, list) and value:
                return value
    return None


def normalize(payload: Any) -> dict[str, Any]:
    """Return ``{"kind": ..., ...}`` describing how to display ``payload``."""
    if isinstance(payload, str):
        return {"kind": "text", "text": payload}

    messages = _find_messages(payload)
    if messages is not None:
        meta = {}
        if isinstance(payload, dict):
            meta = {
                k: v
                for k, v in payload.items()
                if k
                not in ("messages", "trace", "conversation", "history", "turns", "steps")
                and isinstance(v, (str, int, float, bool))
            }
        return {
            "kind": "messages",
            "messages": [_normalize_message(m) for m in messages],
            "meta": meta,
        }

    return {"kind": "json", "text": _pretty(payload)}


def summarize(payload: Any) -> str:
    """One-line description for the page title / header."""
    messages = _find_messages(payload)
    if messages is not None:
        return f"{len(messages)} message{'s' if len(messages) != 1 else ''}"
    if isinstance(payload, str):
        lines = payload.count("\n") + 1
        return f"{lines} line{'s' if lines != 1 else ''}"
    return "JSON document"