pyfect-dev 0.1.0__py3-none-any.whl

pyfect/ai/_stream_resolution.py

@@ -0,0 +1,231 @@
+"""``stream_text``'s tool-call resolution: the provider's parts as they
+come, each tool call's handler started once the stream has moved past the
+call, and the finish parts after every tool result.
+
+A tool call waits one chunk before its handler starts, so a truncating
+finish (``"length"``, ``"content-filter"`` ...) arriving next means the
+handler never runs. On such a finish the running handlers are interrupted
+and each tool call left without a result gets an ``execution-interrupted``
+failure result.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any, cast
+
+import pyfect as fx
+from pyfect._core import Effect
+from pyfect.ai import _approvals as A
+from pyfect.ai import prompt as P
+from pyfect.ai import response as R
+from pyfect.ai import toolkit as K
+from pyfect.ai.id_generator import IdGenerator
+from pyfect.fiber_set import FiberSet
+from pyfect.queue import Queue
+from pyfect.semaphore import Semaphore
+from pyfect.stream import Stream
+from pyfect.stream import _constructors as StreamC
+from pyfect.stream import _sources as StreamSources
+
+if TYPE_CHECKING:
+    from pyfect._concurrency import Concurrency
+
+type EncodedPart = Any
+type Chunk = Sequence[EncodedPart]
+
+_COMPLETE: tuple[str, ...] = ("stop", "tool-calls", "pause")
+
+
+def incomplete_reason(reasons: Sequence[object]) -> str | None:
+    """The first finish reason that leaves tool calls unanswered."""
+    for reason in reasons:
+        if reason not in _COMPLETE:
+            return reason if isinstance(reason, str) else "unknown"
+    return None
+
+
+def interrupted_result(id: str, name: str, reason: str) -> R.AnyToolResultPart:
+    """The failure result of a tool call a truncated response left
+    unanswered."""
+    result = {
+        "type": "execution-interrupted",
+        "reason": "Tool call execution was interrupted because the response finished "
+        f'with reason "{reason}"',
+    }
+    return R.tool_result_part(
+        id=id, name=name, is_failure=True, result=result, encoded_result=result
+    )
+
+
+def result_part(id: str, name: str, result: K.HandlerResult) -> R.AnyToolResultPart:
+    return R.tool_result_part(
+        id=id,
+        name=name,
+        is_failure=result.is_failure,
+        result=result.result,
+        encoded_result=result.encoded_result,
+        preliminary=result.preliminary,
+    )
+
+
+def resolve_stream(
+    kit: K.WithHandler[Any, Any, Any, Any],
+    provider: Stream[EncodedPart, Any, Any],
+    on_chunk: Callable[[Chunk], Effect[list[Any], Any, Any]],
+    concurrency: Concurrency,
+    *,
+    messages: Sequence[P.Message],
+    emitted: Sequence[Any] = (),
+) -> Effect[Stream[Any, Any, Any], Any, Any]:
+    """The stream of ``emitted``, then ``provider``'s parts (decoded and
+    checked by ``on_chunk``) with the tool results of ``kit``'s tools, or
+    approval requests for the calls needing approval (in the context of
+    ``messages``). Needs a ``Scope``, which the handlers and the provider's
+    run live in."""
+
+    def start(queue: Queue[Any, Any]) -> Effect[Stream[Any, Any, Any], Any, Any]:
+        queue.offer_all_unsafe(emitted)
+
+        def run(fibers: FiberSet[None, Any], semaphore: Semaphore | None) -> Effect[Any, Any, Any]:
+            resolution = _Resolution(kit, queue, fibers, semaphore, messages)
+            return resolution.run(provider, on_chunk).fork_scoped()
+
+        return (
+            FiberSet[None, Any]
+            .make()
+            .flat_map(lambda fibers: _permits(concurrency).flat_map(lambda s: run(fibers, s)))
+            .as_(StreamSources.from_queue(queue))
+        )
+
+    return Queue[Any, Any].unbounded().flat_map(start)
+
+
+def _permits(concurrency: Concurrency) -> Effect[Semaphore | None]:
+    if concurrency is None or concurrency == "unbounded":
+        return fx.succeed(None)
+    return cast("Effect[Semaphore | None]", Semaphore.make(concurrency))
+
+
+class _Resolution:
+    """The state of one streamed response: the tool calls waiting for the
+    next chunk, those without a final result yet, and the finish parts."""
+
+    def __init__(
+        self,
+        kit: K.WithHandler[Any, Any, Any, Any],
+        queue: Queue[Any, Any],
+        fibers: FiberSet[None, Any],
+        semaphore: Semaphore | None,
+        messages: Sequence[P.Message],
+    ) -> None:
+        self.kit = kit
+        self.messages = messages
+        self.queue = queue
+        self.fibers = fibers
+        self.semaphore = semaphore
+        self.buffered: list[EncodedPart] = []
+        self.pending: dict[str, str] = {}
+        self.finishes: list[R.FinishPart] = []
+
+    def _incomplete(self) -> str | None:
+        return incomplete_reason([finish.reason for finish in self.finishes])
+
+    def _handle(self, call: EncodedPart) -> Effect[None, Any, Any]:
+        id, name = cast("str", call["id"]), cast("str", call["name"])
+        tool = self.kit.tools.get(name)
+        if tool is None:
+            return fx.void
+
+        def request(approval_id: str) -> None:
+            part = R.ToolApprovalRequestPart(approval_id=approval_id, tool_call_id=id)
+            self.queue.offer_unsafe(part)
+            self.pending.pop(id, None)
+
+        def run(needed: bool) -> Effect[None, Any, Any]:
+            if needed:
+                ids = IdGenerator().flat_map(lambda generator: generator.generate_id())
+                return ids.map(request)
+            return self._run_handler(id, name, call)
+
+        return A.needs_approval(tool, call, self.messages).flat_map(run)
+
+    def _run_handler(self, id: str, name: str, call: EncodedPart) -> Effect[None, Any, Any]:
+        def offer(result: K.HandlerResult) -> Effect[None]:
+            def run() -> None:
+                self.queue.offer_unsafe(result_part(id, name, result))
+                if not result.preliminary:
+                    self.pending.pop(id, None)
+
+            return fx.sync(run)
+
+        handled = self.kit.handle(name, call.get("params"), id)
+        return StreamC.unwrap(handled).run_for_each(offer)
+
+    def _fork_buffered(self) -> Effect[None, Any, Any]:
+        def run() -> Effect[None, Any, Any]:
+            calls, self.buffered = self.buffered, []
+            forks: list[Effect[object, Any, Any]] = []
+            for call in calls:
+                handle = self._handle(call)
+                if self.semaphore is not None:
+                    handle = self.semaphore.with_permits(handle)
+                forks.append(self.fibers.run(handle))
+            return fx.for_each(forks, lambda fork: fork).as_(None)
+
+        return fx.suspend(run)
+
+    def _chunk(
+        self, on_chunk: Callable[[Chunk], Effect[list[Any], Any, Any]], chunk: Chunk
+    ) -> Effect[None, Any, Any]:
+        def emitted(parts: list[Any]) -> Effect[None, Any, Any]:
+            immediate: list[Any] = []
+            for part in parts:
+                if isinstance(part, R.FinishPart):
+                    self.finishes.append(part)
+                else:
+                    immediate.append(part)
+            offered = self.queue.offer_all(immediate).as_(None) if immediate else fx.void
+            start = self._fork_buffered() if self._incomplete() is None else fx.void
+
+            def buffer() -> None:
+                for part in chunk:
+                    if part.get("type") == "tool-call" and part.get("providerExecuted") is not True:
+                        if part.get("name") in self.kit.tools:
+                            self.pending[cast("str", part["id"])] = cast("str", part["name"])
+                        self.buffered.append(part)
+
+            return offered.and_then(start).and_then(fx.sync(buffer))
+
+        return on_chunk(chunk).flat_map(emitted)
+
+    def _settle(self) -> Effect[None, Any, Any]:
+        def run() -> Effect[None, Any, Any]:
+            reason = self._incomplete()
+            if reason is None:
+                return self._fork_buffered().and_then(
+                    self.fibers.join().race_first(self.fibers.await_empty())
+                )
+            return self.fibers.clear().and_then(
+                fx.suspend(
+                    lambda: self.queue.offer_all(
+                        [interrupted_result(id, name, reason) for id, name in self.pending.items()]
+                    ).as_(None)
+                )
+            )
+
+        return fx.suspend(run)
+
+    def run(
+        self,
+        provider: Stream[EncodedPart, Any, Any],
+        on_chunk: Callable[[Chunk], Effect[list[Any], Any, Any]],
+    ) -> Effect[None, Any, Any]:
+        return (
+            provider.run_for_each_array(lambda chunk: self._chunk(on_chunk, chunk))
+            .and_then(self._settle())
+            .and_then(fx.suspend(lambda: self.queue.offer_all(self.finishes)))
+            .and_then(self.queue.end())
+            .tap_cause(self.queue.fail_cause)
+            .as_(None)
+        )

pyfect/ai/_structured_output_checks.py

@@ -0,0 +1,138 @@
+"""The checks of a structured output's rewritten nodes (Effect's
+``internal/structured-output.ts``): checks without the ``expected`` text a
+JSON Schema type already says, no-op copies of the checks the JSON Schema
+compiler needs, a record's lower size bounds as bounds of its entries, and
+the description of a rewritten shape."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import replace
+from typing import Any, cast
+
+from pyfect.schema import ast as A
+from pyfect.schema import is_min_length
+from pyfect.schema._jsonschema._checks import exporter
+
+_UNDESCRIBED = frozenset({"effect/schema/isFinite", "effect/schema/isInt"})
+"""The checks whose ``expected`` text is left out of descriptions: their
+JSON Schema type says it."""
+_DOCUMENTATION = ("description", "expected", "title")
+
+
+def _without(annotations: A.Annotations | None, keys: tuple[str, ...]) -> A.Annotations:
+    return {k: v for k, v in (annotations or {}).items() if k not in keys}
+
+
+def prepare_checks(checks: A.Checks | None) -> A.Checks | None:
+    """``checks`` without the ``expected`` text of ``is_finite`` / ``is_int``."""
+    if checks is None:
+        return None
+    out = tuple(map(_prepare_check, checks))
+    if all(a is b for a, b in zip(out, checks, strict=True)):
+        return checks
+    return cast("A.Checks", out)
+
+
+def _prepare_check(check: A.Check[Any]) -> A.Check[Any]:
+    if isinstance(check, A.FilterGroup):
+        checks = prepare_checks(check.checks)
+        return check if checks is check.checks else replace(check, checks=checks)
+    id = _representation(check.annotations).get("id")
+    if id in _UNDESCRIBED:
+        return replace(check, annotations=_without(check.annotations, ("expected",)))
+    return check
+
+
+def _representation(annotations: A.Annotations | None) -> Mapping[str, object]:
+    representation = (annotations or {}).get("representation")
+    return (
+        cast("Mapping[str, object]", representation) if isinstance(representation, Mapping) else {}
+    )
+
+
+def _skip(_: object, __: A.AST, ___: object) -> None:
+    return None
+
+
+def compiler_checks(checks: A.Checks | None) -> A.Checks | None:
+    """No-op copies of the checks that have a JSON Schema fragment: a
+    rewritten node gives the compiler its own JSON type without running the
+    original check on the rewritten shape."""
+    if checks is None:
+        return None
+    out = [c for check in checks for c in _compiler_check(check)]
+    return cast("A.Checks", tuple(out)) if out else None
+
+
+def _compiler_check(check: A.Check[Any]) -> list[A.Check[Any]]:
+    annotations = _compiler_annotations(check.annotations)
+    if isinstance(check, A.FilterGroup):
+        checks = compiler_checks(check.checks)
+        if checks is not None:
+            return [A.FilterGroup(checks, annotations)]
+    return [] if annotations is None else [A.Filter(_skip, annotations)]
+
+
+def _compiler_annotations(annotations: A.Annotations | None) -> A.Annotations | None:
+    if exporter(annotations) is None:
+        return None
+    kept = ("representation", "to_json_schema")
+    return {k: v for k, v in (annotations or {}).items() if k in kept}
+
+
+def record_checks(checks: A.Checks | None) -> A.Checks | None:
+    """The lower size bounds of a record, as bounds of its entries. Upper
+    bounds are left out: duplicate keys collapse when decoding, so a
+    ``max_items`` could reject entries the codec accepts."""
+    if checks is None:
+        return None
+    out = [c for check in checks for c in _record_check(check)]
+    return cast("A.Checks", tuple(out)) if out else None
+
+
+def _record_check(check: A.Check[Any]) -> list[A.Check[Any]]:
+    if isinstance(check, A.FilterGroup):
+        return [c for child in check.checks for c in _record_check(child)]
+    representation = _representation(check.annotations)
+    payload = representation.get("payload")
+    if not isinstance(payload, Mapping):
+        return []
+    payload = cast("Mapping[str, object]", payload)
+    match representation.get("id"):
+        case "effect/schema/isMinProperties":
+            minimum = payload.get("minProperties")
+        case "effect/schema/isBetweenProperties":
+            minimum = payload.get("minimum")
+        case _:
+            return []
+    if not isinstance(minimum, int | float):
+        return []
+    filter = is_min_length(minimum)
+    return [replace(filter, annotations=_without(filter.annotations, _DOCUMENTATION))]
+
+
+def structural_annotations(node: A.AST, description: str) -> A.Annotations:
+    descriptions: list[str] = []
+    _append_description(descriptions, node.annotations)
+    for check in node.checks or ():
+        _append_check_descriptions(descriptions, check)
+    if not descriptions:
+        return {"description": description}
+    return {"description": f"{description}; {' and '.join(descriptions)}"}
+
+
+def _append_check_descriptions(descriptions: list[str], check: A.Check[Any]) -> None:
+    _append_description(descriptions, check.annotations)
+    if isinstance(check, A.FilterGroup):
+        for child in check.checks:
+            _append_check_descriptions(descriptions, child)
+
+
+def _append_description(descriptions: list[str], annotations: A.Annotations | None) -> None:
+    annotations = annotations or {}
+    description = annotations.get("description")
+    if description is None:
+        description = annotations.get("expected")
+    if isinstance(description, str) and description not in descriptions:
+        descriptions.append(description)

