weakincentives 0.9.0__py3-none-any.whl

weakincentives/runtime/session/reducers.py

@@ -0,0 +1,82 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Built-in reducers for Session state slices."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import cast
+
+from ...dbc import pure
+from ...prompt._types import SupportsDataclass
+from ._types import ReducerContextProtocol, ReducerEvent, TypedReducer
+
+
+@pure
+def append[T: SupportsDataclass](
+    slice_values: tuple[T, ...],
+    event: ReducerEvent,
+    *,
+    context: ReducerContextProtocol,
+) -> tuple[T, ...]:
+    """Append the event value if it is not already present."""
+
+    del context
+    value = cast(T, event.value)
+    if value in slice_values:
+        return slice_values
+    return (*slice_values, value)
+
+
+def upsert_by[T: SupportsDataclass, K](key_fn: Callable[[T], K]) -> TypedReducer[T]:
+    """Return a reducer that upserts items sharing the same derived key."""
+
+    def reducer(
+        slice_values: tuple[T, ...],
+        event: ReducerEvent,
+        *,
+        context: ReducerContextProtocol,
+    ) -> tuple[T, ...]:
+        del context
+        value = cast(T, event.value)
+        key = key_fn(value)
+        updated: list[T] = []
+        replaced = False
+        for existing in slice_values:
+            if key_fn(existing) == key:
+                if not replaced:
+                    updated.append(value)
+                    replaced = True
+                continue
+            updated.append(existing)
+        if not replaced:
+            updated.append(value)
+        return tuple(updated)
+
+    return pure(reducer)
+
+
+@pure
+def replace_latest[T: SupportsDataclass](
+    slice_values: tuple[T, ...],
+    event: ReducerEvent,
+    *,
+    context: ReducerContextProtocol,
+) -> tuple[T, ...]:
+    """Keep only the most recent event value."""
+
+    del context
+    return (cast(T, event.value),)
+
+
+__all__ = ["append", "replace_latest", "upsert_by"]

weakincentives/runtime/session/selectors.py

@@ -0,0 +1,56 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helpers for querying Session slices."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from ...dbc import pure
+from ...prompt._types import SupportsDataclass
+from .session import Session
+
+
+@pure
+def select_all[T: SupportsDataclass](
+    session: Session, slice_type: type[T]
+) -> tuple[T, ...]:
+    """Return the entire slice for the provided type."""
+
+    return session.select_all(slice_type)
+
+
+@pure
+def select_latest[T: SupportsDataclass](
+    session: Session, slice_type: type[T]
+) -> T | None:
+    """Return the most recent item in the slice, if any."""
+
+    values = session.select_all(slice_type)
+    if not values:
+        return None
+    return values[-1]
+
+
+@pure
+def select_where[T: SupportsDataclass](
+    session: Session,
+    slice_type: type[T],
+    predicate: Callable[[T], bool],
+) -> tuple[T, ...]:
+    """Return items that satisfy the predicate."""
+
+    return tuple(value for value in session.select_all(slice_type) if predicate(value))
+
+
+__all__ = ["select_all", "select_latest", "select_where"]

weakincentives/runtime/session/session.py

