composeai 0.1.0__py3-none-any.whl

composeai/__init__.py

@@ -0,0 +1,33 @@
+"""composeai: a radically simple, functional framework for multi-agent AI workflows.
+
+Typed agent functions, pipe/aggregate/map composition with build-time type
+checks, always-on local tracing, and durable, resumable flows.
+"""
+
+from .agentfn import agent, prompt
+from .combinators import aggregate, map, pipe
+from .flow import flow, resume, task
+from .hitl import Interrupt, approve, ask_human
+from .models.compatible import openai_compatible
+from .runs import Budget, Run
+from .tools import tool
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Budget",
+    "Interrupt",
+    "Run",
+    "agent",
+    "aggregate",
+    "approve",
+    "ask_human",
+    "flow",
+    "map",
+    "openai_compatible",
+    "pipe",
+    "prompt",
+    "resume",
+    "task",
+    "tool",
+]

composeai/_encoding.py

@@ -0,0 +1,240 @@
+"""JSON-safe encode/decode for journal and trace persistence.
+
+``to_jsonable``/``from_jsonable`` turn a restricted set of Python values
+into plain ``dict``/``list``/scalar JSON trees (and back), tagging
+non-trivial types with a reserved ``"$kind"`` key so they can be
+rehydrated. Deliberately **no pickle** and **no dynamic import**:
+decoding data that references an unregistered type raises
+:class:`~composeai.errors.SerializationError` naming the missing type
+rather than importing it -- guessing and importing arbitrary dotted
+paths found in persisted data is a code-execution hole (this exact
+issue was found and fixed in a prior audit). Callers must import the
+module that defines a type -- or call :func:`register_serializable` --
+before decoding data that references it.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import datetime
+import enum
+from typing import Any
+
+from pydantic import BaseModel
+
+from .errors import SerializationError
+
+_KIND = "$kind"
+_TYPE = "$type"
+_VALUE = "value"
+
+# Type registry: never used to import anything, only to look up classes
+# that some part of this process has already imported and registered.
+_REGISTRY: dict[str, type] = {}
+
+
+def register_serializable(cls: type) -> type:
+    """Register ``cls`` so :func:`from_jsonable` can rehydrate it by tag.
+
+    Encoding a pydantic model, dataclass, or enum auto-registers its
+    class. Call this explicitly in a process that decodes data it never
+    encoded (e.g. a fresh journal reader started in a new process).
+    Returns ``cls`` unchanged, so it can also be used as a decorator.
+    """
+    _REGISTRY[_type_tag(cls)] = cls
+    return cls
+
+
+def _type_tag(cls: type) -> str:
+    return f"{cls.__module__}:{cls.__qualname__}"
+
+
+def _child_path(path: str, key: str) -> str:
+    return f"{path}.{key}" if path else key
+
+
+def _index_path(path: str, index: int) -> str:
+    return f"{path}[{index}]"
+
+
+def _is_namedtuple(obj: Any) -> bool:
+    """Best-effort, standard-library-idiomatic NamedTuple detection.
+
+    ``isinstance(obj, tuple)`` is true for NamedTuple instances too (they're
+    tuple subclasses) -- this must be checked *before* the generic tuple
+    branch, or a NamedTuple silently degrades to a plain tuple.
+    """
+    return (
+        isinstance(obj, tuple)
+        and type(obj) is not tuple
+        and hasattr(type(obj), "_fields")
+        and hasattr(type(obj), "_asdict")
+    )
+
+
+def to_jsonable(obj: Any) -> Any:
+    """Convert ``obj`` into a plain JSON-safe tree.
+
+    Raises :class:`SerializationError` naming the offending path if any
+    value nested in ``obj`` isn't supported.
+    """
+    return _encode(obj, "")
+
+
+def _encode(obj: Any, path: str) -> Any:
+    if obj is None or isinstance(obj, (bool, int, float, str)):
+        return obj
+
+    if isinstance(obj, BaseModel):
+        register_serializable(type(obj))
+        return {
+            _KIND: "pydantic",
+            _TYPE: _type_tag(type(obj)),
+            _VALUE: obj.model_dump(mode="json"),
+        }
+
+    if isinstance(obj, enum.Enum):
+        register_serializable(type(obj))
+        return {_KIND: "enum", _TYPE: _type_tag(type(obj)), _VALUE: obj.name}
+
+    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
+        register_serializable(type(obj))
+        non_init = [field.name for field in dataclasses.fields(obj) if not field.init]
+        if non_init:
+            raise SerializationError(
+                f"{path or '<root>'}: dataclass {type(obj).__name__!r} has non-init "
+                f"field(s) {non_init!r} (declared with field(init=False), typically "
+                "computed in __post_init__) -- these can't round-trip the journal: "
+                "encoding would capture the computed value, but decoding calls "
+                "cls(**fields), which rejects any keyword the real __init__ doesn't "
+                "accept. Exclude non-init fields from what you journal (e.g. store a "
+                "plain dict/dataclass of only the init fields), or give this "
+                "dataclass a classmethod constructor you call explicitly on decode."
+            )
+        value = {
+            field.name: _encode(getattr(obj, field.name), _child_path(path, field.name))
+            for field in dataclasses.fields(obj)
+        }
+        return {_KIND: "dataclass", _TYPE: _type_tag(type(obj)), _VALUE: value}
+
+    if isinstance(obj, datetime.datetime):
+        return {_KIND: "datetime", _VALUE: obj.isoformat()}
+    if isinstance(obj, datetime.date):
+        return {_KIND: "date", _VALUE: obj.isoformat()}
+    if isinstance(obj, datetime.time):
+        return {_KIND: "time", _VALUE: obj.isoformat()}
+
+    if _is_namedtuple(obj):
+        raise SerializationError(
+            f"{path or '<root>'}: {type(obj).__name__!r} is a NamedTuple, which the "
+            "journal encoder doesn't support -- it's a tuple subclass, so encoding it "
+            "as a plain tuple would silently lose field-name access on decode (an "
+            "AttributeError on the replayed value where the original worked). Convert "
+            "it to a dataclass, a dict, or a plain tuple before journaling it."
+        )
+    if isinstance(obj, tuple):
+        return {
+            _KIND: "tuple",
+            _VALUE: [_encode(v, _index_path(path, i)) for i, v in enumerate(obj)],
+        }
+    if isinstance(obj, frozenset):
+        return {
+            _KIND: "frozenset",
+            _VALUE: [_encode(v, _index_path(path, i)) for i, v in enumerate(obj)],
+        }
+    if isinstance(obj, set):
+        return {
+            _KIND: "set",
+            _VALUE: [_encode(v, _index_path(path, i)) for i, v in enumerate(obj)],
+        }
+
+    if isinstance(obj, list):
+        return [_encode(v, _index_path(path, i)) for i, v in enumerate(obj)]
+
+    if isinstance(obj, dict):
+        for key in obj:
+            if not isinstance(key, str):
+                raise SerializationError(
+                    f"{path or '<root>'}: dict keys must be str, got {key!r}"
+                )
+        encoded = {k: _encode(v, _child_path(path, k)) for k, v in obj.items()}
+        if _KIND in obj:
+            # A plain dict that happens to use our reserved key: escape it
+            # so decoding is unambiguous.
+            return {_KIND: "dict", _VALUE: encoded}
+        return encoded
+
+    raise SerializationError(f"{path or '<root>'}: {type(obj)!r} is not JSON-serializable")
+
+
+def from_jsonable(data: Any) -> Any:
+    """Inverse of :func:`to_jsonable`.
+
+    Raises :class:`SerializationError` if a ``$type`` tag isn't in the
+    registry (see :func:`register_serializable`). Never imports modules
+    to resolve a type.
+    """
+    if isinstance(data, dict):
+        if _KIND in data:
+            return _decode_tagged(data)
+        return {k: from_jsonable(v) for k, v in data.items()}
+    if isinstance(data, list):
+        return [from_jsonable(v) for v in data]
+    return data
+
+
+def _decode_tagged(data: dict[str, Any]) -> Any:
+    kind = data.get(_KIND)
+    try:
+        if kind == "dict":
+            return {k: from_jsonable(v) for k, v in data[_VALUE].items()}
+        if kind == "pydantic":
+            cls = _resolve_type(data[_TYPE])
+            value = {k: from_jsonable(v) for k, v in data[_VALUE].items()}
+            return cls.model_validate(value)
+        if kind == "dataclass":
+            cls = _resolve_type(data[_TYPE])
+            value = {k: from_jsonable(v) for k, v in data[_VALUE].items()}
+            return cls(**value)
+        if kind == "enum":
+            cls = _resolve_type(data[_TYPE])
+            return cls[data[_VALUE]]
+        if kind == "datetime":
+            return datetime.datetime.fromisoformat(data[_VALUE])
+        if kind == "date":
+            return datetime.date.fromisoformat(data[_VALUE])
+        if kind == "time":
+            return datetime.time.fromisoformat(data[_VALUE])
+        if kind == "tuple":
+            return tuple(from_jsonable(v) for v in data[_VALUE])
+        if kind == "set":
+            return {from_jsonable(v) for v in data[_VALUE]}
+        if kind == "frozenset":
+            return frozenset(from_jsonable(v) for v in data[_VALUE])
+    except SerializationError:
+        raise  # already the right type (e.g. _resolve_type's "unregistered type") -- pass through
+    except (KeyError, TypeError) as exc:
+        # A tagged value missing an expected key (e.g. `{"$kind": "dataclass",
+        # "value": {...}}` with no "$type" -- DB corruption, a manual edit, a
+        # future schema/version mismatch, ...) used to raise a raw KeyError/
+        # TypeError here instead of this module's own SerializationError, so
+        # callers that specifically catch SerializationError (the module's
+        # documented contract) didn't catch this.
+        raise SerializationError(
+            f"corrupted journal/state row (kind={kind!r}): {type(exc).__name__}: {exc}"
+        ) from exc
+
+    raise SerializationError(f"Unknown {_KIND!r} tag: {kind!r}")
+
+
+def _resolve_type(tag: str) -> Any:
+    cls = _REGISTRY.get(tag)
+    if cls is not None:
+        return cls
+    module_hint = tag.split(":", 1)[0]
+    raise SerializationError(
+        f"Cannot decode unregistered type {tag!r}: import the module that "
+        f"defines it (e.g. `import {module_hint}`) or call "
+        f"register_serializable(...) before decoding this data. "
+        f"Types are never imported automatically, by design (security)."
+    )

composeai/_ids.py

@@ -0,0 +1,41 @@
+"""Sortable unique identifiers (ULID-style) for composeai.
+
+A ``new_ulid()`` result is a 26-character Crockford Base32 string encoding
+128 bits: a 48-bit Unix-millisecond timestamp followed by 80 bits of
+cryptographically random data from :func:`os.urandom`. Because the
+timestamp occupies the most significant bits, two IDs generated at least
+1 millisecond apart sort lexicographically in time order -- useful for
+run/span/message identifiers that should also read as roughly
+chronological when listed.
+
+Stdlib only; no third-party dependency.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+
+# Crockford's Base32 alphabet: excludes I, L, O, U to avoid visual ambiguity.
+_ALPHABET = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
+_TIMESTAMP_BYTES = 6  # 48 bits
+_RANDOM_BYTES = 10  # 80 bits
+_ENCODED_LENGTH = 26  # 128 bits / 5 bits-per-char (rounded up)
+
+
+def new_ulid() -> str:
+    """Return a new 26-character, lexicographically time-sortable ID."""
+    timestamp_ms = int(time.time() * 1000)
+    ts_bytes = timestamp_ms.to_bytes(_TIMESTAMP_BYTES, byteorder="big")
+    rand_bytes = os.urandom(_RANDOM_BYTES)
+    return _encode(ts_bytes + rand_bytes)
+
+
+def _encode(data: bytes) -> str:
+    """Encode 16 raw bytes (128 bits) as 26 Crockford Base32 characters."""
+    value = int.from_bytes(data, byteorder="big")
+    chars = [""] * _ENCODED_LENGTH
+    for i in range(_ENCODED_LENGTH - 1, -1, -1):
+        chars[i] = _ALPHABET[value & 0x1F]
+        value >>= 5
+    return "".join(chars)