pyfect/ai/_structured_output_codec.py

@@ -0,0 +1,229 @@
+"""The codec of a structured output: a schema's canonical JSON form with
+the shapes providers' strict JSON Schema modes cannot express rewritten
+(Effect's ``internal/structured-output.ts``, ``toCodec``).
+
+A tuple becomes an object with keys ``"0"``, ``"1"``, ... (and
+``"__rest__"`` / ``"__tail_N__"`` for its rest), an object with index
+signatures an array of ``[key, value]`` pairs, an optional key a required
+nullable one, a ``one_of`` union an ``any_of`` one. Each rewrite decodes back
+to the original shape, so the original schema's checks still run on the
+decoded value; a check that has a JSON Schema fragment is kept, as a no-op,
+on the rewritten node for the JSON Schema compiler.
+"""
+
+from __future__ import annotations
+
+import operator
+from collections.abc import Mapping, Sequence
+from dataclasses import replace
+from typing import Any, NamedTuple, cast
+
+from pyfect.ai import _structured_output_checks as C
+from pyfect.json_schema import JsonSchema
+from pyfect.option import Option, Some
+from pyfect.schema import Bottom, Codec, to_codec_json
+from pyfect.schema import ast as A
+from pyfect.schema.transformation import (
+    Transformation,
+    passthrough,
+    transform,
+    transform_optional,
+)
+
+_PASSTHROUGH: Transformation[object, object] = passthrough()
+_REST = "__rest__"
+_RECORD_DESCRIPTION = (
+    "Object encoded as array of [key, value] pairs. Apply object constraints to the decoded object"
+)
+_TUPLE_DESCRIPTION = (
+    "Tuple encoded as an object with numeric string keys ('0', '1', ...). "
+    "If present, '__rest__' contains remaining elements"
+)
+_TUPLE_TAIL_DESCRIPTION = (
+    f"{_TUPLE_DESCRIPTION}. Post-rest elements use '__tail_0__', '__tail_1__', and so on"
+)
+
+
+class StructuredOutputCodec[T, RD, RE](NamedTuple):
+    """What a provider's structured-output transformer gives: the codec
+    decoding the model's output into ``T`` and the JSON Schema to send. The
+    codec stays the authority: the JSON Schema may accept more than it."""
+
+    codec: Codec[T, object, RD, RE]
+    json_schema: JsonSchema
+
+
+def _tail(index: int) -> str:
+    return f"__tail_{index}__"
+
+
+def to_codec[T, E, RD, RE](schema: Codec[T, E, RD, RE]) -> Codec[T, object, RD, RE]:
+    """``schema`` decoding from the provider-safe shape of its JSON form;
+    ``schema`` itself when nothing is rewritten."""
+    original = schema.__codec__().ast
+    json_codec = to_codec_json(schema)
+    encoded = A.to_encoded(json_codec.ast)
+    rewritten = _Rewrite()(encoded)
+    if rewritten is encoded:
+        return schema if json_codec.ast is original else json_codec
+    return Bottom[T, object, RD, RE](A.decode_to(rewritten, json_codec.ast, _PASSTHROUGH))
+
+
+class _Rewrite:
+    def __init__(self) -> None:
+        self._suspended: dict[int, A.Suspend] = {}
+
+    def __call__(self, node: A.AST) -> A.AST:
+        match node:
+            case A.Union():
+                mode = "any_of" if node.mode == "one_of" else node.mode
+                return _rebuild(node, types=tuple(map(self, node.types)), mode=mode)
+            case A.Arrays():
+                if node.elements or len(node.rest) > 1:
+                    return self._tuple_to_object(node)
+                return _rebuild(node, rest=tuple(map(self, node.rest)))
+            case A.Objects():
+                return self._objects(node)
+            case A.Suspend():
+                return self._suspend(node)
+            case _:
+                return _rebuild(node)
+
+    def _objects(self, node: A.Objects) -> A.AST:
+        if _is_empty_params(node):
+            return node
+        if node.index_signatures:
+            return self._object_to_entries(node)
+        return _rebuild(node, properties=tuple(map(self._property, node.properties)))
+
+    def _suspend(self, node: A.Suspend) -> A.Suspend:
+        """One rewritten copy per node, so a recursive schema stays a
+        finite graph."""
+        cached = self._suspended.get(id(node))
+        if cached is None:
+
+            def thunk() -> A.AST:
+                return self(node.target)
+
+            cached = replace(node, thunk=thunk)
+            self._suspended[id(node)] = cached
+        return cached
+
+    def _property(self, prop: A.PropertySignature) -> A.PropertySignature:
+        type = self(prop.type)
+        if prop.optional:
+            return A.PropertySignature(prop.name, _optional_to_nullable(type), True, None, False)
+        return prop if type is prop.type else replace(prop, type=type)
+
+    def _tuple_to_object(self, node: A.Arrays) -> A.AST:
+        elements, rest = node.elements, node.rest
+        properties = [
+            A.PropertySignature(str(i), element.type, element.optional)
+            for i, element in enumerate(elements)
+        ]
+        if rest:
+            properties.append(A.PropertySignature(_REST, _list(rest[0])))
+            properties.extend(
+                A.PropertySignature(_tail(i), rest[i + 1]) for i in range(len(rest) - 1)
+            )
+        description = _TUPLE_TAIL_DESCRIPTION if len(rest) > 1 else _TUPLE_DESCRIPTION
+        from_ = self(
+            A.Objects(
+                properties=tuple(properties),
+                annotations=C.structural_annotations(node, description),
+                checks=C.compiler_checks(node.checks),
+            )
+        )
+        tail_length = max(len(rest) - 1, 0)
+
+        def decode(obj: Mapping[str, object]) -> list[object]:
+            out = [obj[str(i)] for i in range(len(elements)) if str(i) in obj]
+            if _REST in obj:
+                out.extend(cast("Sequence[object]", obj[_REST]))
+            out.extend(obj[_tail(i)] for i in range(tail_length))
+            return out
+
+        def encode(values: Sequence[object]) -> dict[str, object]:
+            out: dict[str, object] = {
+                str(i): values[i] for i in range(min(len(elements), len(values)))
+            }
+            if rest:
+                end = max(len(elements), len(values) - tail_length)
+                out[_REST] = list(values[len(elements) : end])
+                for i in range(tail_length):
+                    out[_tail(i)] = values[end + i]
+            return out
+
+        return A.decode_to(from_, node, transform(decode=decode, encode=encode))
+
+    def _object_to_entries(self, node: A.Objects) -> A.AST:
+        key = _union_or_single(
+            [A.Literal(cast("str", p.name)) for p in node.properties]
+            + [i.parameter for i in node.index_signatures]
+        )
+        value = _union_or_single(
+            [p.type for p in node.properties] + [i.type for i in node.index_signatures]
+        )
+        from_ = self(
+            A.Arrays(
+                rest=(A.Arrays(elements=(A.Element(key), A.Element(value)), mutable=True),),
+                mutable=True,
+                annotations=C.structural_annotations(node, _RECORD_DESCRIPTION),
+                checks=A.join_checks(C.record_checks(node.checks), C.compiler_checks(node.checks)),
+            )
+        )
+
+        def decode(entries: Sequence[Sequence[object]]) -> dict[object, object]:
+            return {entry[0]: entry[1] for entry in entries}
+
+        def encode(obj: Mapping[object, object]) -> list[list[object]]:
+            return [[k, v] for k, v in obj.items()]
+
+        return A.decode_to(from_, node, transform(decode=decode, encode=encode))
+
+
+def _rebuild[N: A.AST](node: N, **parts: Any) -> N:
+    """``node`` with ``parts`` and its checks prepared; ``node`` itself when
+    each part is the one it has."""
+    parts["checks"] = C.prepare_checks(node.checks)
+    if all(_same(getattr(node, name), part) for name, part in parts.items()):
+        return node
+    return replace(node, **parts)
+
+
+def _same(a: object, b: object) -> bool:
+    if isinstance(a, tuple) and isinstance(b, tuple):
+        items_a, items_b = cast("tuple[object, ...]", a), cast("tuple[object, ...]", b)
+        return len(items_a) == len(items_b) and all(map(operator.is_, items_a, items_b))
+    return a is b
+
+
+def _list(item: A.AST) -> A.Arrays:
+    return A.Arrays(rest=(item,), mutable=True)
+
+
+def _is_empty_params(node: A.Objects) -> bool:
+    """Effect's ``Tool.isEmptyParamsRecord``: ``Record(String, Never)``."""
+    if node.properties or len(node.index_signatures) != 1:
+        return False
+    index = node.index_signatures[0]
+    return index.parameter is A.string and isinstance(index.type, A.Never)
+
+
+def _union_or_single(types: list[A.AST]) -> A.AST:
+    unique = list({id(t): t for t in types}.values())
+    return unique[0] if len(unique) == 1 else A.Union(tuple(unique))
+
+
+def _optional_to_nullable(type: A.AST) -> A.AST:
+    """``type`` as an optional key decoded from a required nullable one."""
+
+    def decode(value: Option[object]) -> Option[object]:
+        return None if value is None or value.value is None else value
+
+    def encode(value: Option[object]) -> Option[object]:
+        return Some(None) if value is None else value
+
+    return A.decode_to(
+        A.Union((type, A.null)), type, transform_optional(decode=decode, encode=encode)
+    )