Blame for tree_decoder.py

#!/usr/bin/env python3
"""SML's path-preserving conversation decoder. Python 3.9+, standard library only.

Follows recorded branches; never generates a response, merges equal sentences,
or rewrites the source. JSON pointers use RFC 6901, indices are zero-based.
"""
import argparse
import hashlib
import json
import sys
from dataclasses import dataclass
from pathlib import Path

VERSION = "1.0.0"


class DecodeError(ValueError):
    """The requested path or source structure cannot be decoded faithfully."""


def pointer(base, *parts):
    return base + "".join("/" + str(p).replace("~", "~0").replace("/", "~1") for p in parts)


def resolve(document, path):
    if path == "":
        return document
    if not path.startswith("/"):
        raise DecodeError("A JSON pointer must start with /.")
    value = document
    for part in path[1:].split("/"):
        part = part.replace("~1", "/").replace("~0", "~")
        value = value[int(part)] if isinstance(value, list) else value[part]
    return value


def sentence(value):
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        for key in ("sentence", "content", "text", "words"):
            if isinstance(value.get(key), str):
                return value[key]
    return ""


@dataclass(frozen=True)
class Schema:
    name: str
    prompt_keys: tuple
    branches: str
    answer: str
    prompt_role: str = "Partner"
    answer_role: str = "Response"


QA = Schema("question / answer", ("question",), "branches", "answer")
AAC = Schema("AAC / continuation", ("initial_statement", "other_person_says", "initial_question"),
             "response_branches", "aac_response")
GOLDEN = Schema("doctor / patient", ("root", "doctor"), "branches", "patient", "Doctor", "Patient")
CALL = Schema("caller / dispatcher", ("initial_call", "caller_says"),
              "response_branches", "dispatcher_response", "Caller", "Dispatcher")
DISPATCH = Schema("dispatcher / AAC", ("initial_dispatcher_message", "dispatcher_says"),
                  "response_branches", "aac_response", "Dispatcher", "Caller")
SCHEMAS = (QA, AAC, GOLDEN, CALL, DISPATCH)


def schema_for(value):
    if not isinstance(value, dict):
        return None
    for schema in SCHEMAS:
        if any(k in value for k in schema.prompt_keys) and schema.branches in value:
            return schema
    return None


def tile(value, path, role):
    text = sentence(value)
    if not text:
        raise DecodeError("Missing recorded text at " + (path or "/"))
    metadata = value if isinstance(value, dict) else {}
    return {"text": text, "role": metadata.get("role") or role,
            "pointer": path, "emoji": metadata.get("emoji", ""),
            "tile_id": metadata.get("tile_id"), "metadata": metadata}


class Conversation:
    def __init__(self, value, path, schema=None):
        self.value, self.pointer, self.schema = value, path, schema
        self.linear = schema is None

    def summary(self, index):
        v = self.value
        if self.linear:
            prompt = sentence(v["messages"][0]) if v["messages"] else ""
        else:
            prompt = sentence(next((v[k] for k in self.schema.prompt_keys if k in v), None))
        return {"index": index, "conversation_id": v.get("conversation_id"),
                "title": v.get("topic") or v.get("emergency_type") or
                (v.get("metadata") or {}).get("topic") or prompt or "Untitled conversation",
                "prompt": prompt, "pointer": self.pointer,
                "schema": "linear messages" if self.linear else self.schema.name}

    def _state(self, node, path):
        if node is None:
            return None, [], None
        if isinstance(node, str):
            return tile(node, path, self.schema.prompt_role), [], None
        if not isinstance(node, dict):
            raise DecodeError("Unsupported continuation at " + path)
        key = next((k for k in self.schema.prompt_keys if k in node), None)
        if key is None:
            raise DecodeError("Missing recorded prompt at " + path)
        prompt = tile(node[key], pointer(path, key), self.schema.prompt_role)
        branches = node.get(self.schema.branches, [])
        if not isinstance(branches, list):
            raise DecodeError("Branches must be an array at " + path)
        options = []
        for index, branch in enumerate(branches):
            bp = pointer(path, self.schema.branches, index)
            if not isinstance(branch, dict) or self.schema.answer not in branch:
                raise DecodeError("Missing recorded response at " + bp)
            response = tile(branch[self.schema.answer], pointer(bp, self.schema.answer),
                            self.schema.answer_role)
            options.append({**response, "index": index,
                            "has_continuation": branch.get("continuation") is not None})
        return prompt, options, branches

    def trace(self, choices=()):
        choices = list(choices)
        if any(type(i) is not int or i < 0 for i in choices):
            raise DecodeError("Branch indices must be nonnegative integers.")
        if self.linear:
            return self._linear_trace(choices)
        history = []
        node, path = self.value, self.pointer
        prompt, options, branches = self._state(node, path)
        if prompt:
            history.append({**prompt, "kind": "prompt", "step": 0})
        for step, choice in enumerate(choices, 1):
            if choice >= len(options):
                raise DecodeError("Branch %d does not exist at step %d." % (choice, step))
            response = options[choice]
            history.append({**response, "kind": "response", "step": step})
            branch = branches[choice]
            path = pointer(path, self.schema.branches, choice, "continuation")
            node = branch.get("continuation")
            prompt, options, branches = self._state(node, path)
            if prompt:
                history.append({**prompt, "kind": "prompt", "step": step})
        return self._result(choices, history, options, path if node is not None else None)

    def _linear_trace(self, choices):
        messages = self.value["messages"]
        if any(i != 0 for i in choices) or len(choices) >= len(messages):
            raise DecodeError("Linear conversations have one recorded next message (branch 0).")
        history = []
        for i, message in enumerate(messages[:len(choices) + 1]):
            history.append({**tile(message, pointer(self.pointer, "messages", i), "Message"),
                            "kind": "message", "step": i})
        next_index = len(history)
        options = []
        if next_index < len(messages):
            options = [{**tile(messages[next_index], pointer(self.pointer, "messages", next_index),
                               "Message"), "index": 0, "has_continuation": next_index + 1 < len(messages)}]
        return self._result(choices, history, options, history[-1]["pointer"] if history else None)

    def _result(self, choices, history, options, path):
        return {"decoder_version": VERSION, "conversation_id": self.value.get("conversation_id"),
                "root_pointer": self.pointer, "node_pointer": path,
                "schema": "linear messages" if self.linear else self.schema.name,
                "branch_path": choices, "depth": len(choices), "history": history,
                "options": options, "terminal": not options, "linear": self.linear}

    def walk(self):
        """Emit every reachable state, retaining full history and original pointers."""
        pending = [()]
        while pending:
            path = pending.pop()
            result = self.trace(path)
            yield result
            pending.extend(path + (o["index"],) for o in reversed(result["options"]))


