Skip to content

[FIX]: Enforce PromptTemplate construction invariants #132

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Spencer Schoenberg (spencrr) wants to merge 1 commit into from
Closed

[FIX]: Enforce PromptTemplate construction invariants #132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 43 additions & 8 deletions rampart/common/templates.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

from collections.abc import Hashable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Annotated, TypeAlias, TypeVar
from typing import TYPE_CHECKING, Annotated, TypeAlias, TypeVar, final

import yaml
from jinja2 import (
Expand Down Expand Up @@ -98,14 +98,30 @@ def __init__(
super().__init__(msg)


@dataclass(frozen=True, kw_only=True, slots=True)
@final
@dataclass(
frozen=True,
kw_only=True,
slots=True,
init=False,
eq=False,
)
class PromptTemplate:
"""Compiled prompt template with metadata and an explicit render contract."""

name: str
description: str | None
parameter_keys: tuple[str, ...]
_template: Template = field(repr=False, compare=False)
_template: Template = field(repr=False)

def __init__(self) -> None:
"""Reject construction that bypasses template validation.

Raises:
TypeError: Always. Use :meth:`from_yaml` to construct an instance.
"""
msg = "Use PromptTemplate.from_yaml()"
raise TypeError(msg)

@classmethod
def from_yaml(cls, path: Path) -> Self:
Expand All @@ -127,12 +143,31 @@ def from_yaml(cls, path: Path) -> Self:
a valid prompt template.
"""
definition = _load_yaml_definition(path)
return cls(
name=definition.name,
description=definition.description,
parameter_keys=tuple(definition.parameters),
_template=_compile_template(definition, path=path),
return cls._from_validated_definition(definition=definition, path=path)

@classmethod
def _from_validated_definition(
cls,
*,
definition: _PromptTemplateYaml,
path: Path,
) -> Self:
compiled = _compile_template(definition, path=path)

instance = object.__new__(cls)
object.__setattr__(instance, "name", definition.name) # ruff:ignore[unnecessary-dunder-call]
Comment thread
spencrr marked this conversation as resolved.
object.__setattr__( # ruff:ignore[unnecessary-dunder-call]
instance,
"description",
definition.description,
)
object.__setattr__( # ruff:ignore[unnecessary-dunder-call]
instance,
"parameter_keys",
tuple(definition.parameters),
)
object.__setattr__(instance, "_template", compiled) # ruff:ignore[unnecessary-dunder-call]
return instance
Comment on lines +157 to +170

@nina-msft Nina Chikanov (nina-msft) Jul 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was considering your comment on whether dataclass is still necessary since we are limiting some of its core usage...what do you think about this approach?

The object.__new__ + four object.__setattr__ reconstruction here only exists to work around init=False on a frozen, slotted dataclass — the class opts into the dataclass machinery and then fights it. You can keep the dataclass (its generated __init__, __repr__, frozen guard, and slots) and just gate construction with a module-private sentinel:

_FROM_YAML: Final = object()

@final
@dataclass(frozen=True, kw_only=True, slots=True, eq=False)
class PromptTemplate:
    name: str
    description: str | None
    parameter_keys: tuple[str, ...]
    _template: Template = field(repr=False)
    _token: InitVar[object]

    def __post_init__(self, _token: object) -> None:
        if _token is not _FROM_YAML:
            raise TypeError("Use PromptTemplate.from_yaml()")

from_yaml then just calls cls(name=..., description=..., parameter_keys=tuple(definition.parameters), _template=_compile_template(definition, path=path), _token=_FROM_YAML). That deletes _from_validated_definition, both object.* calls, the custom raising __init__, and all four lint suppressions — for the same guarantees.

I checked this on 3.11: direct construction, bare PromptTemplate(), and a wrong/absent token all raise TypeError; instances stay frozen (FrozenInstanceError); equality stays identity-based; the repr still hides _template; and because _token is an InitVar it never lands in __slots__ or on the instance. If you'd rather not carry a sentinel field, the other option is dropping @dataclass for a plain @final + __slots__ class — but then you hand-roll frozen and __repr__, so this is the smaller change.

@spencrr Spencer Schoenberg (spencrr) Aug 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nina Chikanov (@nina-msft) Yeah I am not a big fan of the token/sentinel thing. Can you PTAL at #144 - maybe a cleaner approach?

@nina-msft Nina Chikanov (nina-msft) Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we go with #144 - can you close this out as stale?

@spencrr Spencer Schoenberg (spencrr) Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Closing both in favor of your comment! Stay tuned!


def render(self, **kwargs: object) -> str:
"""Render with exactly the declared keyword arguments.
Expand Down
32 changes: 31 additions & 1 deletion tests/unit/common/test_templates.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,21 @@

from pathlib import Path
from textwrap import dedent
from typing import TYPE_CHECKING, cast

import pytest
import yaml
from jinja2 import TemplateError
from jinja2 import Template, TemplateError

from rampart.common.templates import (
PromptTemplate,
PromptTemplateDefinitionError,
TemplateParameterError,
)

if TYPE_CHECKING:
from collections.abc import Callable


def _write_yaml(tmp_path: Path, data: object) -> Path:
path = tmp_path / "prompt.yaml"
Expand All @@ -37,6 +41,32 @@ def _write_template(tmp_path: Path, **overrides: object) -> Path:
return _write_yaml(tmp_path, definition)


class TestPromptTemplateInitialization:
def test_rejects_direct_construction(self) -> None:
with pytest.raises(TypeError, match=r"Use PromptTemplate\.from_yaml\(\)"):
PromptTemplate()

def test_rejects_generated_constructor_arguments(self) -> None:
constructor = cast("Callable[..., PromptTemplate]", PromptTemplate)

with pytest.raises(TypeError):
constructor(
name="x",
description=None,
parameter_keys=("declared",),
_template=Template("{{ actual }}"),
)

def test_uses_identity_equality_and_hashing(self, tmp_path: Path) -> None:
path = _write_template(tmp_path)

first = PromptTemplate.from_yaml(path)
second = PromptTemplate.from_yaml(path)

assert first != second
assert len({first, second}) == 2


class TestPromptTemplateFromYaml:
def test_loads_raw_yaml_block_scalars(self, tmp_path: Path) -> None:
path = _write_raw_yaml(
Expand Down