composeai/_schema.py

@@ -0,0 +1,147 @@
+"""Shared JSON Schema post-processing for ``@tool`` and ``@agent`` structured output.
+
+Both derive their schemas from pydantic (``create_model`` / ``TypeAdapter``) and
+both need the same treatment before handing a schema to a model: strict
+``additionalProperties: false`` everywhere -- including nested ``$defs`` --
+and no pydantic-generated ``title`` noise (models rely on ``description``,
+not ``title``, and a stray auto-title only wastes tokens).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import enum
+import inspect
+from collections.abc import Callable
+from typing import Any, get_args, get_type_hints
+
+from pydantic import BaseModel
+
+from ._encoding import register_serializable
+
+
+def resolve_annotations(fn: Callable[..., Any], *, include_extras: bool = True) -> dict[str, Any]:
+    """Best-effort ``typing.get_type_hints`` for ``fn``.
+
+    ``get_type_hints`` resolves string-form annotations (e.g. under
+    ``from __future__ import annotations``) against ``fn.__globals__``, but
+    it cannot see names from an enclosing function's local/closure scope --
+    a common pattern in tests that define a small model right next to the
+    function using it. When that lookup fails, fall back to each
+    parameter's raw annotation from :func:`inspect.signature`: on every
+    supported Python version, an annotation that's already a real object
+    (the common case without the ``__future__`` import, and -- since
+    Python 3.14's lazy annotation evaluation -- even local/closure names
+    under it) survives untouched; only an unresolved string annotation
+    combining *both* the ``__future__`` import *and* a local/closure name
+    would still come back as a plain string here.
+    """
+    try:
+        return get_type_hints(fn, include_extras=include_extras)
+    except NameError:
+        pass
+
+    sig = inspect.signature(fn)
+    raw: dict[str, Any] = {
+        name: param.annotation
+        for name, param in sig.parameters.items()
+        if param.annotation is not inspect.Parameter.empty
+    }
+    if sig.return_annotation is not inspect.Signature.empty:
+        raw["return"] = sig.return_annotation
+    return raw
+
+
+def register_annotation_types(fn: Callable[..., Any]) -> None:
+    """Register every pydantic model / dataclass / enum reachable from ``fn``'s annotations.
+
+    Walks :func:`resolve_annotations` (every parameter plus the return
+    type), recursing into generic args (``typing.get_args`` -- e.g. the
+    ``Inner`` in ``list[Inner]`` or ``dict[str, Outer]``) and, for
+    pydantic models and dataclasses, into their own field types. Called by
+    ``@task``/``@flow``/``@agent`` decoration so a *fresh process* (e.g. a
+    resumed flow started via a new ``python`` invocation) can decode
+    journaled values referencing these types via
+    :func:`~composeai._encoding.from_jsonable` without composeai ever
+    importing a type dynamically -- that would be a code-execution hole,
+    so decoding an unregistered type always raises instead (see
+    :mod:`composeai._encoding`).
+    """
+    hints = resolve_annotations(fn, include_extras=True)
+    seen: set[type] = set()
+    for value in hints.values():
+        _register_type_recursive(value, seen)
+
+
+def _register_type_recursive(annotation: Any, seen: set[type]) -> None:
+    if isinstance(annotation, type):
+        if annotation in seen:
+            return
+        seen.add(annotation)
+        if issubclass(annotation, BaseModel):
+            register_serializable(annotation)
+            for field_info in annotation.model_fields.values():
+                _register_type_recursive(field_info.annotation, seen)
+        elif dataclasses.is_dataclass(annotation):
+            register_serializable(annotation)
+            # `dataclasses.fields(annotation)[i].type` is the field's *raw*
+            # annotation -- under `from __future__ import annotations` (or
+            # any other postponed-evaluation source), that's an unresolved
+            # string (e.g. `"Inner"`), not the actual class object. Recursing
+            # on it directly means `isinstance(..., type)` is always False
+            # and `get_args(...)` always empty, so a dataclass field typed
+            # with another custom dataclass/enum/pydantic model silently
+            # never gets registered -- exactly the gap the pydantic branch
+            # above doesn't have, since `model_fields[...].annotation` is
+            # always pydantic-resolved to a real type object already.
+            # `typing.get_type_hints` resolves the same way `model_fields`
+            # does; fall back to the raw (possibly-string) annotation if
+            # resolution itself fails (e.g. a name only in a caller's local
+            # scope -- see `resolve_annotations`'s own docstring for the
+            # same trade-off).
+            try:
+                hints = get_type_hints(annotation, include_extras=True)
+            except NameError:
+                hints = {}
+            for field in dataclasses.fields(annotation):
+                _register_type_recursive(hints.get(field.name, field.type), seen)
+        elif issubclass(annotation, enum.Enum):
+            register_serializable(annotation)
+        return
+
+    for arg in get_args(annotation):
+        _register_type_recursive(arg, seen)
+
+
+def seal_schema(schema: Any) -> None:
+    """Recursively strip auto-generated ``title`` keys and seal every object schema.
+
+    Mutates ``schema`` in place. "Object schema" means any nested dict with
+    ``"type": "object"`` -- including entries under ``$defs`` -- which gets
+    ``additionalProperties: false`` set, *unless* it already has a
+    dict-valued ``additionalProperties`` (a real schema, not a boolean --
+    what pydantic generates for an open mapping like ``dict[str, V]``/
+    ``Mapping[str, V]``, describing the permitted *value* type). Overwriting
+    that with ``False`` would turn a schema that legitimately accepts any
+    string-keyed dict of ``V`` into one that only ever matches ``{}`` --
+    silently forbidding the model from ever populating the field/argument
+    it's meant to constrain, not forbid. A boolean (or absent)
+    ``additionalProperties`` is still always forced to ``False`` (the
+    fixed-shape-object case this was designed for -- pydantic models and
+    dataclasses). Non-dict values (e.g. entries of a ``"required"`` list)
+    are left untouched.
+    """
+    if not isinstance(schema, dict):
+        return
+    # Auto-generated titles are always strings; a dict under a "title" key is
+    # a *property named title* (inside a "properties" map) and must survive.
+    if isinstance(schema.get("title"), str):
+        schema.pop("title")
+    if schema.get("type") == "object" and not isinstance(schema.get("additionalProperties"), dict):
+        schema["additionalProperties"] = False
+    for value in schema.values():
+        if isinstance(value, dict):
+            seal_schema(value)
+        elif isinstance(value, list):
+            for item in value:
+                seal_schema(item)