weakincentives 0.9.0__py3-none-any.whl

weakincentives/prompt/rendering.py

@@ -0,0 +1,288 @@
+# 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.
+
+"""Rendering helpers for :mod:`weakincentives.prompt`."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping, MutableMapping
+from dataclasses import dataclass, field, is_dataclass, replace
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast, override
+
+from ..deadlines import Deadline
+from ._types import SupportsDataclass, SupportsToolResult
+from .errors import PromptRenderError, PromptValidationError, SectionPath
+from .registry import RegistrySnapshot, SectionNode
+from .response_format import ResponseFormatSection
+from .structured_output import StructuredOutputConfig
+from .tool import Tool
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from .overrides import ToolOverride
+
+
+_EMPTY_TOOL_PARAM_DESCRIPTIONS: Mapping[str, Mapping[str, str]] = MappingProxyType({})
+
+
+OutputT_co = TypeVar("OutputT_co", covariant=True)
+
+
+@dataclass(frozen=True, slots=True)
+class RenderedPrompt[OutputT_co]:
+    """Rendered prompt text paired with structured output metadata."""
+
+    text: str
+    structured_output: StructuredOutputConfig[SupportsDataclass] | None = None
+    deadline: Deadline | None = None
+    _tools: tuple[Tool[SupportsDataclass, SupportsToolResult], ...] = field(
+        default_factory=tuple
+    )
+    _tool_param_descriptions: Mapping[str, Mapping[str, str]] = field(
+        default=_EMPTY_TOOL_PARAM_DESCRIPTIONS
+    )
+
+    @override
+    def __str__(self) -> str:  # pragma: no cover - delegated behaviour
+        return self.text
+
+    @property
+    def tools(self) -> tuple[Tool[SupportsDataclass, SupportsToolResult], ...]:
+        """Tools contributed by enabled sections in traversal order."""
+
+        return self._tools
+
+    @property
+    def tool_param_descriptions(
+        self,
+    ) -> Mapping[str, Mapping[str, str]]:
+        """Description patches keyed by tool name."""
+
+        return self._tool_param_descriptions
+
+    @property
+    def output_type(self) -> type[SupportsDataclass] | None:
+        """Return the declared dataclass type for structured output."""
+
+        if self.structured_output is None:
+            return None
+        return self.structured_output.dataclass_type
+
+    @property
+    def container(self) -> Literal["object", "array"] | None:
+        """Return the declared container for structured output."""
+
+        if self.structured_output is None:
+            return None
+        return self.structured_output.container
+
+    @property
+    def allow_extra_keys(self) -> bool | None:
+        """Return whether extra keys are allowed in structured output."""
+
+        if self.structured_output is None:
+            return None
+        return self.structured_output.allow_extra_keys
+
+
+def _freeze_tool_param_descriptions(
+    descriptions: Mapping[str, dict[str, str]],
+) -> Mapping[str, Mapping[str, str]]:
+    if not descriptions:
+        return MappingProxyType({})
+    frozen: dict[str, Mapping[str, str]] = {}
+    for name, field_mapping in descriptions.items():
+        frozen[name] = MappingProxyType(dict(field_mapping))
+    return MappingProxyType(frozen)
+
+
+OutputT = TypeVar("OutputT")
+
+
+class PromptRenderer[OutputT]:
+    """Render prompts using a registry snapshot."""
+
+    def __init__(
+        self,
+        *,
+        registry: RegistrySnapshot,
+        structured_output: StructuredOutputConfig[SupportsDataclass] | None,
+        response_section: ResponseFormatSection | None,
+    ) -> None:
+        super().__init__()
+        self._registry = registry
+        self._structured_output = structured_output
+        self._response_section: ResponseFormatSection | None = response_section
+
+    def build_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._registry.param_types:
+                raise PromptValidationError(
+                    "Unexpected params type supplied to prompt.",
+                    dataclass_type=params_type,
+                )
+            lookup[params_type] = value
+        return lookup
+
+    def render(
+        self,
+        param_lookup: Mapping[type[SupportsDataclass], SupportsDataclass],
+        overrides: Mapping[SectionPath, str] | None = None,
+        tool_overrides: Mapping[str, ToolOverride] | None = None,
+        *,
+        inject_output_instructions: bool | None = None,
+    ) -> RenderedPrompt[OutputT]:
+        rendered_sections: list[str] = []
+        collected_tools: list[Tool[SupportsDataclass, SupportsToolResult]] = []
+        override_lookup = dict(overrides or {})
+        tool_override_lookup = dict(tool_overrides or {})
+        field_description_patches: dict[str, dict[str, str]] = {}
+
+        for node, section_params in self._iter_enabled_sections(
+            dict(param_lookup),
+            inject_output_instructions=inject_output_instructions,
+        ):
+            override_body = (
+                override_lookup.get(node.path)
+                if getattr(node.section, "accepts_overrides", True)
+                else None
+            )
+            rendered = self._render_section(node, section_params, override_body)
+
+            section_tools = node.section.tools()
+            if section_tools:
+                for tool in section_tools:
+                    override = (
+                        tool_override_lookup.get(tool.name)
+                        if tool.accepts_overrides
+                        else None
+                    )
+                    patched_tool = tool
+                    if override is not None:
+                        if (
+                            override.description is not None
+                            and override.description != tool.description
+                        ):
+                            patched_tool = replace(
+                                tool, description=override.description
+                            )
+                        if override.param_descriptions:
+                            field_description_patches[tool.name] = dict(
+                                override.param_descriptions
+                            )
+                    collected_tools.append(patched_tool)
+
+            if rendered:
+                rendered_sections.append(rendered)
+
+        text = "\n\n".join(rendered_sections)
+
+        return RenderedPrompt[OutputT](
+            text=text,
+            structured_output=self._structured_output,
+            _tools=tuple(collected_tools),
+            _tool_param_descriptions=_freeze_tool_param_descriptions(
+                field_description_patches
+            ),
+        )
+
+    def _iter_enabled_sections(
+        self,
+        param_lookup: MutableMapping[type[SupportsDataclass], SupportsDataclass],
+        *,
+        inject_output_instructions: bool | None = None,
+    ) -> Iterator[tuple[SectionNode[SupportsDataclass], SupportsDataclass | None]]:
+        skip_depth: int | None = None
+
+        for node in self._registry.sections:
+            if skip_depth is not None:
+                if node.depth > skip_depth:
+                    continue
+                skip_depth = None
+
+            section_params = self._registry.resolve_section_params(node, param_lookup)
+
+            if node.section is self._response_section and (
+                inject_output_instructions is not None
+            ):
+                enabled = inject_output_instructions
+            else:
+                try:
+                    enabled = node.section.is_enabled(section_params)
+                except Exception as error:  # pragma: no cover - defensive
+                    raise PromptRenderError(
+                        "Section enabled predicate failed.",
+                        section_path=node.path,
+                        dataclass_type=node.section.param_type,
+                    ) from error
+
+            if not enabled:
+                skip_depth = node.depth
+                continue
+
+            yield node, section_params
+
+    def _render_section(
+        self,
+        node: SectionNode[SupportsDataclass],
+        section_params: SupportsDataclass | None,
+        override_body: str | None,
+    ) -> str:
+        params_type = node.section.param_type
+        try:
+            render_override = getattr(node.section, "render_with_template", None)
+            if override_body is not None and callable(render_override):
+                override_renderer = cast(
+                    Callable[[str, SupportsDataclass | None, int], str],
+                    render_override,
+                )
+                rendered = override_renderer(override_body, section_params, node.depth)
+            else:
+                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
+            raise PromptRenderError(
+                "Section rendering failed.",
+                section_path=node.path,
+                dataclass_type=params_type,
+            ) from error
+
+        return rendered
+
+
+__all__ = ["PromptRenderer", "RenderedPrompt"]

weakincentives/prompt/response_format.py

@@ -0,0 +1,60 @@
+# 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 Final, Literal
+
+from ._types import SupportsDataclass
+from .markdown import MarkdownSection
+
+__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: Final[
+    str
+] = """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(MarkdownSection[ResponseFormatParams]):
+    """Internal section that appends JSON-only response instructions."""
+
+    def __init__(
+        self,
+        *,
+        params: ResponseFormatParams,
+        enabled: Callable[[SupportsDataclass], bool] | None = None,
+        accepts_overrides: bool = False,
+    ) -> None:
+        super().__init__(
+            title="Response Format",
+            key="response-format",
+            template=_RESPONSE_FORMAT_BODY,
+            default_params=params,
+            enabled=enabled,
+            accepts_overrides=accepts_overrides,
+        )

