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

aid_runtime/pipeline.py

@@ -0,0 +1,168 @@
+"""A plain-async pipeline engine.
+
+This is the Phase 0 answer to "the loop". It is a sequential reducer over an
+immutable context: each step receives the context and returns the context,
+optionally enriched. No graph, no scheduler, no orchestration dependency.
+
+That is a deliberate ceiling. When Phase 2 needs branching, retries and
+human-in-the-loop interrupts an orchestration library is the right tool -- but
+adopting one before those requirements exist would place an unobservable graph
+between the user and their own trace.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+from aid_runtime.errors import PipelineError
+from aid_runtime.tracing import get_tracer, record_exception
+
+__all__ = ["STEPS_KEY", "Pipeline", "PipelineContext", "Step", "StopPipeline"]
+
+#: Context key under which :meth:`Pipeline.run` records the steps that ran, in order.
+STEPS_KEY = "steps"
+
+
+class StopPipeline(Exception):  # noqa: N818 - control-flow signal, not an error condition
+    """Raise from a step to end the pipeline successfully and early."""
+
+    def __init__(self, reason: str = "") -> None:
+        super().__init__(reason or "pipeline stopped")
+        self.reason = reason
+
+
+@dataclass(frozen=True, slots=True)
+class PipelineContext:
+    """The value threaded through a pipeline.
+
+    Frozen, so a step cannot mutate what an earlier step observed; it returns a
+    new context instead. That keeps a failed run's state inspectable and makes
+    replaying a pipeline reproducible.
+    """
+
+    values: Mapping[str, Any] = field(default_factory=dict)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self.values.get(key, default)
+
+    def require(self, key: str) -> Any:
+        try:
+            return self.values[key]
+        except KeyError:
+            raise PipelineError(f"pipeline context has no {key!r}") from None
+
+    def with_values(self, **values: Any) -> PipelineContext:
+        return replace(self, values={**self.values, **values})
+
+    def without(self, *keys: str) -> PipelineContext:
+        return replace(self, values={k: v for k, v in self.values.items() if k not in keys})
+
+
+#: What a step may produce: a new context, a mapping to merge in, or nothing.
+StepResult = PipelineContext | Mapping[str, Any] | None
+
+#: A step: an async (or plain) callable from context to a step result.
+Step = Callable[[PipelineContext], StepResult | Awaitable[StepResult]]
+
+
+class Pipeline:
+    """Runs named steps in order, tracing each one."""
+
+    def __init__(
+        self,
+        steps: Sequence[tuple[str, Step]],
+        *,
+        tracer_name: str = "aid_runtime",
+        record_steps: bool = True,
+    ) -> None:
+        if not steps:
+            raise PipelineError("a pipeline needs at least one step")
+        seen: set[str] = set()
+        for name, _ in steps:
+            if name in seen:
+                raise PipelineError(f"duplicate step name {name!r}")
+            seen.add(name)
+        self._steps = tuple(steps)
+        self._tracer_name = tracer_name
+        self._record_steps = record_steps
+
+    @property
+    def step_names(self) -> tuple[str, ...]:
+        return tuple(name for name, _ in self._steps)
+
+    @classmethod
+    def of(cls, *steps: Step, tracer_name: str = "aid_runtime") -> Pipeline:
+        """Build a pipeline from bare callables, naming each by its function."""
+        return cls([(_step_name(step), step) for step in steps], tracer_name=tracer_name)
+
+    def then(self, name: str, step: Step) -> Pipeline:
+        """A new pipeline with ``step`` appended. This instance is unchanged."""
+        return Pipeline(
+            [*self._steps, (name, step)],
+            tracer_name=self._tracer_name,
+            record_steps=self._record_steps,
+        )
+
+    async def run(
+        self, initial: PipelineContext | Mapping[str, Any] | None = None
+    ) -> PipelineContext:
+        """Execute every step, stopping early if one raises :class:`StopPipeline`.
+
+        A step that raises anything else aborts the run and the original exception
+        is chained, so the traceback shows both the pipeline coordinate and the
+        real cause.
+        """
+        context = _as_context(initial)
+        ran: list[str] = []
+        tracer = get_tracer(self._tracer_name)
+
+        for name, step in self._steps:
+            with tracer.start_as_current_span(
+                f"pipeline.{name}",
+                # The failing step's own exception is recorded below; the automatic
+                # recording would add a second event for the PipelineError wrapper
+                # and would hide the original exception type.
+                record_exception=False,
+                set_status_on_exception=False,
+            ) as span:
+                span.set_attribute("aid.pipeline.step", name)
+                try:
+                    context = _as_context(await _resolve(step(context)))
+                except StopPipeline as stop:
+                    ran.append(name)
+                    if stop.reason:
+                        span.set_attribute("aid.pipeline.stop_reason", stop.reason)
+                    return self._finish(context, ran)
+                except Exception as exc:
+                    record_exception(span, exc)
+                    raise PipelineError(f"step {name!r} failed: {exc}") from exc
+            ran.append(name)
+
+        return self._finish(context, ran)
+
+    def _finish(self, context: PipelineContext, ran: Sequence[str]) -> PipelineContext:
+        return context.with_values(**{STEPS_KEY: tuple(ran)}) if self._record_steps else context
+
+
+async def _resolve(value: StepResult | Awaitable[StepResult]) -> StepResult:
+    return await value if inspect.isawaitable(value) else value
+
+
+def _as_context(value: Any) -> PipelineContext:
+    if isinstance(value, PipelineContext):
+        return value
+    if value is None:
+        return PipelineContext()
+    if isinstance(value, Mapping):
+        return PipelineContext(values=dict(value))
+    raise PipelineError(
+        f"a pipeline step must return a PipelineContext, a mapping or None, "
+        f"got {type(value).__name__}"
+    )
+
+
+def _step_name(step: Step) -> str:
+    return getattr(step, "__name__", None) or type(step).__name__

aid_runtime/prompts.py

@@ -0,0 +1,185 @@
+"""Versioned prompt registry with strict typed rendering.
+
+Rendering is *strict*: a template that references a variable nobody supplied, or
+a caller that supplies a variable the template never uses, is an error. In a
+framework whose promise is that regenerating an app is safe, a silently-blank
+prompt variable is exactly the class of bug that must not survive to runtime.
+"""
+
+from __future__ import annotations
+
+import string
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Self
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from aid_runtime.errors import PromptNotFoundError, PromptRenderError
+
+__all__ = ["Example", "PromptRegistry", "PromptTemplate", "RenderedPrompt"]
+
+
+class Example(BaseModel):
+    """One few-shot example. Rendered into the prompt, not interpolated as text."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    input: str
+    output: str
+
+
+class PromptTemplate(BaseModel):
+    """A versioned prompt with declared variables and optional few-shot examples."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    name: str = Field(min_length=1)
+    version: str = Field(min_length=1)
+    template: str = Field(min_length=1)
+    description: str = ""
+    examples: tuple[Example, ...] = ()
+    #: Variable names that must be supplied even when the template body omits them.
+    #: Empty by default: the template body is the source of truth.
+    required: tuple[str, ...] = ()
+
+    @property
+    def variables(self) -> frozenset[str]:
+        """Every field name the template body interpolates."""
+        try:
+            parsed = string.Formatter().parse(self.template)
+            return frozenset(field_name for _, field_name, _, _ in parsed if field_name is not None)
+        except ValueError as exc:
+            # A brace the formatter cannot parse is a broken template, not a bad
+            # variable, so it is reported as the render error the caller handles
+            # rather than escaping as a bare ValueError.
+            raise PromptRenderError(
+                f"prompt {self.name!r} v{self.version} has a malformed template: {exc}"
+            ) from exc
+
+    def render(self, **values: str) -> RenderedPrompt:
+        """Interpolate ``values``, rejecting missing and unused variables."""
+        expected = self.variables
+        required = expected | frozenset(self.required)
+        supplied = set(values)
+
+        missing = sorted(name for name in required if name not in supplied)
+        if missing:
+            raise PromptRenderError(
+                f"prompt {self.name!r} v{self.version} is missing variables: {', '.join(missing)}"
+            )
+        # ``required`` is subtracted here as well as added above: a variable the
+        # caller was told to supply without an interpolation site is supplied
+        # deliberately, not "unused".
+        unused = sorted(supplied - required)
+        if unused:
+            raise PromptRenderError(
+                f"prompt {self.name!r} v{self.version} received unused variables: "
+                f"{', '.join(unused)}; template uses {', '.join(sorted(required)) or '<none>'}"
+            )
+
+        try:
+            body = self.template.format(**values)
+        except (KeyError, IndexError, ValueError) as exc:
+            raise PromptRenderError(
+                f"prompt {self.name!r} v{self.version} failed to render: {exc}"
+            ) from exc
+
+        return RenderedPrompt(
+            name=self.name,
+            version=self.version,
+            text=body,
+            examples=self.examples,
+        )
+
+
+class RenderedPrompt(BaseModel):
+    """A prompt ready to send, carrying the version that produced it.
+
+    The version travels with the text so a trace can answer "which prompt was
+    this?" -- without it, a baseline diff cannot distinguish a model change from
+    a prompt change.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    name: str
+    version: str
+    text: str
+    examples: tuple[Example, ...] = ()
+
+    def as_messages(self, *, system: str | None = None) -> tuple[tuple[str, str], ...]:
+        """Flatten to ``(role, content)`` pairs, examples inlined as few-shot turns."""
+        messages: list[tuple[str, str]] = []
+        if system is not None:
+            messages.append(("system", system))
+        for example in self.examples:
+            messages.append(("user", example.input))
+            messages.append(("assistant", example.output))
+        messages.append(("user", self.text))
+        return tuple(messages)
+
+
+class PromptRegistry:
+    """A registry keyed by prompt name, with per-name version pinning."""
+
+    def __init__(self, templates: Iterable[PromptTemplate] = ()) -> None:
+        self._by_name: dict[str, dict[str, PromptTemplate]] = {}
+        for template in templates:
+            self.register(template)
+
+    def register(self, template: PromptTemplate) -> Self:
+        """Add a template. Re-registering the same name+version is an error."""
+        versions = self._by_name.setdefault(template.name, {})
+        if template.version in versions:
+            raise PromptRenderError(
+                f"prompt {template.name!r} version {template.version!r} is already registered"
+            )
+        versions[template.version] = template
+        return self
+
+    def names(self) -> tuple[str, ...]:
+        return tuple(sorted(self._by_name))
+
+    def versions(self, name: str) -> tuple[str, ...]:
+        return tuple(sorted(self._by_name.get(name, {})))
+
+    def latest(self, name: str) -> PromptTemplate:
+        """The highest registered version, compared numerically per segment."""
+        versions = ordered_versions(self.versions(name))
+        if not versions:
+            raise PromptNotFoundError(name, None, self.names())
+        return self._by_name[name][versions[-1]]
+
+    def get(self, name: str, version: str | None = None) -> PromptTemplate:
+        if version is None:
+            return self.latest(name)
+        try:
+            return self._by_name[name][version]
+        except KeyError:
+            raise PromptNotFoundError(name, version, self.versions(name)) from None
+
+    def render(
+        self,
+        name: str,
+        *,
+        version: str | None = None,
+        values: Mapping[str, str] | None = None,
+    ) -> RenderedPrompt:
+        return self.get(name, version).render(**(dict(values) if values else {}))
+
+
+def ordered_versions(versions: Sequence[str]) -> tuple[str, ...]:
+    """Sort version strings numerically per dotted segment.
+
+    String sorting puts ``v10`` before ``v9``; a prompt registry that silently
+    resolves "latest" to the wrong version is worse than one that refuses.
+    """
+
+    def key(version: str) -> tuple[tuple[int, int, str], ...]:
+        segments: list[tuple[int, int, str]] = []
+        for segment in version.lstrip("v").split("."):
+            digits = "".join(char for char in segment if char.isdigit())
+            segments.append((0, int(digits), "") if digits else (1, 0, segment))
+        return tuple(segments)
+
+    return tuple(sorted(versions, key=key))

