pi-durable-python 1.3.0__py3-none-any.whl

pi_durable/harness/inbox.py

@@ -0,0 +1,255 @@
+"""Built-in ``pi.inbox`` and the submission boundary (spec section 6).
+
+Queued submissions wait in one conversation-scoped document, in ID order. A boundary reads the
+document on the Session line, places the selected writes and user items, and removes them
+positionally; run control itself lives in ``pi.live``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal, TypedDict, cast
+
+from ..documents import define_doc
+from ..entries import UserEntry
+from ..ids import ConversationId, EntryId, SubmissionId
+from ..json import JsonObject, JsonValue, copy_json
+from ..session.drafts import DocumentDraft, DraftList, DraftObject
+from ..session.transaction import Transaction
+from ..types import EntryDraft, EntryRecord
+from .types import InputEntry, QueueMode, UserInput
+
+
+class InboxItem(TypedDict, total=False):
+    """One queued submission: an input with its content, or a passive entry write."""
+
+    id: SubmissionId
+    mode: Literal["steer", "followUp", "write"]
+    content: UserInput
+    entry: JsonObject
+
+
+class InboxState(TypedDict):
+    """The inbox document value: queued submissions in submission-ID order."""
+
+    items: list[InboxItem]
+
+
+InboxDoc = define_doc(
+    {
+        "kind": "pi.inbox",
+        "version": 1,
+        "scope": "conversation",
+        "history": "latest",
+        "fork": "initial",
+        "initial": lambda: {"items": []},
+        "checkpoint_when": lambda value, _ops, _info: len(value.get("items") or []) == 0,
+    }
+)
+
+
+@dataclass(slots=True)
+class QueueModes:
+    """The queue modes a boundary reads on the Session line."""
+
+    steering_mode: QueueMode
+    follow_up_mode: QueueMode
+
+
+@dataclass(slots=True)
+class Boundary:
+    """What a boundary reads before the commit's first table write."""
+
+    conversation_id: ConversationId
+    inbox: DocumentDraft
+    steering_mode: QueueMode
+    follow_up_mode: QueueMode
+    head: EntryId | None
+
+
+@dataclass(frozen=True, slots=True)
+class BoundaryResult:
+    """Selected user submission IDs and whether a reset was among the placed writes."""
+
+    users: list[SubmissionId]
+    reset: bool
+
+
+def _plain(value: Any) -> Any:
+    """Return a plain JSON value from a document draft member or a record."""
+
+    if isinstance(value, DraftObject):
+        return value.snapshot()
+    if isinstance(value, DraftList):
+        return value.snapshot()
+    if isinstance(value, dict):
+        return {key: _plain(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_plain(item) for item in value]
+    return value
+
+
+def _entry_draft(item: Any) -> EntryDraft:
+    return cast(EntryDraft, _plain(item["entry"]))
+
+
+async def prepare_boundary(
+    tx: Transaction, conversation_id: ConversationId, modes: QueueModes
+) -> Boundary:
+    """Read the active-range start and the inbox before the commit's first table write."""
+
+    marker = await tx.latest_head_marker(conversation_id)
+    inbox = await tx.doc(InboxDoc, conversation_id)
+    return Boundary(
+        conversation_id=conversation_id,
+        inbox=inbox,
+        steering_mode=modes.steering_mode,
+        follow_up_mode=modes.follow_up_mode,
+        head=None if marker is None else marker.get("head"),
+    )
+
+
+async def apply_boundary(
+    tx: Transaction, boundary: Boundary, at: Literal["postTools", "final"], now: int
+) -> BoundaryResult:
+    """Place the queued items a boundary selects, writes first and then user items, each in ID order.
+
+    A selected reset turns a ``postTools`` boundary into ``final``. A head write whose target is
+    older than the active range, including one started earlier in this commit, is stale. Selected
+    and stale items are removed positionally.
+    """
+
+    conversation_id = boundary.conversation_id
+    items = boundary.inbox["items"]
+    reset = has_queued_reset(items)
+    final = at == "final" or reset
+
+    def pick(mode: Literal["steer", "followUp"], queue_mode: QueueMode) -> list[int]:
+        indexes = [index for index, item in enumerate(items) if item["mode"] == mode]
+        return indexes if queue_mode == "all" else indexes[:1]
+
+    writes = [index for index, item in enumerate(items) if item["mode"] == "write"]
+    users = [*pick("steer", boundary.steering_mode)]
+    if final:
+        users.extend(pick("followUp", boundary.follow_up_mode))
+    users.sort()
+
+    placed: list[SubmissionId] = []
+    for index in writes:
+        item = items[index]
+        draft = _entry_draft(item)
+        if is_stale(boundary, draft):
+            tx.settle_submission(cast(SubmissionId, item["id"]), {"status": "unanswered", "reason": "stale"})
+            continue
+        entry = await tx.append_entry(conversation_id, draft)
+        head = draft.get("head")
+        if head is not None:
+            boundary.head = entry["id"] if head == "self" else head
+        tx.place_submission(cast(SubmissionId, item["id"]), entry["id"])
+
+    for index in users:
+        item = items[index]
+        entry = await append_input(
+            tx,
+            conversation_id,
+            cast(UserInput, _plain(item.get("content"))),
+            _plain(item["entry"]) if "entry" in item else None,
+            now,
+        )
+        tx.place_submission(cast(SubmissionId, item["id"]), entry["id"])
+        placed.append(cast(SubmissionId, item["id"]))
+
+    for index in sorted([*writes, *users], reverse=True):
+        del items[index]
+    return BoundaryResult(users=placed, reset=reset)
+
+
+async def append_input(
+    tx: Transaction,
+    conversation_id: ConversationId,
+    content: UserInput,
+    entry: InputEntry | None,
+    now: int,
+) -> EntryRecord:
+    """Append the entry an input is placed as: ``entry``'s kind and data, or ``pi.user``."""
+
+    message = {"role": "user", "content": content, "timestamp": now}
+    if entry is None:
+        draft = cast(EntryDraft, {"model": [message]})
+        return await tx.append_entry(UserEntry, conversation_id, draft)
+    draft = cast(EntryDraft, {**pick_input_entry(entry), "model": [message]})
+    return await tx.append_entry(conversation_id, draft)
+
+
+def pick_input_entry(entry: InputEntry) -> InputEntry:
+    """Only ``kind`` and ``data`` of an input's entry, as strict JSON."""
+
+    picked: dict[str, JsonValue] = {"kind": entry.get("kind", "")}
+    if "data" in entry:
+        picked["data"] = copy_json(entry["data"])
+    return cast(InputEntry, picked)
+
+
+def is_display_only(entry: EntryDraft) -> bool:
+    """Whether a write changes nothing the model sees: no ``model``, ``head``, or ``edits``."""
+
+    return entry.get("model") is None and entry.get("head") is None and entry.get("edits") is None
+
+
+def has_queued_reset(items: Any) -> bool:
+    """Whether a reset (``head: "self"`` write) is queued."""
+
+    return any(item["mode"] == "write" and item["entry"].get("head") == "self" for item in items)
+
+
+def is_stale(boundary: Boundary, entry: EntryDraft) -> bool:
+    """Whether a head write targets an entry before the active range."""
+
+    head = entry.get("head")
+    return (
+        isinstance(head, int)
+        and not isinstance(head, bool)
+        and boundary.head is not None
+        and head < boundary.head
+    )
+
+
+async def remove_inbox_item(tx: Transaction, conversation_id: ConversationId, id: SubmissionId) -> None:
+    """Remove a withdrawn submission's item; the caller settles the submission."""
+
+    items = (await tx.doc(InboxDoc, conversation_id))["items"]
+    for index in range(len(items)):
+        if items[index]["id"] == id:
+            del items[index]
+            return
+
+
+async def withdraw_queued_inputs(tx: Transaction, conversation_id: ConversationId) -> None:
+    """Withdraw every queued input of one conversation; queued writes stay for later placement."""
+
+    items = (await tx.doc(InboxDoc, conversation_id))["items"]
+    for index in range(len(items) - 1, -1, -1):
+        item = items[index]
+        if item["mode"] == "write":
+            continue
+        tx.settle_submission(cast(SubmissionId, item["id"]), {"status": "unanswered", "reason": "aborted"})
+        del items[index]
+
+
+__all__ = [
+    "Boundary",
+    "BoundaryResult",
+    "InboxDoc",
+    "InboxItem",
+    "InboxState",
+    "QueueModes",
+    "append_input",
+    "apply_boundary",
+    "has_queued_reset",
+    "is_display_only",
+    "is_stale",
+    "pick_input_entry",
+    "prepare_boundary",
+    "remove_inbox_item",
+    "withdraw_queued_inputs",
+]

pi_durable/harness/live.py

@@ -0,0 +1,374 @@
+"""Built-in ``pi.live``: run control and presentation of current work (spec section 8.2).
+
+``run`` is present exactly while a run owns the conversation. Generation publishes its attempt,
+partial, retry, and deferred state here; a tool call publishes throttled running output into its
+slot, model-issued or nested. A compaction task lists its status here while it selects, summarizes,
+or waits out a retry backoff, and the newest failure stays for a client to show.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, NotRequired, TypedDict, cast
+
+from ..documents import define_doc
+from ..errors import error_message
+from ..ids import ConversationId, EntryId, SubmissionId, TaskId
+from ..json import JsonValue
+from ..session.drafts import DocumentDraft
+from ..session.transaction import Transaction
+from ..types import SubmissionSettlement, TaskRecord
+from .types import CompactionReason, ToolDiagnostic
+
+
+class RunState(TypedDict):
+    """The task that settles the run's inputs and those inputs; present exactly while busy."""
+
+    taskId: TaskId
+    inputs: list[SubmissionId]
+
+
+class SlotProgress(TypedDict, total=False):
+    """What a running tool call publishes, shared by model-issued and nested calls."""
+
+    callId: str
+    name: str
+    status: str
+    output: str
+    droppedBytes: int
+    droppedLines: int
+    details: JsonValue
+    diagnostics: list[ToolDiagnostic]
+
+
+class ToolSlot(SlotProgress):
+    """Presentation of one tool call of the current round.
+
+    ``taskId`` is absent for a call not started yet (a sequential round) and for a call its request
+    did not offer, which starts ``done`` with the ``entry`` generation wrote. ``entry`` is absent
+    when the tool task faulted or was orphaned.
+    """
+
+    taskId: NotRequired[TaskId]
+    entry: NotRequired[EntryId]
+
+
+class NestedToolSummary(TypedDict, total=False):
+    """How a nested call ended: an error flag, its execution time, spend, and bounded error text."""
+
+    isError: bool
+    durationMs: int
+    usage: dict[str, JsonValue]
+    error: str
+
+
+class NestedToolSlot(SlotProgress):
+    """Presentation of one nested call: a call a running tool made through ``executeTool()``."""
+
+    parentCallId: str
+    parentTaskId: TaskId
+    taskId: TaskId
+    arguments: dict[str, JsonValue]
+    summary: NotRequired[NestedToolSummary]
+
+
+class CompactionStatus(TypedDict, total=False):
+    """Presentation of one live compaction task (spec section 8.7)."""
+
+    taskId: TaskId
+    reason: CompactionReason
+    blocking: bool
+    attempt: int
+    retry: dict[str, JsonValue]
+
+
+class CompactionFailure(TypedDict):
+    """The newest compaction that failed since a compaction last produced a summary."""
+
+    taskId: TaskId
+    reason: CompactionReason
+    error: str
+
+
+class LiveState(TypedDict, total=False):
+    """Built-in live conversation state."""
+
+    run: RunState
+    generation: dict[str, JsonValue]
+    tools: list[ToolSlot]
+    nestedTools: list[NestedToolSlot]
+    compactions: list[CompactionStatus]
+    compactionFailure: CompactionFailure
+
+
+def _checkpoint_when(value: dict[str, Any], _ops: Any, _info: Any) -> bool:
+    """A complete base whenever nothing runs: no generation and no running tool slots."""
+
+    if value.get("generation") is not None:
+        return False
+    for key in ("tools", "nestedTools"):
+        slots = value.get(key) or []
+        if any(slot.get("status") == "running" for slot in slots):
+            return False
+    return True
+
+
+LiveDoc = define_doc(
+    {
+        "kind": "pi.live",
+        "version": 1,
+        "scope": "conversation",
+        "history": "latest",
+        "fork": "initial",
+        "initial": lambda: {},
+        "checkpoint_when": _checkpoint_when,
+    }
+)
+
+GENERATION_TASK_NAME = "pi.generation"
+TOOL_TASK_KIND = "pi.tool"
+COMPACTION_TASK_KIND = "pi.compaction"
+RUN_TASK_KINDS = frozenset({GENERATION_TASK_NAME})
+
+
+async def start_generation(
+    tx: Transaction,
+    conversation_id: ConversationId,
+    live: Any,
+    inputs: list[SubmissionId],
+    registry: Any,
+) -> None:
+    """Create the registry's generation task and give it run control of ``inputs``."""
+
+    task = registry.task(GENERATION_TASK_NAME)
+    if task is None:
+        raise RuntimeError("Registry lacks the pi.generation task")
+    await start_run(tx, conversation_id, live, inputs, task)
+
+
+async def start_run(
+    tx: Transaction,
+    conversation_id: ConversationId,
+    live: DocumentDraft,
+    inputs: list[SubmissionId],
+    task: Any,
+) -> None:
+    """Create the generation task and give it run control of ``inputs``."""
+
+    task_id = await tx.create_task(
+        task,
+        {},
+        {"ownership": {"kind": "conversation"}, "conversationId": conversation_id},
+    )
+    live["run"] = {"taskId": task_id, "inputs": list(inputs)}
+
+
+def end_run(
+    tx: Transaction,
+    live: DocumentDraft,
+    task_id: TaskId,
+    settlement: SubmissionSettlement,
+) -> None:
+    """End the run owned by ``task_id``: settle its inputs and remove run control."""
+
+    run = live.get("run")
+    if run is not None and run["taskId"] == task_id:
+        for id in run["inputs"]:
+            tx.settle_submission(cast(SubmissionId, id), settlement)
+        del live["run"]
+    for key in ("generation", "tools", "nestedTools"):
+        if key in live:
+            del live[key]
+
+
+def add_compaction_status(live: Any, status: Mapping[str, Any]) -> None:
+    """Add the status of a compaction task created in this commit; statuses stay in task ID order."""
+
+    statuses = live.get("compactions")
+    if statuses is None:
+        live["compactions"] = [dict(status)]
+    else:
+        statuses.append(dict(status))
+
+
+def compaction_status(live: Any, task_id: TaskId) -> Any:
+    """The status draft of compaction task ``task_id``, if listed."""
+
+    statuses = live.get("compactions")
+    if statuses is None:
+        return None
+    return next((status for status in statuses if status.get("taskId") == task_id), None)
+
+
+def remove_compaction_status(live: Any, task_id: TaskId) -> None:
+    """Remove the status of compaction task ``task_id``, and the list once empty."""
+
+    statuses = live.get("compactions")
+    if statuses is None:
+        return
+    for index in range(len(statuses)):
+        if statuses[index].get("taskId") == task_id:
+            del statuses[index]
+            break
+    if len(statuses) == 0:
+        del live["compactions"]
+
+
+def clear_compaction_failure(live: Any) -> None:
+    """Remove the recorded compaction failure, once a compaction produces a summary."""
+
+    if "compactionFailure" in live:
+        del live["compactionFailure"]
+
+
+def fail_compaction_status(live: Any, task_id: TaskId, error: str) -> None:
+    """Remove the status of compaction task ``task_id`` and record that it failed with ``error``."""
+
+    status = compaction_status(live, task_id)
+    if status is not None:
+        live["compactionFailure"] = {
+            "taskId": status["taskId"],
+            "reason": status["reason"],
+            "error": error,
+        }
+    remove_compaction_status(live, task_id)
+
+
+def tool_slot(live: Any, task_id: TaskId) -> Any:
+    """The slot of tool task ``task_id``, model-issued or nested, if ``pi.live`` still lists it.
+
+    Nested slots are found by binary search: a tool may make many nested calls, each of which looks
+    up its slot in several commits, and their task IDs ascend.
+    """
+
+    slots = live.get("tools") if hasattr(live, "get") else None
+    for slot in slots or ():
+        if slot.get("taskId") == task_id:
+            return slot
+    nested = live.get("nestedTools") if hasattr(live, "get") else None
+    if nested is None:
+        return None
+    low = 0
+    high = len(nested) - 1
+    while low <= high:
+        middle = (low + high) >> 1
+        candidate = nested[middle]
+        candidate_id = candidate.get("taskId")
+        if candidate_id == task_id:
+            return candidate
+        if candidate_id is not None and candidate_id < task_id:
+            low = middle + 1
+        else:
+            high = middle - 1
+    return None
+
+
+def finish_slot(slot: Any) -> None:
+    """Mark a slot done and clear its running output; the caller sets ``entry`` or ``summary``."""
+
+    slot["status"] = "done"
+    clear_progress(slot)
+
+
+def remove_nested_slots(live: Any, task_id: TaskId) -> None:
+    """Remove the nested slots below tool task ``task_id``, transitively; parents precede children."""
+
+    nested = live.get("nestedTools") if hasattr(live, "get") else None
+    if nested is None:
+        return
+    removed = {task_id}
+    for slot in nested:
+        if slot.get("parentTaskId") in removed:
+            removed.add(slot.get("taskId"))
+    for index in range(len(nested) - 1, -1, -1):
+        if nested[index].get("parentTaskId") in removed:
+            del nested[index]
+    if len(nested) == 0:
+        del live["nestedTools"]
+
+
+def clear_progress(slot: Any) -> None:
+    """Remove what a tool published while running; its result entry or a rerun replaces it."""
+
+    for key in ("output", "droppedBytes", "droppedLines", "details", "diagnostics"):
+        if key in slot:
+            del slot[key]
+
+
+async def settle_scheduler_outcome(tx: Transaction, record: TaskRecord, outcome: dict[str, Any]) -> None:
+    """Harness cleanup for a terminal outcome the scheduler writes itself (faulted or orphaned)."""
+
+    if record["kind"] == TOOL_TASK_KIND:
+        live = await tx.doc(LiveDoc, record["conversationId"])
+        slot = tool_slot(live, record["id"])
+        if slot is not None:
+            finish_slot(slot)
+        remove_nested_slots(live, record["id"])
+        return
+    if record["kind"] == COMPACTION_TASK_KIND:
+        failure_text = (
+            str(cast(Mapping[str, Any], outcome.get("error") or {}).get("message", ""))
+            if outcome.get("status") == "faulted"
+            else str(outcome.get("reason", ""))
+        )
+        fail_compaction_status(
+            await tx.doc(LiveDoc, record["conversationId"]), record["id"], failure_text
+        )
+        return
+    if record["kind"] not in RUN_TASK_KINDS:
+        return
+    live = await tx.doc(LiveDoc, record["conversationId"])
+    run = live.get("run")
+    if run is None or run["taskId"] != record["id"]:
+        return
+    from .generation import convert_partial
+
+    await convert_partial(tx, live, record["conversationId"])
+    if outcome.get("status") == "faulted":
+        error = outcome.get("error") or {}
+        end_run(
+            tx,
+            live,
+            record["id"],
+            {
+                "status": "unanswered",
+                "reason": "faulted",
+                "detail": error.get("message", error_message(error)),
+            },
+        )
+        return
+    end_run(
+        tx,
+        live,
+        record["id"],
+        {"status": "unanswered", "reason": outcome.get("reason", "aborted")},
+    )
+
+
+__all__ = [
+    "COMPACTION_TASK_KIND",
+    "GENERATION_TASK_NAME",
+    "CompactionFailure",
+    "CompactionStatus",
+    "LiveDoc",
+    "LiveState",
+    "NestedToolSlot",
+    "NestedToolSummary",
+    "RUN_TASK_KINDS",
+    "RunState",
+    "TOOL_TASK_KIND",
+    "ToolSlot",
+    "add_compaction_status",
+    "clear_compaction_failure",
+    "clear_progress",
+    "compaction_status",
+    "end_run",
+    "fail_compaction_status",
+    "finish_slot",
+    "remove_compaction_status",
+    "remove_nested_slots",
+    "settle_scheduler_outcome",
+    "start_generation",
+    "start_run",
+    "tool_slot",
+]