diff --git a/mellea/core/__init__.py b/mellea/core/__init__.py index ec874d7937..c1d5148978 100644 --- a/mellea/core/__init__.py +++ b/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", diff --git a/mellea/core/base.py b/mellea/core/base.py index 1080ef94b3..73a8952adb 100644 --- a/mellea/core/base.py +++ b/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) ... diff --git a/mellea/stdlib/context/chat.py b/mellea/stdlib/context/chat.py index 37554d14c6..5ebcd1018f 100644 --- a/mellea/stdlib/context/chat.py +++ b/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,9 +299,12 @@ 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: @@ -301,11 +312,11 @@ def _configure(node: ChatContext) -> None: 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 diff --git a/mellea/stdlib/context/compactor.py b/mellea/stdlib/context/compactor.py index 85867d75d2..8ac170e01f 100644 --- a/mellea/stdlib/context/compactor.py +++ b/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), ) diff --git a/mellea/stdlib/context/simple.py b/mellea/stdlib/context/simple.py index 6803caa841..9de808440c 100644 --- a/mellea/stdlib/context/simple.py +++ b/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. diff --git a/mellea/stdlib/functional.py b/mellea/stdlib/functional.py index c5fa4ba266..40f58001b2 100644 --- a/mellea/stdlib/functional.py +++ b/mellea/stdlib/functional.py @@ -14,7 +14,7 @@ import time import uuid from collections.abc import Coroutine, Iterable -from typing import Any, Literal, overload +from typing import Any, Literal, TypeVar, cast, overload from PIL import Image as PILImage @@ -28,6 +28,7 @@ Component, ComputedModelOutputThunk, Context, + ContextTypeMismatchError, GenerateLog, ImageBlock, ImageUrlBlock, @@ -56,11 +57,44 @@ from .context import SimpleContext from .sampling import RejectionSamplingStrategy +# Bound to Context so functions can return the same subtype they were given +# (issue #1522): a `ChatContext` in yields a `ChatContext` out, statically. +ContextT = TypeVar("ContextT", bound=Context) + + +def _enforce_context_type( + input_ctx: ContextT, output_ctx: Context, *, allow_context_type_change: bool +) -> ContextT: + """Enforce the input==output context-type convention (issue #1522). + + Mellea functions return the same `Context` subtype they were given. This + checks that invariant at runtime and returns `output_ctx` narrowed to the + input type. The runtime type of `output_ctx` is unchanged; only the static + type is narrowed, which is sound because the types are asserted equal here. + + Args: + input_ctx (ContextT): The context passed into the function. + output_ctx (Context): The context the function produced. + allow_context_type_change (bool): When `True`, a differing output type is + permitted (the deliberate exception, e.g. switching a session's + context type) and returned as-is. + + Returns: + ContextT: `output_ctx`, narrowed to the input context's type. + + Raises: + ContextTypeMismatchError: If the output context type differs from the + input context type and `allow_context_type_change` is `False`. + """ + if type(output_ctx) is type(input_ctx) or allow_context_type_change: + return cast(ContextT, output_ctx) + raise ContextTypeMismatchError(type(input_ctx), type(output_ctx)) + @overload def act( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -69,13 +103,14 @@ def act( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, -) -> tuple[ComputedModelOutputThunk[S], Context]: ... + allow_context_type_change: bool = False, +) -> tuple[ComputedModelOutputThunk[S], ContextT]: ... @overload def act( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -84,12 +119,13 @@ def act( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, ) -> SamplingResult[S]: ... def act( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -98,7 +134,8 @@ def act( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, -) -> tuple[ComputedModelOutputThunk[S], Context] | SamplingResult[S]: + allow_context_type_change: bool = False, +) -> tuple[ComputedModelOutputThunk[S], ContextT] | SamplingResult[S]: """Runs a generic action, and adds both the action and the result to the context. Args: @@ -114,12 +151,14 @@ def act( `MyModel.model_validate_json(str(result))` to get a typed instance. model_options: additional model options, which will upsert into the model/backend's defaults. tool_calls: if true, tool calling is enabled. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. Raises: ValueError: if `return_sampling_results=True` without a `strategy`, or if `requirements` are provided without a `strategy` to validate them. + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. Returns: - A (ComputedModelOutputThunk, Context) if `return_sampling_results` is `False`, else returns a `SamplingResult`. + A (ComputedModelOutputThunk, Context) with the same context subtype as the input if `return_sampling_results` is `False`, else returns a `SamplingResult`. Always returns ComputedModelOutputThunk since sync functions must await completion. """ out = _run_async_in_thread( @@ -134,6 +173,7 @@ def act( model_options=model_options, tool_calls=tool_calls, silence_context_type_warning=True, # We can safely silence this here since it's in a sync function. + allow_context_type_change=allow_context_type_change, await_result=True, # Sync functions must always await ) # type: ignore[call-overload, misc] ) @@ -152,7 +192,7 @@ def act( @overload def instruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -168,13 +208,14 @@ def instruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, -) -> tuple[ComputedModelOutputThunk[str], Context]: ... + allow_context_type_change: bool = False, +) -> tuple[ComputedModelOutputThunk[str], ContextT]: ... @overload def instruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -190,12 +231,13 @@ def instruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, ) -> SamplingResult[str]: ... def instruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -211,7 +253,8 @@ def instruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, -) -> tuple[ComputedModelOutputThunk[str], Context] | SamplingResult[str]: + allow_context_type_change: bool = False, +) -> tuple[ComputedModelOutputThunk[str], ContextT] | SamplingResult[str]: """Generates from an instruction. Args: @@ -234,9 +277,13 @@ def instruct( tool_calls: If true, tool calling is enabled. images: A list of images to be used in the instruction or None if none. audio: A list of audio blocks to be used in the instruction or None if none. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. Returns: - A (ComputedModelOutputThunk, Context) if `return_sampling_results` is `False`, else returns a `SamplingResult`. + A (ComputedModelOutputThunk, Context) with the same context subtype as the input if `return_sampling_results` is `False`, else returns a `SamplingResult`. Always returns ComputedModelOutputThunk since sync functions must await completion. """ requirements = [] if requirements is None else requirements @@ -271,12 +318,13 @@ def instruct( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=allow_context_type_change, ) # type: ignore[call-overload] def chat( content: str, - context: Context, + context: ContextT, backend: Backend, *, role: Message.Role = "user", @@ -287,7 +335,8 @@ def chat( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, -) -> tuple[Message, Context]: + allow_context_type_change: bool = False, +) -> tuple[Message, ContextT]: """Sends a simple chat message and returns the response. Adds both messages to the Context. Args: @@ -303,9 +352,13 @@ def chat( format: Optional Pydantic model for constrained decoding of the response. model_options: Additional model options to merge with backend defaults. tool_calls: If true, tool calling is enabled. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. Returns: - Tuple of the assistant `Message` and the updated `Context`. + Tuple of the assistant `Message` and the updated `Context` (same subtype as the input). """ if user_variables is not None: content_resolved = Instruction.apply_user_dict_from_jinja( @@ -331,6 +384,7 @@ def chat( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=allow_context_type_change, ) parsed_assistant_message = result.parsed_repr assert isinstance(parsed_assistant_message, Message) @@ -387,13 +441,14 @@ def validate( def query( obj: Any, query: str, - context: Context, + context: ContextT, backend: Backend, *, format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, -) -> tuple[ComputedModelOutputThunk, Context]: + allow_context_type_change: bool = False, +) -> tuple[ComputedModelOutputThunk, ContextT]: """Query method for retrieving information from an object. Args: @@ -404,9 +459,13 @@ def query( format: format for output parsing. model_options: Model options to pass to the backend. tool_calls: If true, the model may make tool calls. Defaults to False. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. Returns: - tuple[ComputedModelOutputThunk, Context]: The result of the query and updated context. + tuple[ComputedModelOutputThunk, Context]: The result of the query and updated context (same subtype as the input). """ if not isinstance(obj, MObjectProtocol): obj = mify(obj) @@ -423,6 +482,7 @@ def query( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=allow_context_type_change, ) return answer @@ -430,12 +490,13 @@ def query( def transform( obj: Any, transformation: str, - context: Context, + context: ContextT, backend: Backend, *, format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, -) -> tuple[ModelOutputThunk | Any, Context]: + allow_context_type_change: bool = False, +) -> tuple[ModelOutputThunk | Any, ContextT]: """Transform method for creating a new object with the transformation applied. Args: @@ -445,9 +506,13 @@ def transform( backend: the backend used to generate the response. format: format for output parsing; usually not needed with transform. model_options: Model options to pass to the backend. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. Returns: - (ModelOutputThunk | Any, Context): The result of the transformation as processed by the backend. If no tools were called, + (ModelOutputThunk | Any, Context): The result of the transformation as processed by the backend, with the same context subtype as the input. If no tools were called, the return type will be always be (ModelOutputThunk, Context). If a tool was called, the return type will be the return type of the function called, usually the type of the object passed in. """ @@ -468,6 +533,7 @@ def transform( format=format, model_options=model_options, tool_calls=True, + allow_context_type_change=allow_context_type_change, ) tools = call_tools(transformed, backend) @@ -512,7 +578,7 @@ def transform( @overload async def aact( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -522,14 +588,15 @@ async def aact( model_options: dict | None = None, tool_calls: bool = False, silence_context_type_warning: bool = False, + allow_context_type_change: bool = False, await_result: Literal[True], -) -> tuple[ComputedModelOutputThunk[S], Context]: ... +) -> tuple[ComputedModelOutputThunk[S], ContextT]: ... @overload async def aact( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -539,14 +606,15 @@ async def aact( model_options: dict | None = None, tool_calls: bool = False, silence_context_type_warning: bool = False, + allow_context_type_change: bool = False, await_result: bool = False, -) -> tuple[ComputedModelOutputThunk[S], Context]: ... +) -> tuple[ComputedModelOutputThunk[S], ContextT]: ... @overload async def aact( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -556,14 +624,15 @@ async def aact( model_options: dict | None = None, tool_calls: bool = False, silence_context_type_warning: bool = False, + allow_context_type_change: bool = False, await_result: Literal[False] = False, -) -> tuple[ModelOutputThunk[S], Context]: ... +) -> tuple[ModelOutputThunk[S], ContextT]: ... @overload async def aact( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -573,13 +642,14 @@ async def aact( model_options: dict | None = None, tool_calls: bool = False, silence_context_type_warning: bool = False, + allow_context_type_change: bool = False, await_result: bool = False, ) -> SamplingResult[S]: ... async def aact( action: Component[S] | CBlock | ModelOutputThunk, - context: Context, + context: ContextT, backend: Backend, *, requirements: list[Requirement] | None = None, @@ -589,8 +659,9 @@ async def aact( model_options: dict | None = None, tool_calls: bool = False, silence_context_type_warning: bool = False, + allow_context_type_change: bool = False, await_result: bool = False, -) -> tuple[ModelOutputThunk[S], Context] | SamplingResult: +) -> tuple[ModelOutputThunk[S], ContextT] | SamplingResult: """Asynchronous version of .act; runs a generic action, and adds both the action and the result to the context. Args: @@ -607,13 +678,15 @@ async def aact( model_options: additional model options, which will upsert into the model/backend's defaults. tool_calls: if true, tool calling is enabled. silence_context_type_warning: if called directly from an asynchronous function, will log a warning if not using a SimpleContext + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError` to enforce the input==output context-type convention. await_result: if False and strategy is None, returns uncomputed ModelOutputThunk for streaming. If True or strategy is not None, awaits and returns ComputedModelOutputThunk. Default is False. Raises: ValueError: if `return_sampling_results=True` without a `strategy`, or if `requirements` are provided without a `strategy` to validate them. + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. Returns: - A (ModelOutputThunk, Context) if `return_sampling_results` is `False`, else returns a `SamplingResult`. + A (ModelOutputThunk, Context) with the same context subtype as the input if `return_sampling_results` is `False`, else returns a `SamplingResult`. """ import time import traceback @@ -753,9 +826,23 @@ async def aact( assert ( sampling_result is not None ) # Needed for the type checker but should never happen. + # `SamplingResult` does not statically track its context type, so the + # input==output convention (issue #1522) can only be enforced at + # runtime here. Check every sample context, not just the chosen one, + # so a strategy that produces a mismatched context for any attempt + # is caught rather than silently returned. + for sample_ctx in sampling_result.sample_contexts: + _enforce_context_type( + context, + sample_ctx, + allow_context_type_change=allow_context_type_change, + ) return sampling_result else: - return result, new_ctx + checked_ctx = _enforce_context_type( + context, new_ctx, allow_context_type_change=allow_context_type_change + ) + return result, checked_ctx except BaseException as exc: # --- component_post_error hook --- @@ -781,7 +868,7 @@ async def aact( @overload async def ainstruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -797,14 +884,15 @@ async def ainstruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: Literal[True], -) -> tuple[ComputedModelOutputThunk[str], Context]: ... +) -> tuple[ComputedModelOutputThunk[str], ContextT]: ... @overload async def ainstruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -820,14 +908,15 @@ async def ainstruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: bool = False, -) -> tuple[ComputedModelOutputThunk[str], Context]: ... +) -> tuple[ComputedModelOutputThunk[str], ContextT]: ... @overload async def ainstruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -843,14 +932,15 @@ async def ainstruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: Literal[False] = False, -) -> tuple[ModelOutputThunk[str], Context]: ... +) -> tuple[ModelOutputThunk[str], ContextT]: ... @overload async def ainstruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -866,13 +956,14 @@ async def ainstruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: bool = False, ) -> SamplingResult[str]: ... async def ainstruct( description: str, - context: Context, + context: ContextT, backend: Backend, *, images: list[ImageBlock | ImageUrlBlock] | list[PILImage.Image] | None = None, @@ -888,8 +979,9 @@ async def ainstruct( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: bool = False, -) -> tuple[ModelOutputThunk[str], Context] | SamplingResult: +) -> tuple[ModelOutputThunk[str], ContextT] | SamplingResult: """Generates from an instruction. Args: @@ -912,10 +1004,14 @@ async def ainstruct( tool_calls: If true, tool calling is enabled. images: A list of images to be used in the instruction or None if none. audio: A list of audio blocks to be used in the instruction or None if none. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. await_result: if False and strategy is None, returns uncomputed ModelOutputThunk for streaming. If True or strategy is not None, awaits and returns ComputedModelOutputThunk. Default is False. Returns: A (ModelOutputThunk, Context) if `return_sampling_results` is `False`, else returns a `SamplingResult`. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. """ requirements = [] if requirements is None else requirements icl_examples = [] if icl_examples is None else icl_examples @@ -949,13 +1045,14 @@ async def ainstruct( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=allow_context_type_change, await_result=await_result, ) # type: ignore[call-overload] async def achat( content: str, - context: Context, + context: ContextT, backend: Backend, *, role: Message.Role = "user", @@ -966,7 +1063,8 @@ async def achat( format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, -) -> tuple[Message, Context]: + allow_context_type_change: bool = False, +) -> tuple[Message, ContextT]: """Sends a simple chat message and returns the response. Adds both messages to the Context. Args: @@ -982,9 +1080,13 @@ async def achat( format: Optional Pydantic model for constrained decoding of the response. model_options: Additional model options to merge with backend defaults. tool_calls: If true, tool calling is enabled. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. Returns: Tuple of the assistant `Message` and the updated `Context`. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. """ if user_variables is not None: content_resolved = Instruction.apply_user_dict_from_jinja( @@ -1010,6 +1112,7 @@ async def achat( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=allow_context_type_change, await_result=True, # Must compute for Message parsing below. ) parsed_assistant_message = result.parsed_repr @@ -1135,41 +1238,44 @@ async def avalidate( async def aquery( obj: Any, query: str, - context: Context, + context: ContextT, backend: Backend, *, format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: Literal[True], -) -> tuple[ComputedModelOutputThunk, Context]: ... +) -> tuple[ComputedModelOutputThunk, ContextT]: ... @overload async def aquery( obj: Any, query: str, - context: Context, + context: ContextT, backend: Backend, *, format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: Literal[False] = False, -) -> tuple[ModelOutputThunk, Context]: ... +) -> tuple[ModelOutputThunk, ContextT]: ... async def aquery( obj: Any, query: str, - context: Context, + context: ContextT, backend: Backend, *, format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, tool_calls: bool = False, + allow_context_type_change: bool = False, await_result: bool = False, -) -> tuple[ModelOutputThunk, Context]: +) -> tuple[ModelOutputThunk, ContextT]: """Query method for retrieving information from an object. Args: @@ -1180,10 +1286,14 @@ async def aquery( format: format for output parsing. model_options: Model options to pass to the backend. tool_calls: If true, the model may make tool calls. Defaults to False. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. await_result: if False (default), returns uncomputed ModelOutputThunk. If True, awaits and returns ComputedModelOutputThunk. Returns: tuple[ModelOutputThunk, Context]: The result of the query and updated context. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. """ if not isinstance(obj, MObjectProtocol): obj = mify(obj) @@ -1200,6 +1310,7 @@ async def aquery( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=allow_context_type_change, await_result=await_result, # type: ignore[call-overload] ) return answer @@ -1208,12 +1319,13 @@ async def aquery( async def atransform( obj: Any, transformation: str, - context: Context, + context: ContextT, backend: Backend, *, format: type[BaseModelSubclass] | None = None, model_options: dict | None = None, -) -> tuple[ModelOutputThunk | Any, Context]: + allow_context_type_change: bool = False, +) -> tuple[ModelOutputThunk | Any, ContextT]: """Transform method for creating a new object with the transformation applied. Args: @@ -1223,11 +1335,15 @@ async def atransform( backend: the backend used to generate the response. format: format for output parsing; usually not needed with transform. model_options: Model options to pass to the backend. + allow_context_type_change: if True, permits the returned context to be a different `Context` subtype than `context`. By default (False), a differing type raises `ContextTypeMismatchError`. Returns: tuple[ModelOutputThunk | Any, Context]: The result of the transformation and updated context. If no tools were called, the first element will always be ModelOutputThunk. If a tool was called, the first element will be the return type of the function called, usually the type of the object passed in. + + Raises: + ContextTypeMismatchError: if the returned context type differs from the input context type and `allow_context_type_change` is `False`. """ if not isinstance(obj, MObjectProtocol): obj = mify(obj) @@ -1246,6 +1362,7 @@ async def atransform( format=format, model_options=model_options, tool_calls=True, + allow_context_type_change=allow_context_type_change, await_result=True, # Must be computed for tool calls. ) diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index bf3571f963..103f2a94c2 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -18,7 +18,7 @@ import contextvars import inspect from copy import copy -from typing import Any, Literal, overload +from typing import Any, Generic, Literal, TypeVar, cast, overload from PIL import Image as PILImage @@ -69,6 +69,12 @@ backend_name_to_class, ) +# Bound to Context so a session remembers the concrete subtype it was built +# with and exposes it statically via `ctx` (issue #1522): a session created with +# a `ChatContext` is a `MelleaSession[ChatContext]`, so `session.ctx` narrows to +# `ChatContext` rather than the base `Context`. +ContextT = TypeVar("ContextT", bound=Context) + # Global context variable for the context session _context_session: contextvars.ContextVar[MelleaSession | None] = contextvars.ContextVar( "context_session", default=None @@ -92,6 +98,62 @@ def get_session() -> MelleaSession: return session +@overload +def start_session( + backend_name: Literal["ollama", "hf", "openai", "watsonx", "litellm"], + model_id: str | ModelIdentifier, + ctx: ContextT, + *, + context_type: None = None, + model_options: dict | None = ..., + plugins: list[Any] | None = ..., + allow_context_type_change: bool = ..., + **backend_kwargs: Any, +) -> MelleaSession[ContextT]: ... + + +@overload +def start_session( + backend_name: Literal["ollama", "hf", "openai", "watsonx", "litellm"] = ..., + model_id: str | ModelIdentifier = ..., + *, + ctx: ContextT, + context_type: None = None, + model_options: dict | None = ..., + plugins: list[Any] | None = ..., + allow_context_type_change: bool = ..., + **backend_kwargs: Any, +) -> MelleaSession[ContextT]: ... + + +@overload +def start_session( + backend_name: Literal["ollama", "hf", "openai", "watsonx", "litellm"] = ..., + model_id: str | ModelIdentifier = ..., + ctx: None = None, + *, + context_type: Literal["chat"], + model_options: dict | None = ..., + plugins: list[Any] | None = ..., + allow_context_type_change: bool = ..., + **backend_kwargs: Any, +) -> MelleaSession[ChatContext]: ... + + +@overload +def start_session( + backend_name: Literal["ollama", "hf", "openai", "watsonx", "litellm"] = ..., + model_id: str | ModelIdentifier = ..., + ctx: None = None, + *, + context_type: Literal["simple"] | None = None, + model_options: dict | None = ..., + plugins: list[Any] | None = ..., + allow_context_type_change: bool = ..., + **backend_kwargs: Any, +) -> MelleaSession[SimpleContext]: ... + + def start_session( backend_name: Literal["ollama", "hf", "openai", "watsonx", "litellm"] = "ollama", model_id: str | ModelIdentifier = IBM_GRANITE_4_1_3B, @@ -100,6 +162,7 @@ def start_session( context_type: Literal["simple", "chat"] | None = None, model_options: dict | None = None, plugins: list[Any] | None = None, + allow_context_type_change: bool = False, **backend_kwargs: Any, ) -> MelleaSession: """Start a new Mellea session. Can be used as a context manager or called directly. @@ -129,6 +192,10 @@ def start_session( plugins: Optional list of plugins scoped to this session. Accepts `@hook`-decorated functions, `@plugin`-decorated class instances, `MelleaPlugin` instances, or `PluginSet` instances. + allow_context_type_change: When `True`, the session's model-interaction + methods permit a call to return a different `Context` subtype than + the session currently holds. By default (`False`), a differing type + raises `ContextTypeMismatchError` (issue #1522). **backend_kwargs: Additional keyword arguments passed to the backend constructor. Returns: @@ -231,7 +298,12 @@ def start_session( + (f", model_options={model_options}" if model_options else "") ) - session = MelleaSession(backend, resolved_ctx, session_id=session_id) + session = MelleaSession( + backend, + resolved_ctx, + session_id=session_id, + allow_context_type_change=allow_context_type_change, + ) # Register session-scoped plugins if plugins: @@ -258,7 +330,7 @@ def start_session( return session -class MelleaSession: +class MelleaSession(Generic[ContextT]): """Mellea sessions are a THIN wrapper around `m` convenience functions with NO special semantics. Using a Mellea session is not required, but it does represent the "happy path" of Mellea programming. Some nice things about ussing a `MelleaSession`: @@ -276,22 +348,58 @@ class MelleaSession: session. ctx (Context | None): The conversation context. Defaults to a new `SimpleContext` if `None`. + allow_context_type_change (bool): When `True`, model-interaction methods + permit a call to return a different `Context` subtype than the + session currently holds. By default (`False`), a differing type + raises `ContextTypeMismatchError`, enforcing the input==output + context-type convention (issue #1522). Set this when a session + deliberately switches context types mid-run. + + The session is generic in its context type (issue #1522): the concrete + `Context` subtype passed at construction is remembered as the type parameter, + so `session.ctx` narrows to that subtype statically. `MelleaSession(backend, + ChatContext())` is a `MelleaSession[ChatContext]` whose `ctx` is typed + `ChatContext`; omitting `ctx` yields a `MelleaSession[SimpleContext]`. Attributes: ctx (Context): The active conversation context; never `None` (defaults to a fresh `SimpleContext` when `None` is passed). Updated after - every call that produces model output. + every call that produces model output. Statically typed as the + session's context type parameter. id (str): Unique session UUID assigned at construction. + allow_context_type_change (bool): Whether model-interaction methods + permit the returned context to change subtype (see Args). """ # ``ctx`` is exposed as a property below; backing field is ``_ctx``. + @overload + def __init__( + self: MelleaSession[SimpleContext], + backend: Backend, + ctx: None = None, + *, + session_id: str | None = None, + allow_context_type_change: bool = False, + ) -> None: ... + + @overload + def __init__( + self: MelleaSession[ContextT], + backend: Backend, + ctx: ContextT, + *, + session_id: str | None = None, + allow_context_type_change: bool = False, + ) -> None: ... + def __init__( self, backend: Backend, ctx: Context | None = None, *, session_id: str | None = None, + allow_context_type_change: bool = False, ): """Initialize MelleaSession with a backend and optional conversation context. @@ -303,20 +411,33 @@ def __init__( """ import uuid + # The overloads above bind ContextT (to the ctx's type, or SimpleContext + # when ctx is None); the implementation body is unparameterized, so cast + # the resolved context back to ContextT for the typed _init_fields. + resolved_ctx = cast(ContextT, ctx if ctx is not None else SimpleContext()) self._init_fields( session_id if session_id is not None else str(uuid.uuid4()), backend, - ctx if ctx is not None else SimpleContext(), + resolved_ctx, + allow_context_type_change=allow_context_type_change, ) self._auto_bind_model() - def _init_fields(self, session_id: str, backend: Backend, ctx: Context) -> None: + def _init_fields( + self, + session_id: str, + backend: Backend, + ctx: ContextT, + *, + allow_context_type_change: bool = False, + ) -> None: """Set all instance fields. Shared by __init__ and __copy__ so neither diverges.""" self.id = session_id - self.backend = backend + self.backend: Backend = backend + self.allow_context_type_change = allow_context_type_change # Bypass the ctx setter so this initial assignment doesn't count as an # interaction. - self._ctx: Context = ctx + self._ctx: ContextT = ctx self._interaction_count: int = 0 self._session_logger = MelleaLogger.get_logger() self._context_token = None @@ -324,12 +445,12 @@ def _init_fields(self, session_id: str, backend: Backend, ctx: Context) -> None: self._exit_stack: contextlib.ExitStack | None = None @property - def ctx(self) -> Context: - """The session's current conversation context.""" + def ctx(self) -> ContextT: + """The session's current conversation context, typed as the session's context subtype.""" return self._ctx @ctx.setter - def ctx(self, value: Context) -> None: + def ctx(self, value: ContextT) -> None: """Replace the context and count this as one interaction. Every model-interaction code path in this class assigns to `self.ctx` @@ -350,7 +471,11 @@ def _auto_bind_model(self) -> None: and getattr(self.backend, "model_id", None) is not None ): # Bypass the setter — binding at construction is not an interaction. - self._ctx = self.ctx._bind_model(getattr(self.backend, "model_id")) + # `_bind_model` returns a `ChatContext`; the `isinstance` guard above + # means `ContextT` is `ChatContext` here, so the cast is sound. + self._ctx = cast( + ContextT, self.ctx._bind_model(getattr(self.backend, "model_id")) + ) def __enter__(self): """Enter context manager and set this session as the current global session.""" @@ -396,16 +521,23 @@ def __exit__(self, exc_type, exc_val, exc_tb): self._exit_stack = None finish_session_span(self.id, exception=exc_val) - def __copy__(self): + def __copy__(self) -> MelleaSession[ContextT]: """Use self.clone. Copies the current session but keeps references to the backend and context.""" import uuid - new = object.__new__(MelleaSession) - new._init_fields(str(uuid.uuid4()), self.backend, self.ctx) + new: MelleaSession[ContextT] = object.__new__( + cast("type[MelleaSession[ContextT]]", MelleaSession) + ) + new._init_fields( + str(uuid.uuid4()), + self.backend, + self.ctx, + allow_context_type_change=self.allow_context_type_change, + ) # Do not call _auto_bind_model: the session state is already settled. return new - def clone(self) -> MelleaSession: + def clone(self) -> MelleaSession[ContextT]: """Useful for running multiple generation requests while keeping the context at a given point in time. Returns: @@ -546,10 +678,14 @@ def act( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=self.allow_context_type_change, ) # type: ignore if isinstance(r, SamplingResult): - self.ctx = r.result_ctx + # SamplingResult doesn't track its context type statically; the + # runtime guard in mfuncs already enforced input==output type + # (unless allow_context_type_change), so this cast is sound. + self.ctx = cast(ContextT, r.result_ctx) return r else: result, context = r @@ -656,10 +792,14 @@ def instruct( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=self.allow_context_type_change, ) if isinstance(r, SamplingResult): - self.ctx = r.result_ctx + # SamplingResult doesn't track its context type statically; the + # runtime guard in mfuncs already enforced input==output type + # (unless allow_context_type_change), so this cast is sound. + self.ctx = cast(ContextT, r.result_ctx) return r else: # It's a tuple[ModelOutputThunk, Context]. @@ -709,6 +849,7 @@ def chat( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=self.allow_context_type_change, ) self.ctx = context @@ -777,6 +918,7 @@ def query( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=self.allow_context_type_change, ) self.ctx = context return result @@ -809,6 +951,7 @@ def transform( backend=self.backend, format=format, model_options=model_options, + allow_context_type_change=self.allow_context_type_change, ) self.ctx = context return result @@ -914,10 +1057,14 @@ async def aact( model_options=model_options, tool_calls=tool_calls, await_result=await_result, + allow_context_type_change=self.allow_context_type_change, ) # type: ignore if isinstance(r, SamplingResult): - self.ctx = r.result_ctx + # SamplingResult doesn't track its context type statically; the + # runtime guard in mfuncs already enforced input==output type + # (unless allow_context_type_change), so this cast is sound. + self.ctx = cast(ContextT, r.result_ctx) return r else: result, context = r @@ -1072,10 +1219,14 @@ async def ainstruct( model_options=model_options, tool_calls=tool_calls, await_result=await_result, + allow_context_type_change=self.allow_context_type_change, ) if isinstance(r, SamplingResult): - self.ctx = r.result_ctx + # SamplingResult doesn't track its context type statically; the + # runtime guard in mfuncs already enforced input==output type + # (unless allow_context_type_change), so this cast is sound. + self.ctx = cast(ContextT, r.result_ctx) return r else: # It's a tuple[ModelOutputThunk, Context]. @@ -1125,6 +1276,7 @@ async def achat( format=format, model_options=model_options, tool_calls=tool_calls, + allow_context_type_change=self.allow_context_type_change, ) self.ctx = context @@ -1220,6 +1372,7 @@ async def aquery( model_options=model_options, tool_calls=tool_calls, await_result=await_result, # type: ignore[call-overload] + allow_context_type_change=self.allow_context_type_change, ) self.ctx = context return result @@ -1252,6 +1405,7 @@ async def atransform( backend=self.backend, format=format, model_options=model_options, + allow_context_type_change=self.allow_context_type_change, ) self.ctx = context return result diff --git a/test/stdlib/test_context_type_enforcement.py b/test/stdlib/test_context_type_enforcement.py new file mode 100644 index 0000000000..77e0b67ae6 --- /dev/null +++ b/test/stdlib/test_context_type_enforcement.py @@ -0,0 +1,372 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for context-type enforcement in the functional layer (issue #1522). + +Mellea functions must return the same `Context` subtype they were given. These +tests use `DummyBackend` (no LLM) to verify the runtime guard: the type is +preserved on the happy path, a mismatch raises `ContextTypeMismatchError`, and +the `allow_context_type_change` escape hatch permits a deliberate change. +""" + +import pytest + +from mellea.backends.dummy import DummyBackend +from mellea.core import ( + BaseModelSubclass, + C, + CBlock, + Component, + ComputedModelOutputThunk, + Context, + ContextTypeMismatchError, + GenerateLog, + ModelOutputThunk, + SamplingResult, + SamplingStrategy, + Span, +) +from mellea.stdlib.components import Message +from mellea.stdlib.context import ChatContext, SimpleContext +from mellea.stdlib.functional import aact + + +class _MinimalContext(Context): + """A minimal `Context` subclass that overrides only the two abstract methods. + + Used to prove that the inherited, `Self`-typed helpers (`new_instance`, + `reset_to_new`) return the *subclass* type rather than the base `Context`, + and that a straightforward `add` override returns its own type too. + """ + + def add(self, c: Span) -> "_MinimalContext": + """Return a new `_MinimalContext` node with `c` appended.""" + return _MinimalContext.from_previous(self, c) + + def view_for_generation(self) -> list[Span] | None: + """Return the full linear history.""" + return self.as_list() + + +class _TypeChangingBackend(DummyBackend): + """Backend that always returns a `SimpleContext`, ignoring the input type. + + Used to force the input/output context-type mismatch the guard must catch. + """ + + async def _generate_from_context( + self, + action: Component[C] | CBlock | ModelOutputThunk, + ctx: Context, + *, + format: type[BaseModelSubclass] | None = None, + model_options: dict | None = None, + tool_calls: bool = False, + ) -> tuple[ModelOutputThunk[C], Context]: + """Return an output whose context is a fresh `SimpleContext`.""" + mot: ModelOutputThunk = ModelOutputThunk(value="dummy") + new_ctx = SimpleContext().add(action).add(mot) + return mot, new_ctx # type: ignore[return-value] + + +def _action() -> Message: + return Message(role="user", content="hello") + + +async def test_chat_context_type_preserved(): + """A `ChatContext` in yields a `ChatContext` out; no error is raised.""" + backend = DummyBackend(responses=None) + ctx_in = ChatContext() + + _, ctx_out = await aact( + _action(), ctx_in, backend, silence_context_type_warning=True + ) + + assert type(ctx_out) is type(ctx_in) + assert isinstance(ctx_out, ChatContext) + + +async def test_simple_context_type_preserved(): + """A `SimpleContext` in yields a `SimpleContext` out; no error is raised.""" + backend = DummyBackend(responses=None) + ctx_in = SimpleContext() + + _, ctx_out = await aact(_action(), ctx_in, backend) + + assert type(ctx_out) is type(ctx_in) + assert isinstance(ctx_out, SimpleContext) + + +async def test_type_mismatch_raises(): + """A backend that changes the context type trips the guard by default.""" + backend = _TypeChangingBackend(responses=None) + ctx_in = ChatContext() + + with pytest.raises(ContextTypeMismatchError): + await aact(_action(), ctx_in, backend, silence_context_type_warning=True) + + +async def test_escape_hatch_allows_type_change(): + """`allow_context_type_change=True` permits a deliberate type change.""" + backend = _TypeChangingBackend(responses=None) + ctx_in = ChatContext() + + _, ctx_out = await aact( + _action(), + ctx_in, + backend, + silence_context_type_warning=True, + allow_context_type_change=True, + ) + + assert isinstance(ctx_out, SimpleContext) + + +# --- self-returning Context functions preserve the concrete subtype (#1522) --- +# +# The runtime guard in `functional.py` only holds if the context helpers that +# claim to return `Self` (`add`) or the subclass type (`new_instance`, +# `reset_to_new`) actually do so when subclassed. These tests pin that +# invariant directly on the context types. + + +class _ChatSubclass(ChatContext): + """A bare subclass of `ChatContext` that adds no behaviour of its own. + + Inherits `add` / `new_instance` unchanged. If those helpers construct the + hard-coded `ChatContext` type instead of `type(self)`, an instance of this + class silently degrades to `ChatContext` on the first `add` — the exact + regression these tests guard against. + """ + + +class _SimpleSubclass(SimpleContext): + """A bare subclass of `SimpleContext`; see `_ChatSubclass` for the rationale.""" + + +# The stdlib contexts, a direct `Context` subclass, and bare subclasses of the +# stdlib contexts. The subclasses are what catch the "named constructor instead +# of `type(self)`" bug: they inherit `add` verbatim, so they only stay their own +# type if the inherited helper builds `type(self)`. +_CONTEXT_TYPES = [ + ChatContext, + SimpleContext, + _MinimalContext, + _ChatSubclass, + _SimpleSubclass, +] + + +@pytest.mark.parametrize("ctx_cls", _CONTEXT_TYPES) +def test_add_returns_same_subtype(ctx_cls): + """`ctx.add(...)` returns a context of the same subtype it was called on.""" + ctx = ctx_cls() + added = ctx.add(_action()) + assert type(added) is ctx_cls + + +@pytest.mark.parametrize("ctx_cls", _CONTEXT_TYPES) +def test_add_twice_preserves_subtype(ctx_cls): + """Chaining `add` keeps the subtype, so history nodes never demote.""" + ctx = ctx_cls().add(_action()).add(Message(role="assistant", content="hi")) + assert type(ctx) is ctx_cls + + +@pytest.mark.parametrize("ctx_cls", _CONTEXT_TYPES) +def test_new_instance_returns_same_subtype(ctx_cls): + """`ctx.new_instance()` returns a fresh root context of the same subtype.""" + ctx = ctx_cls().add(_action()) + fresh = ctx.new_instance() + assert type(fresh) is ctx_cls + assert fresh.is_root_node + + +@pytest.mark.parametrize("ctx_cls", _CONTEXT_TYPES) +def test_reset_to_new_returns_same_subtype(ctx_cls): + """The `reset_to_new()` classmethod returns an instance of the class it is called on.""" + fresh = ctx_cls.reset_to_new() + assert type(fresh) is ctx_cls + assert fresh.is_root_node + + +def test_chat_subclass_survives_compaction(): + """A `ChatContext` subclass stays its own type after the compactor fires. + + Compaction rebuilds the linked list via `_rebuild_chat_context`, which must + reconstruct the concrete subtype (`type(ctx)`) rather than a plain + `ChatContext`. A `window_size` of 1 forces the `WindowCompactor` to run on + the third `add`. + """ + ctx = _ChatSubclass(window_size=1) + ctx = ( + ctx.add(_action()).add(Message(role="assistant", content="one")).add(_action()) + ) + assert type(ctx) is _ChatSubclass + + +async def test_chat_subclass_passes_functional_guard(): + """A subclassed `ChatContext` round-trips through `aact` without tripping the guard. + + This is the end-to-end payoff of the `type(self)` fix: because `add` + preserves the subtype, input type == output type and + `_enforce_context_type` is satisfied with no escape hatch. + """ + backend = DummyBackend(responses=None) + ctx_in = _ChatSubclass() + + _, ctx_out = await aact( + _action(), ctx_in, backend, silence_context_type_warning=True + ) + + assert type(ctx_out) is _ChatSubclass + + +# --- session-level allow_context_type_change flag (#1522) --- + + +async def test_session_enforces_context_type_by_default(): + """A session trips the guard when a backend changes the context type.""" + from mellea import MelleaSession + + session = MelleaSession(_TypeChangingBackend(responses=None), ChatContext()) + assert session.allow_context_type_change is False + + with pytest.raises(ContextTypeMismatchError): + await session.aact(_action()) + + +async def test_session_flag_allows_context_type_change(): + """`allow_context_type_change=True` lets a session switch context types.""" + from mellea import MelleaSession + + session = MelleaSession( + _TypeChangingBackend(responses=None), + ChatContext(), + allow_context_type_change=True, + ) + assert session.allow_context_type_change is True + + await session.aact(_action()) + assert isinstance(session.ctx, SimpleContext) + + +def test_session_flag_preserved_across_clone(): + """`clone()` carries the `allow_context_type_change` flag to the copy.""" + from mellea import MelleaSession + + session = MelleaSession( + DummyBackend(responses=None), ChatContext(), allow_context_type_change=True + ) + assert session.clone().allow_context_type_change is True + + +# --- context-type enforcement over ALL sample contexts (#1522) --- +# +# `aact(..., return_sampling_results=True)` returns a `SamplingResult` whose +# type is not statically tracked, so the input==output convention is enforced at +# runtime against *every* `sample_contexts` entry — not just the chosen one — so +# a strategy that produces a mismatched context for any attempt is caught rather +# than silently returned. These tests drive that loop with a stub strategy that +# lets each test dictate the exact contexts attached to the result. + + +def _computed(value: str) -> ComputedModelOutputThunk: + """Build a computed thunk with a final-marked `GenerateLog`, as `aact` requires.""" + mot: ModelOutputThunk = ModelOutputThunk(value=value) + mot._generate_log = GenerateLog(is_final_result=True) + return ComputedModelOutputThunk(mot) + + +class _StubStrategy(SamplingStrategy): + """A sampling strategy that returns a caller-supplied list of sample contexts. + + Bypasses real generation entirely: `sample` ignores the backend and returns + a `SamplingResult` whose `sample_contexts` are exactly `self._contexts`, with + one computed generation per context. This lets a test attach a deliberately + mismatched context to any slot and assert whether the functional guard fires. + """ + + def __init__(self, contexts: list[Context]) -> None: + """Store the contexts to attach, one per generation, to the result.""" + self._contexts = contexts + + async def sample( + self, + action, + context, + backend, + requirements, + *, + validation_ctx=None, + format=None, + model_options=None, + tool_calls=False, + ) -> SamplingResult: + """Return a `SamplingResult` wrapping `self._contexts`; the last is chosen.""" + gens = [_computed(f"gen{i}") for i in range(len(self._contexts))] + return SamplingResult( + result_index=len(gens) - 1, + success=True, + sample_generations=gens, + sample_contexts=list(self._contexts), + ) + + +async def test_sampling_result_all_contexts_type_checked(): + """A mismatched context in a *non-chosen* sample slot trips the guard. + + The chosen (last) context matches the input type, but an earlier attempt is a + `SimpleContext`. The guard must still catch it — enforcement covers the whole + `sample_contexts` list, not just `result_ctx`. + """ + backend = DummyBackend(responses=None) + ctx_in = ChatContext() + strategy = _StubStrategy([SimpleContext(), ChatContext()]) + + with pytest.raises(ContextTypeMismatchError): + await aact( + _action(), + ctx_in, + backend, + strategy=strategy, + return_sampling_results=True, + silence_context_type_warning=True, + ) + + +async def test_sampling_result_matching_contexts_pass(): + """When every sample context matches the input type, the result is returned.""" + backend = DummyBackend(responses=None) + ctx_in = ChatContext() + strategy = _StubStrategy([ChatContext(), ChatContext()]) + + result = await aact( + _action(), + ctx_in, + backend, + strategy=strategy, + return_sampling_results=True, + silence_context_type_warning=True, + ) + + assert isinstance(result, SamplingResult) + assert all(type(c) is ChatContext for c in result.sample_contexts) + + +async def test_sampling_result_escape_hatch_allows_all_mismatches(): + """`allow_context_type_change=True` skips the per-sample check entirely.""" + backend = DummyBackend(responses=None) + ctx_in = ChatContext() + strategy = _StubStrategy([SimpleContext(), SimpleContext()]) + + result = await aact( + _action(), + ctx_in, + backend, + strategy=strategy, + return_sampling_results=True, + silence_context_type_warning=True, + allow_context_type_change=True, + ) + + assert isinstance(result, SamplingResult) diff --git a/test/stdlib/test_session_audio_unit.py b/test/stdlib/test_session_audio_unit.py index 3aa2c2b184..f68b876a7e 100644 --- a/test/stdlib/test_session_audio_unit.py +++ b/test/stdlib/test_session_audio_unit.py @@ -30,6 +30,10 @@ def _make_session() -> MagicMock: session = MagicMock(spec=MelleaSession) session.ctx = SimpleContext() session.backend = MagicMock() + # Real session methods read this instance attribute when forwarding to + # mfuncs; a spec'd mock doesn't provide it automatically (it's set in + # _init_fields, not on the class). + session.allow_context_type_change = False return session diff --git a/test/typing/check_context_propagation.py b/test/typing/check_context_propagation.py new file mode 100644 index 0000000000..71a1d304ba --- /dev/null +++ b/test/typing/check_context_propagation.py @@ -0,0 +1,92 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mypy checks for context-type propagation (issue #1522). + +Mellea functions must return the *same* `Context` subtype they were given, not +a widened base `Context`. These `assert_type` checks fail under mypy until the +functional layer threads the input context subtype through to its return type. +""" + +from typing import Any, assert_type, cast + +from mellea.core import Backend, ComputedModelOutputThunk, ModelOutputThunk +from mellea.stdlib.components import Instruction, Message +from mellea.stdlib.context import ChatContext, SimpleContext +from mellea.stdlib.functional import ( + aact, + achat, + act, + ainstruct, + aquery, + chat, + instruct, + query, + transform, +) + +backend = cast(Backend, None) +action: Instruction = cast(Instruction, None) +chat_ctx = cast(ChatContext, None) +simple_ctx = cast(SimpleContext, None) + + +# --- sync functions preserve the concrete context subtype --- + + +def check_act_chat_ctx() -> None: + r = act(action, chat_ctx, backend) + assert_type(r, tuple[ComputedModelOutputThunk[str], ChatContext]) + + +def check_act_simple_ctx() -> None: + r = act(action, simple_ctx, backend) + assert_type(r, tuple[ComputedModelOutputThunk[str], SimpleContext]) + + +def check_instruct_chat_ctx() -> None: + r = instruct("test", chat_ctx, backend) + assert_type(r, tuple[ComputedModelOutputThunk[str], ChatContext]) + + +def check_query_chat_ctx() -> None: + r = query(object(), "q", chat_ctx, backend) + assert_type(r, tuple[ComputedModelOutputThunk[Any], ChatContext]) + + +def check_transform_simple_ctx() -> None: + r = transform(object(), "t", simple_ctx, backend) + assert_type(r, tuple[ModelOutputThunk[Any] | Any, SimpleContext]) + + +def check_chat_chat_ctx() -> None: + r = chat("hi", chat_ctx, backend) + assert_type(r, tuple[Message, ChatContext]) + + +# --- async functions preserve the concrete context subtype --- + + +async def check_aact_chat_ctx() -> None: + r = await aact(action, chat_ctx, backend, strategy=None, await_result=True) + assert_type(r, tuple[ComputedModelOutputThunk[str], ChatContext]) + + +async def check_aact_simple_ctx() -> None: + r = await aact(action, simple_ctx, backend, strategy=None) + assert_type(r, tuple[ModelOutputThunk[str], SimpleContext]) + + +async def check_ainstruct_chat_ctx() -> None: + r = await ainstruct("test", chat_ctx, backend, strategy=None, await_result=True) + assert_type(r, tuple[ComputedModelOutputThunk[str], ChatContext]) + + +async def check_aquery_simple_ctx() -> None: + r = await aquery(object(), "q", simple_ctx, backend) + assert_type(r, tuple[ModelOutputThunk[Any], SimpleContext]) + + +async def check_achat_chat_ctx() -> None: + r = await achat("hi", chat_ctx, backend) + assert_type(r, tuple[Message, ChatContext]) diff --git a/test/typing/check_session_context_type.py b/test/typing/check_session_context_type.py new file mode 100644 index 0000000000..d2a17fc106 --- /dev/null +++ b/test/typing/check_session_context_type.py @@ -0,0 +1,77 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mypy checks for MelleaSession context-type parameterization (issue #1522). + +`MelleaSession` is generic in its context type, so `session.ctx` must narrow to +the concrete `Context` subtype the session was built with rather than widening +to the base `Context`. These `assert_type` checks verify that the constructor +overloads, the `start_session` overloads, and `clone()` all preserve the +parameter. Verification is via `uv run mypy .`; the functions never execute. +""" + +from typing import assert_type, cast + +from mellea import MelleaSession, start_session +from mellea.core import Backend +from mellea.stdlib.context import ChatContext, SimpleContext + +backend = cast(Backend, None) +chat_ctx = cast(ChatContext, None) +simple_ctx = cast(SimpleContext, None) + + +# --- constructor infers the context type parameter --- + + +def check_ctor_chat_ctx() -> None: + session = MelleaSession(backend, chat_ctx) + assert_type(session, MelleaSession[ChatContext]) + assert_type(session.ctx, ChatContext) + + +def check_ctor_simple_ctx() -> None: + session = MelleaSession(backend, simple_ctx) + assert_type(session, MelleaSession[SimpleContext]) + assert_type(session.ctx, SimpleContext) + + +def check_ctor_no_ctx_defaults_simple() -> None: + session = MelleaSession(backend) + assert_type(session, MelleaSession[SimpleContext]) + assert_type(session.ctx, SimpleContext) + + +# --- start_session infers the context type parameter --- + + +def check_start_session_context_type_chat() -> None: + session = start_session(context_type="chat") + assert_type(session, MelleaSession[ChatContext]) + assert_type(session.ctx, ChatContext) + + +def check_start_session_context_type_simple() -> None: + session = start_session(context_type="simple") + assert_type(session, MelleaSession[SimpleContext]) + assert_type(session.ctx, SimpleContext) + + +def check_start_session_default_is_simple() -> None: + session = start_session() + assert_type(session, MelleaSession[SimpleContext]) + assert_type(session.ctx, SimpleContext) + + +def check_start_session_explicit_ctx() -> None: + session = start_session(ctx=chat_ctx) + assert_type(session, MelleaSession[ChatContext]) + assert_type(session.ctx, ChatContext) + + +# --- clone preserves the context type parameter --- + + +def check_clone_preserves_type() -> None: + session = MelleaSession(backend, chat_ctx) + assert_type(session.clone(), MelleaSession[ChatContext])