stratum-py 0.1.0__py3-none-any.whl

stratum/__init__.py

@@ -0,0 +1,120 @@
+"""
+Stratum — LLM calls that behave like the rest of your code.
+
+Public API surface (v1):
+
+    Decorators:   @contract, @infer, @compute, @flow, @refine
+    Types:        Budget, opaque, Probabilistic, HumanDecision, HumanReviewContext
+    HITL:         await_human
+    Concurrency:  parallel, debate, race
+    Utilities:    configure, run
+    Errors:       StratumError and all subclasses
+    Trace:        TraceRecord, all_records, clear_traces
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from .contracts import contract, opaque, is_registered
+from .budget import Budget
+from .exceptions import (
+    StratumError,
+    StratumCompileError,
+    PreconditionFailed,
+    PostconditionFailed,
+    ParseFailure,
+    BudgetExceeded,
+    ConvergenceFailure,
+    ConsensusFailure,
+    ParallelValidationFailed,
+    HITLTimeoutError,
+    StabilityAssertionError,
+)
+from .decorators import infer, compute, flow, refine
+from .trace import TraceRecord, all_records, clear as clear_traces
+from ._config import configure
+from .types import Probabilistic, HumanDecision, HumanReviewContext, Success, Failure
+from .hitl import await_human, ReviewSink, ConsoleReviewSink, PendingReview
+from .concurrency import parallel, debate, race
+from .flow_scope import FlowScope
+from . import exporters
+
+
+def run(coro: Any) -> Any:
+    """
+    Synchronous shim for non-async contexts (scripts, notebooks).
+
+    Manages an event loop internally. MUST NOT be called from inside an
+    already-running event loop.
+    """
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            raise RuntimeError(
+                "stratum.run() must not be called from inside a running event loop. "
+                "Use 'await' directly instead."
+            )
+        return loop.run_until_complete(coro)
+    except RuntimeError as exc:
+        if "no current event loop" in str(exc).lower():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                return loop.run_until_complete(coro)
+            finally:
+                loop.close()
+        raise
+
+
+__all__ = [
+    # Decorators
+    "contract",
+    "infer",
+    "compute",
+    "flow",
+    "refine",
+    # Types
+    "Budget",
+    "opaque",
+    "Probabilistic",
+    "Success",
+    "Failure",
+    "HumanDecision",
+    "HumanReviewContext",
+    # HITL
+    "await_human",
+    "ReviewSink",
+    "ConsoleReviewSink",
+    "PendingReview",
+    # Concurrency
+    "parallel",
+    "debate",
+    "race",
+    # Flow context
+    "FlowScope",
+    # Configuration
+    "configure",
+    "run",
+    # Trace
+    "TraceRecord",
+    "all_records",
+    "clear_traces",
+    # Errors
+    "StratumError",
+    "StratumCompileError",
+    "PreconditionFailed",
+    "PostconditionFailed",
+    "ParseFailure",
+    "BudgetExceeded",
+    "ConvergenceFailure",
+    "ConsensusFailure",
+    "ParallelValidationFailed",
+    "HITLTimeoutError",
+    "StabilityAssertionError",
+    # Registry
+    "is_registered",
+    # Exporters
+    "exporters",
+]

stratum/_config.py

@@ -0,0 +1,48 @@
+"""Global Stratum configuration."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+_config: dict[str, Any] = {
+    "client": None,           # uses litellm directly if None
+    "review_sink": None,      # ConsoleReviewSink if None
+    "tracer": None,           # None = no OTel export
+    "default_model": "claude-sonnet-4-6",
+    "test_mode": False,       # True → sample sample_n times for Probabilistic[T]
+    "sample_n": 5,            # samples per @infer call in test_mode
+}
+
+
+def configure(
+    client: Any = None,
+    review_sink: Any = None,
+    tracer: Any = None,
+    default_model: str | None = None,
+    test_mode: bool | None = None,
+    sample_n: int | None = None,
+) -> None:
+    """
+    Set global Stratum configuration.
+
+    Configuration is global and set once at startup. Per-function decorator
+    annotations take precedence over global defaults.
+    """
+    if client is not None:
+        _config["client"] = client
+    if review_sink is not None:
+        _config["review_sink"] = review_sink
+    if tracer is not None:
+        _config["tracer"] = tracer
+    if default_model is not None:
+        _config["default_model"] = default_model
+    if test_mode is not None:
+        _config["test_mode"] = test_mode
+    if sample_n is not None:
+        _config["sample_n"] = sample_n
+
+
+def get_config() -> dict[str, Any]:
+    """Return the current configuration dict (mutable reference)."""
+    return _config

