aid-runtime-python 0.1.39__py3-none-any.whl

aid_runtime/agents/memory.py

@@ -0,0 +1,241 @@
+"""Agent memory: thread scope, and one durable adapter.
+
+Ratified decision D5 draws the line this module implements: in-process memory is
+the default, exactly one store-backed adapter ships, and there is no plugin zoo.
+A ``backend`` naming a declared store is served by whatever store the caller
+registered under that id; with none registered, the in-process store answers.
+
+``thread`` scope is the checkpointer unit -- the conversation the loop is already
+carrying -- so it needs no store and no key. ``user`` and ``global`` scope are the
+cross-run facts, keyed by a caller-supplied key and stored behind
+:class:`MemoryStore`. The IR vocabulary is ``thread``/``user``/``global``; the
+``session`` spelling L1 removed is not reintroduced here.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import tempfile
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Protocol, runtime_checkable
+
+from aid_runtime.cost import estimate_tokens
+from aid_runtime.errors import ConfigurationError, FeatureDeferredError
+from aid_runtime.types import Message
+
+__all__ = [
+    "IMPLEMENTED_MEMORY_KINDS",
+    "MEMORY_SCOPES",
+    "FileMemoryStore",
+    "InProcessMemoryStore",
+    "MemoryStore",
+    "ThreadMemory",
+    "check_memory_supported",
+    "render_context",
+]
+
+#: The IR's scope vocabulary. ``session`` is deliberately absent.
+MEMORY_SCOPES = ("thread", "user", "global")
+
+#: The kinds served without a summarizer or an embedding model. The rest are
+#: declared in the IR and refused by name, never silently downgraded.
+IMPLEMENTED_MEMORY_KINDS = ("buffer", "window")
+
+
+@runtime_checkable
+class MemoryStore(Protocol):
+    """Append-only, scoped storage behind the memory scopes that outlive a run."""
+
+    async def read(self, *, scope: str, key: str) -> tuple[str, ...]:
+        """Every value recorded for ``(scope, key)``, oldest first."""
+        ...
+
+    async def append(self, *, scope: str, key: str, value: str) -> None:
+        """Record one value. Append-only is what makes replay deterministic."""
+        ...
+
+
+class InProcessMemoryStore:
+    """The default: a dict, lost when the process ends.
+
+    That is the honest default for a library that must not decide where a host
+    application keeps its data. A durable adapter is opt-in, and there is one.
+    """
+
+    def __init__(self) -> None:
+        self._values: dict[tuple[str, str], list[str]] = {}
+
+    async def read(self, *, scope: str, key: str) -> tuple[str, ...]:
+        return tuple(self._values.get((scope, key), ()))
+
+    async def append(self, *, scope: str, key: str, value: str) -> None:
+        self._values.setdefault((scope, key), []).append(value)
+
+    def clear(self) -> None:
+        self._values.clear()
+
+
+class FileMemoryStore:
+    """One JSON file, one lock-free writer.
+
+    The file is rewritten through a temporary file and an atomic replace so a crash
+    mid-write cannot leave a half-file that a later read would parse as truth.
+    Writes are serialized by an ``asyncio.Lock``; the file I/O itself runs in a
+    worker thread so the loop is not blocked by disk.
+    """
+
+    def __init__(self, path: str | Path) -> None:
+        self._path = Path(path)
+        self._lock = asyncio.Lock()
+
+    @property
+    def path(self) -> Path:
+        return self._path
+
+    async def read(self, *, scope: str, key: str) -> tuple[str, ...]:
+        data = await asyncio.to_thread(self._load)
+        return tuple(data.get(scope, {}).get(key, ()))
+
+    async def append(self, *, scope: str, key: str, value: str) -> None:
+        async with self._lock:
+            data = await asyncio.to_thread(self._load)
+            data.setdefault(scope, {}).setdefault(key, []).append(value)
+            await asyncio.to_thread(self._save, data)
+
+    def _load(self) -> dict[str, dict[str, list[str]]]:
+        if not self._path.exists():
+            return {}
+        try:
+            raw = json.loads(self._path.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError) as exc:
+            raise ConfigurationError(f"memory store {self._path} is unreadable: {exc}") from exc
+        if not isinstance(raw, dict):
+            raise ConfigurationError(f"memory store {self._path} is not a JSON object")
+        return {
+            str(scope): {str(key): [str(value) for value in values] for key, values in keys.items()}
+            for scope, keys in raw.items()
+            if isinstance(keys, dict)
+        }
+
+    def _save(self, data: dict[str, dict[str, list[str]]]) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        handle, temporary = tempfile.mkstemp(
+            dir=self._path.parent, prefix=f".{self._path.name}.", suffix=".tmp"
+        )
+        try:
+            with os.fdopen(handle, "w", encoding="utf-8") as stream:
+                json.dump(data, stream, sort_keys=True, separators=(",", ":"))
+            Path(temporary).replace(self._path)
+        except OSError:
+            Path(temporary).unlink(missing_ok=True)
+            raise
+
+
+def check_memory_supported(
+    *,
+    memory_id: str,
+    kind: str,
+    summarize_model: str | None,
+    ttl: str | None,
+    redact: Sequence[str],
+) -> None:
+    """Refuse a declared memory the runtime cannot honour exactly.
+
+    ``summary``, ``semantic``, ``episodic`` and ``vector`` would each need a model
+    or an index that L2 does not have, and so would ``summarizeModel``; ``ttl``
+    needs a clock and ``redact`` needs redaction on the tracing path. Serving them
+    approximately would make the declaration lie, so they are refused by name.
+    """
+    if kind not in IMPLEMENTED_MEMORY_KINDS:
+        raise FeatureDeferredError(
+            f"memory {memory_id!r} asks for kind {kind!r}; L2 implements "
+            f"{', '.join(IMPLEMENTED_MEMORY_KINDS)} and refuses the rest by name"
+        )
+    if summarize_model is not None:
+        raise FeatureDeferredError(
+            f"memory {memory_id!r} names summarizeModel {summarize_model!r}; "
+            "no summarizer is implemented, so an unbounded buffer is served instead"
+        )
+    if ttl is not None:
+        raise FeatureDeferredError(
+            f"memory {memory_id!r} declares ttl {ttl!r}; expiry enforcement is not implemented"
+        )
+    if redact:
+        raise FeatureDeferredError(
+            f"memory {memory_id!r} declares redaction rules; memory redaction is not implemented"
+        )
+
+
+class ThreadMemory:
+    """Bounds the transcript a single agent turn may replay.
+
+    ``thread`` scope is the conversation the loop is already carrying, so this is
+    a trimmer rather than a store: when the transcript exceeds ``token_budget``,
+    whole tool-call groups are dropped oldest-first. A group is an assistant turn
+    with its tool observations, because dropping one without the other would
+    replay a tool result whose call is gone.
+
+    A message larger than the whole budget is kept rather than truncated: half a
+    tool observation is worse than an over-budget one, and the budget is a target
+    for what to *replay*, not a hard accounting ceiling.
+    """
+
+    def __init__(self, token_budget: int | None = None) -> None:
+        if token_budget is not None and token_budget < 1:
+            raise ConfigurationError("thread memory tokenBudget must be positive")
+        self._token_budget = token_budget
+
+    @property
+    def token_budget(self) -> int | None:
+        return self._token_budget
+
+    def tokens(self, messages: Sequence[Message]) -> int:
+        return sum(estimate_tokens(message.content) for message in messages)
+
+    def trim(self, messages: Sequence[Message]) -> tuple[Message, ...]:
+        if self._token_budget is None:
+            return tuple(messages)
+        kept = list(messages)
+        total = self.tokens(kept)
+        start = next(
+            (index for index, message in enumerate(kept) if message.tool_calls),
+            len(kept),
+        )
+        while total > self._token_budget and start < len(kept):
+            end = start + 1
+            while end < len(kept) and kept[end].role == "tool":
+                end += 1
+            total -= self.tokens(kept[start:end])
+            del kept[start:end]
+        while total > self._token_budget:
+            victim = next(
+                (
+                    index
+                    for index, message in enumerate(kept)
+                    if message.role != "system" and index < len(kept) - 1
+                ),
+                None,
+            )
+            if victim is None:
+                break
+            total -= estimate_tokens(kept[victim].content)
+            del kept[victim]
+        return tuple(kept)
+
+
+def render_context(entries: Sequence[str], token_budget: int) -> tuple[str, int]:
+    """Fit stored entries into ``token_budget``, newest first. Returns text, tokens."""
+    chosen: list[str] = []
+    tokens = 0
+    for entry in reversed(entries):
+        cost = estimate_tokens(entry) + 1
+        if chosen and tokens + cost > token_budget:
+            break
+        chosen.append(entry)
+        tokens += cost
+        if tokens >= token_budget:
+            break
+    return "\n".join(reversed(chosen)), tokens

aid_runtime/agents/policy.py

@@ -0,0 +1,163 @@
+"""Evaluation of the IR's ``PolicyExpr`` AST.
+
+The grammar is the one ``packages/ir/src/policy`` parses, and the AST is plain
+JSON precisely so a target does not need a second parser:
+
+    expression := or
+    or         := and ('or' and)*
+    and        := comparison ('and' comparison)*
+    comparison := unary (('==' | '!=' | '<' | '<=' | '>' | '>=') unary)?
+    unary      := 'not' unary | primary
+    primary    := 'true' | 'false' | string | number | call | '(' expression ')'
+    call       := builtin '(' (operand (',' operand)*)? ')'
+
+The builtins are the closed set ``hasRole``, ``hasScope``, ``isAuthenticated``, so
+the context this evaluates against is a :class:`Principal` -- who is asking --
+and not loop state such as a step counter. Anything the evaluator cannot decide
+raises: for a human-in-the-loop gate, "could not decide" must never degrade into
+"do not ask".
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from aid_runtime.errors import PolicyEvaluationError
+
+__all__ = ["Principal", "evaluate_policy", "policy_source"]
+
+_COMPARISONS = frozenset({"==", "!=", "<", "<=", ">", ">="})
+_LOGICAL = frozenset({"and", "or"})
+
+
+@dataclass(frozen=True, slots=True)
+class Principal:
+    """Who a policy is evaluated for.
+
+    Defaults are the anonymous-but-authenticated caller an app with no auth layer
+    has: ``isAuthenticated()`` is true and both role and scope sets are empty, so
+    ``hasRole('admin')`` is false rather than undefined.
+    """
+
+    roles: frozenset[str] = frozenset()
+    scopes: frozenset[str] = frozenset()
+    authenticated: bool = True
+
+
+def evaluate_policy(node: Mapping[str, Any], principal: Principal) -> bool:
+    """Evaluate a parsed policy against ``principal``. Never returns "maybe"."""
+    value = _evaluate(node, principal)
+    if not isinstance(value, bool):
+        raise PolicyEvaluationError(
+            f"policy did not evaluate to a boolean, got {type(value).__name__}: "
+            f"{policy_source(node)}"
+        )
+    return value
+
+
+def policy_source(node: Mapping[str, Any]) -> str:
+    """Render an AST back to source, so a trace or error can quote the policy."""
+    kind = node.get("kind")
+    if kind in {"string", "number", "boolean"}:
+        return _literal_source(node.get("value"))
+    if kind == "call":
+        arguments = ", ".join(policy_source(argument) for argument in node.get("args", ()))
+        return f"{node.get('name')}({arguments})"
+    if kind == "comparison":
+        return f"{policy_source(node['left'])} {node.get('op')} {policy_source(node['right'])}"
+    if kind in _LOGICAL:
+        return f"({policy_source(node['left'])} {kind} {policy_source(node['right'])})"
+    if kind == "not":
+        return f"not {policy_source(node['operand'])}"
+    raise PolicyEvaluationError(f"unsupported policy node kind {kind!r}")
+
+
+def _literal_source(value: Any) -> str:
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, str):
+        return f'"{value}"'
+    return str(value)
+
+
+def _evaluate(node: Mapping[str, Any], principal: Principal) -> Any:
+    kind = node.get("kind")
+    if kind in {"string", "number", "boolean"}:
+        return node.get("value")
+    if kind == "call":
+        return _call(node, principal)
+    if kind == "comparison":
+        return _compare(node, principal)
+    if kind in _LOGICAL:
+        return _logic(kind, node, principal)
+    if kind == "not":
+        return not _as_bool(_evaluate(node["operand"], principal), node)
+    raise PolicyEvaluationError(f"unsupported policy node kind {kind!r}")
+
+
+def _call(node: Mapping[str, Any], principal: Principal) -> bool:
+    name = node.get("name")
+    arguments = [_evaluate(argument, principal) for argument in node.get("args", ())]
+    if name == "isAuthenticated":
+        if arguments:
+            raise PolicyEvaluationError("isAuthenticated() takes no arguments")
+        return principal.authenticated
+    if name == "hasRole":
+        return _membership(name, arguments, principal.roles)
+    if name == "hasScope":
+        return _membership(name, arguments, principal.scopes)
+    raise PolicyEvaluationError(f"unknown policy builtin {name!r}")
+
+
+def _membership(name: str, arguments: list[Any], held: frozenset[str]) -> bool:
+    if len(arguments) != 1:
+        raise PolicyEvaluationError(f"{name}() takes exactly one argument")
+    (wanted,) = arguments
+    if not isinstance(wanted, str):
+        raise PolicyEvaluationError(f"{name}() argument must be a string, got {wanted!r}")
+    return wanted in held
+
+
+def _compare(node: Mapping[str, Any], principal: Principal) -> bool:
+    operator = node.get("op")
+    if operator not in _COMPARISONS:
+        raise PolicyEvaluationError(f"unknown comparison operator {operator!r}")
+    left = _evaluate(node["left"], principal)
+    right = _evaluate(node["right"], principal)
+    if operator == "==":
+        return bool(left == right)
+    if operator == "!=":
+        return bool(left != right)
+    try:
+        if operator == "<":
+            return bool(left < right)
+        if operator == "<=":
+            return bool(left <= right)
+        if operator == ">":
+            return bool(left > right)
+        return bool(left >= right)
+    except TypeError as exc:
+        raise PolicyEvaluationError(
+            f"cannot compare {type(left).__name__} {operator} {type(right).__name__}: "
+            f"{policy_source(node)}"
+        ) from exc
+
+
+def _logic(kind: str, node: Mapping[str, Any], principal: Principal) -> bool:
+    left = _as_bool(_evaluate(node["left"], principal), node)
+    # ``and``/``or`` short-circuit on the left operand's truth value.
+    if kind == "and" and not left:
+        return False
+    if kind == "or" and left:
+        return True
+    return _as_bool(_evaluate(node["right"], principal), node)
+
+
+def _as_bool(value: Any, node: Mapping[str, Any]) -> bool:
+    if not isinstance(value, bool):
+        raise PolicyEvaluationError(
+            f"expected a boolean operand in {policy_source(node)}, got {value!r}"
+        )
+    return value

