weakincentives 0.9.0__py3-none-any.whl

weakincentives/prompt/protocols.py

@@ -0,0 +1,103 @@
+# 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.
+
+"""Structural typing primitives shared across prompt tooling."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar
+
+from ..deadlines import Deadline
+from ._overrides_protocols import PromptOverridesStoreProtocol
+from ._types import SupportsDataclass
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from ..runtime.events._types import EventBus
+    from ..runtime.session.protocols import SessionProtocol
+
+PromptOutputT = TypeVar("PromptOutputT", covariant=True)
+RenderedOutputT = TypeVar("RenderedOutputT", covariant=True)
+AdapterOutputT = TypeVar("AdapterOutputT")
+
+
+class PromptResponseProtocol(Protocol[AdapterOutputT]):
+    prompt_name: str
+    text: str | None
+    output: AdapterOutputT | None
+    tool_results: tuple[object, ...]
+    provider_payload: Mapping[str, Any] | None
+
+
+class RenderedPromptProtocol(Protocol[RenderedOutputT]):
+    """Interface satisfied by rendered prompt snapshots."""
+
+    @property
+    def text(self) -> str: ...
+
+    @property
+    def output_type(self) -> type[Any] | None: ...
+
+    @property
+    def container(self) -> Literal["object", "array"] | None: ...
+
+    @property
+    def allow_extra_keys(self) -> bool | None: ...
+
+    @property
+    def deadline(self) -> Deadline | None: ...
+
+    @property
+    def tools(self) -> tuple[object, ...]: ...
+
+    @property
+    def tool_param_descriptions(self) -> Mapping[str, Mapping[str, str]]: ...
+
+
+class PromptProtocol(Protocol[PromptOutputT]):
+    """Interface describing the subset of Prompt state exposed to tools."""
+
+    ns: str
+    key: str
+    name: str | None
+
+    def render(
+        self,
+        *params: SupportsDataclass,
+        overrides_store: PromptOverridesStoreProtocol | None = None,
+        tag: str = "latest",
+        inject_output_instructions: bool | None = None,
+    ) -> RenderedPromptProtocol[PromptOutputT]: ...
+
+
+class ProviderAdapterProtocol(Protocol[AdapterOutputT]):
+    """Interface describing the subset of adapter behaviour required by tools."""
+
+    def evaluate(
+        self,
+        prompt: PromptProtocol[AdapterOutputT],
+        *params: SupportsDataclass,
+        parse_output: bool = True,
+        bus: EventBus,
+        session: SessionProtocol,
+        deadline: Deadline | None = None,
+        overrides_store: PromptOverridesStoreProtocol | None = None,
+        overrides_tag: str = "latest",
+    ) -> PromptResponseProtocol[AdapterOutputT]: ...
+
+
+__all__ = [
+    "PromptProtocol",
+    "PromptResponseProtocol",
+    "ProviderAdapterProtocol",
+    "RenderedPromptProtocol",
+]

weakincentives/prompt/registry.py

@@ -0,0 +1,375 @@
+# 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.
+
+"""Section registration helpers for :mod:`weakincentives.prompt`."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
+from dataclasses import dataclass, fields, is_dataclass, replace
+from types import MappingProxyType
+from typing import Any, cast
+
+from ..dbc import invariant
+from ._types import SupportsDataclass, SupportsToolResult
+from .errors import PromptRenderError, PromptValidationError, SectionPath
+from .section import Section
+from .tool import Tool
+
+
+@dataclass(frozen=True, slots=True)
+class SectionNode[ParamsT: SupportsDataclass]:
+    """Flattened view of a section within a prompt."""
+
+    section: Section[ParamsT]
+    depth: int
+    path: SectionPath
+
+
+@dataclass(frozen=True)
+class RegistrySnapshot:
+    """Immutable view over registered prompt sections."""
+
+    sections: tuple[SectionNode[SupportsDataclass], ...]
+    params_registry: Mapping[
+        type[SupportsDataclass], tuple[SectionNode[SupportsDataclass], ...]
+    ]
+    defaults_by_path: Mapping[SectionPath, SupportsDataclass]
+    defaults_by_type: Mapping[type[SupportsDataclass], SupportsDataclass]
+    placeholders: Mapping[SectionPath, frozenset[str]]
+    tool_name_registry: Mapping[str, SectionPath]
+
+    def resolve_section_params(
+        self,
+        node: SectionNode[SupportsDataclass],
+        param_lookup: MutableMapping[type[SupportsDataclass], SupportsDataclass],
+    ) -> SupportsDataclass | None:
+        """Return parameters for a section, applying defaults when necessary."""
+
+        params_type = node.section.param_type
+        if params_type is None:
+            return None
+
+        section_params: SupportsDataclass | None = param_lookup.get(params_type)
+
+        if section_params is None:
+            default_value = self.defaults_by_path.get(node.path)
+            if default_value is not None:
+                section_params = clone_dataclass(default_value)
+            else:
+                type_default = self.defaults_by_type.get(params_type)
+                if type_default is not None:
+                    section_params = clone_dataclass(type_default)
+                else:
+                    try:
+                        constructor = cast(
+                            Callable[[], SupportsDataclass | None], params_type
+                        )
+                        section_params = constructor()
+                    except TypeError as error:
+                        raise PromptRenderError(
+                            "Missing parameters for section.",
+                            section_path=node.path,
+                            dataclass_type=params_type,
+                        ) from error
+
+        result: SupportsDataclass | None = section_params
+        if result is None:
+            raise PromptRenderError(
+                "Section constructor must return a dataclass instance.",
+                section_path=node.path,
+                dataclass_type=params_type,
+            )
+
+        if not is_dataclass(result):
+            raise PromptRenderError(
+                "Section constructor must return a dataclass instance.",
+                section_path=node.path,
+                dataclass_type=params_type,
+            )
+
+        return result
+
+    @property
+    def param_types(self) -> set[type[SupportsDataclass]]:
+        """Return the set of parameter dataclasses registered for sections."""
+
+        return set(self.params_registry.keys())
+
+
+def _registry_paths_are_registered(
+    registry: PromptRegistry,
+) -> tuple[bool, str] | bool:
+    """Ensure internal registries only reference known section nodes."""
+
+    section_nodes = registry._section_nodes  # pyright: ignore[reportPrivateUsage]
+    node_by_path = {node.path: node for node in section_nodes}
+    defaults_by_path = registry._defaults_by_path  # pyright: ignore[reportPrivateUsage]
+    placeholders = registry._placeholders  # pyright: ignore[reportPrivateUsage]
+    tool_name_registry = registry._tool_name_registry  # pyright: ignore[reportPrivateUsage]
+    defaults_by_type = registry._defaults_by_type  # pyright: ignore[reportPrivateUsage]
+
+    unknown_default_paths = [
+        path for path in defaults_by_path if path not in node_by_path
+    ]
+    if unknown_default_paths:
+        return (
+            False,
+            f"defaults reference unknown paths: {sorted(unknown_default_paths)!r}",
+        )
+
+    for path, default in defaults_by_path.items():
+        node = node_by_path[path]
+        params_type = node.section.param_type
+        if params_type is None:
+            return False, f"section at {path!r} does not accept params but has defaults"
+        if type(default) is not params_type:
+            return False, (
+                "default params type mismatch for path "
+                f"{path!r}: expected {params_type.__name__}, got {type(default).__name__}"
+            )
+
+    unknown_placeholder_paths = [
+        path for path in placeholders if path not in node_by_path
+    ]
+    if unknown_placeholder_paths:
+        return False, (
+            "placeholders reference unknown paths: "
+            f"{sorted(unknown_placeholder_paths)!r}"
+        )
+
+    unknown_tool_paths = [
+        path for path in tool_name_registry.values() if path not in node_by_path
+    ]
+    if unknown_tool_paths:
+        return False, f"tools reference unknown paths: {sorted(unknown_tool_paths)!r}"
+
+    for params_type, default in defaults_by_type.items():
+        if type(default) is not params_type:
+            return False, (
+                "default by type mismatch for "
+                f"{params_type.__name__}: got {type(default).__name__}"
+            )
+
+    return True
+
+
+def _params_registry_is_consistent(
+    registry: PromptRegistry,
+) -> tuple[bool, str] | bool:
+    """Ensure params registry entries point at known nodes with matching types."""
+
+    section_nodes = list(registry._section_nodes)  # pyright: ignore[reportPrivateUsage]
+    params_registry = registry._params_registry  # pyright: ignore[reportPrivateUsage]
+    for params_type, nodes in params_registry.items():
+        for node in nodes:
+            if node not in section_nodes:
+                return False, (
+                    "params registry references unknown node at path "
+                    f"{node.path!r} for {params_type.__name__}"
+                )
+            node_params_type = node.section.param_type
+            if node_params_type is None:
+                return False, (
+                    "params registry references section without params at path "
+                    f"{node.path!r}"
+                )
+            if node_params_type is not params_type:
+                return False, (
+                    "params registry type mismatch for path "
+                    f"{node.path!r}: expected {params_type.__name__}, "
+                    f"found {node_params_type.__name__}"
+                )
+
+    return True
+
+
+@invariant(
+    _registry_paths_are_registered,
+    _params_registry_is_consistent,
+)
+class PromptRegistry:
+    """Collect and validate prompt sections prior to rendering."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._section_nodes: list[SectionNode[SupportsDataclass]] = []
+        self._params_registry: dict[
+            type[SupportsDataclass], list[SectionNode[SupportsDataclass]]
+        ] = {}
+        self._defaults_by_path: dict[SectionPath, SupportsDataclass] = {}
+        self._defaults_by_type: dict[type[SupportsDataclass], SupportsDataclass] = {}
+        self._placeholders: dict[SectionPath, set[str]] = {}
+        self._tool_name_registry: dict[str, SectionPath] = {}
+
+    def register_sections(self, sections: Sequence[Section[SupportsDataclass]]) -> None:
+        """Register the provided root sections."""
+
+        for section in sections:
+            self._register_section(section, path=(section.key,), depth=0)
+
+    def register_section(
+        self,
+        section: Section[SupportsDataclass],
+        *,
+        path: SectionPath,
+        depth: int,
+    ) -> None:
+        """Register a single section at the supplied path and depth."""
+
+        self._register_section(section, path=path, depth=depth)
+
+    def _register_section(
+        self,
+        section: Section[SupportsDataclass],
+        *,
+        path: SectionPath,
+        depth: int,
+    ) -> None:
+        params_type = section.param_type
+        if params_type is not None and not is_dataclass(params_type):
+            raise PromptValidationError(
+                "Section params must be a dataclass.",
+                section_path=path,
+                dataclass_type=params_type,
+            )
+
+        node: SectionNode[SupportsDataclass] = SectionNode(
+            section=section, depth=depth, path=path
+        )
+        self._section_nodes.append(node)
+
+        if params_type is not None:
+            self._params_registry.setdefault(params_type, []).append(node)
+
+        if params_type is not None and section.default_params is not None:
+            default_value = section.default_params
+            if isinstance(default_value, type) or not is_dataclass(default_value):
+                raise PromptValidationError(
+                    "Section defaults must be dataclass instances.",
+                    section_path=path,
+                    dataclass_type=params_type,
+                )
+            if type(default_value) is not params_type:
+                raise PromptValidationError(
+                    "Section defaults must match section params type.",
+                    section_path=path,
+                    dataclass_type=params_type,
+                )
+            self._defaults_by_path[path] = default_value
+            _ = self._defaults_by_type.setdefault(params_type, default_value)
+
+        section_placeholders = section.placeholder_names()
+        self._placeholders[path] = set(section_placeholders)
+        if params_type is None:
+            if section_placeholders:
+                placeholder = sorted(section_placeholders)[0]
+                raise PromptValidationError(
+                    "Section does not accept parameters but declares placeholders.",
+                    section_path=path,
+                    placeholder=placeholder,
+                )
+        else:
+            param_fields = {field.name for field in fields(params_type)}
+            unknown_placeholders = section_placeholders - param_fields
+            if unknown_placeholders:
+                placeholder = sorted(unknown_placeholders)[0]
+                raise PromptValidationError(
+                    "Template references unknown placeholder.",
+                    section_path=path,
+                    dataclass_type=params_type,
+                    placeholder=placeholder,
+                )
+
+        section_tools = cast(tuple[object, ...], section.tools())
+        if section_tools:
+            for tool in section_tools:
+                if not isinstance(tool, Tool):
+                    raise PromptValidationError(
+                        "Section tools must be Tool instances.",
+                        section_path=path,
+                        dataclass_type=params_type,
+                    )
+                typed_tool = cast(Tool[SupportsDataclass, SupportsToolResult], tool)
+                self._register_section_tools(
+                    typed_tool,
+                    path,
+                )
+
+        for child in section.children:
+            child_path = (*path, child.key)
+            self._register_section(child, path=child_path, depth=depth + 1)
+
+    def _register_section_tools[
+        ParamsT: SupportsDataclass,
+        ResultT: SupportsToolResult,
+    ](
+        self,
+        tool: Tool[ParamsT, ResultT],
+        path: SectionPath,
+    ) -> None:
+        params_type = tool.params_type
+        if not is_dataclass(params_type):
+            raise PromptValidationError(
+                "Tool parameters must be dataclass types.",
+                section_path=path,
+                dataclass_type=params_type,
+            )
+
+        existing_path = self._tool_name_registry.get(tool.name)
+        if existing_path is not None:
+            raise PromptValidationError(
+                "Duplicate tool name registered for prompt.",
+                section_path=path,
+                dataclass_type=tool.params_type,
+            )
+
+        self._tool_name_registry[tool.name] = path
+
+    def snapshot(self) -> RegistrySnapshot:
+        """Return an immutable snapshot of the registered sections."""
+
+        params_registry: dict[
+            type[SupportsDataclass], tuple[SectionNode[SupportsDataclass], ...]
+        ] = {
+            params_type: tuple(nodes)
+            for params_type, nodes in self._params_registry.items()
+        }
+        defaults_by_path = MappingProxyType(dict(self._defaults_by_path))
+        defaults_by_type = MappingProxyType(dict(self._defaults_by_type))
+        placeholders = MappingProxyType(
+            {path: frozenset(names) for path, names in self._placeholders.items()}
+        )
+        tool_name_registry = MappingProxyType(dict(self._tool_name_registry))
+
+        return RegistrySnapshot(
+            sections=tuple(self._section_nodes),
+            params_registry=MappingProxyType(params_registry),
+            defaults_by_path=defaults_by_path,
+            defaults_by_type=defaults_by_type,
+            placeholders=placeholders,
+            tool_name_registry=tool_name_registry,
+        )
+
+
+def clone_dataclass(instance: SupportsDataclass) -> SupportsDataclass:
+    """Return a shallow copy of the provided dataclass instance."""
+
+    return cast(SupportsDataclass, replace(cast(Any, instance)))
+
+
+__all__ = [
+    "PromptRegistry",
+    "RegistrySnapshot",
+    "SectionNode",
+    "clone_dataclass",
+]