stratum/budget.py

@@ -0,0 +1,61 @@
+"""Budget dataclass — time and cost envelope for @infer and @flow calls."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Budget:
+    """
+    Declares a time and/or cost budget for an @infer or @flow invocation.
+
+    Either or both of `ms` and `usd` may be specified. Unspecified axes are
+    unbounded.
+    """
+
+    ms: int | None = None       # wall-clock milliseconds
+    usd: float | None = None    # cost ceiling in USD
+
+    # Runtime tracking — not part of the public API, not shown in repr
+    _start_ms: float = field(
+        default_factory=lambda: time.monotonic() * 1000,
+        init=False,
+        repr=False,
+        compare=False,
+    )
+    _spent_usd: float = field(
+        default=0.0,
+        init=False,
+        repr=False,
+        compare=False,
+    )
+
+    def remaining_seconds(self) -> float | None:
+        """
+        Return remaining wall-clock time in seconds, or None if no ms limit.
+        Returns 0.0 if the budget is already exhausted.
+        """
+        if self.ms is None:
+            return None
+        elapsed_ms = (time.monotonic() * 1000) - self._start_ms
+        remaining_ms = self.ms - elapsed_ms
+        return max(0.0, remaining_ms / 1000.0)
+
+    def record_cost(self, usd: float) -> None:
+        """Accumulate a cost charge against this budget."""
+        self._spent_usd += usd
+
+    def is_cost_exceeded(self) -> bool:
+        """Return True if cumulative cost has reached or exceeded the usd ceiling."""
+        if self.usd is None:
+            return False
+        return self._spent_usd >= self.usd
+
+    def clone(self) -> Budget:
+        """
+        Create a fresh budget with the same limits, resetting the elapsed clock
+        and spent cost to zero. Used by @flow to create a per-execution envelope.
+        """
+        return Budget(ms=self.ms, usd=self.usd)

stratum/compiler.py

@@ -0,0 +1,160 @@
+"""Prompt compiler — deterministic assembly and SHA-256 hash."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any
+
+
+def _format_value(value: Any) -> str:
+    """
+    Format a value for inline prompt display.
+
+    Contract instances are rendered as a dict of their public attributes.
+    """
+    if isinstance(value, str):
+        return value
+    if hasattr(value, "__dict__"):
+        public = {k: v for k, v in value.__dict__.items() if not k.startswith("_")}
+        return str(public)
+    return str(value)
+
+
+def compile_prompt(
+    intent: str,
+    context: list[str],
+    inputs: dict[str, Any],
+    opaque_fields: set[str],
+    retry_reasons: list[str],
+) -> str:
+    """
+    Assemble the prompt string sent to the LLM.
+
+    Assembly order (per spec §4.1):
+      1. intent
+      2. context annotations (in declaration order)
+      3. non-opaque input bindings
+      4. retry context (only when retry_reasons is non-empty)
+      5. opaque data reference (only when opaque_fields is non-empty)
+
+    The output schema is NOT included — it is enforced via the structured
+    outputs API (tool_choice).
+
+    Spec §4.2: raises StratumCompileError if an opaque field name appears as
+    an inline {field} reference in intent or context strings.
+    """
+    # Spec §4.2: raise StratumCompileError if an opaque field is referenced
+    # inline in intent or context strings.
+    if opaque_fields:
+        from .exceptions import StratumCompileError
+        for text in [intent, *context]:
+            for field_name in opaque_fields:
+                if f"{{{field_name}}}" in text:
+                    raise StratumCompileError(
+                        f"opaque field '{field_name}' must not appear in inline "
+                        "string interpolation (intent or context). "
+                        "Opaque fields are passed as structured attachments only."
+                    )
+    parts: list[str] = []
+
+    # 1. Intent
+    parts.append(intent)
+
+    # 2. Context annotations
+    for ctx in context:
+        if ctx:
+            parts.append(ctx)
+
+    # 3. Non-opaque input bindings
+    non_opaque = {k: v for k, v in inputs.items() if k not in opaque_fields}
+    if non_opaque:
+        parts.append("Inputs:")
+        for key, value in non_opaque.items():
+            parts.append(f"  {key}: {_format_value(value)}")
+
+    # 4. Retry context — only on retries (attempt > 0)
+    if retry_reasons:
+        parts.append("Previous attempt failed:")
+        for reason in retry_reasons:
+            parts.append(f"  - {reason}")
+        parts.append("Fix these issues specifically.")
+
+    # 5. Opaque field reference
+    if opaque_fields:
+        names = ", ".join(sorted(opaque_fields))
+        parts.append(f"See attached data for: {names}")
+
+    return "\n".join(parts)
+
+
+def compile_prompt_stable(
+    intent: str,
+    context: list[str],
+    opaque_fields: set[str],
+) -> str:
+    """
+    Return the stable, cacheable prefix of the compiled prompt: intent + context.
+
+    Runs the opaque-field inline-reference check (spec §4.2).
+    Used by the executor to build Anthropic prompt-cache content blocks.
+    """
+    if opaque_fields:
+        from .exceptions import StratumCompileError
+        for text in [intent, *context]:
+            for field_name in opaque_fields:
+                if f"{{{field_name}}}" in text:
+                    raise StratumCompileError(
+                        f"opaque field '{field_name}' must not appear in inline "
+                        "string interpolation (intent or context). "
+                        "Opaque fields are passed as structured attachments only."
+                    )
+    parts = [intent]
+    for ctx in context:
+        if ctx:
+            parts.append(ctx)
+    return "\n".join(parts)
+
+
+def compile_prompt_variable(
+    inputs: dict[str, Any],
+    opaque_fields: set[str],
+    retry_reasons: list[str],
+) -> str:
+    """
+    Return the variable suffix of the compiled prompt: inputs + retry + opaque reference.
+
+    Used by the executor to build the uncached portion of Anthropic content blocks.
+    """
+    parts: list[str] = []
+    non_opaque = {k: v for k, v in inputs.items() if k not in opaque_fields}
+    if non_opaque:
+        parts.append("Inputs:")
+        for key, value in non_opaque.items():
+            parts.append(f"  {key}: {_format_value(value)}")
+    if retry_reasons:
+        parts.append("Previous attempt failed:")
+        for reason in retry_reasons:
+            parts.append(f"  - {reason}")
+        parts.append("Fix these issues specifically.")
+    if opaque_fields:
+        names = ", ".join(sorted(opaque_fields))
+        parts.append(f"See attached data for: {names}")
+    return "\n".join(parts)
+
+
+def prompt_hash(prompt: str) -> str:
+    """Return the first 12 hex characters of the SHA-256 of the prompt string."""
+    return hashlib.sha256(prompt.encode()).hexdigest()[:12]
+
+
+def build_opaque_attachment(
+    inputs: dict[str, Any],
+    opaque_fields: set[str],
+) -> dict[str, Any] | None:
+    """
+    Return a dict of {field_name: value} for all opaque fields, or None if
+    there are no opaque fields.
+    """
+    result = {k: v for k, v in inputs.items() if k in opaque_fields}
+    return result if result else None