aid_runtime/agents/state.py

@@ -0,0 +1,182 @@
+"""The value types of a run: its steps, its result, and its pause.
+
+These live together because they are one shape: a checkpoint *is* a result that
+stopped early, and a resumed result is a checkpoint plus more steps. Keeping them
+in one module is also what lets ``agents.loop`` and a caller share the types
+without an import cycle.
+
+:class:`AgentCheckpoint` carries the steps completed so far, so a run that paused
+and resumed reports the whole transcript of steps rather than only the part that
+happened after the resume -- a trace that forgets the pre-pause steps is not a
+trace of the run.
+
+``caller`` is the frame that was awaiting a *delegated* result when the pause
+happened inside the delegate. Transfer has no caller to return to, which is the
+distinction the loop is built on; a transfer checkpoint never has one.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from aid_runtime.agents.guardrails import GuardrailEscalation, GuardrailRecord
+from aid_runtime.types import Message, ToolCall
+
+__all__ = [
+    "AgentCheckpoint",
+    "HumanDecision",
+    "PendingApproval",
+    "RunResult",
+    "RunStatus",
+    "StepKind",
+    "StepRecord",
+]
+
+
+class RunStatus(StrEnum):
+    """How a run ended.
+
+    ``PAUSED`` is an outcome, not an error. So is ``BLOCKED``: a human rejected an
+    escalation the guardrails raised, which is a decision the run respected rather
+    than a failure it suffered.
+    """
+
+    COMPLETED = "completed"
+    PAUSED = "paused"
+    MAX_STEPS = "max-steps"
+    BLOCKED = "blocked"
+
+
+class StepKind(StrEnum):
+    """What one model step did with its output."""
+
+    ANSWER = "answer"
+    TOOLS = "tools"
+    TRANSFER = "transfer"
+    DELEGATE = "delegate"
+
+
+class StepRecord(BaseModel):
+    """One model call, and what was done with its output.
+
+    ``degraded_from`` is set only on the step that stepped down the alias ladder,
+    so "did this run degrade?" is a question the trace answers without arithmetic.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    index: int = Field(ge=1)
+    agent_id: str
+    kind: StepKind
+    alias: str
+    model: str
+    cost_usd: float = Field(ge=0.0)
+    total_tokens: int = Field(ge=0)
+    tool_names: tuple[str, ...] = ()
+    target: str | None = None
+    degraded_from: str | None = None
+
+
+class HumanDecision(StrEnum):
+    """What a human answered. Rejection is a normal answer, not an error."""
+
+    APPROVE = "approve"
+    REJECT = "reject"
+
+
+class AgentCheckpoint(BaseModel):
+    """Everything a paused agent needs to continue, and nothing that is derived."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    agent_id: str = Field(min_length=1)
+    step: int = Field(ge=0)
+    messages: tuple[Message, ...]
+    steps: tuple[StepRecord, ...] = ()
+    calls: tuple[ToolCall, ...] = ()
+    channel: str = ""
+    when: dict[str, Any] = Field(default_factory=dict)
+    # `approve`: `calls` are awaiting a person. `delegate`: `calls` holds the one
+    # control call whose delegate paused, and the resumed frame must append the
+    # delegate's answer rather than execute anything.
+    resume_kind: Literal["approve", "delegate"] = "approve"
+    cost_anchor_usd: float = Field(ge=0.0)
+    token_anchor: int = Field(ge=0)
+    spent_usd: float = Field(ge=0.0)
+    spent_tokens: int = Field(ge=0)
+    degraded: bool = False
+    # The alias the paused frame was actually calling. A resumed run has to keep
+    # it: `degraded` records that this run already stepped down, and resuming on
+    # the dear alias would silently undo a bounded decision the trace still shows.
+    alias: str = ""
+    path: tuple[str, ...] = ()
+    memory_key: str = ""
+    caller: AgentCheckpoint | None = None
+    # Set when a guardrail escalated. Its `stage` tells the resume path what to do:
+    # an `input` escalation has already been checked, an `output` one holds the
+    # answer to release on approval, a `tool-call` one holds `calls` to execute.
+    escalation: GuardrailEscalation | None = None
+    guardrails: tuple[GuardrailRecord, ...] = ()
+
+
+class PendingApproval(BaseModel):
+    """The question a pause is asking, separated from the resume state."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    agent_id: str
+    channel: str
+    step: int
+    calls: tuple[ToolCall, ...]
+    when: dict[str, Any] = Field(default_factory=dict)
+    path: tuple[str, ...] = ()
+    escalation: GuardrailEscalation | None = None
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint: AgentCheckpoint) -> PendingApproval:
+        return cls(
+            agent_id=checkpoint.agent_id,
+            channel=checkpoint.channel,
+            step=checkpoint.step,
+            calls=checkpoint.calls,
+            when=checkpoint.when,
+            path=checkpoint.path,
+            escalation=checkpoint.escalation,
+        )
+
+
+class RunResult(BaseModel):
+    """What a run produced -- including when it produced no answer."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    status: RunStatus
+    agent_id: str
+    answer: str = ""
+    steps: tuple[StepRecord, ...] = ()
+    path: tuple[str, ...] = ()
+    cost_usd: float = Field(default=0.0, ge=0.0)
+    total_tokens: int = Field(default=0, ge=0)
+    checkpoint: AgentCheckpoint | None = None
+    guardrails: tuple[GuardrailRecord, ...] = ()
+
+    @property
+    def completed(self) -> bool:
+        return self.status is RunStatus.COMPLETED
+
+    @property
+    def degraded(self) -> tuple[tuple[str, str], ...]:
+        """``(from, to)`` for every step that stepped down the alias ladder."""
+        pairs: list[tuple[str, str]] = []
+        for step in self.steps:
+            if step.degraded_from is not None:
+                pairs.append((step.degraded_from, step.alias))
+        return tuple(pairs)
+
+    @property
+    def approval(self) -> PendingApproval | None:
+        """The question a paused run is asking, or ``None`` when it did not pause."""
+        return None if self.checkpoint is None else PendingApproval.from_checkpoint(self.checkpoint)