weakincentives/prompt/section.py

@@ -0,0 +1,166 @@
+# 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
+
+import inspect
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, ClassVar, TypeVar, cast
+
+if TYPE_CHECKING:
+    from .tool import Tool
+
+from ._generic_params_specializer import GenericParamsSpecializer
+from ._normalization import normalize_component_key
+from ._types import SupportsDataclass, SupportsToolResult
+
+SectionParamsT = TypeVar("SectionParamsT", bound=SupportsDataclass, covariant=True)
+
+EnabledPredicate = Callable[[SupportsDataclass], bool] | Callable[[], bool]
+
+
+class Section(GenericParamsSpecializer[SectionParamsT], ABC):
+    """Abstract building block for prompt content."""
+
+    _generic_owner_name: ClassVar[str | None] = "Section"
+
+    def __init__(
+        self,
+        *,
+        title: str,
+        key: str,
+        default_params: SectionParamsT | None = None,
+        children: Sequence[Section[SupportsDataclass]] | None = None,
+        enabled: EnabledPredicate | None = None,
+        tools: Sequence[object] | None = None,
+        accepts_overrides: bool = True,
+    ) -> None:
+        super().__init__()
+        params_candidate = getattr(self.__class__, "_params_type", None)
+        candidate_type = (
+            params_candidate if isinstance(params_candidate, type) else None
+        )
+        params_type = cast(type[SupportsDataclass] | None, candidate_type)
+
+        self.params_type: type[SectionParamsT] | None = cast(
+            type[SectionParamsT] | None, params_type
+        )
+        self.param_type: type[SectionParamsT] | None = self.params_type
+        self.title = title
+        self.key = self._normalize_key(key)
+        self.default_params = default_params
+        self.accepts_overrides = accepts_overrides
+
+        if self.params_type is None and self.default_params is not None:
+            raise TypeError("Section without parameters cannot define default_params.")
+
+        normalized_children: list[Section[SupportsDataclass]] = []
+        raw_children: Sequence[object] = cast(Sequence[object], children or ())
+        for child in raw_children:
+            if not isinstance(child, Section):
+                raise TypeError("Section children must be Section instances.")
+            normalized_children.append(cast(Section[SupportsDataclass], child))
+        self.children = tuple(normalized_children)
+        self._enabled: Callable[[SupportsDataclass | None], bool] | None = (
+            self._normalize_enabled(enabled, params_type)
+        )
+        self._tools = self._normalize_tools(tools)
+
+    def is_enabled(self, params: SupportsDataclass | None) -> 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: SupportsDataclass | None, 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, SupportsToolResult], ...]:
+        """Return the tools exposed by this section."""
+
+        return self._tools
+
+    def original_body_template(self) -> str | None:
+        """Return the template text that participates in hashing, when available."""
+
+        return None
+
+    @staticmethod
+    def _normalize_key(key: str) -> str:
+        return normalize_component_key(key, owner="Section")
+
+    @staticmethod
+    def _normalize_tools(
+        tools: Sequence[object] | None,
+    ) -> tuple[Tool[SupportsDataclass, SupportsToolResult], ...]:
+        if not tools:
+            return ()
+
+        from .tool import Tool
+
+        normalized: list[Tool[SupportsDataclass, SupportsToolResult]] = []
+        for tool in tools:
+            if not isinstance(tool, Tool):
+                raise TypeError("Section tools must be Tool instances.")
+            normalized.append(cast(Tool[SupportsDataclass, SupportsToolResult], tool))
+        return tuple(normalized)
+
+    @staticmethod
+    def _normalize_enabled(
+        enabled: EnabledPredicate | None,
+        params_type: type[SupportsDataclass] | None,
+    ) -> Callable[[SupportsDataclass | None], bool] | None:
+        if enabled is None:
+            return None
+        if params_type is None and not _callable_requires_positional_argument(enabled):
+            zero_arg = cast(Callable[[], bool], enabled)
+
+            def _without_params(_: SupportsDataclass | None) -> bool:
+                return bool(zero_arg())
+
+            return _without_params
+
+        coerced = cast(Callable[[SupportsDataclass], bool], enabled)
+
+        def _with_params(value: SupportsDataclass | None) -> bool:
+            return bool(coerced(cast(SupportsDataclass, value)))
+
+        return _with_params
+
+
+def _callable_requires_positional_argument(callback: EnabledPredicate) -> bool:
+    try:
+        signature = inspect.signature(callback)
+    except (TypeError, ValueError):
+        return True
+    for parameter in signature.parameters.values():
+        if (
+            parameter.kind
+            in (
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            )
+            and parameter.default is inspect.Signature.empty
+        ):
+            return True
+    return False
+
+
+__all__ = ["Section"]