@@ -0,0 +1,387 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Session state container synchronized with the event bus."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, replace
+from datetime import UTC, datetime
+from threading import RLock
+from typing import Any, cast, override
+from uuid import UUID, uuid4
+
+from ...dbc import invariant
+from ...prompt._types import SupportsDataclass
+from ..events import EventBus, PromptExecuted, PromptRendered, ToolInvoked
+from ..logging import StructuredLogger, get_logger
+from ._slice_types import SessionSlice, SessionSliceType
+from ._types import ReducerContextProtocol, ReducerEvent, TypedReducer
+from .dataclasses import is_dataclass_instance
+from .protocols import SessionProtocol, SnapshotProtocol
+from .reducers import append
+from .snapshots import (
+    Snapshot,
+    SnapshotRestoreError,
+    SnapshotSerializationError,
+    SnapshotState,
+    normalize_snapshot_state,
+)
+
+logger: StructuredLogger = get_logger(__name__, context={"component": "session"})
+
+
+type DataEvent = PromptExecuted | PromptRendered | ToolInvoked
+
+
+_PROMPT_RENDERED_TYPE: type[SupportsDataclass] = cast(
+    type[SupportsDataclass], PromptRendered
+)
+_TOOL_INVOKED_TYPE: type[SupportsDataclass] = cast(type[SupportsDataclass], ToolInvoked)
+_PROMPT_EXECUTED_TYPE: type[SupportsDataclass] = cast(
+    type[SupportsDataclass], PromptExecuted
+)
+
+EMPTY_SLICE: SessionSlice = ()
+
+
+def _append_event(
+    slice_values: tuple[SupportsDataclass, ...],
+    event: ReducerEvent,
+    *,
+    context: ReducerContextProtocol,
+) -> tuple[SupportsDataclass, ...]:
+    del context
+    appended = cast(SupportsDataclass, event)
+    return (*slice_values, appended)
+
+
+@dataclass(slots=True)
+class _ReducerRegistration:
+    reducer: TypedReducer[Any]
+    slice_type: SessionSliceType
+
+
+def _session_id_is_well_formed(session: "Session") -> bool:  # noqa: UP037
+    return len(session.session_id.bytes) == 16
+
+
+def _created_at_has_tz(session: "Session") -> bool:  # noqa: UP037
+    return session.created_at.tzinfo is not None
+
+
+def _created_at_is_utc(session: "Session") -> bool:  # noqa: UP037
+    return session.created_at.tzinfo == UTC
+
+
+@invariant(
+    _session_id_is_well_formed,
+    _created_at_has_tz,
+    _created_at_is_utc,
+)
+class Session(SessionProtocol):
+    """Collect dataclass payloads from prompt executions and tool invocations."""
+
+    def __init__(
+        self,
+        *,
+        bus: EventBus,
+        session_id: UUID | None = None,
+        created_at: datetime | None = None,
+    ) -> None:
+        super().__init__()
+        resolved_session_id = session_id if session_id is not None else uuid4()
+        resolved_created_at = (
+            created_at if created_at is not None else datetime.now(UTC)
+        )
+        if resolved_created_at.tzinfo is None:
+            msg = "Session created_at must be timezone-aware."
+            raise ValueError(msg)
+
+        self.session_id: UUID = resolved_session_id
+        self.created_at: datetime = resolved_created_at.astimezone(UTC)
+        self._bus: EventBus = bus
+        self._reducers: dict[SessionSliceType, list[_ReducerRegistration]] = {}
+        self._state: dict[SessionSliceType, SessionSlice] = {}
+        self._lock = RLock()
+        self._subscriptions_attached = False
+        self._attach_to_bus(bus)
+
+    def clone(
+        self,
+        *,
+        bus: EventBus,
+        session_id: UUID | None = None,
+        created_at: datetime | None = None,
+    ) -> Session:
+        """Return a new session that mirrors the current state and reducers."""
+
+        with self._lock:
+            reducer_snapshot = [
+                (data_type, tuple(registrations))
+                for data_type, registrations in self._reducers.items()
+            ]
+            state_snapshot = dict(self._state)
+
+        clone = Session(
+            bus=bus,
+            session_id=session_id if session_id is not None else self.session_id,
+            created_at=created_at if created_at is not None else self.created_at,
+        )
+
+        for data_type, registrations in reducer_snapshot:
+            for registration in registrations:
+                clone.register_reducer(
+                    data_type,
+                    registration.reducer,
+                    slice_type=registration.slice_type,
+                )
+
+        with clone._lock:
+            clone._state = state_snapshot
+
+        return clone
+
+    def register_reducer[S: SupportsDataclass](
+        self,
+        data_type: SessionSliceType,
+        reducer: TypedReducer[S],
+        *,
+        slice_type: type[S] | None = None,
+    ) -> None:
+        """Register a reducer for the provided data type."""
+
+        target_slice_type: SessionSliceType = (
+            data_type if slice_type is None else slice_type
+        )
+        registration = _ReducerRegistration(
+            reducer=cast(TypedReducer[Any], reducer),
+            slice_type=target_slice_type,
+        )
+        with self._lock:
+            bucket = self._reducers.setdefault(data_type, [])
+            bucket.append(registration)
+            _ = self._state.setdefault(target_slice_type, EMPTY_SLICE)
+
+    def select_all[S: SupportsDataclass](self, slice_type: type[S]) -> tuple[S, ...]:
+        """Return the tuple slice maintained for the provided type."""
+
+        with self._lock:
+            return cast(tuple[S, ...], self._state.get(slice_type, EMPTY_SLICE))
+
+    def seed_slice[S: SupportsDataclass](
+        self, slice_type: type[S], values: Iterable[S]
+    ) -> None:
+        """Initialize or replace the stored tuple for the provided type."""
+
+        with self._lock:
+            self._state[slice_type] = tuple(values)
+
+    @override
+    def reset(self) -> None:
+        """Clear all stored slices while preserving reducer registrations."""
+
+        with self._lock:
+            slice_types: set[SessionSliceType] = set(self._state)
+            for registrations in self._reducers.values():
+                for registration in registrations:
+                    slice_types.add(registration.slice_type)
+
+            self._state = dict.fromkeys(slice_types, EMPTY_SLICE)
+
+    @property
+    def event_bus(self) -> EventBus:
+        """Return the event bus backing this session."""
+
+        return self._bus
+
+    @override
+    def snapshot(self) -> SnapshotProtocol:
+        """Capture an immutable snapshot of the current session state."""
+
+        with self._lock:
+            state_snapshot: dict[SessionSliceType, SessionSlice] = dict(self._state)
+        for ephemeral in (
+            _TOOL_INVOKED_TYPE,
+            _PROMPT_EXECUTED_TYPE,
+            _PROMPT_RENDERED_TYPE,
+        ):
+            _ = state_snapshot.pop(ephemeral, None)
+        try:
+            normalized: SnapshotState = normalize_snapshot_state(state_snapshot)
+        except ValueError as error:
+            msg = "Unable to serialize session slices"
+            raise SnapshotSerializationError(msg) from error
+
+        created_at = datetime.now(UTC)
+        return Snapshot(created_at=created_at, slices=normalized)
+
+    @override
+    def rollback(self, snapshot: SnapshotProtocol) -> None:
+        """Restore session slices from the provided snapshot."""
+
+        registered_slices = self._registered_slice_types()
+        missing = [
+            slice_type
+            for slice_type in snapshot.slices
+            if slice_type not in registered_slices
+        ]
+        if missing:
+            missing_names = ", ".join(sorted(cls.__qualname__ for cls in missing))
+            msg = f"Slice types not registered: {missing_names}"
+            raise SnapshotRestoreError(msg)
+
+        with self._lock:
+            new_state: dict[SessionSliceType, SessionSlice] = dict(self._state)
+            for slice_type in registered_slices:
+                new_state[slice_type] = snapshot.slices.get(slice_type, EMPTY_SLICE)
+
+            self._state = new_state
+
+    def _registered_slice_types(self) -> set[SessionSliceType]:
+        with self._lock:
+            types: set[SessionSliceType] = set(self._state)
+            for registrations in self._reducers.values():
+                for registration in registrations:
+                    types.add(registration.slice_type)
+            return types
+
+    def _on_tool_invoked(self, event: object) -> None:
+        tool_event = cast(ToolInvoked, event)
+        self._handle_tool_invoked(tool_event)
+
+    def _on_prompt_executed(self, event: object) -> None:
+        prompt_event = cast(PromptExecuted, event)
+        self._handle_prompt_executed(prompt_event)
+
+    def _on_prompt_rendered(self, event: object) -> None:
+        start_event = cast(PromptRendered, event)
+        self._handle_prompt_rendered(start_event)
+
+    def _handle_tool_invoked(self, event: ToolInvoked) -> None:
+        normalized_event = event
+        payload = event.value if event.value is not None else event.result.value
+        if event.value is None and is_dataclass_instance(payload):
+            normalized_event = replace(event, value=payload)
+
+        self._dispatch_data_event(
+            _TOOL_INVOKED_TYPE,
+            cast(ReducerEvent, normalized_event),
+        )
+
+        if normalized_event.value is not None:
+            self._dispatch_data_event(
+                type(normalized_event.value),
+                cast(ReducerEvent, normalized_event),
+            )
+
+    def _handle_prompt_executed(self, event: PromptExecuted) -> None:
+        normalized_event = event
+        output = event.result.output
+        if event.value is None and is_dataclass_instance(output):
+            normalized_event = replace(event, value=output)
+
+        self._dispatch_data_event(
+            _PROMPT_EXECUTED_TYPE,
+            cast(ReducerEvent, normalized_event),
+        )
+
+        if normalized_event.value is not None:
+            self._dispatch_data_event(
+                type(normalized_event.value),
+                cast(ReducerEvent, normalized_event),
+            )
+            return
+
+        if isinstance(output, Iterable) and not isinstance(output, (str, bytes)):
+            for item in cast(Iterable[object], output):
+                if is_dataclass_instance(item):
+                    enriched_event = replace(normalized_event, value=item)
+                    self._dispatch_data_event(
+                        type(item),
+                        cast(ReducerEvent, enriched_event),
+                    )
+
+    def _handle_prompt_rendered(self, event: PromptRendered) -> None:
+        self._dispatch_data_event(
+            _PROMPT_RENDERED_TYPE,
+            cast(ReducerEvent, event),
+        )
+
+    def _dispatch_data_event(
+        self, data_type: SessionSliceType, event: ReducerEvent
+    ) -> None:
+        from .reducer_context import build_reducer_context
+
+        with self._lock:
+            registrations = list(self._reducers.get(data_type, ()))
+            if not registrations:
+                default_reducer: TypedReducer[Any]
+                if data_type in {_TOOL_INVOKED_TYPE, _PROMPT_EXECUTED_TYPE}:
+                    default_reducer = cast(TypedReducer[Any], _append_event)
+                else:
+                    default_reducer = cast(TypedReducer[Any], append)
+                registrations = [
+                    _ReducerRegistration(
+                        reducer=default_reducer,
+                        slice_type=data_type,
+                    )
+                ]
+            event_bus = self._bus
+
+        context = build_reducer_context(session=self, event_bus=event_bus)
+
+        for registration in registrations:
+            slice_type = registration.slice_type
+            while True:
+                with self._lock:
+                    previous = self._state.get(slice_type, EMPTY_SLICE)
+                try:
+                    result = registration.reducer(previous, event, context=context)
+                except Exception:  # log and continue
+                    reducer_name = getattr(
+                        registration.reducer, "__qualname__", repr(registration.reducer)
+                    )
+                    logger.exception(
+                        "Reducer application failed.",
+                        event="session_reducer_failed",
+                        context={
+                            "reducer": reducer_name,
+                            "data_type": data_type.__qualname__,
+                            "slice_type": slice_type.__qualname__,
+                        },
+                    )
+                    break
+                normalized = tuple(result)
+                with self._lock:
+                    current = self._state.get(slice_type, EMPTY_SLICE)
+                    if current is previous or current == normalized:
+                        self._state[slice_type] = normalized
+                        break
+
+    def _attach_to_bus(self, bus: EventBus) -> None:
+        with self._lock:
+            if self._subscriptions_attached and self._bus is bus:
+                return
+            self._bus = bus
+            self._subscriptions_attached = True
+            bus.subscribe(ToolInvoked, self._on_tool_invoked)
+            bus.subscribe(PromptExecuted, self._on_prompt_executed)
+            bus.subscribe(PromptRendered, self._on_prompt_rendered)
+
+
+__all__ = [
+    "DataEvent",
+    "Session",
+    "TypedReducer",
+]