class Document:
    def __init__(self, value):
        self.value = value
        self.conversations = []
        self._discover(value)

    def _discover(self, value):
        pending = [(value, "")]
        while pending:
            item, path = pending.pop()
            schema = schema_for(item)
            if schema:
                self.conversations.append(Conversation(item, path, schema))
                continue  # Continuations belong to this root, never independent conversations.
            if isinstance(item, dict) and isinstance(item.get("messages"), list):
                if item["messages"] and all(sentence(m) for m in item["messages"]):
                    self.conversations.append(Conversation(item, path))
                continue
            if isinstance(item, dict):
                pending.extend((v, pointer(path, k)) for k, v in reversed(list(item.items()))
                               if isinstance(v, (dict, list)))
            elif isinstance(item, list):
                pending.extend((item[i], pointer(path, i)) for i in reversed(range(len(item)))
                               if isinstance(item[i], (dict, list)))

    def summaries(self):
        return [c.summary(i) for i, c in enumerate(self.conversations)]

    def trace(self, conversation=0, choices=()):
        if type(conversation) is not int or not 0 <= conversation < len(self.conversations):
            raise DecodeError("Conversation index is out of range, or this file contains no supported conversations.")
        return {**self.conversations[conversation].trace(choices), "conversation_index": conversation}

    def walk(self):
        for index, conversation in enumerate(self.conversations):
            for state in conversation.walk():
                yield {**state, "conversation_index": index}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", type=Path, help="Original tree JSON file")
    parser.add_argument("--list", action="store_true", help="List conversation roots")
    parser.add_argument("--conversation", type=int, default=0, help="Zero-based conversation index")
    parser.add_argument("--path", default="", help="Zero-based branch indices, e.g. 0,2,1")
    parser.add_argument("--all", action="store_true", help="Export every reachable state as JSONL")
    args = parser.parse_args()
    try:
        raw = args.source.read_bytes()
        doc = Document(json.loads(raw))
        provenance = {"sha256": hashlib.sha256(raw).hexdigest(), "file_name": args.source.name}
        if args.list:
            print(json.dumps(doc.summaries(), ensure_ascii=False, indent=2))
        elif args.all:
            if not doc.conversations:
                raise DecodeError("No supported conversation trees or message sequences in this file.")
            for row in doc.walk():
                print(json.dumps({**row, "source": provenance}, ensure_ascii=False))
        else:
            choices = [int(i) for i in args.path.split(",")] if args.path else []
            print(json.dumps({**doc.trace(args.conversation, choices), "source": provenance},
                             ensure_ascii=False, indent=2))
    except (OSError, ValueError, RecursionError) as error:
        parser.exit(2, "Decoder error: %s\n" % error)


if __name__ == "__main__":
    main()