weakincentives/prompt/structured_output.py

@@ -0,0 +1,179 @@
+# 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
+
+import json
+import re
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Final, Literal, Protocol, TypeVar, cast
+
+from ..serde.parse import parse as parse_dataclass
+from ..types import JSONValue, ParseableDataclassT
+from ._types import SupportsDataclass
+
+__all__ = [
+    "ARRAY_WRAPPER_KEY",
+    "OutputParseError",
+    "StructuredOutputConfig",
+    "parse_structured_output",
+]
+
+ARRAY_WRAPPER_KEY: Final[str] = "items"
+
+_JSON_FENCE_PATTERN: Final[re.Pattern[str]] = re.compile(
+    r"```json\s*\n(.*?)```", re.IGNORECASE | re.DOTALL
+)
+
+
+DataclassT = TypeVar("DataclassT", bound=SupportsDataclass)
+
+
+@dataclass(frozen=True, slots=True)
+class StructuredOutputConfig[DataclassT]:
+    """Resolved structured output declaration for a prompt."""
+
+    dataclass_type: type[DataclassT]
+    container: Literal["object", "array"]
+    allow_extra_keys: bool
+
+
+class StructuredRenderedPrompt[PayloadT](Protocol):
+    @property
+    def structured_output(self) -> StructuredOutputConfig[SupportsDataclass] | None:
+        """Structured output metadata declared by the prompt."""
+
+
+class OutputParseError(Exception):
+    """Raised when structured output parsing fails."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        dataclass_type: type[SupportsDataclass] | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.dataclass_type = dataclass_type
+
+
+def parse_structured_output[PayloadT](
+    output_text: str, rendered: StructuredRenderedPrompt[PayloadT]
+) -> PayloadT:
+    """Parse a model response into the structured output type declared by the prompt."""
+
+    config = rendered.structured_output
+    if config is None:
+        raise OutputParseError("Prompt does not declare structured output.")
+
+    dataclass_type = config.dataclass_type
+    container = config.container
+    allow_extra_keys = config.allow_extra_keys
+    payload = _extract_json_payload(output_text, dataclass_type)
+    try:
+        parsed = parse_dataclass_payload(
+            dataclass_type,
+            container,
+            payload,
+            allow_extra_keys=allow_extra_keys,
+            object_error="Expected top-level JSON object.",
+            array_error="Expected top-level JSON array.",
+            array_item_error="Array item at index {index} is not an object.",
+        )
+    except (TypeError, ValueError) as error:
+        raise OutputParseError(str(error), dataclass_type=dataclass_type) from error
+
+    return cast(PayloadT, parsed)
+
+
+def _extract_json_payload(
+    text: str, dataclass_type: type[SupportsDataclass]
+) -> JSONValue:
+    fenced_match = _JSON_FENCE_PATTERN.search(text)
+    if fenced_match is not None:
+        block = fenced_match.group(1).strip()
+        try:
+            return json.loads(block)
+        except json.JSONDecodeError as error:
+            raise OutputParseError(
+                "Failed to decode JSON from fenced code block.",
+                dataclass_type=dataclass_type,
+            ) from error
+
+    stripped = text.strip()
+    if stripped:
+        try:
+            return json.loads(stripped)
+        except json.JSONDecodeError:
+            pass
+
+    decoder = json.JSONDecoder()
+    for index, character in enumerate(text):
+        if character not in "{[":
+            continue
+        try:
+            payload, _ = decoder.raw_decode(text, index)
+        except json.JSONDecodeError:
+            continue
+        return cast(JSONValue, payload)
+
+    raise OutputParseError(
+        "No JSON object or array found in assistant message.",
+        dataclass_type=dataclass_type,
+    )
+
+
+def parse_dataclass_payload(
+    dataclass_type: type[ParseableDataclassT],
+    container: Literal["object", "array"],
+    payload: JSONValue,
+    *,
+    allow_extra_keys: bool,
+    object_error: str,
+    array_error: str,
+    array_item_error: str,
+) -> ParseableDataclassT | list[ParseableDataclassT]:
+    if container not in {"object", "array"}:
+        raise TypeError("Unknown output container declared.")
+
+    extra_mode: Literal["ignore", "forbid"] = "ignore" if allow_extra_keys else "forbid"
+
+    if container == "object":
+        if not isinstance(payload, Mapping):
+            raise TypeError(object_error)
+        mapping_payload = cast(Mapping[str, JSONValue], payload)
+        return parse_dataclass(dataclass_type, mapping_payload, extra=extra_mode)
+
+    if isinstance(payload, Mapping):
+        mapping_payload = cast(Mapping[str, JSONValue], payload)
+        if ARRAY_WRAPPER_KEY not in mapping_payload:
+            raise TypeError(array_error)
+        payload = mapping_payload[ARRAY_WRAPPER_KEY]
+    if not isinstance(payload, Sequence) or isinstance(
+        payload, (str, bytes, bytearray)
+    ):
+        raise TypeError(array_error)
+    sequence_payload = cast(Sequence[JSONValue], payload)
+    parsed_items: list[ParseableDataclassT] = []
+    for index, item in enumerate(sequence_payload):
+        if not isinstance(item, Mapping):
+            raise TypeError(array_item_error.format(index=index))
+        mapping_item = cast(Mapping[str, JSONValue], item)
+        parsed_item = parse_dataclass(
+            dataclass_type,
+            mapping_item,
+            extra=extra_mode,
+        )
+        parsed_items.append(parsed_item)
+    return parsed_items