build-small-hackathon/trace-field-notes

1"""Trace parsing and narrative-message extraction."""2 3from __future__ import annotations4 5import json6from pathlib import Path7from typing import Any, Iterable8 9from schemas import AgentType, NarrativeMessage10 11 12TEXT_KEYS = ("text", "message", "summary", "transcript", "output", "body")13TOOLISH_TYPE_FRAGMENTS = (14    "tool",15    "function_call",16    "function_result",17    "command",18    "exec",19    "screenshot",20    "image",21    "patch",22    "diff",23)24TOOLISH_KEYS = (25    "tool_call_id",26    "tool_use_id",27    "tool_calls",28    "tool_results",29    "function_call",30    "arguments",31    "input_json",32    "output_json",33)34 35 36class TraceParseError(ValueError):37    """Raised when an uploaded trace cannot be parsed into narrative messages."""38 39 40def parse_trace(41    path: str | Path,42    *,43    include_user_context: bool = True,44    ignore_tool_calls: bool = True,45) -> tuple[list[NarrativeMessage], AgentType]:46    """Parse an uploaded trace and return visible narrative messages plus agent guess."""47 48    trace_path = Path(path)49    records = load_records(trace_path)50    agent_type = guess_agent_type(records, trace_path)51 52    messages: list[NarrativeMessage] = []53    for raw_index, record in enumerate(records):54        for role, text, timestamp, source in normalize_record(55            record,56            raw_index=raw_index,57            ignore_tool_calls=ignore_tool_calls,58        ):59            cleaned = normalize_whitespace(text)60            if not cleaned:61                continue62            if role == "assistant" or (role == "user" and include_user_context):63                messages.append(64                    NarrativeMessage(65                        index=len(messages),66                        role=role,67                        text=cleaned,68                        timestamp=timestamp,69                        source=source,70                    )71                )72 73    return messages, agent_type74 75 76def load_records(path: Path) -> list[Any]:77    """Load JSONL, JSON, or plain text records from disk."""78 79    try:80        text = path.read_text(encoding="utf-8", errors="replace")81    except OSError as exc:82        raise TraceParseError(f"Could not read uploaded file: {exc}") from exc83 84    if not text.strip():85        raise TraceParseError("The uploaded trace is empty.")86 87    suffix = path.suffix.lower()88    if suffix == ".json":89        try:90            parsed = json.loads(text)91        except json.JSONDecodeError as exc:92            raise TraceParseError(f"Invalid JSON: {exc}") from exc93        return records_from_json(parsed)94 95    if suffix in {".jsonl", ".log", ".txt", ""}:96        records = try_jsonl(text)97        if records:98            return records99        return records_from_plain_text(text)100 101    records = try_jsonl(text)102    return records if records else records_from_plain_text(text)103 104 105def records_from_json(parsed: Any) -> list[Any]:106    if isinstance(parsed, list):107        return parsed108    if isinstance(parsed, dict):109        for key in ("messages", "turns", "events", "records", "items"):110            value = parsed.get(key)111            if isinstance(value, list):112                return value113        return [parsed]114    return [{"type": "text", "role": "assistant", "content": str(parsed)}]115 116 117def try_jsonl(text: str) -> list[Any]:118    records: list[Any] = []119    saw_json = False120    for line in text.splitlines():121        if not line.strip():122            continue123        try:124            records.append(json.loads(line))125            saw_json = True126        except json.JSONDecodeError:127            if saw_json:128                records.append({"type": "text", "role": "assistant", "content": line})129            else:130                return []131    return records if saw_json else []132 133 134def records_from_plain_text(text: str) -> list[Any]:135    records: list[Any] = []136    current_role = "assistant"137    buffer: list[str] = []138 139    def flush() -> None:140        nonlocal buffer141        content = "\n".join(buffer).strip()142        if content:143            records.append({"type": "text", "role": current_role, "content": content})144        buffer = []145 146    for line in text.splitlines():147        lowered = line.strip().lower()148        if lowered.startswith(("assistant:", "agent:")):149            flush()150            current_role = "assistant"151            buffer.append(line.split(":", 1)[1].strip())152        elif lowered.startswith("user:"):153            flush()154            current_role = "user"155            buffer.append(line.split(":", 1)[1].strip())156        else:157            buffer.append(line)158    flush()159 160    if not records:161        records.append({"type": "text", "role": "assistant", "content": text})162    return records163 164 165def guess_agent_type(records: Iterable[Any], path: Path | None = None) -> AgentType:166    path_text = str(path or "").lower()167    if ".codex" in path_text or "/codex/" in path_text:168        return "codex"169    if ".claude" in path_text or "claude" in path_text:170        return "claude_code"171    if ".pi" in path_text or "/pi/" in path_text:172        return "pi"173 174    sample = list(records[:20] if isinstance(records, list) else records)175    for record in sample:176        if not isinstance(record, dict):177            continue178        top_type = str(record.get("type", "")).lower()179        payload = record.get("payload")180        message = record.get("message")181        if top_type in {"session_meta", "turn_context", "response_item", "event_msg"}:182            return "codex"183        if isinstance(payload, dict) and (184            payload.get("originator") == "codex_cli"185            or str(payload.get("type", "")).startswith(("agent_", "user_"))186        ):187            return "codex"188        if "parentUuid" in record or "sessionId" in record or "userType" in record:189            return "claude_code"190        if isinstance(message, dict) and "claude" in str(message.get("model", "")).lower():191            return "claude_code"192        if top_type.startswith("pi_") or "pi agent" in json.dumps(record, default=str).lower()[:1000]:193            return "pi"194    return "unknown"195 196 197def normalize_record(198    record: Any,199    *,200    raw_index: int,201    ignore_tool_calls: bool,202) -> list[tuple[str, str, str | None, str]]:203    """Return zero or more role/text/timestamp/source tuples from one raw record."""204 205    if isinstance(record, str):206        return [("assistant", record, None, "plain_text")]207    if not isinstance(record, dict):208        return [("assistant", str(record), None, "plain_text")]209 210    timestamp = find_timestamp(record)211    candidates: list[tuple[str | None, Any, str]] = []212 213    payload = record.get("payload")214    if isinstance(payload, dict):215        role = normalize_role(payload.get("role"))216        if role is None and str(payload.get("type", "")).lower().startswith("agent"):217            role = "assistant"218        if role is None and str(payload.get("type", "")).lower().startswith("user"):219            role = "user"220        for key in ("content", "message", "summary", "text"):221            if key in payload:222                candidates.append((role, payload[key], f"payload.{key}"))223 224    message = record.get("message")225    if isinstance(message, dict):226        role = normalize_role(message.get("role")) or normalize_role(record.get("type"))227        for key in ("content", "text", "message"):228            if key in message:229                candidates.append((role, message[key], f"message.{key}"))230    elif message is not None:231        role = normalize_role(record.get("role")) or normalize_role(record.get("type"))232        candidates.append((role, message, "message"))233 234    role = normalize_role(record.get("role")) or normalize_role(record.get("type"))235    for key in ("content", "text", "summary", "body"):236        if key in record:237            candidates.append((role, record[key], key))238 239    normalized: list[tuple[str, str, str | None, str]] = []240    seen: set[tuple[str, str]] = set()241    for maybe_role, content, source in candidates:242        role = maybe_role or "assistant"243        if role not in {"assistant", "user"}:244            continue245        text = extract_text(content, ignore_tool_calls=ignore_tool_calls)246        if not text:247            continue248        key = (role, text)249        if key in seen:250            continue251        seen.add(key)252        normalized.append((role, text, timestamp, source))253 254    return normalized255 256 257def normalize_role(value: Any) -> str | None:258    role = str(value or "").lower()259    if role in {"assistant", "agent", "agent_message", "response_item"}:260        return "assistant"261    if role in {"user", "human", "user_message"}:262        return "user"263    return None264 265 266def find_timestamp(record: dict[str, Any]) -> str | None:267    for key in ("timestamp", "created_at", "time", "date"):268        value = record.get(key)269        if isinstance(value, str) and value.strip():270            return value.strip()271    for key in ("payload", "message", "snapshot"):272        value = record.get(key)273        if isinstance(value, dict):274            nested = find_timestamp(value)275            if nested:276                return nested277    return None278 279 280def extract_text(content: Any, *, ignore_tool_calls: bool) -> str:281    """Extract visible prose from known chat content shapes."""282 283    if content is None:284        return ""285    if isinstance(content, str):286        return content287    if isinstance(content, (int, float, bool)):288        return str(content)289    if isinstance(content, list):290        parts = [extract_text(item, ignore_tool_calls=ignore_tool_calls) for item in content]291        return "\n\n".join(part for part in parts if part.strip())292    if isinstance(content, dict):293        if ignore_tool_calls and is_toolish(content):294            return ""295        for key in TEXT_KEYS:296            value = content.get(key)297            if value is not None:298                text = extract_text(value, ignore_tool_calls=ignore_tool_calls)299                if text.strip():300                    return text301        if "content" in content:302            return extract_text(content["content"], ignore_tool_calls=ignore_tool_calls)303    return ""304 305 306def is_toolish(item: dict[str, Any]) -> bool:307    item_type = str(item.get("type", "")).lower()308    role = str(item.get("role", "")).lower()309    name = str(item.get("name", "")).lower()310    if role == "tool":311        return True312    if any(fragment in item_type for fragment in TOOLISH_TYPE_FRAGMENTS):313        return True314    if any(fragment in name for fragment in TOOLISH_TYPE_FRAGMENTS):315        return True316    return any(key in item for key in TOOLISH_KEYS)317 318 319def normalize_whitespace(text: str) -> str:320    lines = [line.rstrip() for line in text.replace("\r\n", "\n").replace("\r", "\n").split("\n")]321    while lines and not lines[0].strip():322        lines.pop(0)323    while lines and not lines[-1].strip():324        lines.pop()325    return "\n".join(lines)326