weakincentives 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

weakincentives/prompts/prompt.py

@@ -0,0 +1,440 @@
+# 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.
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass, field, fields, is_dataclass, replace
+from typing import Any, ClassVar, Literal, cast, get_args, get_origin
+
+from ._types import SupportsDataclass
+from .errors import (
+    PromptRenderError,
+    PromptValidationError,
+    SectionPath,
+)
+from .response_format import ResponseFormatParams, ResponseFormatSection
+from .section import Section
+from .tool import Tool
+
+
+@dataclass(frozen=True, slots=True)
+class RenderedPrompt[OutputT = Any]:
+    """Rendered prompt text paired with structured output metadata."""
+
+    text: str
+    output_type: type[Any] | None
+    output_container: Literal["object", "array"] | None
+    allow_extra_keys: bool | None
+    _tools: tuple[Tool[SupportsDataclass, SupportsDataclass], ...] = field(
+        default_factory=tuple
+    )
+
+    def __str__(self) -> str:  # pragma: no cover - convenience for logging
+        return self.text
+
+    @property
+    def tools(self) -> tuple[Tool[SupportsDataclass, SupportsDataclass], ...]:
+        """Tools contributed by enabled sections in traversal order."""
+
+        return self._tools
+
+
+def _clone_dataclass(instance: SupportsDataclass) -> SupportsDataclass:
+    """Return a shallow copy of the provided dataclass instance."""
+
+    return cast(SupportsDataclass, replace(cast(Any, instance)))
+
+
+def _format_specialization_argument(argument: object | None) -> str:
+    if argument is None:  # pragma: no cover - defensive formatting
+        return "?"
+    if isinstance(argument, type):
+        return argument.__name__
+    return repr(argument)  # pragma: no cover - fallback for debugging
+
+
+@dataclass(frozen=True, slots=True)
+class PromptSectionNode[ParamsT: SupportsDataclass]:
+    """Flattened view of a section within a prompt."""
+
+    section: Section[ParamsT]
+    depth: int
+    path: SectionPath
+
+
+class Prompt[OutputT = Any]:
+    """Coordinate prompt sections and their parameter bindings."""
+
+    _output_container_spec: ClassVar[Literal["object", "array"] | None] = None
+    _output_dataclass_candidate: ClassVar[Any] = None
+
+    def __class_getitem__(cls, item: object) -> type[Prompt[Any]]:
+        origin = get_origin(item)
+        candidate = item
+        container: Literal["object", "array"] | None = "object"
+
+        if origin is list:
+            args = get_args(item)
+            candidate = args[0] if len(args) == 1 else None
+            container = "array"
+            label = f"list[{_format_specialization_argument(candidate)}]"
+        else:
+            container = "object"
+            label = _format_specialization_argument(candidate)
+
+        name = f"{cls.__name__}[{label}]"
+        namespace = {
+            "__module__": cls.__module__,
+            "_output_container_spec": container if candidate is not None else None,
+            "_output_dataclass_candidate": candidate,
+        }
+        return type(name, (cls,), namespace)
+
+    def __init__(
+        self,
+        *,
+        name: str | None = None,
+        sections: Sequence[Section[Any]] | None = None,
+        inject_output_instructions: bool = True,
+        allow_extra_keys: bool = False,
+    ) -> None:
+        self.name = name
+        base_sections: list[Section[SupportsDataclass]] = [
+            cast(Section[SupportsDataclass], section) for section in sections or ()
+        ]
+        self._sections: tuple[Section[SupportsDataclass], ...] = tuple(base_sections)
+        self._section_nodes: list[PromptSectionNode[SupportsDataclass]] = []
+        self._params_registry: dict[
+            type[SupportsDataclass], list[PromptSectionNode[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] = {}
+
+        self._output_type: type[Any] | None
+        self._output_container: Literal["object", "array"] | None
+        self._allow_extra_keys: bool | None
+        (
+            self._output_type,
+            self._output_container,
+            self._allow_extra_keys,
+        ) = self._resolve_output_spec(allow_extra_keys)
+
+        self.inject_output_instructions = inject_output_instructions
+
+        for section in base_sections:
+            self._register_section(section, path=(section.title,), depth=0)
+
+        self._response_section: ResponseFormatSection | None = None
+        if self._output_type is not None and self._output_container is not None:
+            response_params = self._build_response_format_params()
+            response_section = ResponseFormatSection(
+                params=response_params,
+                enabled=lambda _params, prompt=self: prompt.inject_output_instructions,
+            )
+            self._response_section = response_section
+            section_for_registry = cast(Section[SupportsDataclass], response_section)
+            self._sections += (section_for_registry,)
+            self._register_section(
+                section_for_registry,
+                path=(response_section.title,),
+                depth=0,
+            )
+
+    def render(self, *params: SupportsDataclass) -> RenderedPrompt[OutputT]:
+        """Render the prompt using provided parameter dataclass instances."""
+
+        param_lookup = self._collect_param_lookup(params)
+        rendered_sections: list[str] = []
+        collected_tools: list[Tool[SupportsDataclass, SupportsDataclass]] = []
+
+        for node, section_params in self._iter_enabled_sections(param_lookup):
+            params_type = node.section.params
+            try:
+                rendered = node.section.render(section_params, node.depth)
+            except PromptRenderError as error:
+                if error.section_path and error.dataclass_type:
+                    raise
+                raise PromptRenderError(
+                    error.message,
+                    section_path=node.path,
+                    dataclass_type=params_type,
+                    placeholder=error.placeholder,
+                ) from error
+            except Exception as error:  # pragma: no cover - defensive guard
+                raise PromptRenderError(
+                    "Section rendering failed.",
+                    section_path=node.path,
+                    dataclass_type=params_type,
+                ) from error
+
+            section_tools = node.section.tools()
+            if section_tools:
+                collected_tools.extend(section_tools)
+
+            if rendered:
+                rendered_sections.append(rendered)
+
+        text = "\n\n".join(rendered_sections)
+
+        return RenderedPrompt(
+            text=text,
+            output_type=self._output_type,
+            output_container=self._output_container,
+            allow_extra_keys=self._allow_extra_keys,
+            _tools=tuple(collected_tools),
+        )
+
+    def _register_section(
+        self,
+        section: Section[SupportsDataclass],
+        *,
+        path: SectionPath,
+        depth: int,
+    ) -> None:
+        params_type = section.params
+        if not is_dataclass(params_type):
+            raise PromptValidationError(
+                "Section params must be a dataclass.",
+                section_path=path,
+                dataclass_type=params_type,
+            )
+
+        node: PromptSectionNode[SupportsDataclass] = PromptSectionNode(
+            section=section, depth=depth, path=path
+        )
+        self._section_nodes.append(node)
+        self._params_registry.setdefault(params_type, []).append(node)
+
+        if section.defaults is not None:
+            default_value = section.defaults
+            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)
+        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,
+            )
+
+        self._register_section_tools(section, path)
+
+        for child in section.children:
+            child_path = path + (child.title,)
+            self._register_section(child, path=child_path, depth=depth + 1)
+
+    @property
+    def sections(self) -> tuple[PromptSectionNode[SupportsDataclass], ...]:
+        return tuple(self._section_nodes)
+
+    @property
+    def params_types(self) -> set[type[SupportsDataclass]]:
+        return set(self._params_registry.keys())
+
+    def _resolve_output_spec(
+        self, allow_extra_keys: bool
+    ) -> tuple[type[Any] | None, Literal["object", "array"] | None, bool | None]:
+        candidate = getattr(type(self), "_output_dataclass_candidate", None)
+        container = cast(
+            Literal["object", "array"] | None,
+            getattr(type(self), "_output_container_spec", None),
+        )
+
+        if candidate is None or container is None:
+            return None, None, None
+
+        if not isinstance(candidate, type):  # pragma: no cover - defensive guard
+            candidate_type = cast(type[Any], type(candidate))
+            raise PromptValidationError(
+                "Prompt output type must be a dataclass.",
+                dataclass_type=candidate_type,
+            )
+
+        if not is_dataclass(candidate):
+            bad_dataclass = cast(type[Any], candidate)
+            raise PromptValidationError(
+                "Prompt output type must be a dataclass.",
+                dataclass_type=bad_dataclass,
+            )
+
+        dataclass_type = cast(type[Any], candidate)
+        return dataclass_type, container, allow_extra_keys
+
+    def _build_response_format_params(self) -> ResponseFormatParams:
+        container = self._output_container
+        if container is None:  # pragma: no cover - defensive guard
+            raise RuntimeError(
+                "Output container missing during response format construction."
+            )
+
+        article: Literal["a", "an"] = (
+            "an" if container.startswith(("a", "e", "i", "o", "u")) else "a"
+        )
+        extra_clause = ". Do not add extra keys." if not self._allow_extra_keys else "."
+        return ResponseFormatParams(
+            article=article,
+            container=container,
+            extra_clause=extra_clause,
+        )
+
+    def _collect_param_lookup(
+        self, params: tuple[SupportsDataclass, ...]
+    ) -> dict[type[SupportsDataclass], SupportsDataclass]:
+        lookup: dict[type[SupportsDataclass], SupportsDataclass] = {}
+        for value in params:
+            if isinstance(value, type):
+                provided_type: type[Any] = value
+            else:
+                provided_type = type(value)
+            if isinstance(value, type) or not is_dataclass(value):
+                raise PromptValidationError(
+                    "Prompt expects dataclass instances.",
+                    dataclass_type=provided_type,
+                )
+            params_type = cast(type[SupportsDataclass], provided_type)
+            if params_type in lookup:
+                raise PromptValidationError(
+                    "Duplicate params type supplied to prompt.",
+                    dataclass_type=params_type,
+                )
+            if params_type not in self._params_registry:
+                raise PromptValidationError(
+                    "Unexpected params type supplied to prompt.",
+                    dataclass_type=params_type,
+                )
+            lookup[params_type] = value
+        return lookup
+
+    def _resolve_section_params(
+        self,
+        node: PromptSectionNode[SupportsDataclass],
+        param_lookup: dict[type[SupportsDataclass], SupportsDataclass],
+    ) -> SupportsDataclass:
+        params_type = node.section.params
+        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], 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
+
+        return section_params
+
+    def _iter_enabled_sections(
+        self,
+        param_lookup: dict[type[SupportsDataclass], SupportsDataclass],
+    ) -> Iterator[tuple[PromptSectionNode[SupportsDataclass], SupportsDataclass]]:
+        skip_depth: int | None = None
+
+        for node in self._section_nodes:
+            if skip_depth is not None:
+                if node.depth > skip_depth:
+                    continue
+                skip_depth = None
+
+            section_params = self._resolve_section_params(node, param_lookup)
+
+            try:
+                enabled = node.section.is_enabled(section_params)
+            except Exception as error:  # pragma: no cover - defensive guard
+                raise PromptRenderError(
+                    "Section enabled predicate failed.",
+                    section_path=node.path,
+                    dataclass_type=node.section.params,
+                ) from error
+
+            if not enabled:
+                skip_depth = node.depth
+                continue
+
+            yield node, section_params
+
+    def _register_section_tools(
+        self,
+        section: Section[SupportsDataclass],
+        path: SectionPath,
+    ) -> None:
+        section_tools = section.tools()
+        if not section_tools:
+            return
+
+        tools_iterable = cast(Sequence[object], section_tools)
+        for tool_candidate in tools_iterable:
+            if not isinstance(tool_candidate, Tool):
+                raise PromptValidationError(
+                    "Section tools() must return Tool instances.",
+                    section_path=path,
+                    dataclass_type=section.params,
+                )
+            tool: Tool[SupportsDataclass, SupportsDataclass] = cast(
+                Tool[SupportsDataclass, SupportsDataclass], tool_candidate
+            )
+            params_type = cast(
+                type[SupportsDataclass] | None, getattr(tool, "params_type", None)
+            )
+            if not isinstance(params_type, type) or not is_dataclass(params_type):
+                raise PromptValidationError(
+                    "Tool params_type must be a dataclass type.",
+                    section_path=path,
+                    dataclass_type=(
+                        params_type
+                        if isinstance(params_type, type)
+                        else 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
+
+
+__all__ = ["Prompt", "PromptSectionNode", "RenderedPrompt"]

weakincentives/prompts/response_format.py

@@ -0,0 +1,54 @@
+# 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.
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Literal
+
+from .text import TextSection
+
+__all__ = ["ResponseFormatParams", "ResponseFormatSection"]
+
+
+@dataclass(slots=True)
+class ResponseFormatParams:
+    """Parameter payload for the auto-generated response format section."""
+
+    article: Literal["a", "an"]
+    container: Literal["object", "array"]
+    extra_clause: str
+
+
+_RESPONSE_FORMAT_BODY = """Return ONLY a single fenced JSON code block. Do not include any text
+before or after the block.
+
+The top-level JSON value MUST be ${article} ${container} that matches the fields
+of the expected schema${extra_clause}"""
+
+
+class ResponseFormatSection(TextSection[ResponseFormatParams]):
+    """Internal section that appends JSON-only response instructions."""
+
+    def __init__(
+        self,
+        *,
+        params: ResponseFormatParams,
+        enabled: Callable[[ResponseFormatParams], bool] | None = None,
+    ) -> None:
+        super().__init__(
+            title="Response Format",
+            body=_RESPONSE_FORMAT_BODY,
+            defaults=params,
+            enabled=enabled,
+        )

weakincentives/prompts/section.py

@@ -0,0 +1,120 @@
+# 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.
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, ClassVar, cast
+
+if TYPE_CHECKING:
+    from .tool import Tool
+
+from ._types import SupportsDataclass
+
+
+class Section[ParamsT: SupportsDataclass](ABC):
+    """Abstract building block for prompt content."""
+
+    _params_type: ClassVar[type[SupportsDataclass] | None] = None
+
+    def __init__(
+        self,
+        *,
+        title: str,
+        defaults: ParamsT | None = None,
+        children: Sequence[object] | None = None,
+        enabled: Callable[[ParamsT], bool] | None = None,
+        tools: Sequence[object] | None = None,
+    ) -> None:
+        params_type = cast(
+            type[ParamsT] | None, getattr(self.__class__, "_params_type", None)
+        )
+        if params_type is None:
+            raise TypeError(
+                "Section must be instantiated with a concrete ParamsT type."
+            )
+
+        self.params_type: type[ParamsT] = params_type
+        self.params: type[ParamsT] = params_type
+        self.title = title
+        self.defaults = defaults
+
+        normalized_children: list[Section[SupportsDataclass]] = []
+        for child in children or ():
+            if not isinstance(child, Section):
+                raise TypeError("Section children must be Section instances.")
+            normalized_children.append(cast(Section[SupportsDataclass], child))
+        self.children: tuple[Section[SupportsDataclass], ...] = tuple(
+            normalized_children
+        )
+        self._enabled = enabled
+        self._tools = self._normalize_tools(tools)
+
+    def is_enabled(self, params: ParamsT) -> bool:
+        """Return True when the section should render for the given params."""
+
+        if self._enabled is None:
+            return True
+        return bool(self._enabled(params))
+
+    @abstractmethod
+    def render(self, params: ParamsT, depth: int) -> str:
+        """Produce markdown output for the section at the supplied depth."""
+
+    def placeholder_names(self) -> set[str]:
+        """Return placeholder identifiers used by the section template."""
+
+        return set()
+
+    def tools(self) -> tuple[Tool[SupportsDataclass, SupportsDataclass], ...]:
+        """Return the tools exposed by this section."""
+
+        return self._tools
+
+    @classmethod
+    def __class_getitem__(cls, item: object) -> type[Section[SupportsDataclass]]:
+        params_type = cls._normalize_generic_argument(item)
+
+        class _SpecializedSection(cls):  # type: ignore[misc]
+            pass
+
+        _SpecializedSection.__name__ = cls.__name__
+        _SpecializedSection.__qualname__ = cls.__qualname__
+        _SpecializedSection.__module__ = cls.__module__
+        _SpecializedSection._params_type = cast(type[SupportsDataclass], params_type)
+        return _SpecializedSection  # type: ignore[return-value]
+
+    @staticmethod
+    def _normalize_generic_argument(item: object) -> object:
+        if isinstance(item, tuple):
+            raise TypeError("Section[...] expects a single type argument.")
+        return item
+
+    @staticmethod
+    def _normalize_tools(
+        tools: Sequence[object] | None,
+    ) -> tuple[Tool[SupportsDataclass, SupportsDataclass], ...]:
+        if not tools:
+            return ()
+
+        from .tool import Tool
+
+        normalized: list[Tool[SupportsDataclass, SupportsDataclass]] = []
+        for tool in tools:
+            if not isinstance(tool, Tool):
+                raise TypeError("Section tools must be Tool instances.")
+            normalized.append(cast(Tool[SupportsDataclass, SupportsDataclass], tool))
+        return tuple(normalized)
+
+
+__all__ = ["Section"]