stratum/concurrency.py

@@ -0,0 +1,215 @@
+"""Concurrency primitives: parallel, race, debate."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable
+
+from .exceptions import ConsensusFailure, ParallelValidationFailed
+from .types import Failure, Success
+
+
+async def parallel(
+    *coros: Any,
+    require: str | int = "all",
+    validate: Callable[[list], bool] | None = None,
+) -> Any:
+    """
+    Run coroutines concurrently with configurable success semantics.
+
+    require="all"  → all must succeed; any failure cancels rest and re-raises.
+                     Returns a tuple matching input order.
+    require="any"  → first success wins, rest cancelled. Returns single result.
+    require=N: int → at least N must succeed. Returns list of N results.
+    require=0      → collect all regardless of failure. Returns list[Success|Failure].
+
+    validate       → optional callable on collected results; False → ParallelValidationFailed.
+    """
+    if require == "all":
+        async with asyncio.TaskGroup() as tg:
+            tasks = [tg.create_task(c) for c in coros]
+        results = [t.result() for t in tasks]
+
+        if validate is not None and not validate(results):
+            raise ParallelValidationFailed()
+
+        return tuple(results)
+
+    if require == "any":
+        if not coros:
+            raise ValueError("parallel(require='any'): requires at least one coroutine")
+        tasks = [asyncio.create_task(c) for c in coros]
+        pending: set = set(tasks)
+        last_exc: Exception | None = None
+
+        while pending:
+            done, pending = await asyncio.wait(
+                pending, return_when=asyncio.FIRST_COMPLETED
+            )
+            winner = None
+            for d in done:
+                if d.exception() is None:
+                    winner = d
+                    break
+                last_exc = d.exception()
+
+            if winner is not None:
+                # Cancel remaining pending tasks
+                for p in pending:
+                    p.cancel()
+                    try:
+                        await p
+                    except (asyncio.CancelledError, Exception):
+                        pass
+                # Drain other done tasks so their exceptions are retrieved
+                for d in done:
+                    if d is not winner:
+                        try:
+                            d.exception()
+                        except (asyncio.CancelledError, Exception):
+                            pass
+                result = winner.result()
+                if validate is not None and not validate([result]):
+                    raise ParallelValidationFailed()
+                return result
+
+        if last_exc is not None:
+            raise last_exc
+        raise RuntimeError("parallel: all coroutines failed with no exception recorded")
+
+    if isinstance(require, int) and require == 0:
+        # Collect all regardless of failure — wrap in Success/Failure
+        raw = await asyncio.gather(*coros, return_exceptions=True)
+        results = [
+            Success(r) if not isinstance(r, Exception) else Failure(r)
+            for r in raw
+        ]
+        if validate is not None and not validate(results):
+            raise ParallelValidationFailed()
+        return results
+
+    if isinstance(require, int) and require > 0:
+        # At least require many must succeed
+        all_results = await asyncio.gather(*coros, return_exceptions=True)
+        successes = [r for r in all_results if not isinstance(r, Exception)]
+        failures = [r for r in all_results if isinstance(r, Exception)]
+
+        if len(successes) < require:
+            if failures:
+                raise failures[0]
+            raise RuntimeError(
+                f"parallel: needed {require} successes, got {len(successes)}"
+            )
+
+        results = successes[:require]
+        if validate is not None and not validate(results):
+            raise ParallelValidationFailed()
+        return results
+
+    raise ValueError(f"parallel: invalid require value: {require!r}")
+
+
+async def race(*coros: Any) -> Any:
+    """
+    Submit all coroutines concurrently. First to complete without raising wins.
+    Remaining coroutines are cancelled. If all raise, re-raises the last error.
+    """
+    if not coros:
+        raise ValueError("race: requires at least one coroutine")
+    tasks = [asyncio.create_task(c) for c in coros]
+    pending: set = set(tasks)
+    last_exc: Exception | None = None
+
+    while pending:
+        done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+        winner = None
+        for d in done:
+            exc = d.exception()
+            if exc is None:
+                winner = d
+                break
+            last_exc = exc
+
+        if winner is not None:
+            # Cancel the rest
+            for p in pending:
+                p.cancel()
+                try:
+                    await p
+                except (asyncio.CancelledError, Exception):
+                    pass
+            # Drain other done tasks so their exceptions are retrieved
+            for d in done:
+                if d is not winner:
+                    try:
+                        d.exception()
+                    except (asyncio.CancelledError, Exception):
+                        pass
+            return winner.result()
+
+    if last_exc is not None:
+        raise last_exc
+    raise RuntimeError("race: all coroutines failed")
+
+
+async def debate(
+    agents: list[Callable],
+    topic: Any,
+    rounds: int = 2,
+    *,
+    synthesize: Callable,
+) -> Any:
+    """
+    Multi-agent debate protocol.
+
+    Round 1: all agents invoked concurrently with topic.
+    Rounds 2..N: each agent invoked concurrently with topic + other agents' previous arguments.
+    After all rounds, convergence is computed and synthesize is called.
+
+    synthesize is required.
+    """
+    if not agents:
+        raise ValueError("debate: agents list must not be empty")
+
+    # Round 1 — initial arguments (concurrent)
+    initial_results = await asyncio.gather(*[agent(topic=topic) for agent in agents])
+    arguments: list[Any] = list(initial_results)
+    history: list[list[Any]] = [arguments]
+
+    # Rebuttal rounds — each round all agents run concurrently
+    for _round in range(1, rounds):
+        rebuttal_coros = [
+            agents[i](
+                topic=topic,
+                previous_arguments=[arguments[j] for j in range(len(arguments)) if j != i],
+            )
+            for i in range(len(agents))
+        ]
+        new_args = list(await asyncio.gather(*rebuttal_coros))
+        arguments = new_args
+        history.append(list(arguments))
+
+    # Compute convergence — use agree_on field from agents if declared
+    last_round = history[-1]
+    agree_on_field: str | None = None
+    for agent in agents:
+        spec = getattr(agent, "_stratum_spec", None)
+        if spec is not None and getattr(spec, "agree_on", None):
+            agree_on_field = spec.agree_on
+            break
+
+    def _get_field(obj: Any, name: str) -> Any:
+        if hasattr(obj, name):
+            return getattr(obj, name)
+        if isinstance(obj, dict):
+            return obj.get(name)
+        return obj
+
+    if agree_on_field is not None:
+        comparison_values = {str(_get_field(a, agree_on_field)) for a in last_round}
+    else:
+        comparison_values = {str(a) for a in last_round}
+
+    converged = len(comparison_values) == 1
+
+    return await synthesize(topic=topic, arguments=history, converged=converged)