feat: preserve Context subtype and enforce input==output context type

mellea/core/__init__.py

@@ -23,6 +23,7 @@
     ComputedModelOutputThunk,
     Context,
     ContextTurn,
+    ContextTypeMismatchError,
     GenerateLog,
     GenerateType,
     GenerationMetadata,
@@ -76,6 +77,7 @@ def __getattr__(name: str) -> object:
     "ComputedModelOutputThunk",
     "Context",
     "ContextTurn",
+    "ContextTypeMismatchError",
     "Formatter",
     "GenerateLog",
     "GenerateType",

mellea/core/base.py

@@ -37,6 +37,7 @@
     Literal,
     ParamSpec,
     Protocol,
+    Self,
     TypeVar,
     runtime_checkable,
 )
@@ -1679,6 +1680,32 @@ class ContextTurn:
 ContextT = TypeVar("ContextT", bound="Context")
 
 
+class ContextTypeMismatchError(TypeError):
+    """Raised when a function returns a different `Context` subtype than it was given.
+
+    Mellea's convention is that the context type flowing out of a function equals
+    the context type flowing in (see issue #1522). This error enforces that
+    invariant. It is raised by the functional layer when a backend or sampling
+    strategy produces a context whose type differs from the input context's type,
+    unless the caller opted in to a deliberate type change via
+    `allow_context_type_change=True`.
+
+    Args:
+        input_type (type): The type of the context passed into the function.
+        output_type (type): The type of the context the function produced.
+    """
+
+    def __init__(self, input_type: type, output_type: type) -> None:
+        """Build the error message from the mismatched input and output context types."""
+        super().__init__(
+            f"Context type changed during generation: input was "
+            f"{input_type.__name__!r} but output is {output_type.__name__!r}. "
+            "Mellea functions must return the same Context subtype they were "
+            "given. If this change is deliberate (e.g. switching the context type "
+            "associated with a session), pass allow_context_type_change=True."
+        )
+
+
 class Context(abc.ABC):
     """A `Context` is used to track the state of a `MelleaSession`.
 
@@ -1738,7 +1765,7 @@ def reset_to_new(cls: type[ContextT]) -> ContextT:
         """
         return cls()
 
-    def new_instance(self) -> Context:
+    def new_instance(self) -> Self:
         """Return a new empty root context, preserving any subclass configuration.
 
         The base implementation calls `reset_to_new()`, which returns a bare
@@ -1747,7 +1774,7 @@ def new_instance(self) -> Context:
         should override this to propagate their config into the fresh instance.
 
         Returns:
-            Context: A freshly initialised root context of the same type.
+            Self: A freshly initialised root context of the same type.
         """
         return self.reset_to_new()
 
@@ -1877,14 +1904,15 @@ def last_turn(self) -> ContextTurn | None:
     # Abstract methods below this line.
 
     @abc.abstractmethod
-    def add(self, c: Span) -> Context:
+    def add(self, c: Span) -> Self:
         """Returns a new context obtained by appending `c` to this context.
 
         Args:
             c (Span): The component, content block, or model output to add to the context.
 
         Returns:
-            Context: A new context node with `c` as its data and this context as its previous node.
+            Self: A new context node of the same type with `c` as its data and this
+            context as its previous node.
         """
         # something along ....from_previous(self, c)
         ...

mellea/stdlib/context/chat.py

@@ -117,8 +117,12 @@ def model_id(self) -> str | ModelIdentifier | None:
         return self._model_id
 
     def _make_root(self, model_id: str | ModelIdentifier | None) -> ChatContext:
-        """Return a new empty root `ChatContext`, propagating all `_propagated_fields` then binding `model_id`."""
-        new = ChatContext()
+        """Return a new empty root `ChatContext`, propagating all `_propagated_fields` then binding `model_id`.
+
+        Uses `type(self)`, not `ChatContext`, so a subclass gets an instance of
+        itself back rather than being silently demoted to `ChatContext`.
+        """
+        new = type(self)()
         for field in self._propagated_fields:
             setattr(new, field, getattr(self, field))
         # Override whatever _propagated_fields copied for _model_id: the caller
@@ -175,9 +179,12 @@ def add(self, c: Span) -> ChatContext:
                 block, or model output to append.
 
         Returns:
-            ChatContext: A new `ChatContext` carrying the same configuration.
+            ChatContext: A new context of the same concrete subtype carrying the
+            same configuration.
         """
-        new = ChatContext.from_previous(self, c)
+        # `type(self)`, not `ChatContext`, so a subclass gets an instance of
+        # itself back rather than being silently demoted to `ChatContext`.
+        new = type(self).from_previous(self, c)
         for field in self._propagated_fields:
             setattr(new, field, getattr(self, field))
         if self._compactor is not None:
@@ -278,6 +285,7 @@ def _rebuild_chat_context(
     compactor: InlineCompactor | None = None,
     token_context_length_limit: int | None = None,
     model_id: str | ModelIdentifier | None = None,
+    cls: type[ChatContext] = ChatContext,
 ) -> ChatContext:
     """Build a fresh `ChatContext` linked-list without triggering compaction.
 
@@ -291,21 +299,24 @@ def _rebuild_chat_context(
         compactor: Compactor to attach to every node of the rebuilt context.
         token_context_length_limit: Token budget to attach to every node.
         model_id: Model identifier to attach to every node.
+        cls: The concrete `ChatContext` subtype to construct. Compactors pass
+            `type(ctx)` so a subclassed context is rebuilt as its own type
+            rather than being demoted to `ChatContext`.
 
     Returns:
-        A new `ChatContext` whose linear history is exactly `components`.
+        A new context of type `cls` whose linear history is exactly `components`.
     """
 
     def _configure(node: ChatContext) -> None:
         node._compactor = compactor
         node._token_context_length_limit = token_context_length_limit
         node._model_id = model_id
 
-    ctx: ChatContext = ChatContext.__new__(ChatContext)
+    ctx: ChatContext = cls.__new__(cls)
     Context.__init__(ctx)
     _configure(ctx)
     for c in components:
-        new: ChatContext = ChatContext.__new__(ChatContext)
+        new: ChatContext = cls.__new__(cls)
         new._previous = ctx
         new._data = c
         new._is_root = False

mellea/stdlib/context/compactor.py

@@ -273,6 +273,7 @@ def compact(
             compactor=ctx._compactor,
             token_context_length_limit=ctx._token_context_length_limit,
             model_id=ctx._model_id,
+            cls=type(ctx),
         )
 
 
@@ -619,4 +620,5 @@ async def _async_compact(
             compactor=ctx._compactor,
             token_context_length_limit=ctx._token_context_length_limit,
             model_id=ctx._model_id,
+            cls=type(ctx),
         )

mellea/stdlib/context/simple.py

@@ -19,10 +19,12 @@ def add(self, c: Span) -> SimpleContext:
                 block, or model output to record.
 
         Returns:
-            SimpleContext: A new `SimpleContext` containing only the added entry;
-            prior history is not retained.
+            SimpleContext: A new context of the same concrete subtype containing
+            only the added entry; prior history is not retained.
         """
-        return SimpleContext.from_previous(self, c)
+        # `type(self)`, not `SimpleContext`, so a subclass gets an instance of
+        # itself back rather than being silently demoted to `SimpleContext`.
+        return type(self).from_previous(self, c)
 
     def view_for_generation(self) -> list[Span] | None:
         """Return an empty list, since `SimpleContext` does not pass history to the model.