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

pi_durable/harness/usage.py

@@ -0,0 +1,88 @@
+"""Built-in ``pi.usage``: the ledger of one conversation's own spend.
+
+Port of v1.3.0 ``src/harness/usage.ts``. Every built-in writer of an assistant or tool-result entry
+adds that message's usage in the same commit; compaction summarization attempts add theirs without
+an entry. ``Harness.usage()`` sums every conversation's document into the Session total.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, MutableMapping
+from typing import Any, TypedDict, cast
+
+from ..documents import define_doc
+from ..ids import ConversationId
+from ..json import JsonValue, copy_json
+from ..session.transaction import Transaction
+
+
+class UsageState(TypedDict):
+    """Spend keyed by ``provider/modelId`` and by tool name."""
+
+    models: dict[str, dict[str, JsonValue]]
+    tools: dict[str, dict[str, JsonValue]]
+
+
+UsageDoc = define_doc(
+    {
+        "kind": "pi.usage",
+        "version": 1,
+        "scope": "conversation",
+        "history": "latest",
+        "fork": "initial",
+        "initial": lambda: {"models": {}, "tools": {}},
+        "checkpoint_when": lambda _state, _ops, _info: True,
+    }
+)
+
+
+async def record_usage(
+    tx: Transaction,
+    conversation_id: ConversationId,
+    bucket: str,
+    key: str,
+    usage: Mapping[str, Any],
+) -> None:
+    """Add ``usage`` to one bucket of the conversation's ``pi.usage``."""
+
+    totals = (await tx.doc(UsageDoc, conversation_id))[bucket]
+    current = totals.get(key) if hasattr(totals, "get") else None
+    if current is None:
+        totals[key] = _strict_usage(usage)
+    else:
+        add_usage(cast(MutableMapping[str, Any], current), usage)
+
+
+def add_usage(total: MutableMapping[str, Any], usage: Mapping[str, Any]) -> None:
+    """Add every counter of ``usage`` to ``total``; optional counters add once either side reports them."""
+
+    for key in ("input", "output", "cacheRead", "cacheWrite", "totalTokens"):
+        total[key] = int(total.get(key, 0)) + int(usage.get(key, 0))
+    for key in ("cacheWrite1h", "reasoning"):
+        if usage.get(key) is not None:
+            total[key] = int(total.get(key, 0)) + int(usage[key])
+    if "cost" not in total:
+        total["cost"] = {}
+    total_cost = cast(MutableMapping[str, Any], total["cost"])
+    usage_cost = cast(Mapping[str, Any], usage.get("cost") or {})
+    for key in ("input", "output", "cacheRead", "cacheWrite", "total"):
+        total_cost[key] = total_cost.get(key, 0) + usage_cost.get(key, 0)
+
+
+def add_usage_state(sum_value: UsageState, state: Mapping[str, Any]) -> None:
+    """Add every bucket of ``state`` into ``sum``."""
+
+    for bucket in ("models", "tools"):
+        destination = sum_value[bucket]
+        for key, value in cast(Mapping[str, Any], state.get(bucket) or {}).items():
+            if key in destination:
+                add_usage(cast(MutableMapping[str, Any], destination[key]), value)
+            else:
+                destination[key] = cast(dict[str, JsonValue], copy_json(value))
+
+
+def _strict_usage(usage: Mapping[str, Any]) -> dict[str, JsonValue]:
+    return cast(dict[str, JsonValue], copy_json(dict(usage)))
+
+
+__all__ = ["UsageDoc", "UsageState", "add_usage", "add_usage_state", "record_usage"]

pi_durable/harness/util.py