aid_runtime/providers/__init__.py

@@ -0,0 +1,20 @@
+"""Provider adapters.
+
+Exactly two ship in Phase 0: one real HTTP adapter and one deterministic fake for
+tests and the eval-gate fixture. Provider breadth (Anthropic, Azure, Bedrock) is
+deliberately out of scope.
+"""
+
+from __future__ import annotations
+
+from aid_runtime.providers.base import Provider
+from aid_runtime.providers.fake import FakeProvider, ScriptedTurn
+from aid_runtime.providers.http import DEFAULT_BASE_URL, OpenAICompatibleProvider
+
+__all__ = [
+    "DEFAULT_BASE_URL",
+    "FakeProvider",
+    "OpenAICompatibleProvider",
+    "Provider",
+    "ScriptedTurn",
+]

aid_runtime/providers/base.py

@@ -0,0 +1,57 @@
+"""Provider adapter contract.
+
+A provider translates the runtime's shapes into one wire protocol. It knows
+nothing about aliases, budgets, or tracing -- those live in the gateway, so an
+adapter stays a thin, independently testable translation layer.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Protocol, runtime_checkable
+
+from aid_runtime.types import (
+    Capability,
+    CompletionRequest,
+    CompletionResponse,
+    ModelSpec,
+    StreamChunk,
+)
+
+__all__ = ["Provider"]
+
+
+@runtime_checkable
+class Provider(Protocol):
+    """A single upstream model vendor endpoint."""
+
+    @property
+    def name(self) -> str:
+        """Identifier referenced by :attr:`ModelSpec.provider`."""
+        ...
+
+    @property
+    def capabilities(self) -> frozenset[Capability]:
+        """What this endpoint can do at all, before per-model negotiation."""
+        ...
+
+    async def complete(self, request: CompletionRequest, *, model: ModelSpec) -> CompletionResponse:
+        """Perform one request/response call.
+
+        Implementations must not retry in a way that hides which model answered:
+        the gateway owns fallback and reports every attempt it made.
+        """
+        ...
+
+    def stream(self, request: CompletionRequest, *, model: ModelSpec) -> AsyncIterator[StreamChunk]:
+        """Stream one call as text deltas, ending with a usage-bearing chunk.
+
+        Only adapters that advertise :attr:`Capability.STREAMING` implement this.
+        The first delta must be yielded before the call completes, and the final
+        chunk carries ``finish_reason`` and ``usage`` so the gateway can account
+        for the stream exactly as it accounts for :meth:`complete`.
+        """
+
+    async def aclose(self) -> None:
+        """Release transport resources. Must be idempotent."""
+        ...

aid_runtime/providers/fake.py

@@ -0,0 +1,256 @@
+"""A deterministic provider for tests.
+
+This is not a mock in the "returns a canned value" sense. It is a *seeded
+simulation* of a probabilistic model: given the same seed and the same request it
+produces byte-identical output, but it deliberately answers correctly only a
+configurable fraction of the time. That is the whole point -- it lets the eval
+gate be tested against genuine non-determinism, so a test can prove that run-to-run
+variation is tolerated while a real regression is caught.
+
+A *script* is the second way to be deterministic, and the one an agent loop needs:
+a real model's tool calls are the part of its output that a canned text answer
+cannot express, so :class:`ScriptedTurn` scripts them turn by turn. An off-by-one
+in a loop's step accounting is invisible against a provider that can only answer.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import random
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from aid_runtime.errors import ProviderError
+from aid_runtime.types import (
+    Capability,
+    CompletionRequest,
+    CompletionResponse,
+    FinishReason,
+    ModelSpec,
+    StreamChunk,
+    TokenUsage,
+    ToolCall,
+)
+
+__all__ = ["FakeProvider", "ScriptedTurn"]
+
+#: The fixed chunk size :meth:`FakeProvider.stream` splits text into, so a test can
+#: assert that a long answer streams as several deltas rather than one blob.
+_STREAM_CHUNK = 4
+
+# A responder sees the request and returns the model's text. Receiving the index
+# of the call lets a responder model "the third call fails" without mutable state.
+Responder = Callable[[CompletionRequest, int], str]
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptedTurn:
+    """One scripted model turn: some text, some tool calls, or both.
+
+    ``calls`` is ``(tool name, arguments)`` pairs rather than :class:`ToolCall`
+    because a script should not have to invent unique call ids -- the provider
+    derives ``call_<turn>_<position>`` from the turn's position in the script, so
+    two turns calling the same tool still produce distinguishable ids.
+    """
+
+    text: str = ""
+    calls: tuple[tuple[str, Mapping[str, Any]], ...] = ()
+
+    @classmethod
+    def answer(cls, text: str) -> ScriptedTurn:
+        return cls(text=text)
+
+    @classmethod
+    def tool(cls, name: str, arguments: Mapping[str, Any] | None = None) -> ScriptedTurn:
+        return cls(calls=((name, dict(arguments or {})),))
+
+    @classmethod
+    def tools(cls, *calls: tuple[str, Mapping[str, Any]]) -> ScriptedTurn:
+        return cls(calls=tuple(calls))
+
+    @property
+    def finish_reason(self) -> FinishReason:
+        return "tool-calls" if self.calls else "stop"
+
+    def tool_calls(self, turn: int) -> tuple[ToolCall, ...]:
+        return tuple(
+            ToolCall(id=f"call_{turn}_{position}", name=name, arguments=dict(arguments))
+            for position, (name, arguments) in enumerate(self.calls)
+        )
+
+
+class FakeProvider:
+    """Deterministic, seedable, and optionally inaccurate."""
+
+    def __init__(
+        self,
+        *,
+        responder: Responder | None = None,
+        answers: Mapping[str, str] | None = None,
+        script: Sequence[ScriptedTurn] = (),
+        accuracy: float = 1.0,
+        seed: int = 0,
+        name: str = "fake",
+        fail_with: Sequence[str] = (),
+        tokens_per_char: float = 0.25,
+        overhead_tokens: int = 8,
+    ) -> None:
+        if not 0.0 <= accuracy <= 1.0:
+            raise ValueError("accuracy must be between 0.0 and 1.0")
+        self._responder = responder
+        # `answers` maps the *user message* to the ideal response, which lets a
+        # test wire a dataset to a provider without writing a callable.
+        self._answers = dict(answers or {})
+        # A script is consumed one turn per call, then the provider behaves
+        # normally -- so a scripted loop that overshoots its script is visible
+        # rather than silently answered.
+        self._script = tuple(script)
+        self._accuracy = accuracy
+        self._seed = seed
+        self._name = name
+        # Each entry is a message for a successive call. Once exhausted the
+        # provider behaves normally, so a test can prove retry-then-succeed.
+        self._failures = list(fail_with)
+        self._served = 0
+        self._tokens_per_char = tokens_per_char
+        self._overhead_tokens = overhead_tokens
+        self.calls: list[CompletionRequest] = []
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def capabilities(self) -> frozenset[Capability]:
+        return frozenset(
+            {
+                Capability.TOOL_CALLING,
+                Capability.STRUCTURED_OUTPUT,
+                Capability.LONG_CONTEXT,
+                Capability.STREAMING,
+            }
+        )
+
+    async def complete(self, request: CompletionRequest, *, model: ModelSpec) -> CompletionResponse:
+        self.calls.append(request)
+        index = len(self.calls) - 1
+
+        if self._failures:
+            raise ProviderError(str(self._failures.pop(0)))
+
+        input_tokens = self._count_tokens(_prompt_text(request))
+        if self._served < len(self._script):
+            turn = self._script[self._served]
+            self._served += 1
+            return CompletionResponse(
+                text=turn.text,
+                model=model.name,
+                provider=self._name,
+                usage=TokenUsage(
+                    input_tokens=input_tokens,
+                    output_tokens=self._count_tokens(turn.text),
+                ),
+                tool_calls=turn.tool_calls(self._served - 1),
+                finish_reason=turn.finish_reason,
+            )
+
+        text = self._respond_to(request, index)
+        return CompletionResponse(
+            text=text,
+            model=model.name,
+            provider=self._name,
+            usage=TokenUsage(
+                input_tokens=input_tokens,
+                output_tokens=self._count_tokens(text),
+            ),
+            finish_reason="stop",
+        )
+
+    async def stream(
+        self, request: CompletionRequest, *, model: ModelSpec
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream the same answer :meth:`complete` would return, in fixed chunks.
+
+        The script and failure queue are consumed identically to ``complete``, so a
+        streaming call and a request/response call observe the same deterministic
+        sequence.
+        """
+        self.calls.append(request)
+        index = len(self.calls) - 1
+
+        if self._failures:
+            raise ProviderError(str(self._failures.pop(0)))
+
+        input_tokens = self._count_tokens(_prompt_text(request))
+        if self._served < len(self._script):
+            turn = self._script[self._served]
+            self._served += 1
+            text = turn.text
+            finish_reason = turn.finish_reason
+        else:
+            text = self._respond_to(request, index)
+            finish_reason = "stop"
+
+        usage = TokenUsage(
+            input_tokens=input_tokens,
+            output_tokens=self._count_tokens(text),
+        )
+        for start in range(0, len(text), _STREAM_CHUNK):
+            yield StreamChunk(delta=text[start : start + _STREAM_CHUNK])
+        yield StreamChunk(finish_reason=finish_reason, usage=usage)
+
+    async def aclose(self) -> None:
+        return None
+
+    def _respond_to(self, request: CompletionRequest, index: int) -> str:
+        if self._responder is not None:
+            return self._responder(request, index)
+
+        prompt = _prompt_text(request)
+        ideal = self._answers.get(prompt)
+        if ideal is None:
+            for key, value in self._answers.items():
+                if key in prompt:
+                    ideal = value
+                    break
+
+        if ideal is None:
+            # No expectation to hit, so accuracy does not apply: echo deterministically.
+            return f"echo:{_digest(prompt, self._seed)}"
+
+        if self._accuracy >= 1.0:
+            return ideal
+        if self._draw(prompt, index) < self._accuracy:
+            return ideal
+        # A wrong answer must still be deterministic and visibly wrong.
+        return f"{ideal}#{_digest(prompt, self._seed + index)}"
+
+    def _draw(self, prompt: str, index: int) -> float:
+        """A stable pseudo-random draw for this (seed, call index, prompt) triple.
+
+        Deriving the draw from a hash rather than a sequential PRNG stream means a
+        recorded baseline reproduces exactly: the same call in the same position
+        always draws the same number. It is deliberately *not* stable across
+        positions -- the call index is part of the input -- so re-running a dataset
+        re-draws every case and an N-run sample genuinely varies. That variation is
+        what the eval gate has to tolerate, so it has to exist here.
+        """
+        rng = random.Random(f"{self._seed}:{index}:{prompt}")
+        return rng.random()
+
+    def _count_tokens(self, text: str) -> int:
+        return int(len(text) * self._tokens_per_char) + self._overhead_tokens
+
+
+def _prompt_text(request: CompletionRequest) -> str:
+    return "\n".join(message.content for message in request.messages)
+
+
+def _digest(prompt: str, seed: int) -> str:
+    return hashlib.sha256(f"{seed}:{prompt}".encode()).hexdigest()[:8]
+
+
+def with_accuracy(accuracy: float, *, seed: int = 0) -> FakeProvider:
+    """Convenience factory used by the eval-gate tests."""
+    return FakeProvider(accuracy=accuracy, seed=seed)