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

aid_runtime/cache.py

@@ -0,0 +1,185 @@
+"""A completion cache with exact (prompt-keyed) and opt-in semantic reuse.
+
+The exact key is ``(model, prompt transcript hash, tool identity hash, temperature,
+max output tokens, stop sequences)``. Semantic reuse is strictly opt-in (an explicit
+threshold and TTL; there is no default) and is gated by the *same* structural key: a
+stored entry may only answer a new request whose model, tool identity, temperature,
+output cap and stop sequences all match, and then only when the prompt text is
+similar enough and young enough.
+
+The invalidation rule, stated: an entry is reusable if and only if (a) the model,
+tool-identity hash, temperature, output cap and stop sequences are identical, (b)
+the Jaccard similarity of the two prompt transcripts is at least ``threshold``, and
+(c) the entry is younger than ``ttl_seconds``. Text similarity is the *match* test,
+never the storage key: the key stays a structural hash so a prompt change can never
+resurrect a response that was computed with different tools, a different model, a
+different cap, or different stop sequences.
+
+The similarity measure is lexical (token Jaccard), not an embedding cosine. A true
+embedding similarity would need an embedder at cache-construction time; wiring that
+in is deferred rather than implied.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import time
+from dataclasses import dataclass, field
+
+from aid_runtime.types import CompletionRequest, CompletionResponse, ModelSpec
+
+__all__ = ["CompletionCache", "cache_key_for"]
+
+
+@dataclass(frozen=True, slots=True)
+class CacheKey:
+    """The structural key: everything in a request that can change its output.
+
+    Model, transcript hash, full tool identity, temperature, output cap, and stop
+    sequences. ``tools_hash`` covers each tool's name, description, and parameters
+    in request order, since a provider preserves that order. The cap and stop
+    sequences are stored verbatim rather than omitted: a cached body must never
+    answer a request that asked for a different cap or stop set.
+    """
+
+    model: str
+    prompt_hash: str
+    tools_hash: str
+    temperature: str
+    max_output_tokens: str
+    stop: str
+
+    def same_structure(self, other: CacheKey) -> bool:
+        """Whether two keys may share a completion, ignoring prompt similarity."""
+        return (
+            self.model == other.model
+            and self.tools_hash == other.tools_hash
+            and self.temperature == other.temperature
+            and self.max_output_tokens == other.max_output_tokens
+            and self.stop == other.stop
+        )
+
+
+def cache_key_for(model: str, request: CompletionRequest) -> CacheKey:
+    """Derive the structural key for a request.
+
+    "Prompt hash" hashes the full transcript (roles, content, and tool context),
+    not concatenated text, so a system/user/tool distinction is part of the key.
+    The output cap and stop sequences are part of the key too, so a cached body
+    cannot answer a request that asked for a different cap or stop set.
+    """
+    return CacheKey(
+        model=model,
+        prompt_hash=_transcript_hash(request),
+        tools_hash=_tools_hash(request),
+        temperature=repr(request.temperature),
+        max_output_tokens=repr(request.max_output_tokens),
+        stop=repr(request.stop),
+    )
+
+
+@dataclass(slots=True)
+class CacheEntry:
+    response: CompletionResponse
+    prompt: str
+    stored_at: float = field(default_factory=time.monotonic)
+
+
+class CompletionCache:
+    """An in-process, async-safe cache of provider completions.
+
+    Async-safe by construction: lookups and inserts are synchronous and complete
+    before the first ``await`` in the caller, and a caller never observes a
+    half-populated entry because an entry is written whole, once, after the call
+    completes. In-flight identical requests are *not* coalesced: each miss calls
+    the provider. That is the explicit choice here -- coalescing would add a
+    per-key condition variable and change failure semantics -- and a test pins the
+    duplicate-work behaviour.
+    """
+
+    def __init__(self) -> None:
+        self._entries: dict[CacheKey, CacheEntry] = {}
+
+    def get(self, key: CacheKey) -> CompletionResponse | None:
+        entry = self._entries.get(key)
+        return None if entry is None else entry.response
+
+    def semantic_get(
+        self,
+        spec: ModelSpec,
+        request: CompletionRequest,
+        *,
+        threshold: float,
+        ttl_seconds: int,
+    ) -> CompletionResponse | None:
+        """The opt-in semantic lookup: same structure, similar prompt, young enough."""
+        key = cache_key_for(spec.name, request)
+        prompt = _transcript_text(request)
+        now = time.monotonic()
+        for existing_key, entry in self._entries.items():
+            if not existing_key.same_structure(key):
+                continue
+            if now - entry.stored_at > ttl_seconds:
+                continue
+            if _jaccard(prompt, entry.prompt) < threshold:
+                continue
+            return entry.response
+        return None
+
+    def put(self, key: CacheKey, response: CompletionResponse, request: CompletionRequest) -> None:
+        self._entries[key] = CacheEntry(response=response, prompt=_transcript_text(request))
+
+    def clear(self) -> None:
+        self._entries.clear()
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+
+def _transcript_text(request: CompletionRequest) -> str:
+    return "\n".join(message.content for message in request.messages)
+
+
+def _transcript_hash(request: CompletionRequest) -> str:
+    messages = [
+        {
+            "role": message.role,
+            "content": message.content,
+            "tool_call_id": message.tool_call_id,
+            "name": message.name,
+            "tool_calls": [call.model_dump(mode="json") for call in message.tool_calls],
+        }
+        for message in request.messages
+    ]
+    blob = json.dumps(messages, sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()
+
+
+def _tools_hash(request: CompletionRequest) -> str:
+    blob = json.dumps(
+        [
+            {
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": tool.parameters,
+            }
+            for tool in request.tools
+        ],
+        sort_keys=True,
+        separators=(",", ":"),
+    )
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()
+
+
+def _jaccard(left: str, right: str) -> float:
+    """Token Jaccard similarity in [0, 1]; two empty transcripts are identical."""
+    a = set(_tokens(left))
+    b = set(_tokens(right))
+    if not a and not b:
+        return 1.0
+    return len(a & b) / len(a | b)
+
+
+def _tokens(text: str) -> list[str]:
+    return [word for word in text.lower().replace("?", " ").split() if word]

aid_runtime/concurrency.py

@@ -0,0 +1,70 @@
+"""A single shared concurrency limiter for a generated application.
+
+``deployments.concurrency`` produces exactly one :class:`ConcurrencyLimiter`. The
+model gateway, the tool executor, and the agent engine's delegation all acquire it,
+so there is one limiter -- not one per component -- and a single cap bounds every
+unit of work the app performs concurrently.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+
+from aid_runtime.errors import ConfigurationError
+
+__all__ = ["ConcurrencyLimiter", "holds_slot"]
+
+#: Whether the current task is already inside :meth:`ConcurrencyLimiter.run`.
+#:
+#: A fan-out worker holds a slot for its whole body, and that body's model and tool
+#: calls acquire the same limiter. Acquiring twice would deadlock at a cap of one,
+#: so an acquisition made while a slot is already held is a no-op: the work is
+#: already counted against the cap above it. The flag is a ``ContextVar`` because
+#: each worker runs in its own task, and a task's copy must not leak to a sibling.
+_HELD: ContextVar[bool] = ContextVar("aid_runtime_holds_slot", default=False)
+
+
+def holds_slot() -> bool:
+    """Whether this task already holds a slot of the limiter it would acquire."""
+    return _HELD.get()
+
+
+class ConcurrencyLimiter:
+    """An ``asyncio.Semaphore`` with a name and an explicit, validated cap.
+
+    Sharing the same instance between the gateway and the executor is what makes
+    "one limiter, not two" true: both acquire the same counter, so tool calls and
+    model calls draw from the same budget of concurrent slots.
+    """
+
+    def __init__(self, limit: int) -> None:
+        if limit < 1:
+            raise ConfigurationError(f"concurrency must be at least 1, got {limit}")
+        self.limit = limit
+        self._semaphore = asyncio.Semaphore(limit)
+
+    @asynccontextmanager
+    async def run(self) -> AsyncIterator[None]:
+        """Hold one slot for the duration of a single unit of work.
+
+        Cancellation and errors release the slot: ``asyncio.Semaphore`` is an async
+        context manager whose ``__aexit__`` releases on both the normal and
+        exceptional paths, so a cancelled or failed call cannot leak capacity.
+
+        Nested acquisition inside an already-held slot yields at once rather than
+        waiting: the enclosing unit of work is already counted, and waiting would
+        deadlock at a cap of one, where the fan-out worker holding the only slot is
+        the caller of the gateway that wants it.
+        """
+        if _HELD.get():
+            yield
+            return
+        async with self._semaphore:
+            token = _HELD.set(True)
+            try:
+                yield
+            finally:
+                _HELD.reset(token)