@@ -0,0 +1,204 @@
+"""Scheduler support: pending waits, paginated scans, and closed-Harness errors."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from collections.abc import Awaitable, Callable, Hashable, Mapping
+from typing import Any, Generic, TypeVar, cast
+
+from ..context import AbortError, AbortSignal, Context
+from ..errors import SessionFailed
+from ..json import JsonObject
+from ..types import Cursor, Page, TaskRecord
+
+T = TypeVar("T")
+K = TypeVar("K", bound=Hashable)
+
+
+def cancellation_error(signal: AbortSignal) -> BaseException:
+    """Return what an aborted context rejects with: its reason if it is an error, else ``AbortError``."""
+
+    reason = signal.reason
+    if isinstance(reason, BaseException):
+        return reason
+    return AbortError(reason)
+
+
+class Waiters(Generic[K, T]):
+    """Pending waits by key. Each settles once: through ``resolve``, ``reject_all``, or its cancellation."""
+
+    def __init__(self) -> None:
+        self._sets: dict[K, set[asyncio.Future[T]]] = {}
+
+    def add(self, key: K, context: Context) -> Awaitable[T]:
+        signal = context.abort_signal
+        if signal is not None and signal.aborted:
+            return rejected(cancellation_error(signal))
+        loop = asyncio.get_running_loop()
+        waiter: asyncio.Future[T] = loop.create_future()
+        waiter.add_done_callback(lambda future: future.exception() if not future.cancelled() else None)
+        self._sets.setdefault(key, set()).add(waiter)
+        if signal is not None:
+
+            def cancel() -> None:
+                values = self._sets.get(key)
+                if values is not None:
+                    values.discard(waiter)
+                    if not values and self._sets.get(key) is values:
+                        del self._sets[key]
+                if not waiter.done():
+                    waiter.set_exception(cancellation_error(signal))
+
+            remove = signal.add_listener(cancel)
+            waiter.add_done_callback(lambda _: remove())
+        return waiter
+
+    def keys(self) -> list[K]:
+        return list(self._sets)
+
+    def resolve(self, key: K, value: T) -> None:
+        values = self._sets.pop(key, None)
+        for waiter in values or ():
+            if not waiter.done():
+                waiter.set_result(value)
+
+    def reject_all(self, error: BaseException) -> None:
+        sets = list(self._sets.values())
+        self._sets.clear()
+        for values in sets:
+            for waiter in values:
+                if not waiter.done():
+                    waiter.set_exception(error)
+
+
+def rejected(error: BaseException) -> asyncio.Future[Any]:
+    future: asyncio.Future[Any] = asyncio.get_running_loop().create_future()
+    future.set_exception(error)
+    future.add_done_callback(lambda f: f.exception() if not f.cancelled() else None)
+    return future
+
+
+def resolved(value: T) -> asyncio.Future[T]:
+    future: asyncio.Future[T] = asyncio.get_running_loop().create_future()
+    future.set_result(value)
+    return future
+
+
+async def scan_all(scan: Callable[[Cursor | None], Awaitable[Page[Any, Cursor]]]) -> list[Any]:
+    """Every item of a paginated scan, in page order."""
+
+    items: list[Any] = []
+    cursor: Cursor | None = None
+    while True:
+        page = await scan(cursor)
+        items.extend(page["items"])
+        cursor = page.get("next")
+        if cursor is None:
+            return items
+
+
+def closed_error(session: Any) -> BaseException:
+    """What a closed Harness rejects with: ``SessionFailed`` when it closed because the Session failed."""
+
+    failure = session.failure
+    return RuntimeError("Harness is closed") if failure is None else SessionFailed(failure)
+
+
+def call_with_args(handler: Callable[..., Any], *args: Any) -> Any:
+    """Call ``handler`` with only the arguments it declares, as JavaScript ignores extras."""
+
+    try:
+        parameters = list(inspect.signature(handler).parameters.values())
+    except (TypeError, ValueError):
+        return handler(*args)
+    if any(parameter.kind is inspect.Parameter.VAR_POSITIONAL for parameter in parameters):
+        return handler(*args)
+    positional = [
+        parameter
+        for parameter in parameters
+        if parameter.kind
+        in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+    ]
+    return handler(*args[: len(positional)])
+
+
+def assign_json(target: Any, key: Any, value: Any) -> None:
+    """Assign ``value`` at ``target[key]`` leaf by leaf, so a string leaf becomes an append."""
+
+    slots = target
+    current = slots.get(key) if hasattr(slots, "get") else slots[key]
+    if _is_record(current) and _is_record(value):
+        for name in list(current.keys()):
+            if name not in value:
+                del current[name]
+        for name, child in value.items():
+            assign_json(current, name, child)
+        return
+    if _is_array(current) and _is_array(value) and len(current) <= len(value):
+        for index in range(len(value)):
+            if index < len(current):
+                assign_json(current, index, value[index])
+            else:
+                current.append(value[index])
+        return
+    if current != value:
+        slots[key] = value
+
+
+def _is_record(value: Any) -> bool:
+    from ..session.drafts import DraftObject
+
+    return isinstance(value, (dict, DraftObject))
+
+
+def _is_array(value: Any) -> bool:
+    from ..session.drafts import DraftList
+
+    return isinstance(value, (list, DraftList))
+
+
+#: Host-facing settings use PEP 8 names; provider options and persisted checkpoints use the
+#: camelCase pi-ai request shape.
+_STREAM_OPTION_KEYS = {
+    "timeout_ms": "timeoutMs",
+    "max_retries": "maxRetries",
+    "max_retry_delay_ms": "maxRetryDelayMs",
+    "cache_retention": "cacheRetention",
+}
+
+
+def task_checkpoint(task: TaskRecord) -> dict[str, Any]:
+    """The checkpoint a task's current phase reads."""
+
+    return cast(dict[str, Any], task["state"]["checkpoint"])
+
+
+async def call_hook(handler: Callable[..., Any], *args: Any) -> Any:
+    """Call a hook with only the arguments it declares and await a returned awaitable."""
+
+    result = call_with_args(handler, *args)
+    if hasattr(result, "__await__"):
+        return await cast(Awaitable[Any], result)
+    return result
+
+
+def request_stream_options(stream: Mapping[str, Any]) -> JsonObject:
+    """Map Python-facing stream settings to the camelCase provider request shape."""
+
+    return cast(JsonObject, {_STREAM_OPTION_KEYS.get(key, key): value for key, value in stream.items()})
+
+
+__all__ = [
+    "Waiters",
+    "assign_json",
+    "call_hook",
+    "call_with_args",
+    "cancellation_error",
+    "closed_error",
+    "rejected",
+    "request_stream_options",
+    "resolved",
+    "scan_all",
+    "task_checkpoint",
+]