aid_runtime/control.py

@@ -0,0 +1,329 @@
+"""Composite pipeline steps: ``branch``, ``parallel``, ``map`` and ``reduce``.
+
+:class:`~aid_runtime.pipeline.Pipeline` is a sequential reducer, and it stays one.
+Each factory here returns an ordinary :data:`~aid_runtime.pipeline.Step`, so a
+composite composes with the flat steps by the same rule and replays unchanged on
+the streaming route, which inlines the pipeline instead of driving it step by step.
+
+Six context keys are reserved for the family, and no step may bind a name equal to
+one of them: ``principal``, ``limiter``, ``item``, ``index``, ``reduced`` and
+``depth``. The first two are seeded by the application; the rest belong to the
+composites -- ``item``, ``index`` and ``depth`` are written and removed around a
+fan-out, while ``reduced`` is the binding a ``reduce`` body must produce for the
+composite to collect -- so an inner step can read ``item`` without a matching
+binding leaking to a later sibling.
+
+Two bounds are enforced here and both fail the run rather than degrade it:
+
+* nesting, at :data:`MAX_STEP_DEPTH`. The spec layer rejects a deeper pipeline
+  before generation, so reaching this bound at run time means a hand-built
+  pipeline, and the reason is named in the error.
+* concurrency, through the application's one :class:`ConcurrencyLimiter`. A
+  fan-out with no limiter in the context raises rather than running unbounded,
+  because the cap is what makes a fan-out safe and its absence is a wiring bug.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from typing import Any, TypeVar
+
+from aid_runtime.agents.policy import Principal, evaluate_policy, policy_source
+from aid_runtime.concurrency import ConcurrencyLimiter, holds_slot
+from aid_runtime.errors import PipelineError
+from aid_runtime.pipeline import STEPS_KEY, Pipeline, PipelineContext
+
+__all__ = [
+    "DEPTH_KEY",
+    "INDEX_KEY",
+    "ITEM_KEY",
+    "LIMITER_KEY",
+    "MAX_STEP_DEPTH",
+    "PRINCIPAL_KEY",
+    "REDUCED_KEY",
+    "RESERVED_KEYS",
+    "branch",
+    "map_step",
+    "parallel",
+    "reduce_step",
+]
+
+#: Deepest composite nesting a pipeline may declare.
+#:
+#: The spec layer owns this number and the generator *emits* it into the generated
+#: package, so a generated app passes its own build's value; this default is what a
+#: hand-written pipeline gets. ``test_max_step_depth_matches_the_emitted_constant``
+#: holds the two together, so a drift fails a test instead of diverging silently.
+MAX_STEP_DEPTH = 4
+
+#: Who the caller is, for ``branch``. Seeded by the route.
+PRINCIPAL_KEY = "principal"
+#: The application's single limiter, from ``deployments.concurrency``.
+LIMITER_KEY = "limiter"
+#: The current item and its position, written by ``map``/``reduce`` for one iteration.
+ITEM_KEY = "item"
+INDEX_KEY = "index"
+#: How many composites deep the current step is.
+DEPTH_KEY = "depth"
+#: The binding ``reduce`` collects from each iteration. Reserved so a reduce body's
+#: own binding cannot collide with the name the composite reads.
+REDUCED_KEY = "reduced"
+
+#: Names a composite owns. A spec binding equal to one of them is refused upstream.
+RESERVED_KEYS = frozenset({PRINCIPAL_KEY, LIMITER_KEY, ITEM_KEY, INDEX_KEY, DEPTH_KEY, REDUCED_KEY})
+
+#: A composite: ``async`` by construction, so ``await step(context)`` is well-typed.
+#: ``Step`` also allows a sync return, which is what makes a generated ``await`` of one
+#: untypeable; the factories below return this narrower alias instead.
+CompositeStep = Callable[[PipelineContext], Awaitable[PipelineContext]]
+
+_StepT = TypeVar("_StepT", bound=CompositeStep)
+
+#: Values a fan-out child produced that its parent did not already hold.
+_NewValues = dict[str, Any]
+
+
+def _named(step: _StepT, name: str | None) -> _StepT:
+    """Name a composite's callable, so a trace and ``Pipeline.of`` identify it.
+
+    ``Pipeline`` names a step by the callable's ``__name__`` and refuses two steps
+    with the same name, so a pipeline holding two ``branch`` composites must give
+    them distinct names. The generator passes one; the default is the kind.
+    """
+    if name is not None:
+        step.__name__ = name
+    return step
+
+
+def _depth(context: PipelineContext) -> int:
+    value = context.get(DEPTH_KEY, 0)
+    return 0 if value is None else int(value)
+
+
+def _enter(context: PipelineContext, max_depth: int, where: str) -> PipelineContext:
+    """The context a nested step list runs in, or a named refusal past the bound."""
+    depth = _depth(context)
+    if depth >= max_depth:
+        raise PipelineError(
+            f"composite step {where} is nested deeper than {max_depth} levels; "
+            f"nesting is bounded so a self-similar spec cannot recurse without end"
+        )
+    return context.with_values(**{DEPTH_KEY: depth + 1})
+
+
+def _limiter(context: PipelineContext, where: str) -> ConcurrencyLimiter:
+    limiter = context.get(LIMITER_KEY)
+    if not isinstance(limiter, ConcurrencyLimiter):
+        raise PipelineError(
+            f"step {where} fans out but the pipeline context has no concurrency "
+            f"limiter; declare `concurrency` under a deployment in the spec so the "
+            f"generator seeds one"
+        )
+    return limiter
+
+
+def _escape(context: PipelineContext) -> PipelineContext:
+    """Drop the iteration keys, so nothing a composite added reaches a later sibling."""
+    return context.without(DEPTH_KEY, ITEM_KEY, INDEX_KEY, REDUCED_KEY)
+
+
+async def _run_body(body: Pipeline, context: PipelineContext) -> tuple[PipelineContext, _NewValues]:
+    """Run one iteration and return its context with only the bindings it added.
+
+    Only the added names come back, so a composite hands its siblings the bindings
+    the spec declared and nothing else. The inner run records its own ``steps``; it
+    is excluded here because the outer ``Pipeline.run`` writes that key itself when
+    the composite returns, and an inner list would otherwise be read as a binding.
+    """
+    result = await body.run(context)
+    added = {
+        key: value
+        for key, value in result.values.items()
+        if key not in context.values and key not in RESERVED_KEYS and key != STEPS_KEY
+    }
+    return result, added
+
+
+async def _fan_out(
+    workers: Sequence[Awaitable[tuple[PipelineContext, _NewValues]]],
+) -> list[tuple[PipelineContext, _NewValues]]:
+    """Run every worker, in the order given, whatever order they finish in.
+
+    ``asyncio.gather`` returns results positionally, so the caller's ordering rule is
+    input order even when the last worker finishes first.
+
+    A worker inside a slot this task already holds runs its siblings one at a time
+    instead of acquiring a second slot: the limiter is shared with the gateway and
+    the tool executor, so a worker that acquired again while its own body waited for
+    the gateway would deadlock at a cap of one. Serialising keeps the cap and
+    cannot deadlock, which is the trade this family makes at nesting depth two.
+    """
+    if holds_slot():
+        return [await worker for worker in workers]
+    gathered = await asyncio.gather(*workers, return_exceptions=True)
+    # Every worker is awaited even after one fails, so no task is left running and the
+    # error that surfaces is the first in input order rather than the first to finish.
+    for entry in gathered:
+        if isinstance(entry, BaseException):
+            raise entry
+    return [entry for entry in gathered if not isinstance(entry, BaseException)]
+
+
+def branch(
+    *,
+    when: Mapping[str, Any],
+    then: Pipeline,
+    otherwise: Pipeline | None = None,
+    max_depth: int = MAX_STEP_DEPTH,
+    name: str | None = None,
+) -> CompositeStep:
+    """Run ``then`` when the caller's principal satisfies ``when``, else ``otherwise``.
+
+    ``when`` is a ``PolicyExpr`` AST, evaluated against the principal and nothing
+    else: the grammar has no operand for a context value, so a branch cannot read
+    what an earlier step produced. It is a ``Mapping`` rather than a parsed object
+    so the generated code passes the AST the IR already carries, and
+    :func:`evaluate_policy` stays the only evaluator.
+
+    No principal in the context is an error, not a false. Defaulting would decide an
+    authorization question for a caller nobody identified.
+    """
+
+    async def _branch(context: PipelineContext) -> PipelineContext:
+        principal = context.get(PRINCIPAL_KEY)
+        if not isinstance(principal, Principal):
+            raise PipelineError(
+                "branch needs a principal in the pipeline context to evaluate "
+                f"{policy_source(when)!r}; none was seeded"
+            )
+        nested = _enter(context, max_depth, f"branch ({policy_source(when)})")
+        chosen = then if evaluate_policy(when, principal) else otherwise
+        if chosen is None:
+            # No alternative declared: the branch binds nothing at all.
+            return _escape(context)
+        _, added = await _run_body(chosen, nested)
+        return _escape(context.with_values(**added))
+
+    return _named(_branch, name)
+
+
+def parallel(
+    *,
+    branches: Sequence[Pipeline],
+    max_depth: int = MAX_STEP_DEPTH,
+    name: str | None = None,
+) -> CompositeStep:
+    """Run every arm concurrently under the shared cap, then merge what they bound.
+
+    A failing arm fails the whole run: the results are inspected in input order and
+    the first failure is raised, so which error surfaces does not depend on which
+    arm happened to fail first. Two arms binding the same new name is an error
+    rather than a silent last-writer-wins, because the winning value would depend on
+    completion order and the ordering rule exists to remove that.
+    """
+
+    async def _parallel(context: PipelineContext) -> PipelineContext:
+        if not branches:
+            return _escape(context)
+        limiter = _limiter(context, "parallel")
+        nested = _enter(context, max_depth, "parallel")
+
+        async def run_arm(arm: Pipeline) -> tuple[PipelineContext, _NewValues]:
+            async with limiter.run():
+                return await _run_body(arm, nested)
+
+        outcomes = await _fan_out([run_arm(arm) for arm in branches])
+        merged: dict[str, Any] = {}
+        for _, added in outcomes:
+            for key, value in added.items():
+                if key in merged:
+                    raise PipelineError(
+                        f"parallel arms both bind {key!r}; a name may be bound by one arm"
+                    )
+                merged[key] = value
+        return _escape(context.with_values(**merged))
+
+    return _named(_parallel, name)
+
+
+def map_step(
+    *,
+    over: str,
+    as_: str,
+    step: Pipeline,
+    max_depth: int = MAX_STEP_DEPTH,
+    name: str | None = None,
+) -> CompositeStep:
+    """Run ``step`` once per item of ``over``, binding the results in input order.
+
+    ``as_`` holds one mapping of the names the body bound for that item, so an item
+    that bound nothing contributes ``{}`` rather than vanishing: the list keeps the
+    length of the input. An empty or absent-over value binds an empty list.
+    """
+
+    async def _map(context: PipelineContext) -> PipelineContext:
+        items = _items(context, over, "map")
+        limiter = _limiter(context, "map")
+        nested = _enter(context, max_depth, f"map over {over!r}")
+
+        async def run_item(index: int, item: Any) -> tuple[PipelineContext, _NewValues]:
+            iteration = nested.with_values(**{ITEM_KEY: item, INDEX_KEY: index})
+            async with limiter.run():
+                return await _run_body(step, iteration)
+
+        outcomes = await _fan_out(
+            [run_item(index, item) for index, item in enumerate(items)],
+        )
+        results = [added for _, added in outcomes]
+        return _escape(context.with_values(**{as_: results}))
+
+    return _named(_map, name)
+
+
+def reduce_step(
+    *,
+    over: str,
+    as_: str,
+    step: Pipeline,
+    max_depth: int = MAX_STEP_DEPTH,
+    name: str | None = None,
+) -> CompositeStep:
+    """Run ``step`` once per item of ``over``, binding its :data:`REDUCED_KEY` results.
+
+    ``step`` is a one-step pipeline that renders the reduce prompt for one item and
+    binds :data:`REDUCED_KEY`; the composite collects that binding into ``as_`` in
+    input order. Combining the results is a following ``generate`` step, so this
+    layer never invents a merge strategy.
+    """
+
+    async def _reduce(context: PipelineContext) -> PipelineContext:
+        items = _items(context, over, "reduce")
+        limiter = _limiter(context, "reduce")
+        nested = _enter(context, max_depth, f"reduce over {over!r}")
+
+        async def run_item(index: int, item: Any) -> tuple[PipelineContext, _NewValues]:
+            iteration = nested.with_values(**{ITEM_KEY: item, INDEX_KEY: index})
+            async with limiter.run():
+                result = await step.run(iteration)
+            if REDUCED_KEY not in result.values:
+                raise PipelineError(
+                    f"reduce over {over!r} produced no {REDUCED_KEY!r} binding for item {index}"
+                )
+            return result, {REDUCED_KEY: result.require(REDUCED_KEY)}
+
+        outcomes = await _fan_out(
+            [run_item(index, item) for index, item in enumerate(items)],
+        )
+        results = [added[REDUCED_KEY] for _, added in outcomes]
+        return _escape(context.with_values(**{as_: results}))
+
+    return _named(_reduce, name)
+
+
+def _items(context: PipelineContext, over: str, where: str) -> list[Any]:
+    """The sequence a ``map``/``reduce`` iterates, refusing a missing or odd value."""
+    value = context.require(over)
+    if isinstance(value, str) or not isinstance(value, Sequence):
+        raise PipelineError(f"{where} over {over!r} needs a list, got {type(value).__name__}")
+    return list(value)