From ebd6b7a051d58e507cb81b4ecc478a7ce289c642 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Mon, 15 Jun 2026 14:55:39 -0400 Subject: [PATCH 01/11] first pass at model aware sliding context window Signed-off-by: AngeloDanducci --- mellea/backends/context_lengths.py | 112 +++++++++++++++ mellea/backends/model_ids.py | 36 ++++- mellea/stdlib/context.py | 67 +++++++-- mellea/stdlib/session.py | 21 +++ test/stdlib/test_base_context.py | 222 +++++++++++++++++++++++++++++ 5 files changed, 445 insertions(+), 13 deletions(-) create mode 100644 mellea/backends/context_lengths.py diff --git a/mellea/backends/context_lengths.py b/mellea/backends/context_lengths.py new file mode 100644 index 0000000000..06e889a6c9 --- /dev/null +++ b/mellea/backends/context_lengths.py @@ -0,0 +1,112 @@ +"""Model context-length lookup table. + +Maps known model identifiers to their maximum context window in tokens. +The primary lookup uses ``ModelIdentifier.context_length`` when the caller +passes a ``ModelIdentifier`` object. String lookups fall back to a name +table keyed on ``hf_model_name`` and ``ollama_name``. +""" + +from __future__ import annotations + +from .model_ids import ModelIdentifier + + +def get_context_length(model_id: str | ModelIdentifier) -> int | None: + """Return the maximum context length in tokens for a known model, or ``None``. + + Priority: + + 1. If *model_id* is a ``ModelIdentifier`` with a non-``None`` + ``context_length`` field, return that value directly. + 2. Otherwise perform a name-based lookup in the built-in table, + checking ``hf_model_name`` then ``ollama_name`` (for + ``ModelIdentifier`` inputs) or the raw string (for string inputs). + + Args: + model_id: A ``ModelIdentifier`` instance or a raw model-name string. + + Returns: + The context length in tokens, or ``None`` if the model is unknown or + its context length is not recorded in the catalog. + """ + if isinstance(model_id, ModelIdentifier): + if model_id.context_length is not None: + return model_id.context_length + names_to_check = [model_id.hf_model_name, model_id.ollama_name] + else: + names_to_check = [model_id] + + for name in names_to_check: + if name and name in _CONTEXT_LENGTH_TABLE: + return _CONTEXT_LENGTH_TABLE[name] + return None + + +# Fallback name-keyed table for raw strings and ModelIdentifiers constructed +# without a context_length field. Keyed on both hf_model_name and ollama_name +# so either variant resolves correctly. +_CONTEXT_LENGTH_TABLE: dict[str, int] = { + # IBM Granite 4.x dense + "ibm-granite/granite-4.1-3b": 131072, + "granite4.1:3b": 131072, + "ibm-granite/granite-4.1-8b": 131072, + "granite4.1:8b": 131072, + "ibm-granite/granite-4.1-30b": 131072, + "granite4.1:30b": 131072, + # IBM Granite 4.x hybrid + "ibm-granite/granite-4.0-h-micro": 131072, + "granite4:micro-h": 131072, + "ibm-granite/granite-4.0-h-tiny": 131072, + "granite4:tiny-h": 131072, + "ibm-granite/granite-4.0-h-small": 131072, + "granite4:small-h": 131072, + "ibm-granite/granite-4.0-h-1b": 131072, + "granite4:1b-h": 131072, + "ibm-granite/granite-4.0-h-350m": 32768, + "granite4:350m-h": 32768, + # IBM Granite 4.0 micro (pre-built) + "ibm-granite/granite-4.0-micro": 131072, + "granite4:micro": 131072, + # IBM Granite 3.x + "ibm-granite/granite-3.3-8b-instruct": 131072, + "granite3.3:8b": 131072, + "ibm-granite/granite-3.2-8b-instruct": 131072, + "granite3.2:8b": 131072, + # Meta Llama 4 + "unsloth/Llama-4-Scout-17B-16E-Instruct": 10485760, + "llama4:scout": 10485760, + "unsloth/Llama-4-Maverick-17B-128E-Instruct": 1048576, + "llama4:maverick": 1048576, + # Meta Llama 3 + "unsloth/Llama-3.3-70B-Instruct": 131072, + "llama3.3:70b": 131072, + "unsloth/Llama-3.2-3B-Instruct": 131072, + "llama3.2:3b": 131072, + "unsloth/Llama-3.2-1B": 131072, + "llama3.2:1b": 131072, + # Mistral + "mistralai/Mistral-7B-Instruct-v0.3": 32768, + "mistral:7b": 32768, + "mistralai/Mistral-Small-3.1-24B-Instruct-2503": 131072, + "mistral-small:latest": 131072, + "mistralai/Mistral-Large-Instruct-2411": 131072, + "mistral-large:latest": 131072, + # Qwen 3 + "Qwen/Qwen3-0.6B": 32768, + "qwen3:0.6b": 32768, + "Qwen/Qwen3-1.7B": 32768, + "qwen3:1.7b": 32768, + "Qwen/Qwen3-8B": 40960, + "qwen3:8b": 40960, + "Qwen/Qwen3-14B": 40960, + "qwen3:14b": 40960, + # Microsoft Phi + "microsoft/phi-4": 16384, + "phi4:14b": 16384, + # DeepSeek + "deepseek-ai/DeepSeek-R1-Distill-Llama-8B": 131072, + "deepseek-r1:8b": 131072, + # SmolLM2 + "HuggingFaceTB/SmolLM2-1.7B-Instruct": 8192, + "smollm2:1.7b": 8192, +} diff --git a/mellea/backends/model_ids.py b/mellea/backends/model_ids.py index 230e0da0b4..526a1eb4d0 100644 --- a/mellea/backends/model_ids.py +++ b/mellea/backends/model_ids.py @@ -37,6 +37,7 @@ class ModelIdentifier: bedrock_name: str | None = None hf_tokenizer_name: str | None = None # if None, is the same as hf_model_name + context_length: int | None = None #################### @@ -48,30 +49,35 @@ class ModelIdentifier: hf_model_name="ibm-granite/granite-4.0-h-micro", ollama_name="granite4:micro-h", watsonx_name=None, # Only h-small available on Watsonx + context_length=131072, ) IBM_GRANITE_4_HYBRID_TINY = ModelIdentifier( hf_model_name="ibm-granite/granite-4.0-h-tiny", ollama_name="granite4:tiny-h", watsonx_name=None, # Only h-small available on Watsonx + context_length=131072, ) IBM_GRANITE_4_HYBRID_SMALL = ModelIdentifier( hf_model_name="ibm-granite/granite-4.0-h-small", ollama_name="granite4:small-h", watsonx_name="ibm/granite-4-h-small", + context_length=131072, ) IBM_GRANITE_4_HYBRID_1B = ModelIdentifier( hf_model_name="ibm-granite/granite-4.0-h-1b", ollama_name="granite4:1b-h", watsonx_name=None, + context_length=131072, ) IBM_GRANITE_4_HYBRID_350m = ModelIdentifier( hf_model_name="ibm-granite/granite-4.0-h-350m", ollama_name="granite4:350m-h", watsonx_name=None, + context_length=32768, ) # Granite 4.1 Dense Models @@ -79,14 +85,19 @@ class ModelIdentifier: hf_model_name="ibm-granite/granite-4.1-3b", ollama_name="granite4.1:3b", watsonx_name=None, + context_length=131072, ) IBM_GRANITE_4_1_8B = ModelIdentifier( - hf_model_name="ibm-granite/granite-4.1-8b", ollama_name="granite4.1:8b" + hf_model_name="ibm-granite/granite-4.1-8b", + ollama_name="granite4.1:8b", + context_length=131072, ) IBM_GRANITE_4_1_30B = ModelIdentifier( - hf_model_name="ibm-granite/granite-4.1-30b", ollama_name="granite4.1:30b" + hf_model_name="ibm-granite/granite-4.1-30b", + ollama_name="granite4.1:30b", + context_length=131072, ) IBM_GRANITE_GUARDIAN_4_1_8B = ModelIdentifier( @@ -100,18 +111,21 @@ class ModelIdentifier: hf_model_name="ibm-granite/granite-3.2-8b-instruct", ollama_name="granite3.2:8b", watsonx_name="ibm/granite-3-2b-instruct", + context_length=131072, ) IBM_GRANITE_3_3_8B = ModelIdentifier( hf_model_name="ibm-granite/granite-3.3-8b-instruct", ollama_name="granite3.3:8b", watsonx_name="ibm/granite-3-3-8b-instruct", + context_length=131072, ) IBM_GRANITE_4_MICRO_3B = ModelIdentifier( hf_model_name="ibm-granite/granite-4.0-micro", ollama_name="granite4:micro", watsonx_name="ibm/granite-4-h-small", # Keeping hybrid version here for backwards compatibility. + context_length=131072, ) # Granite 3.3 Vision Model (2B) @@ -160,6 +174,7 @@ class ModelIdentifier: ollama_name="llama4:scout", hf_tokenizer_name="unsloth/Llama-4-Scout-17B-16E-Instruct", mlx_name="mlx-community/Llama-4-Scout-17B-16E-Instruct-4bit", + context_length=10485760, ) META_LLAMA_4_MAVERICK_17B_128E_INSTRUCT = ModelIdentifier( @@ -168,6 +183,7 @@ class ModelIdentifier: watsonx_name=None, # NOTE: we do have a fp8 model in watsonx (meta-llama/llama-4-maverick-17b-128e-instruct-fp8) Not sure if we want to include it here. hf_tokenizer_name="unsloth/Llama-4-Maverick-17B-128E-Instruct", mlx_name="mlx-community/Llama-4-Maverick-17B-128E-Instruct-4bit", + context_length=1048576, ) #### LLAMA 3 models #### @@ -177,12 +193,14 @@ class ModelIdentifier: watsonx_name="meta-llama/llama-3-3-70b-instruct", hf_tokenizer_name="unsloth/Llama-3.3-70B-Instruct", mlx_name="mlx-community/Llama-3.3-70B-Instruct-4bit", + context_length=131072, ) META_LLAMA_3_2_3B = ModelIdentifier( hf_model_name="unsloth/Llama-3.2-3B-Instruct", ollama_name="llama3.2:3b", watsonx_name="meta-llama/llama-3-2-3b-instruct", + context_length=131072, ) META_LLAMA_GUARD3_1B = ModelIdentifier( @@ -190,7 +208,9 @@ class ModelIdentifier: ) META_LLAMA_3_2_1B = ModelIdentifier( - ollama_name="llama3.2:1b", hf_model_name="unsloth/Llama-3.2-1B" + ollama_name="llama3.2:1b", + hf_model_name="unsloth/Llama-3.2-1B", + context_length=131072, ) ######################## @@ -200,18 +220,21 @@ class ModelIdentifier: MISTRALAI_MISTRAL_0_3_7B = ModelIdentifier( hf_model_name="mistralai/Mistral-7B-Instruct-v0.3", # Mistral 7B v0.3 ollama_name="mistral:7b", # Ollama + context_length=32768, ) MISTRALAI_MISTRAL_SMALL_24B = ModelIdentifier( hf_model_name="mistralai/Mistral-Small-3.1-24B-Instruct-2503", ollama_name="mistral-small:latest", watsonx_name="mistralai/mistral-small-3-1-24b-instruct-2503", + context_length=131072, ) MISTRALAI_MISTRAL_LARGE_123B = ModelIdentifier( hf_model_name="mistralai/Mistral-Large-Instruct-2411", ollama_name="mistral-large:latest", watsonx_name="mistralai/mistral-large", + context_length=131072, ) @@ -222,21 +245,25 @@ class ModelIdentifier: QWEN3_0_6B = ModelIdentifier( hf_model_name="Qwen/Qwen3-0.6B", # Qwen 0.6B ollama_name="qwen3:0.6b", # Ollama + context_length=32768, ) QWEN3_1_7B = ModelIdentifier( hf_model_name="Qwen/Qwen3-1.7B", # Qwen 1.7B ollama_name="qwen3:1.7b", # Ollama + context_length=32768, ) QWEN3_8B = ModelIdentifier( hf_model_name="Qwen/Qwen3-8B", # Qwen 8B ollama_name="qwen3:8b", # Ollama + context_length=40960, ) QWEN3_14B = ModelIdentifier( hf_model_name="Qwen/Qwen3-14B", # Qwen 14B ollama_name="qwen3:14b", # Ollama + context_length=40960, ) ########################### @@ -274,6 +301,7 @@ class ModelIdentifier: MS_PHI_4_14B = ModelIdentifier( hf_model_name="microsoft/phi-4", # Microsoft Phi-4 14B ollama_name="phi4:14b", # Ollama + context_length=16384, ) MS_PHI_4_MINI_REASONING_4B = ModelIdentifier( @@ -285,6 +313,7 @@ class ModelIdentifier: DEEPSEEK_R1_8B = ModelIdentifier( hf_model_name="deepseek-ai/DeepSeek-R1-Distill-Llama-8B", ollama_name="deepseek-r1:8b", + context_length=131072, ) @@ -292,6 +321,7 @@ class ModelIdentifier: ollama_name="smollm2:1.7b", hf_model_name="HuggingFaceTB/SmolLM2-1.7B-Instruct", mlx_name="mlx-community/SmolLM2-1.7B-Instruct", + context_length=8192, ) HF_SMOLLM3_3B_no_ollama = ModelIdentifier( diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index b8a748fabc..078bfb5312 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -9,6 +9,8 @@ from __future__ import annotations +from ..backends.model_ids import ModelIdentifier + # Leave unused `ContextTurn` import for import ergonomics. from ..core import CBlock, Component, Context, ContextTurn @@ -19,13 +21,49 @@ class ChatContext(Context): Args: window_size (int | None): Maximum number of context turns to include when calling `view_for_generation`. `None` (the default) means the full - history is always returned. + history is returned, unless a `model_id` is bound and has a known + context length — in that case the model's context length (in tokens) + is used as an upper bound on the number of items returned. + model_id (str | ModelIdentifier | None): Optional model identifier used + for automatic context-window sizing. When set and `window_size` is + `None`, `view_for_generation` looks up the model's known context + length and uses it as the window size. Explicit `window_size` always + takes priority. + + Note: + `context_length` is measured in tokens; `window_size` counts context + items (CBlocks / Components). When deriving the window from a model's + context length, the token count is used as a maximum item count — a + conservative proxy that is correct in practice because the number of + context items is always far smaller than the token budget. """ - def __init__(self, *, window_size: int | None = None): - """Initialize ChatContext with an optional sliding-window size.""" + def __init__( + self, + *, + window_size: int | None = None, + model_id: str | ModelIdentifier | None = None, + ): + """Initialize ChatContext with an optional sliding-window size and model binding.""" super().__init__() self._window_size = window_size + self._model_id: str | ModelIdentifier | None = model_id + + def bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: + """Return a new root `ChatContext` with the given model bound, preserving `window_size`. + + Creates a new root node (not a linked continuation) with the same + `window_size` and the provided `model_id`. Intended to be called on a + freshly-created root context before any items are added; subsequent + `add()` calls propagate the binding automatically. + + Args: + model_id: The model identifier to associate with this context. + + Returns: + ChatContext: A new root `ChatContext` with `_model_id` set. + """ + return ChatContext(window_size=self._window_size, model_id=model_id) def add(self, c: Component | CBlock) -> ChatContext: """Add a new component or CBlock to the context and return the updated context. @@ -34,25 +72,34 @@ def add(self, c: Component | CBlock) -> ChatContext: c (Component | CBlock): The component or content block to append. Returns: - ChatContext: A new `ChatContext` with the added entry, preserving the - current `window_size` setting. + ChatContext: A new `ChatContext` with the added entry, preserving + both `window_size` and any `model_id` binding. """ new = ChatContext.from_previous(self, c) new._window_size = self._window_size + new._model_id = self._model_id return new def view_for_generation(self) -> list[Component | CBlock] | None: """Return the context entries to pass to the model, respecting the configured window. - Uses the `window_size` set during initialisation to limit how many past - turns are included. `None` is returned when the underlying history is - non-linear. + Window size resolution priority: + + 1. Explicit ``window_size`` passed at construction — always wins. + 2. Model-derived context length looked up via ``model_id`` — used when + ``window_size`` is ``None`` and a model binding exists. + 3. ``None`` — return the full history (unbounded). Returns: list[Component | CBlock] | None: Ordered list of context entries up to - `window_size` turns, or `None` if the history is non-linear. + the effective window size, or `None` if the history is non-linear. """ - return self.as_list(self._window_size) + effective_window = self._window_size + if effective_window is None and self._model_id is not None: + from ..backends.context_lengths import get_context_length + + effective_window = get_context_length(self._model_id) + return self.as_list(effective_window) class SimpleContext(Context): diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index 0109ec5552..6c07588aa3 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -298,6 +298,17 @@ def __init__( self._log_context_token = None self._exit_stack: contextlib.ExitStack | None = None + # Bind the backend's model_id to the context when the context is a + # ChatContext without an existing model binding. This enables automatic + # context-window sizing in view_for_generation() when window_size is + # not explicitly set. + if ( + isinstance(self.ctx, ChatContext) + and self.ctx._model_id is None + and hasattr(self.backend, "model_id") + ): + self.ctx = self.ctx.bind_model(self.backend.model_id) + def __enter__(self): """Enter context manager and set this session as the current global session.""" # Called directly, not via hook: OTel Token attach/detach is task-affine, @@ -389,7 +400,17 @@ def reset(self) -> None: _run_async_in_thread( invoke_hook(HookType.SESSION_RESET, payload, backend=self.backend) ) + old_ctx = self.ctx self.ctx = self.ctx.reset_to_new() + # Re-bind the model_id to the fresh context so automatic window sizing + # continues to work after a reset. + if ( + isinstance(self.ctx, ChatContext) + and self.ctx._model_id is None + and isinstance(old_ctx, ChatContext) + and old_ctx._model_id is not None + ): + self.ctx = self.ctx.bind_model(old_ctx._model_id) def cleanup(self, *, exception: BaseException | None = None) -> None: """Clean up session resources and deregister session-scoped plugins. diff --git a/test/stdlib/test_base_context.py b/test/stdlib/test_base_context.py index 2fccb11fd0..0659c8ee07 100644 --- a/test/stdlib/test_base_context.py +++ b/test/stdlib/test_base_context.py @@ -1,5 +1,15 @@ +from unittest.mock import MagicMock + import pytest +from mellea.backends.context_lengths import get_context_length +from mellea.backends.model_ids import ( + IBM_GRANITE_4_1_8B, + META_LLAMA_3_3_70B, + META_LLAMA_4_SCOUT_17B_16E_INSTRUCT, + MISTRALAI_MISTRAL_0_3_7B, + ModelIdentifier, +) from mellea.core import CBlock, Context from mellea.stdlib.context import ChatContext, SimpleContext @@ -68,5 +78,217 @@ def test_actions_for_available_tools(): assert actions[i] == for_generation[i] +# --------------------------------------------------------------------------- +# get_context_length unit tests +# --------------------------------------------------------------------------- + + +def test_get_context_length_from_model_identifier_field(): + assert get_context_length(IBM_GRANITE_4_1_8B) == 131072 + + +def test_get_context_length_returns_none_for_unknown_model_identifier(): + unknown = ModelIdentifier(hf_model_name="org/totally-unknown-model") + assert get_context_length(unknown) is None + + +def test_get_context_length_from_raw_hf_string(): + assert get_context_length("mistralai/Mistral-7B-Instruct-v0.3") == 32768 + + +def test_get_context_length_from_raw_ollama_string(): + assert get_context_length("mistral:7b") == 32768 + + +def test_get_context_length_unknown_string(): + assert get_context_length("not-a-real-model") is None + + +def test_get_context_length_llama4_scout(): + assert get_context_length(META_LLAMA_4_SCOUT_17B_16E_INSTRUCT) == 10485760 + + +# --------------------------------------------------------------------------- +# ChatContext constructor with model_id +# --------------------------------------------------------------------------- + + +def test_chat_context_model_id_constructor(): + ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B) + assert ctx._model_id is IBM_GRANITE_4_1_8B + + +def test_chat_context_existing_window_size_unchanged(): + ctx = ChatContext(window_size=3) + for i in range(5): + ctx = ctx.add(CBlock(f"item {i}")) + assert len(ctx.view_for_generation()) == 3 + + +def test_chat_context_default_no_model_id(): + ctx = ChatContext() + assert ctx._model_id is None + assert ctx._window_size is None + + +# --------------------------------------------------------------------------- +# bind_model +# --------------------------------------------------------------------------- + + +def test_bind_model_returns_new_root_context(): + ctx = ChatContext() + bound = ctx.bind_model(IBM_GRANITE_4_1_8B) + assert bound._model_id is IBM_GRANITE_4_1_8B + assert bound.is_root_node + + +def test_bind_model_does_not_mutate_original(): + ctx = ChatContext() + ctx.bind_model(IBM_GRANITE_4_1_8B) + assert ctx._model_id is None + + +def test_bind_model_preserves_window_size(): + ctx = ChatContext(window_size=5) + bound = ctx.bind_model(IBM_GRANITE_4_1_8B) + assert bound._window_size == 5 + + +# --------------------------------------------------------------------------- +# model_id propagates through add() +# --------------------------------------------------------------------------- + + +def test_model_id_propagates_through_add(): + ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B) + ctx = ctx.add(CBlock("hello")) + ctx = ctx.add(CBlock("world")) + assert ctx._model_id is IBM_GRANITE_4_1_8B + + +def test_model_id_propagates_through_bind_then_add(): + ctx = ChatContext().bind_model(META_LLAMA_3_3_70B) + for i in range(3): + ctx = ctx.add(CBlock(f"msg {i}")) + assert ctx._model_id is META_LLAMA_3_3_70B + + +# --------------------------------------------------------------------------- +# view_for_generation with model_id +# --------------------------------------------------------------------------- + + +def test_view_for_generation_model_context_length_as_upper_bound(): + # Mistral 7B has context_length=32768; we have far fewer items so all are returned. + ctx = ChatContext(model_id=MISTRALAI_MISTRAL_0_3_7B) + for i in range(10): + ctx = ctx.add(CBlock(f"item {i}")) + result = ctx.view_for_generation() + assert len(result) == 10 + + +def test_view_for_generation_explicit_window_size_beats_model(): + ctx = ChatContext(window_size=2, model_id=IBM_GRANITE_4_1_8B) + for i in range(10): + ctx = ctx.add(CBlock(f"item {i}")) + result = ctx.view_for_generation() + assert len(result) == 2 + + +def test_view_for_generation_no_model_id_returns_full_history(): + ctx = ChatContext() + for i in range(8): + ctx = ctx.add(CBlock(f"item {i}")) + result = ctx.view_for_generation() + assert len(result) == 8 + + +def test_view_for_generation_unknown_model_returns_full_history(): + unknown = ModelIdentifier(hf_model_name="org/unknown-model") + ctx = ChatContext(model_id=unknown) + for i in range(5): + ctx = ctx.add(CBlock(f"item {i}")) + result = ctx.view_for_generation() + assert len(result) == 5 + + +# --------------------------------------------------------------------------- +# reset_to_new and session.reset() +# --------------------------------------------------------------------------- + + +def test_reset_to_new_returns_empty_root(): + ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B, window_size=3) + for i in range(3): + ctx = ctx.add(CBlock(f"msg {i}")) + reset_ctx = ctx.reset_to_new() + assert isinstance(reset_ctx, ChatContext) + assert reset_ctx.is_root_node + assert len(reset_ctx.as_list()) == 0 + + +def test_session_reset_rebinds_model_id(): + from mellea.stdlib.session import MelleaSession + + mock_backend = MagicMock() + mock_backend.model_id = IBM_GRANITE_4_1_8B + ctx = ChatContext() + session = MelleaSession(mock_backend, ctx) + assert session.ctx._model_id is IBM_GRANITE_4_1_8B + + session.ctx = session.ctx.add(CBlock("some history")) + session.reset() + + assert isinstance(session.ctx, ChatContext) + assert session.ctx._model_id is IBM_GRANITE_4_1_8B + assert len(session.ctx.as_list()) == 0 + + +# --------------------------------------------------------------------------- +# MelleaSession auto-binding +# --------------------------------------------------------------------------- + + +def test_session_binds_model_id_to_chat_context(): + from mellea.stdlib.session import MelleaSession + + mock_backend = MagicMock() + mock_backend.model_id = IBM_GRANITE_4_1_8B + ctx = ChatContext() + session = MelleaSession(mock_backend, ctx) + assert isinstance(session.ctx, ChatContext) + assert session.ctx._model_id is IBM_GRANITE_4_1_8B + + +def test_session_does_not_override_explicit_model_id(): + from mellea.stdlib.session import MelleaSession + + mock_backend = MagicMock() + mock_backend.model_id = IBM_GRANITE_4_1_8B + ctx = ChatContext(model_id=META_LLAMA_3_3_70B) + session = MelleaSession(mock_backend, ctx) + assert session.ctx._model_id is META_LLAMA_3_3_70B + + +def test_session_does_not_bind_for_simple_context(): + from mellea.stdlib.session import MelleaSession + + mock_backend = MagicMock() + mock_backend.model_id = IBM_GRANITE_4_1_8B + ctx = SimpleContext() + session = MelleaSession(mock_backend, ctx) + assert isinstance(session.ctx, SimpleContext) + + +def test_session_graceful_when_backend_has_no_model_id(): + from mellea.stdlib.session import MelleaSession + + mock_backend = MagicMock(spec=[]) # no attributes at all + ctx = ChatContext() + session = MelleaSession(mock_backend, ctx) + assert session.ctx._model_id is None + + if __name__ == "__main__": pytest.main([__file__]) From 412ac5a4d7a1c4e5a002b517edc8f92d48f6ba6d Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Mon, 15 Jun 2026 17:52:01 -0400 Subject: [PATCH 02/11] additional pass at context length window Signed-off-by: AngeloDanducci --- mellea/backends/context_lengths.py | 110 +++++++++++----------------- mellea/backends/model_ids.py | 24 ++++-- mellea/stdlib/context.py | 61 ++++++++++++---- mellea/stdlib/session.py | 28 +++---- test/stdlib/test_base_context.py | 113 +++++++++++++++++++++++------ 5 files changed, 212 insertions(+), 124 deletions(-) diff --git a/mellea/backends/context_lengths.py b/mellea/backends/context_lengths.py index 06e889a6c9..7e3fc0e6dc 100644 --- a/mellea/backends/context_lengths.py +++ b/mellea/backends/context_lengths.py @@ -3,7 +3,18 @@ Maps known model identifiers to their maximum context window in tokens. The primary lookup uses ``ModelIdentifier.context_length`` when the caller passes a ``ModelIdentifier`` object. String lookups fall back to a name -table keyed on ``hf_model_name`` and ``ollama_name``. +table keyed on ``hf_model_name`` and ``ollama_name`` of every +``ModelIdentifier`` constant defined in ``model_ids``. The table is built +automatically at import time so there is a single source of truth: the +``context_length`` field on each constant. + +Only ``hf_model_name`` and ``ollama_name`` are indexed — matching the two +fields checked in the ``ModelIdentifier`` branch of ``get_context_length``. +Platform-specific IDs (``mlx_name``, ``openai_name``, ``bedrock_name``) and +tokenizer references (``hf_tokenizer_name``) are intentionally excluded: the +former would populate the table with deployment-platform strings that are +not valid serving names, and the latter is a tokenizer reference, not a +model-serving identifier. """ from __future__ import annotations @@ -21,6 +32,9 @@ def get_context_length(model_id: str | ModelIdentifier) -> int | None: 2. Otherwise perform a name-based lookup in the built-in table, checking ``hf_model_name`` then ``ollama_name`` (for ``ModelIdentifier`` inputs) or the raw string (for string inputs). + Only ``hf_model_name`` and ``ollama_name`` are indexed in the table, + so platform-specific names (``mlx_name``, ``openai_name``, + ``bedrock_name``) do not resolve via raw string lookup. Args: model_id: A ``ModelIdentifier`` instance or a raw model-name string. @@ -34,7 +48,11 @@ def get_context_length(model_id: str | ModelIdentifier) -> int | None: return model_id.context_length names_to_check = [model_id.hf_model_name, model_id.ollama_name] else: - names_to_check = [model_id] + # LiteLLM and similar backends prefix model names with a provider + # slug (e.g. "ollama_chat/granite4.1:3b"). Strip the prefix so the + # bare name can resolve against the table. + stripped = model_id.split("/", 1)[-1] if "/" in model_id else model_id + names_to_check = [model_id, stripped] if stripped != model_id else [model_id] for name in names_to_check: if name and name in _CONTEXT_LENGTH_TABLE: @@ -42,71 +60,25 @@ def get_context_length(model_id: str | ModelIdentifier) -> int | None: return None +def _build_table() -> dict[str, int]: + import mellea.backends.model_ids as _m + + table: dict[str, int] = {} + for obj in vars(_m).values(): + if not isinstance(obj, ModelIdentifier) or obj.context_length is None: + continue + # Only index hf_model_name and ollama_name — these are the two fields + # checked by the ModelIdentifier branch of get_context_length, keeping + # the table and the lookup logic symmetric. Platform-specific names + # (mlx_name, openai_name, bedrock_name) and hf_tokenizer_name are + # intentionally excluded; see module docstring for rationale. + for name in (obj.hf_model_name, obj.ollama_name): + if name: + table[name] = obj.context_length + return table + + # Fallback name-keyed table for raw strings and ModelIdentifiers constructed -# without a context_length field. Keyed on both hf_model_name and ollama_name -# so either variant resolves correctly. -_CONTEXT_LENGTH_TABLE: dict[str, int] = { - # IBM Granite 4.x dense - "ibm-granite/granite-4.1-3b": 131072, - "granite4.1:3b": 131072, - "ibm-granite/granite-4.1-8b": 131072, - "granite4.1:8b": 131072, - "ibm-granite/granite-4.1-30b": 131072, - "granite4.1:30b": 131072, - # IBM Granite 4.x hybrid - "ibm-granite/granite-4.0-h-micro": 131072, - "granite4:micro-h": 131072, - "ibm-granite/granite-4.0-h-tiny": 131072, - "granite4:tiny-h": 131072, - "ibm-granite/granite-4.0-h-small": 131072, - "granite4:small-h": 131072, - "ibm-granite/granite-4.0-h-1b": 131072, - "granite4:1b-h": 131072, - "ibm-granite/granite-4.0-h-350m": 32768, - "granite4:350m-h": 32768, - # IBM Granite 4.0 micro (pre-built) - "ibm-granite/granite-4.0-micro": 131072, - "granite4:micro": 131072, - # IBM Granite 3.x - "ibm-granite/granite-3.3-8b-instruct": 131072, - "granite3.3:8b": 131072, - "ibm-granite/granite-3.2-8b-instruct": 131072, - "granite3.2:8b": 131072, - # Meta Llama 4 - "unsloth/Llama-4-Scout-17B-16E-Instruct": 10485760, - "llama4:scout": 10485760, - "unsloth/Llama-4-Maverick-17B-128E-Instruct": 1048576, - "llama4:maverick": 1048576, - # Meta Llama 3 - "unsloth/Llama-3.3-70B-Instruct": 131072, - "llama3.3:70b": 131072, - "unsloth/Llama-3.2-3B-Instruct": 131072, - "llama3.2:3b": 131072, - "unsloth/Llama-3.2-1B": 131072, - "llama3.2:1b": 131072, - # Mistral - "mistralai/Mistral-7B-Instruct-v0.3": 32768, - "mistral:7b": 32768, - "mistralai/Mistral-Small-3.1-24B-Instruct-2503": 131072, - "mistral-small:latest": 131072, - "mistralai/Mistral-Large-Instruct-2411": 131072, - "mistral-large:latest": 131072, - # Qwen 3 - "Qwen/Qwen3-0.6B": 32768, - "qwen3:0.6b": 32768, - "Qwen/Qwen3-1.7B": 32768, - "qwen3:1.7b": 32768, - "Qwen/Qwen3-8B": 40960, - "qwen3:8b": 40960, - "Qwen/Qwen3-14B": 40960, - "qwen3:14b": 40960, - # Microsoft Phi - "microsoft/phi-4": 16384, - "phi4:14b": 16384, - # DeepSeek - "deepseek-ai/DeepSeek-R1-Distill-Llama-8B": 131072, - "deepseek-r1:8b": 131072, - # SmolLM2 - "HuggingFaceTB/SmolLM2-1.7B-Instruct": 8192, - "smollm2:1.7b": 8192, -} +# without a context_length field. Built automatically from the ModelIdentifier +# constants in model_ids so there is a single source of truth. +_CONTEXT_LENGTH_TABLE: dict[str, int] = _build_table() diff --git a/mellea/backends/model_ids.py b/mellea/backends/model_ids.py index 526a1eb4d0..6642f9e753 100644 --- a/mellea/backends/model_ids.py +++ b/mellea/backends/model_ids.py @@ -101,7 +101,7 @@ class ModelIdentifier: ) IBM_GRANITE_GUARDIAN_4_1_8B = ModelIdentifier( - hf_model_name="ibm-granite/granite-guardian-4.1-8b" + hf_model_name="ibm-granite/granite-guardian-4.1-8b", context_length=131072 ) @@ -133,34 +133,36 @@ class ModelIdentifier: hf_model_name="ibm-granite/granite-vision-3.3-2b", ollama_name="ibm/granite3.3-vision:2b", watsonx_name=None, + context_length=131072, ) IBM_GRANITE_GUARDIAN_3_0_2B = ModelIdentifier( hf_model_name="ibm-granite/granite-guardian-3.0-2b", ollama_name="granite3-guardian:2b", + context_length=4096, ) IBM_GRANITE_4_TINY_PREVIEW_7B = ModelIdentifier( - hf_model_name="ibm-granite/granite-4.0-tiny-preview" + hf_model_name="ibm-granite/granite-4.0-tiny-preview", context_length=131072 ) IBM_GRANITE_4_TINY_PREVIEW_BASE_7B = ModelIdentifier( - hf_model_name="ibm-granite/granite-4.0-tiny-base-preview" + hf_model_name="ibm-granite/granite-4.0-tiny-base-preview", context_length=131072 ) # Pre-Built Granite Switch Models IBM_GRANITE_SWITCH_4_1_3B_PREVIEW = ModelIdentifier( - hf_model_name="ibm-granite/granite-switch-4.1-3b-preview" + hf_model_name="ibm-granite/granite-switch-4.1-3b-preview", context_length=131072 ) """Granite Switch Preview Model. Adapters: `citations`, `query_rewrite`, `query_clarification`, `hallucination_detection`, `answerability`, `policy-guardrails`, `guardian-core`, `uncertainty`, `requirement-check`, `context-attribution`, `factuality-detection`, `factuality-correction`.""" # Document what adapters are included by default here. IBM_GRANITE_SWITCH_4_1_8B_PREVIEW = ModelIdentifier( - hf_model_name="ibm-granite/granite-switch-4.1-8b-preview" + hf_model_name="ibm-granite/granite-switch-4.1-8b-preview", context_length=131072 ) """Granite Switch Preview Model. Adapters: `citations`, `query_rewrite`, `query_clarification`, `hallucination_detection`, `answerability`, `policy-guardrails`, `guardian-core`, `uncertainty`, `requirement-check`, `context-attribution`, `factuality-detection`, `factuality-correction`.""" # Document what adapters are included by default here. IBM_GRANITE_SWITCH_4_1_30B_PREVIEW = ModelIdentifier( - hf_model_name="ibm-granite/granite-switch-4.1-30b-preview" + hf_model_name="ibm-granite/granite-switch-4.1-30b-preview", context_length=131072 ) """Granite Switch Preview Model. Adapters: `citations`, `query_rewrite`, `query_clarification`, `hallucination_detection`, `answerability`, `policy-guardrails`, `guardian-core`, `uncertainty`, `requirement-check`, `context-attribution`, `factuality-detection`, `factuality-correction`.""" # Document what adapters are included by default here. @@ -204,7 +206,9 @@ class ModelIdentifier: ) META_LLAMA_GUARD3_1B = ModelIdentifier( - ollama_name="llama-guard3:1b", hf_model_name="meta-llama/Llama-Guard-3-1B" + ollama_name="llama-guard3:1b", + hf_model_name="meta-llama/Llama-Guard-3-1B", + context_length=131072, ) META_LLAMA_3_2_1B = ModelIdentifier( @@ -274,11 +278,13 @@ class ModelIdentifier: hf_model_name="openai/gpt-oss-20b", # OpenAI GPT-OSS 20B ollama_name="gpt-oss:20b", # Ollama bedrock_name="openai.gpt-oss-20b", + context_length=131072, ) OPENAI_GPT_OSS_120B = ModelIdentifier( hf_model_name="openai/gpt-oss-120b", # OpenAI GPT-OSS 120B ollama_name="gpt-oss:120b", # Ollama bedrock_name="openai.gpt-oss-120b", + context_length=131072, ) ########################### @@ -296,6 +302,7 @@ class ModelIdentifier: GOOGLE_GEMMA_3N_E4B = ModelIdentifier( hf_model_name="google/gemma-3n-e4b-it", # Google Gemma 3N E4B ollama_name="gemma3n:e4b", # Ollama + context_length=32768, ) MS_PHI_4_14B = ModelIdentifier( @@ -307,6 +314,7 @@ class ModelIdentifier: MS_PHI_4_MINI_REASONING_4B = ModelIdentifier( hf_model_name="microsoft/Phi-4-mini-flash-reasoning", ollama_name="phi4-mini-reasoning:3.8b", + context_length=131072, ) @@ -325,5 +333,5 @@ class ModelIdentifier: ) HF_SMOLLM3_3B_no_ollama = ModelIdentifier( - hf_model_name="HuggingFaceTB/SmolLM3-3B", ollama_name="" + hf_model_name="HuggingFaceTB/SmolLM3-3B", ollama_name=None, context_length=65536 ) diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index 078bfb5312..4723f7b35f 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -9,6 +9,7 @@ from __future__ import annotations +from ..backends.context_lengths import get_context_length from ..backends.model_ids import ModelIdentifier # Leave unused `ContextTurn` import for import ergonomics. @@ -32,10 +33,13 @@ class ChatContext(Context): Note: `context_length` is measured in tokens; `window_size` counts context - items (CBlocks / Components). When deriving the window from a model's - context length, the token count is used as a maximum item count — a - conservative proxy that is correct in practice because the number of - context items is always far smaller than the token budget. + items (CBlocks / Components). When only a `model_id` is bound (no + explicit `window_size`), the model's context length in tokens is used + as an absolute upper bound on the number of items passed to the model. + In practice this ceiling is never reached — real conversations do not + accumulate hundreds of thousands of items — so the full history is + returned for typical sessions. Set `window_size` explicitly to enforce + a tighter item-count limit. """ def __init__( @@ -49,21 +53,54 @@ def __init__( self._window_size = window_size self._model_id: str | ModelIdentifier | None = model_id - def bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: + @property + def model_id(self) -> str | ModelIdentifier | None: + """The model identifier bound to this context, or ``None`` if unbound.""" + return self._model_id + + def _make_root(self, model_id: str | ModelIdentifier | None) -> ChatContext: + """Return a new empty root ``ChatContext`` with the given ``model_id``, preserving ``window_size``.""" + return ChatContext(window_size=self._window_size, model_id=model_id) + + def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: """Return a new root `ChatContext` with the given model bound, preserving `window_size`. - Creates a new root node (not a linked continuation) with the same - `window_size` and the provided `model_id`. Intended to be called on a - freshly-created root context before any items are added; subsequent - `add()` calls propagate the binding automatically. + Internal use only — called by ``MelleaSession`` to wire the backend's + model identifier into the context at session construction and after + ``reset()``. To bind a model at construction time, pass ``model_id=`` + directly to ``ChatContext()``. Args: model_id: The model identifier to associate with this context. Returns: - ChatContext: A new root `ChatContext` with `_model_id` set. + ChatContext: A new root `ChatContext` with ``_model_id`` set. + + Raises: + ValueError: If called on a non-root (non-empty) context. History + would be silently discarded, so this is disallowed. Create a + new ``ChatContext(model_id=...)`` instead. """ - return ChatContext(window_size=self._window_size, model_id=model_id) + if not self.is_root_node: + raise ValueError( + "_bind_model() must be called on a root (empty) ChatContext. " + "To bind a model on a context that already has history, create a " + "new ChatContext(model_id=...) before adding items." + ) + return self._make_root(model_id) + + def reset_to_new(self) -> ChatContext: # type: ignore[override] + """Return a new empty root ``ChatContext``, preserving ``window_size`` and ``model_id``. + + Overrides ``Context.reset_to_new()`` so that the model binding and + window size set at construction are not lost when a session is reset + or when user code calls this method directly. + + Returns: + ChatContext: A fresh root context with the same ``window_size`` + and ``model_id`` as this context, but no history. + """ + return self._make_root(self._model_id) def add(self, c: Component | CBlock) -> ChatContext: """Add a new component or CBlock to the context and return the updated context. @@ -96,8 +133,6 @@ def view_for_generation(self) -> list[Component | CBlock] | None: """ effective_window = self._window_size if effective_window is None and self._model_id is not None: - from ..backends.context_lengths import get_context_length - effective_window = get_context_length(self._model_id) return self.as_list(effective_window) diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index 6c07588aa3..d17002224f 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -304,10 +304,11 @@ def __init__( # not explicitly set. if ( isinstance(self.ctx, ChatContext) - and self.ctx._model_id is None + and self.ctx.model_id is None + and self.ctx.is_root_node and hasattr(self.backend, "model_id") ): - self.ctx = self.ctx.bind_model(self.backend.model_id) + self.ctx = self.ctx._bind_model(self.backend.model_id) def __enter__(self): """Enter context manager and set this session as the current global session.""" @@ -355,8 +356,19 @@ def __exit__(self, exc_type, exc_val, exc_tb): def __copy__(self): """Use self.clone. Copies the current session but keeps references to the backend and context.""" - new = MelleaSession(backend=self.backend, ctx=self.ctx) + import uuid + + # Bypass __init__'s auto-bind logic: the session state is already + # settled, so we replicate it directly rather than re-deriving it. + new = object.__new__(MelleaSession) + new.id = str(uuid.uuid4()) + new.backend = self.backend + new.ctx = self.ctx new._session_logger = self._session_logger + new._context_token = None + new._log_context_token = None + new._session_span = None + new._exit_stack = None # Explicitly don't copy over the _context_token. return new @@ -400,17 +412,7 @@ def reset(self) -> None: _run_async_in_thread( invoke_hook(HookType.SESSION_RESET, payload, backend=self.backend) ) - old_ctx = self.ctx self.ctx = self.ctx.reset_to_new() - # Re-bind the model_id to the fresh context so automatic window sizing - # continues to work after a reset. - if ( - isinstance(self.ctx, ChatContext) - and self.ctx._model_id is None - and isinstance(old_ctx, ChatContext) - and old_ctx._model_id is not None - ): - self.ctx = self.ctx.bind_model(old_ctx._model_id) def cleanup(self, *, exception: BaseException | None = None) -> None: """Clean up session resources and deregister session-scoped plugins. diff --git a/test/stdlib/test_base_context.py b/test/stdlib/test_base_context.py index 0659c8ee07..13de3e34a9 100644 --- a/test/stdlib/test_base_context.py +++ b/test/stdlib/test_base_context.py @@ -115,7 +115,7 @@ def test_get_context_length_llama4_scout(): def test_chat_context_model_id_constructor(): ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B) - assert ctx._model_id is IBM_GRANITE_4_1_8B + assert ctx.model_id is IBM_GRANITE_4_1_8B def test_chat_context_existing_window_size_unchanged(): @@ -127,34 +127,41 @@ def test_chat_context_existing_window_size_unchanged(): def test_chat_context_default_no_model_id(): ctx = ChatContext() - assert ctx._model_id is None + assert ctx.model_id is None assert ctx._window_size is None # --------------------------------------------------------------------------- -# bind_model +# _bind_model # --------------------------------------------------------------------------- -def test_bind_model_returns_new_root_context(): +def test__bind_model_returns_new_root_context(): ctx = ChatContext() - bound = ctx.bind_model(IBM_GRANITE_4_1_8B) - assert bound._model_id is IBM_GRANITE_4_1_8B + bound = ctx._bind_model(IBM_GRANITE_4_1_8B) + assert bound.model_id is IBM_GRANITE_4_1_8B assert bound.is_root_node -def test_bind_model_does_not_mutate_original(): +def test__bind_model_does_not_mutate_original(): ctx = ChatContext() - ctx.bind_model(IBM_GRANITE_4_1_8B) - assert ctx._model_id is None + ctx._bind_model(IBM_GRANITE_4_1_8B) + assert ctx.model_id is None -def test_bind_model_preserves_window_size(): +def test__bind_model_preserves_window_size(): ctx = ChatContext(window_size=5) - bound = ctx.bind_model(IBM_GRANITE_4_1_8B) + bound = ctx._bind_model(IBM_GRANITE_4_1_8B) assert bound._window_size == 5 +def test__bind_model_raises_on_non_root_context(): + ctx = ChatContext() + ctx = ctx.add(CBlock("some history")) + with pytest.raises(ValueError, match="_bind_model\\(\\) must be called on a root"): + ctx._bind_model(IBM_GRANITE_4_1_8B) + + # --------------------------------------------------------------------------- # model_id propagates through add() # --------------------------------------------------------------------------- @@ -164,14 +171,14 @@ def test_model_id_propagates_through_add(): ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B) ctx = ctx.add(CBlock("hello")) ctx = ctx.add(CBlock("world")) - assert ctx._model_id is IBM_GRANITE_4_1_8B + assert ctx.model_id is IBM_GRANITE_4_1_8B def test_model_id_propagates_through_bind_then_add(): - ctx = ChatContext().bind_model(META_LLAMA_3_3_70B) + ctx = ChatContext()._bind_model(META_LLAMA_3_3_70B) for i in range(3): ctx = ctx.add(CBlock(f"msg {i}")) - assert ctx._model_id is META_LLAMA_3_3_70B + assert ctx.model_id is META_LLAMA_3_3_70B # --------------------------------------------------------------------------- @@ -179,8 +186,8 @@ def test_model_id_propagates_through_bind_then_add(): # --------------------------------------------------------------------------- -def test_view_for_generation_model_context_length_as_upper_bound(): - # Mistral 7B has context_length=32768; we have far fewer items so all are returned. +def test_view_for_generation_model_bound_returns_full_history_when_under_limit(): + # 10 items is far below the 32768-token ceiling — full history returned. ctx = ChatContext(model_id=MISTRALAI_MISTRAL_0_3_7B) for i in range(10): ctx = ctx.add(CBlock(f"item {i}")) @@ -188,6 +195,17 @@ def test_view_for_generation_model_context_length_as_upper_bound(): assert len(result) == 10 +def test_view_for_generation_model_bound_used_as_upper_bound(): + # A tiny context_length=5 triggers truncation; most recent items are retained. + tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=5) + ctx = ChatContext(model_id=tiny) + for i in range(10): + ctx = ctx.add(CBlock(f"item {i}")) + result = ctx.view_for_generation() + assert len(result) == 5 + assert str(result[-1]) == "item 9" + + def test_view_for_generation_explicit_window_size_beats_model(): ctx = ChatContext(window_size=2, model_id=IBM_GRANITE_4_1_8B) for i in range(10): @@ -228,6 +246,15 @@ def test_reset_to_new_returns_empty_root(): assert len(reset_ctx.as_list()) == 0 +def test_reset_to_new_preserves_model_id_and_window_size(): + ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B, window_size=5) + for i in range(3): + ctx = ctx.add(CBlock(f"msg {i}")) + reset_ctx = ctx.reset_to_new() + assert reset_ctx.model_id is IBM_GRANITE_4_1_8B + assert reset_ctx._window_size == 5 + + def test_session_reset_rebinds_model_id(): from mellea.stdlib.session import MelleaSession @@ -235,13 +262,13 @@ def test_session_reset_rebinds_model_id(): mock_backend.model_id = IBM_GRANITE_4_1_8B ctx = ChatContext() session = MelleaSession(mock_backend, ctx) - assert session.ctx._model_id is IBM_GRANITE_4_1_8B + assert session.ctx.model_id is IBM_GRANITE_4_1_8B session.ctx = session.ctx.add(CBlock("some history")) session.reset() assert isinstance(session.ctx, ChatContext) - assert session.ctx._model_id is IBM_GRANITE_4_1_8B + assert session.ctx.model_id is IBM_GRANITE_4_1_8B assert len(session.ctx.as_list()) == 0 @@ -258,7 +285,7 @@ def test_session_binds_model_id_to_chat_context(): ctx = ChatContext() session = MelleaSession(mock_backend, ctx) assert isinstance(session.ctx, ChatContext) - assert session.ctx._model_id is IBM_GRANITE_4_1_8B + assert session.ctx.model_id is IBM_GRANITE_4_1_8B def test_session_does_not_override_explicit_model_id(): @@ -268,7 +295,7 @@ def test_session_does_not_override_explicit_model_id(): mock_backend.model_id = IBM_GRANITE_4_1_8B ctx = ChatContext(model_id=META_LLAMA_3_3_70B) session = MelleaSession(mock_backend, ctx) - assert session.ctx._model_id is META_LLAMA_3_3_70B + assert session.ctx.model_id is META_LLAMA_3_3_70B def test_session_does_not_bind_for_simple_context(): @@ -287,7 +314,51 @@ def test_session_graceful_when_backend_has_no_model_id(): mock_backend = MagicMock(spec=[]) # no attributes at all ctx = ChatContext() session = MelleaSession(mock_backend, ctx) - assert session.ctx._model_id is None + assert session.ctx.model_id is None + + +def test_session_does_not_bind_when_context_has_history(): + from mellea.stdlib.session import MelleaSession + + mock_backend = MagicMock() + mock_backend.model_id = IBM_GRANITE_4_1_8B + ctx = ChatContext() + ctx = ctx.add(CBlock("pre-existing history")) + session = MelleaSession(mock_backend, ctx) + # Non-root context must not be auto-bound (would silently discard history). + assert session.ctx.model_id is None + assert len(session.ctx.as_list()) == 1 + + +def test_get_context_length_litellm_prefixed_string(): + # LiteLLM prefixes model names with a provider slug, e.g. "ollama_chat/granite4.1:3b". + # The stripped bare name should resolve to the correct context length. + assert get_context_length("ollama_chat/granite4.1:3b") == 131072 + assert get_context_length("ollama/granite4.1:8b") == 131072 + # Multi-segment strip still resolves when the remainder is a known HF name. + assert get_context_length("huggingface/ibm-granite/granite-4.1-3b") == 131072 + # Completely unknown prefix+name returns None. + assert get_context_length("someprefix/not-a-real-model") is None + + +def test_clone_does_not_double_bind_replaced_root_context(): + from copy import copy + + from mellea.stdlib.session import MelleaSession + + mock_backend = MagicMock() + mock_backend.model_id = IBM_GRANITE_4_1_8B + session = MelleaSession(mock_backend, ChatContext()) + assert session.ctx.model_id is IBM_GRANITE_4_1_8B + + # User replaces ctx with a fresh unbound root — no model_id. + session.ctx = ChatContext() + assert session.ctx.model_id is None + + cloned = copy(session) + # Clone must preserve the replaced (unbound) context, not re-bind from backend. + assert cloned.ctx.model_id is None + assert cloned.id != session.id # fresh session ID if __name__ == "__main__": From 1f759d579c9ba40b2f7e91c06a4a1409a6a17bfc Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Mon, 15 Jun 2026 18:10:19 -0400 Subject: [PATCH 03/11] remove double backticks Signed-off-by: AngeloDanducci --- mellea/backends/context_lengths.py | 38 +++++++++++++++--------------- mellea/stdlib/context.py | 30 +++++++++++------------ 2 files changed, 34 insertions(+), 34 deletions(-) diff --git a/mellea/backends/context_lengths.py b/mellea/backends/context_lengths.py index 7e3fc0e6dc..a1bfd77807 100644 --- a/mellea/backends/context_lengths.py +++ b/mellea/backends/context_lengths.py @@ -1,17 +1,17 @@ """Model context-length lookup table. Maps known model identifiers to their maximum context window in tokens. -The primary lookup uses ``ModelIdentifier.context_length`` when the caller -passes a ``ModelIdentifier`` object. String lookups fall back to a name -table keyed on ``hf_model_name`` and ``ollama_name`` of every -``ModelIdentifier`` constant defined in ``model_ids``. The table is built +The primary lookup uses `ModelIdentifier.context_length` when the caller +passes a `ModelIdentifier` object. String lookups fall back to a name +table keyed on `hf_model_name` and `ollama_name` of every +`ModelIdentifier` constant defined in `model_ids`. The table is built automatically at import time so there is a single source of truth: the -``context_length`` field on each constant. +`context_length` field on each constant. -Only ``hf_model_name`` and ``ollama_name`` are indexed — matching the two -fields checked in the ``ModelIdentifier`` branch of ``get_context_length``. -Platform-specific IDs (``mlx_name``, ``openai_name``, ``bedrock_name``) and -tokenizer references (``hf_tokenizer_name``) are intentionally excluded: the +Only `hf_model_name` and `ollama_name` are indexed — matching the two +fields checked in the `ModelIdentifier` branch of `get_context_length`. +Platform-specific IDs (`mlx_name`, `openai_name`, `bedrock_name`) and +tokenizer references (`hf_tokenizer_name`) are intentionally excluded: the former would populate the table with deployment-platform strings that are not valid serving names, and the latter is a tokenizer reference, not a model-serving identifier. @@ -23,24 +23,24 @@ def get_context_length(model_id: str | ModelIdentifier) -> int | None: - """Return the maximum context length in tokens for a known model, or ``None``. + """Return the maximum context length in tokens for a known model, or `None`. Priority: - 1. If *model_id* is a ``ModelIdentifier`` with a non-``None`` - ``context_length`` field, return that value directly. + 1. If *model_id* is a `ModelIdentifier` with a non-`None` + `context_length` field, return that value directly. 2. Otherwise perform a name-based lookup in the built-in table, - checking ``hf_model_name`` then ``ollama_name`` (for - ``ModelIdentifier`` inputs) or the raw string (for string inputs). - Only ``hf_model_name`` and ``ollama_name`` are indexed in the table, - so platform-specific names (``mlx_name``, ``openai_name``, - ``bedrock_name``) do not resolve via raw string lookup. + checking `hf_model_name` then `ollama_name` (for + `ModelIdentifier` inputs) or the raw string (for string inputs). + Only `hf_model_name` and `ollama_name` are indexed in the table, + so platform-specific names (`mlx_name`, `openai_name`, + `bedrock_name`) do not resolve via raw string lookup. Args: - model_id: A ``ModelIdentifier`` instance or a raw model-name string. + model_id: A `ModelIdentifier` instance or a raw model-name string. Returns: - The context length in tokens, or ``None`` if the model is unknown or + The context length in tokens, or `None` if the model is unknown or its context length is not recorded in the catalog. """ if isinstance(model_id, ModelIdentifier): diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index 4723f7b35f..22f535302f 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -55,31 +55,31 @@ def __init__( @property def model_id(self) -> str | ModelIdentifier | None: - """The model identifier bound to this context, or ``None`` if unbound.""" + """The model identifier bound to this context, or `None` if unbound.""" return self._model_id def _make_root(self, model_id: str | ModelIdentifier | None) -> ChatContext: - """Return a new empty root ``ChatContext`` with the given ``model_id``, preserving ``window_size``.""" + """Return a new empty root `ChatContext` with the given `model_id`, preserving `window_size`.""" return ChatContext(window_size=self._window_size, model_id=model_id) def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: """Return a new root `ChatContext` with the given model bound, preserving `window_size`. - Internal use only — called by ``MelleaSession`` to wire the backend's + Internal use only — called by `MelleaSession` to wire the backend's model identifier into the context at session construction and after - ``reset()``. To bind a model at construction time, pass ``model_id=`` - directly to ``ChatContext()``. + `reset()`. To bind a model at construction time, pass `model_id=` + directly to `ChatContext()`. Args: model_id: The model identifier to associate with this context. Returns: - ChatContext: A new root `ChatContext` with ``_model_id`` set. + ChatContext: A new root `ChatContext` with `_model_id` set. Raises: ValueError: If called on a non-root (non-empty) context. History would be silently discarded, so this is disallowed. Create a - new ``ChatContext(model_id=...)`` instead. + new `ChatContext(model_id=...)` instead. """ if not self.is_root_node: raise ValueError( @@ -90,15 +90,15 @@ def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: return self._make_root(model_id) def reset_to_new(self) -> ChatContext: # type: ignore[override] - """Return a new empty root ``ChatContext``, preserving ``window_size`` and ``model_id``. + """Return a new empty root `ChatContext`, preserving `window_size` and `model_id`. - Overrides ``Context.reset_to_new()`` so that the model binding and + Overrides `Context.reset_to_new()` so that the model binding and window size set at construction are not lost when a session is reset or when user code calls this method directly. Returns: - ChatContext: A fresh root context with the same ``window_size`` - and ``model_id`` as this context, but no history. + ChatContext: A fresh root context with the same `window_size` + and `model_id` as this context, but no history. """ return self._make_root(self._model_id) @@ -122,10 +122,10 @@ def view_for_generation(self) -> list[Component | CBlock] | None: Window size resolution priority: - 1. Explicit ``window_size`` passed at construction — always wins. - 2. Model-derived context length looked up via ``model_id`` — used when - ``window_size`` is ``None`` and a model binding exists. - 3. ``None`` — return the full history (unbounded). + 1. Explicit `window_size` passed at construction — always wins. + 2. Model-derived context length looked up via `model_id` — used when + `window_size` is `None` and a model binding exists. + 3. `None` — return the full history (unbounded). Returns: list[Component | CBlock] | None: Ordered list of context entries up to From beac933cde826251599d7ad62629dc7cd821be57 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Mon, 15 Jun 2026 22:24:48 -0400 Subject: [PATCH 04/11] add missing batch of changes and review changes Signed-off-by: AngeloDanducci --- mellea/backends/context_lengths.py | 9 +--- mellea/stdlib/context.py | 74 +++++++++++++++++++++--------- mellea/stdlib/session.py | 39 ++++++++-------- test/stdlib/test_base_context.py | 33 +++++++++++-- 4 files changed, 104 insertions(+), 51 deletions(-) diff --git a/mellea/backends/context_lengths.py b/mellea/backends/context_lengths.py index a1bfd77807..570088e1ca 100644 --- a/mellea/backends/context_lengths.py +++ b/mellea/backends/context_lengths.py @@ -19,6 +19,8 @@ from __future__ import annotations +import mellea.backends.model_ids as _m + from .model_ids import ModelIdentifier @@ -61,17 +63,10 @@ def get_context_length(model_id: str | ModelIdentifier) -> int | None: def _build_table() -> dict[str, int]: - import mellea.backends.model_ids as _m - table: dict[str, int] = {} for obj in vars(_m).values(): if not isinstance(obj, ModelIdentifier) or obj.context_length is None: continue - # Only index hf_model_name and ollama_name — these are the two fields - # checked by the ModelIdentifier branch of get_context_length, keeping - # the table and the lookup logic symmetric. Platform-specific names - # (mlx_name, openai_name, bedrock_name) and hf_tokenizer_name are - # intentionally excluded; see module docstring for rationale. for name in (obj.hf_model_name, obj.ollama_name): if name: table[name] = obj.context_length diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index 22f535302f..bf719c4695 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -34,12 +34,11 @@ class ChatContext(Context): Note: `context_length` is measured in tokens; `window_size` counts context items (CBlocks / Components). When only a `model_id` is bound (no - explicit `window_size`), the model's context length in tokens is used - as an absolute upper bound on the number of items passed to the model. - In practice this ceiling is never reached — real conversations do not - accumulate hundreds of thousands of items — so the full history is - returned for typical sessions. Set `window_size` explicitly to enforce - a tighter item-count limit. + explicit `window_size`), `view_for_generation` estimates per-item token + counts (via `len(str(item)) // 4`) and walks history newest-first, + dropping the oldest items until the running sum fits within + `context_length`. Set `window_size` explicitly to enforce an item-count + limit instead of a token budget. """ def __init__( @@ -79,7 +78,9 @@ def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: Raises: ValueError: If called on a non-root (non-empty) context. History would be silently discarded, so this is disallowed. Create a - new `ChatContext(model_id=...)` instead. + new `ChatContext(model_id=...)` instead. Note: `__init__` is + the unrestricted path — it produces the same end state without + this root check. """ if not self.is_root_node: raise ValueError( @@ -89,16 +90,16 @@ def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: ) return self._make_root(model_id) - def reset_to_new(self) -> ChatContext: # type: ignore[override] + def new_instance(self) -> ChatContext: """Return a new empty root `ChatContext`, preserving `window_size` and `model_id`. - Overrides `Context.reset_to_new()` so that the model binding and - window size set at construction are not lost when a session is reset - or when user code calls this method directly. + Use this instead of `reset_to_new()` when you need to preserve the + model binding and window size from an existing instance. `reset_to_new()` + is a classmethod that returns a bare `ChatContext()` with no configuration. Returns: ChatContext: A fresh root context with the same `window_size` - and `model_id` as this context, but no history. + and `model_id` as this instance, but no history. """ return self._make_root(self._model_id) @@ -122,19 +123,50 @@ def view_for_generation(self) -> list[Component | CBlock] | None: Window size resolution priority: - 1. Explicit `window_size` passed at construction — always wins. - 2. Model-derived context length looked up via `model_id` — used when - `window_size` is `None` and a model binding exists. - 3. `None` — return the full history (unbounded). + 1. Explicit `window_size` passed at construction — item-count limit, always wins. + 2. Model-derived context length looked up via `model_id` — token-budget truncation. + Items are added newest-first until adding the next item would exceed the budget. + Token count is estimated as `len(str(item)) // 4`. + 3. No limit — return the full history. Returns: list[Component | CBlock] | None: Ordered list of context entries up to - the effective window size, or `None` if the history is non-linear. + the effective window, or `None` if the history is non-linear. """ - effective_window = self._window_size - if effective_window is None and self._model_id is not None: - effective_window = get_context_length(self._model_id) - return self.as_list(effective_window) + if self._window_size is not None: + return self.as_list(self._window_size) + + if self._model_id is not None: + token_budget = get_context_length(self._model_id) + if token_budget is not None: + return self._as_list_token_budget(token_budget) + + return self.as_list(None) + + def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: + """Return history items that fit within *token_budget*, dropping oldest first. + + Walks the linked list from newest to oldest, accumulating items until + adding the next item would exceed the budget. The returned list is in + chronological order (oldest-first), matching `as_list` behaviour. + Token count per item is estimated as `len(str(item)) // 4`. + """ + collected: list[Component | CBlock] = [] + spent = 0 + current: Context = self + while not current.is_root_node: + item = current.node_data + assert item is not None + cost = max(1, len(str(item)) // 4) + if spent + cost > token_budget: + break + collected.append(item) + spent += cost + prev = current.previous_node + assert prev is not None + current = prev + collected.reverse() + return collected class SimpleContext(Context): diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index d17002224f..cbc4c79fb5 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -290,18 +290,25 @@ def __init__( """Initialize MelleaSession with a backend and optional conversation context.""" import uuid - self.id = session_id if session_id is not None else str(uuid.uuid4()) + self._init_fields( + session_id if session_id is not None else str(uuid.uuid4()), + backend, + ctx if ctx is not None else SimpleContext(), + ) + self._auto_bind_model() + + def _init_fields(self, session_id: str, backend: Backend, ctx: Context) -> None: + """Set all instance fields. Shared by __init__ and __copy__ so neither diverges.""" + self.id = session_id self.backend = backend - self.ctx: Context = ctx if ctx is not None else SimpleContext() + self.ctx: Context = ctx self._session_logger = MelleaLogger.get_logger() self._context_token = None self._log_context_token = None self._exit_stack: contextlib.ExitStack | None = None - # Bind the backend's model_id to the context when the context is a - # ChatContext without an existing model binding. This enables automatic - # context-window sizing in view_for_generation() when window_size is - # not explicitly set. + def _auto_bind_model(self) -> None: + """Bind the backend's model_id to a bare ChatContext, enabling auto context-window sizing.""" if ( isinstance(self.ctx, ChatContext) and self.ctx.model_id is None @@ -358,19 +365,9 @@ def __copy__(self): """Use self.clone. Copies the current session but keeps references to the backend and context.""" import uuid - # Bypass __init__'s auto-bind logic: the session state is already - # settled, so we replicate it directly rather than re-deriving it. new = object.__new__(MelleaSession) - new.id = str(uuid.uuid4()) - new.backend = self.backend - new.ctx = self.ctx - new._session_logger = self._session_logger - new._context_token = None - new._log_context_token = None - new._session_span = None - new._exit_stack = None - # Explicitly don't copy over the _context_token. - + new._init_fields(str(uuid.uuid4()), self.backend, self.ctx) + # Do not call _auto_bind_model: the session state is already settled. return new def clone(self) -> MelleaSession: @@ -412,7 +409,11 @@ def reset(self) -> None: _run_async_in_thread( invoke_hook(HookType.SESSION_RESET, payload, backend=self.backend) ) - self.ctx = self.ctx.reset_to_new() + self.ctx = ( + self.ctx.new_instance() + if isinstance(self.ctx, ChatContext) + else self.ctx.reset_to_new() + ) def cleanup(self, *, exception: BaseException | None = None) -> None: """Clean up session resources and deregister session-scoped plugins. diff --git a/test/stdlib/test_base_context.py b/test/stdlib/test_base_context.py index 13de3e34a9..0a1279bbf3 100644 --- a/test/stdlib/test_base_context.py +++ b/test/stdlib/test_base_context.py @@ -197,6 +197,8 @@ def test_view_for_generation_model_bound_returns_full_history_when_under_limit() def test_view_for_generation_model_bound_used_as_upper_bound(): # A tiny context_length=5 triggers truncation; most recent items are retained. + # Each "item N" is 6 chars → cost = max(1, 6 // 4) = 1 token each. + # Budget 5 fits the 5 most-recent items. tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=5) ctx = ChatContext(model_id=tiny) for i in range(10): @@ -206,6 +208,20 @@ def test_view_for_generation_model_bound_used_as_upper_bound(): assert str(result[-1]) == "item 9" +def test_view_for_generation_token_budget_drops_oldest(): + # Each item is a 100-char string → cost = 100 // 4 = 25 tokens. + # Budget 60: fits 2 items (50 tokens), the 3rd would exceed the budget. + tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=60) + ctx = ChatContext(model_id=tiny) + for i in range(5): + ctx = ctx.add(CBlock("x" * 100 + str(i))) # 101 chars → cost 25 + result = ctx.view_for_generation() + assert len(result) == 2 + # Newest items are kept. + assert str(result[-1]).endswith("4") + assert str(result[-2]).endswith("3") + + def test_view_for_generation_explicit_window_size_beats_model(): ctx = ChatContext(window_size=2, model_id=IBM_GRANITE_4_1_8B) for i in range(10): @@ -246,13 +262,22 @@ def test_reset_to_new_returns_empty_root(): assert len(reset_ctx.as_list()) == 0 -def test_reset_to_new_preserves_model_id_and_window_size(): +def test_reset_to_new_classmethod_does_not_preserve_config(): + # reset_to_new() is a classmethod — it returns a bare ChatContext() with + # no window_size or model_id, regardless of the instance it's called on. + ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B, window_size=5) + reset_ctx = ctx.reset_to_new() + assert reset_ctx.model_id is None + assert reset_ctx._window_size is None + + +def test_new_instance_preserves_model_id_and_window_size(): ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B, window_size=5) for i in range(3): ctx = ctx.add(CBlock(f"msg {i}")) - reset_ctx = ctx.reset_to_new() - assert reset_ctx.model_id is IBM_GRANITE_4_1_8B - assert reset_ctx._window_size == 5 + new_ctx = ctx.new_instance() + assert new_ctx.model_id is IBM_GRANITE_4_1_8B + assert new_ctx._window_size == 5 def test_session_reset_rebinds_model_id(): From d582785963d66fd7c01f6327db5d07d84b8a3924 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Tue, 16 Jun 2026 15:09:45 -0400 Subject: [PATCH 05/11] review feedback Signed-off-by: AngeloDanducci --- mellea/backends/context_lengths.py | 8 ++-- mellea/backends/model_ids.py | 2 +- mellea/stdlib/context.py | 35 +++++++++++++-- mellea/stdlib/session.py | 10 +++-- test/stdlib/test_base_context.py | 70 +++++++++++++++++++++++++----- 5 files changed, 101 insertions(+), 24 deletions(-) diff --git a/mellea/backends/context_lengths.py b/mellea/backends/context_lengths.py index 570088e1ca..f8235c441f 100644 --- a/mellea/backends/context_lengths.py +++ b/mellea/backends/context_lengths.py @@ -34,9 +34,6 @@ def get_context_length(model_id: str | ModelIdentifier) -> int | None: 2. Otherwise perform a name-based lookup in the built-in table, checking `hf_model_name` then `ollama_name` (for `ModelIdentifier` inputs) or the raw string (for string inputs). - Only `hf_model_name` and `ollama_name` are indexed in the table, - so platform-specific names (`mlx_name`, `openai_name`, - `bedrock_name`) do not resolve via raw string lookup. Args: model_id: A `ModelIdentifier` instance or a raw model-name string. @@ -69,6 +66,11 @@ def _build_table() -> dict[str, int]: continue for name in (obj.hf_model_name, obj.ollama_name): if name: + if name in table and table[name] != obj.context_length: + raise ValueError( + f"context_length collision for {name!r}: " + f"{table[name]} vs {obj.context_length}" + ) table[name] = obj.context_length return table diff --git a/mellea/backends/model_ids.py b/mellea/backends/model_ids.py index 6642f9e753..e9b09b6639 100644 --- a/mellea/backends/model_ids.py +++ b/mellea/backends/model_ids.py @@ -261,7 +261,7 @@ class ModelIdentifier: QWEN3_8B = ModelIdentifier( hf_model_name="Qwen/Qwen3-8B", # Qwen 8B ollama_name="qwen3:8b", # Ollama - context_length=40960, + context_length=40960, # 8B+ series; smaller Qwen3 models use 32768 ) QWEN3_14B = ModelIdentifier( diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index bf719c4695..7706a112d5 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -13,7 +13,7 @@ from ..backends.model_ids import ModelIdentifier # Leave unused `ContextTurn` import for import ergonomics. -from ..core import CBlock, Component, Context, ContextTurn +from ..core import CBlock, Component, Context, ContextTurn, MelleaLogger class ChatContext(Context): @@ -112,6 +112,13 @@ def add(self, c: Component | CBlock) -> ChatContext: Returns: ChatContext: A new `ChatContext` with the added entry, preserving both `window_size` and any `model_id` binding. + + Note: + Invariant: every ChatContext-specific field (_window_size, _model_id) + must be copied here AND in _make_root(). from_previous() initialises + them to None; we patch them back manually. Any new field added to + ChatContext must be propagated in both places or it will silently + disappear from child nodes. """ new = ChatContext.from_previous(self, c) new._window_size = self._window_size @@ -149,22 +156,42 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: Walks the linked list from newest to oldest, accumulating items until adding the next item would exceed the budget. The returned list is in chronological order (oldest-first), matching `as_list` behaviour. - Token count per item is estimated as `len(str(item)) // 4`. + Token count per item is estimated as `len(str(item)) // 4`; note that + `str()` falls back to `repr()` for `Component` subclasses, so the + estimate reflects repr boilerplate rather than rendered content. + + A headroom factor of 0.8 is applied so the backend retains capacity for + the current action and the model's generated response. """ + effective_budget = int(token_budget * 0.8) collected: list[Component | CBlock] = [] spent = 0 + total = 0 current: Context = self while not current.is_root_node: item = current.node_data assert item is not None - cost = max(1, len(str(item)) // 4) - if spent + cost > token_budget: + cost = max( + 1, len(str(item)) // 4 + ) # str() falls back to repr() for Components + total += 1 + if spent + cost > effective_budget: break collected.append(item) spent += cost prev = current.previous_node assert prev is not None current = prev + dropped = total - len(collected) + if dropped: + MelleaLogger.get_logger().debug( + "Context truncated: dropped %d item(s) to stay within %d-token budget " + "(effective budget after 0.8 headroom: %d tokens, used: %d tokens).", + dropped, + token_budget, + effective_budget, + spent, + ) collected.reverse() return collected diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index cbc4c79fb5..0466bc4229 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -313,9 +313,9 @@ def _auto_bind_model(self) -> None: isinstance(self.ctx, ChatContext) and self.ctx.model_id is None and self.ctx.is_root_node - and hasattr(self.backend, "model_id") + and getattr(self.backend, "model_id", None) is not None ): - self.ctx = self.ctx._bind_model(self.backend.model_id) + self.ctx = self.ctx._bind_model(getattr(self.backend, "model_id")) def __enter__(self): """Enter context manager and set this session as the current global session.""" @@ -399,8 +399,10 @@ def reset(self) -> None: """Reset the context state to a fresh, empty context of the same type. Fires the `SESSION_RESET` plugin hook if any plugins are registered, then - replaces `self.ctx` with the result of `ctx.reset_to_new()`, discarding - all accumulated conversation history. + replaces `self.ctx` with a fresh empty context, discarding all accumulated + conversation history. For `ChatContext`, uses `new_instance()` so the + `model_id` and `window_size` bindings are preserved; for all other context + types, uses `reset_to_new()`. """ if has_plugins(HookType.SESSION_RESET): from ..plugins.hooks.session import SessionResetPayload diff --git a/test/stdlib/test_base_context.py b/test/stdlib/test_base_context.py index 0a1279bbf3..77ac0e0c6b 100644 --- a/test/stdlib/test_base_context.py +++ b/test/stdlib/test_base_context.py @@ -196,30 +196,31 @@ def test_view_for_generation_model_bound_returns_full_history_when_under_limit() def test_view_for_generation_model_bound_used_as_upper_bound(): - # A tiny context_length=5 triggers truncation; most recent items are retained. + # A tiny context_length=10 triggers truncation; most recent items are retained. # Each "item N" is 6 chars → cost = max(1, 6 // 4) = 1 token each. - # Budget 5 fits the 5 most-recent items. - tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=5) + # Effective budget = int(10 * 0.8) = 8; fits the 8 most-recent items. + tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=10) ctx = ChatContext(model_id=tiny) - for i in range(10): + for i in range(15): ctx = ctx.add(CBlock(f"item {i}")) result = ctx.view_for_generation() - assert len(result) == 5 - assert str(result[-1]) == "item 9" + assert len(result) == 8 + assert str(result[-1]) == "item 14" def test_view_for_generation_token_budget_drops_oldest(): # Each item is a 100-char string → cost = 100 // 4 = 25 tokens. - # Budget 60: fits 2 items (50 tokens), the 3rd would exceed the budget. - tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=60) + # context_length=130: effective budget = int(130 * 0.8) = 104. + # Fits 4 items (100 tokens); the 5th would push spent to 125 > 104. + tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=130) ctx = ChatContext(model_id=tiny) - for i in range(5): + for i in range(7): ctx = ctx.add(CBlock("x" * 100 + str(i))) # 101 chars → cost 25 result = ctx.view_for_generation() - assert len(result) == 2 + assert len(result) == 4 # Newest items are kept. - assert str(result[-1]).endswith("4") - assert str(result[-2]).endswith("3") + assert str(result[-1]).endswith("6") + assert str(result[-2]).endswith("5") def test_view_for_generation_explicit_window_size_beats_model(): @@ -386,5 +387,50 @@ def test_clone_does_not_double_bind_replaced_root_context(): assert cloned.id != session.id # fresh session ID +# --------------------------------------------------------------------------- +# _as_list_token_budget boundary and _build_table collision tests +# --------------------------------------------------------------------------- + + +def test_as_list_token_budget_fits_exactly(): + # Verify > vs >= boundary: an item whose cost equals the remaining budget + # is INCLUDED (condition is `spent + cost > effective_budget`, so equality passes). + # context_length=10 → effective = int(10 * 0.8) = 8. + # Two items, each 8 chars → cost = max(1, 8 // 4) = 2 tokens each; total = 4. + # Four such items = 8 tokens, exactly equal to effective_budget — all four included. + tiny = ModelIdentifier(hf_model_name="org/exact-fit-model", context_length=10) + ctx = ChatContext(model_id=tiny) + for _ in range(4): + ctx = ctx.add(CBlock("abcdefgh")) # 8 chars → cost 2 + result = ctx.view_for_generation() + assert len(result) == 4 + + +def test_build_table_raises_on_collision(): + import mellea.backends.context_lengths as cl + from mellea.backends.context_lengths import _build_table + from mellea.backends.model_ids import ModelIdentifier + + colliding_a = ModelIdentifier( + hf_model_name="shared/model-name", context_length=1024 + ) + colliding_b = ModelIdentifier( + hf_model_name="shared/model-name", context_length=2048 + ) + + # Temporarily inject two colliding constants into the model_ids module that + # _build_table() scans via vars(_m). + import mellea.backends.model_ids as _m_module + + _m_module.COLLIDE_A = colliding_a # type: ignore[attr-defined] + _m_module.COLLIDE_B = colliding_b # type: ignore[attr-defined] + try: + with pytest.raises(ValueError, match="context_length collision"): + _build_table() + finally: + del _m_module.COLLIDE_A # type: ignore[attr-defined] + del _m_module.COLLIDE_B # type: ignore[attr-defined] + + if __name__ == "__main__": pytest.main([__file__]) From 5f77ed14d71c57bd243c4213efde9789b343ea1d Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Tue, 16 Jun 2026 16:49:05 -0400 Subject: [PATCH 06/11] additional review feedback Signed-off-by: AngeloDanducci --- mellea/backends/context_lengths.py | 8 ++++++ mellea/core/base.py | 13 +++++++++ mellea/stdlib/context.py | 43 ++++++++++++++++++++---------- mellea/stdlib/session.py | 15 ++++++----- test/stdlib/test_base_context.py | 39 +++++++++++---------------- 5 files changed, 75 insertions(+), 43 deletions(-) diff --git a/mellea/backends/context_lengths.py b/mellea/backends/context_lengths.py index f8235c441f..a1d509e799 100644 --- a/mellea/backends/context_lengths.py +++ b/mellea/backends/context_lengths.py @@ -41,6 +41,14 @@ def get_context_length(model_id: str | ModelIdentifier) -> int | None: Returns: The context length in tokens, or `None` if the model is unknown or its context length is not recorded in the catalog. + + Warning: + The values returned here reflect each model's theoretical maximum context + window as recorded in the catalog — **not** the context window actually + used by a running server. Ollama in particular defaults to + ``num_ctx=2048`` regardless of the model's rated maximum; a caller that + relies on this lookup to size a sliding context window will silently + overflow Ollama's wire-level limit. """ if isinstance(model_id, ModelIdentifier): if model_id.context_length is not None: diff --git a/mellea/core/base.py b/mellea/core/base.py index ff2202559b..a078e1c30a 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -1028,6 +1028,19 @@ def reset_to_new(cls: type[ContextT]) -> ContextT: """ return cls() + def new_instance(self) -> Context: + """Return a new empty root context, preserving any subclass configuration. + + The base implementation calls `reset_to_new()`, which returns a bare + instance with no history and no config. Subclasses that carry + configuration (e.g. `ChatContext` with `model_id` and `window_size`) + should override this to propagate their config into the fresh instance. + + Returns: + Context: A freshly initialised root context of the same type. + """ + return self.reset_to_new() + # Internal functions below this line. @property diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index 7706a112d5..15cd55a422 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -15,6 +15,8 @@ # Leave unused `ContextTurn` import for import ergonomics. from ..core import CBlock, Component, Context, ContextTurn, MelleaLogger +logger = MelleaLogger.get_logger() + class ChatContext(Context): """Initializes a chat context with unbounded window_size and is_chat=True by default. @@ -31,6 +33,12 @@ class ChatContext(Context): length and uses it as the window size. Explicit `window_size` always takes priority. + Class Attributes: + _propagated_fields: Instance-attribute names copied by `add()` and + `_make_root()` into every descendant node. Add new `ChatContext` + fields here so they propagate automatically without touching both + methods. + Note: `context_length` is measured in tokens; `window_size` counts context items (CBlocks / Components). When only a `model_id` is bound (no @@ -39,8 +47,20 @@ class ChatContext(Context): dropping the oldest items until the running sum fits within `context_length`. Set `window_size` explicitly to enforce an item-count limit instead of a token budget. + + The token estimate is intentionally conservative: `str()` falls back to + `repr()` for `Component` subclasses, so it counts repr boilerplate rather + than rendered prompt content. A 0.55 headroom factor (retaining 55 % of + the model's rated context length as the effective budget) is applied to + absorb this repr-vs-render skew and to leave capacity for the current + action and the model's generated response. The result is a rough + lower-bound — expect the actual token usage to be noticeably below the + budget ceiling. Use `window_size` for precise item-count control, or + wait for a future tokenizer-backed estimate. """ + _propagated_fields: tuple[str, ...] = ("_window_size", "_model_id") + def __init__( self, *, @@ -112,17 +132,10 @@ def add(self, c: Component | CBlock) -> ChatContext: Returns: ChatContext: A new `ChatContext` with the added entry, preserving both `window_size` and any `model_id` binding. - - Note: - Invariant: every ChatContext-specific field (_window_size, _model_id) - must be copied here AND in _make_root(). from_previous() initialises - them to None; we patch them back manually. Any new field added to - ChatContext must be propagated in both places or it will silently - disappear from child nodes. """ new = ChatContext.from_previous(self, c) - new._window_size = self._window_size - new._model_id = self._model_id + for field in self._propagated_fields: + setattr(new, field, getattr(self, field)) return new def view_for_generation(self) -> list[Component | CBlock] | None: @@ -160,10 +173,12 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: `str()` falls back to `repr()` for `Component` subclasses, so the estimate reflects repr boilerplate rather than rendered content. - A headroom factor of 0.8 is applied so the backend retains capacity for - the current action and the model's generated response. + A headroom factor of 0.55 is applied to absorb repr-vs-render skew and + to leave capacity for the current action and the model's generated + response. This is a conservative approximation; a tokenizer-backed + estimate is a known follow-up. """ - effective_budget = int(token_budget * 0.8) + effective_budget = int(token_budget * 0.55) collected: list[Component | CBlock] = [] spent = 0 total = 0 @@ -184,9 +199,9 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: current = prev dropped = total - len(collected) if dropped: - MelleaLogger.get_logger().debug( + logger.debug( "Context truncated: dropped %d item(s) to stay within %d-token budget " - "(effective budget after 0.8 headroom: %d tokens, used: %d tokens).", + "(effective budget after 0.55 headroom: %d tokens, used: %d tokens).", dropped, token_budget, effective_budget, diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index 0466bc4229..e8ead9e216 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -287,7 +287,14 @@ def __init__( *, session_id: str | None = None, ): - """Initialize MelleaSession with a backend and optional conversation context.""" + """Initialize MelleaSession with a backend and optional conversation context. + + Note: + When *ctx* is a bare (unbound, empty) `ChatContext` and *backend* + exposes a `model_id`, the constructor replaces `self.ctx` with a + re-bound copy that carries the model identifier. After construction, + ``session.ctx is ctx`` will therefore be ``False``. + """ import uuid self._init_fields( @@ -411,11 +418,7 @@ def reset(self) -> None: _run_async_in_thread( invoke_hook(HookType.SESSION_RESET, payload, backend=self.backend) ) - self.ctx = ( - self.ctx.new_instance() - if isinstance(self.ctx, ChatContext) - else self.ctx.reset_to_new() - ) + self.ctx = self.ctx.new_instance() def cleanup(self, *, exception: BaseException | None = None) -> None: """Clean up session resources and deregister session-scoped plugins. diff --git a/test/stdlib/test_base_context.py b/test/stdlib/test_base_context.py index 77ac0e0c6b..fb10f18d9e 100644 --- a/test/stdlib/test_base_context.py +++ b/test/stdlib/test_base_context.py @@ -198,26 +198,26 @@ def test_view_for_generation_model_bound_returns_full_history_when_under_limit() def test_view_for_generation_model_bound_used_as_upper_bound(): # A tiny context_length=10 triggers truncation; most recent items are retained. # Each "item N" is 6 chars → cost = max(1, 6 // 4) = 1 token each. - # Effective budget = int(10 * 0.8) = 8; fits the 8 most-recent items. + # Effective budget = int(10 * 0.55) = 5; fits the 5 most-recent items. tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=10) ctx = ChatContext(model_id=tiny) for i in range(15): ctx = ctx.add(CBlock(f"item {i}")) result = ctx.view_for_generation() - assert len(result) == 8 + assert len(result) == 5 assert str(result[-1]) == "item 14" def test_view_for_generation_token_budget_drops_oldest(): # Each item is a 100-char string → cost = 100 // 4 = 25 tokens. - # context_length=130: effective budget = int(130 * 0.8) = 104. - # Fits 4 items (100 tokens); the 5th would push spent to 125 > 104. + # context_length=130: effective budget = int(130 * 0.55) = 71. + # Fits 2 items (50 tokens); the 3rd would push spent to 75 > 71. tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=130) ctx = ChatContext(model_id=tiny) for i in range(7): ctx = ctx.add(CBlock("x" * 100 + str(i))) # 101 chars → cost 25 result = ctx.view_for_generation() - assert len(result) == 4 + assert len(result) == 2 # Newest items are kept. assert str(result[-1]).endswith("6") assert str(result[-2]).endswith("5") @@ -395,19 +395,18 @@ def test_clone_does_not_double_bind_replaced_root_context(): def test_as_list_token_budget_fits_exactly(): # Verify > vs >= boundary: an item whose cost equals the remaining budget # is INCLUDED (condition is `spent + cost > effective_budget`, so equality passes). - # context_length=10 → effective = int(10 * 0.8) = 8. - # Two items, each 8 chars → cost = max(1, 8 // 4) = 2 tokens each; total = 4. - # Four such items = 8 tokens, exactly equal to effective_budget — all four included. - tiny = ModelIdentifier(hf_model_name="org/exact-fit-model", context_length=10) + # context_length=20 → effective = int(20 * 0.55) = 11. + # Items of 8 chars → cost = max(1, 8 // 4) = 2 tokens each. + # Five such items = 10 tokens ≤ 11 (all included); 6th = 12 > 11 (excluded). + tiny = ModelIdentifier(hf_model_name="org/exact-fit-model", context_length=20) ctx = ChatContext(model_id=tiny) - for _ in range(4): + for _ in range(6): ctx = ctx.add(CBlock("abcdefgh")) # 8 chars → cost 2 result = ctx.view_for_generation() - assert len(result) == 4 + assert len(result) == 5 -def test_build_table_raises_on_collision(): - import mellea.backends.context_lengths as cl +def test_build_table_raises_on_collision(monkeypatch): from mellea.backends.context_lengths import _build_table from mellea.backends.model_ids import ModelIdentifier @@ -418,18 +417,12 @@ def test_build_table_raises_on_collision(): hf_model_name="shared/model-name", context_length=2048 ) - # Temporarily inject two colliding constants into the model_ids module that - # _build_table() scans via vars(_m). import mellea.backends.model_ids as _m_module - _m_module.COLLIDE_A = colliding_a # type: ignore[attr-defined] - _m_module.COLLIDE_B = colliding_b # type: ignore[attr-defined] - try: - with pytest.raises(ValueError, match="context_length collision"): - _build_table() - finally: - del _m_module.COLLIDE_A # type: ignore[attr-defined] - del _m_module.COLLIDE_B # type: ignore[attr-defined] + monkeypatch.setattr(_m_module, "COLLIDE_A", colliding_a, raising=False) + monkeypatch.setattr(_m_module, "COLLIDE_B", colliding_b, raising=False) + with pytest.raises(ValueError, match="context_length collision"): + _build_table() if __name__ == "__main__": From 3eb087480c1d322090fab312589942f57953eaf3 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Wed, 17 Jun 2026 13:54:46 -0400 Subject: [PATCH 07/11] make root now iterates propagated fields Signed-off-by: AngeloDanducci --- mellea/stdlib/context.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index 15cd55a422..1534081062 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -78,8 +78,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` with the given `model_id`, preserving `window_size`.""" - return ChatContext(window_size=self._window_size, model_id=model_id) + """Return a new empty root `ChatContext` with the given `model_id`, propagating all `_propagated_fields`.""" + new = ChatContext() + for field in self._propagated_fields: + setattr(new, field, getattr(self, field)) + new._model_id = model_id + return new def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: """Return a new root `ChatContext` with the given model bound, preserving `window_size`. From 0a649e1b7d61b52b6f00841f9837cb231349b5f9 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Mon, 22 Jun 2026 11:56:00 -0400 Subject: [PATCH 08/11] review feedback Signed-off-by: AngeloDanducci --- mellea/stdlib/context.py | 64 +++++++++++++++++++++------------------- mellea/stdlib/session.py | 2 +- 2 files changed, 35 insertions(+), 31 deletions(-) diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index 1534081062..f8891e8d8e 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -41,22 +41,22 @@ class ChatContext(Context): Note: `context_length` is measured in tokens; `window_size` counts context - items (CBlocks / Components). When only a `model_id` is bound (no - explicit `window_size`), `view_for_generation` estimates per-item token - counts (via `len(str(item)) // 4`) and walks history newest-first, - dropping the oldest items until the running sum fits within - `context_length`. Set `window_size` explicitly to enforce an item-count - limit instead of a token budget. - - The token estimate is intentionally conservative: `str()` falls back to - `repr()` for `Component` subclasses, so it counts repr boilerplate rather - than rendered prompt content. A 0.55 headroom factor (retaining 55 % of - the model's rated context length as the effective budget) is applied to - absorb this repr-vs-render skew and to leave capacity for the current - action and the model's generated response. The result is a rough - lower-bound — expect the actual token usage to be noticeably below the - budget ceiling. Use `window_size` for precise item-count control, or - wait for a future tokenizer-backed estimate. + items (CBlocks / Components). When both `window_size` and `model_id` + are set, `window_size` always takes priority — the token-budget path is + skipped entirely. When only `model_id` is bound, `view_for_generation` + estimates per-item token counts (via ``len(rendered) // 4``) and walks + history newest-first, dropping the oldest items until the running sum + fits within `context_length`. Set `window_size` explicitly to enforce + an item-count limit instead of a token budget. + + Per-item token count is estimated as ``len(rendered) // 4`` where + ``rendered`` is the string produced by `TemplateFormatter` (the same + renderer used at generation time), so the estimate reflects actual prompt + content. A 0.75 headroom factor (retaining 75 % of the model's rated + context length) is applied to leave capacity for the system prompt, + injected tool schemas, the current action, and the model's response — + none of which are tracked here. Use ``window_size`` for precise + item-count control. """ _propagated_fields: tuple[str, ...] = ("_window_size", "_model_id") @@ -82,6 +82,8 @@ def _make_root(self, model_id: str | ModelIdentifier | None) -> ChatContext: new = ChatContext() for field in self._propagated_fields: setattr(new, field, getattr(self, field)) + # Override whatever _propagated_fields copied for _model_id: the caller + # explicitly supplies the model_id to bind (e.g. _bind_model changes it). new._model_id = model_id return new @@ -150,7 +152,7 @@ def view_for_generation(self) -> list[Component | CBlock] | None: 1. Explicit `window_size` passed at construction — item-count limit, always wins. 2. Model-derived context length looked up via `model_id` — token-budget truncation. Items are added newest-first until adding the next item would exceed the budget. - Token count is estimated as `len(str(item)) // 4`. + Token count is estimated as ``len(rendered) // 4`` via `TemplateFormatter`. 3. No limit — return the full history. Returns: @@ -173,16 +175,20 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: Walks the linked list from newest to oldest, accumulating items until adding the next item would exceed the budget. The returned list is in chronological order (oldest-first), matching `as_list` behaviour. - Token count per item is estimated as `len(str(item)) // 4`; note that - `str()` falls back to `repr()` for `Component` subclasses, so the - estimate reflects repr boilerplate rather than rendered content. - - A headroom factor of 0.55 is applied to absorb repr-vs-render skew and - to leave capacity for the current action and the model's generated - response. This is a conservative approximation; a tokenizer-backed - estimate is a known follow-up. + + Per-item token count is estimated as ``len(rendered) // 4`` where + ``rendered`` is the string produced by `TemplateFormatter` for the + bound model — the same renderer used at generation time. A 0.75 + headroom factor (retaining 75 % of the rated context length) reserves + capacity for the system prompt, injected tool schemas, the current + action, and the model's response. Use ``window_size`` for precise + item-count control. """ - effective_budget = int(token_budget * 0.55) + assert self._model_id is not None + from ..formatters import TemplateFormatter # deferred to avoid circular import + + formatter = TemplateFormatter(self._model_id) + effective_budget = int(token_budget * 0.75) collected: list[Component | CBlock] = [] spent = 0 total = 0 @@ -190,9 +196,7 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: while not current.is_root_node: item = current.node_data assert item is not None - cost = max( - 1, len(str(item)) // 4 - ) # str() falls back to repr() for Components + cost = max(1, len(formatter.print(item)) // 4) total += 1 if spent + cost > effective_budget: break @@ -205,7 +209,7 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: if dropped: logger.debug( "Context truncated: dropped %d item(s) to stay within %d-token budget " - "(effective budget after 0.55 headroom: %d tokens, used: %d tokens).", + "(effective budget after 0.75 headroom: %d tokens, used: %d tokens).", dropped, token_budget, effective_budget, diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index e8ead9e216..e1a262da5e 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -293,7 +293,7 @@ def __init__( When *ctx* is a bare (unbound, empty) `ChatContext` and *backend* exposes a `model_id`, the constructor replaces `self.ctx` with a re-bound copy that carries the model identifier. After construction, - ``session.ctx is ctx`` will therefore be ``False``. + `session.ctx is ctx` will therefore be `False`. """ import uuid From ff3d36245af149e88ebac06173ad939a83465d83 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Mon, 22 Jun 2026 12:21:11 -0400 Subject: [PATCH 09/11] update token budget tests Signed-off-by: AngeloDanducci --- test/stdlib/test_base_context.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/test/stdlib/test_base_context.py b/test/stdlib/test_base_context.py index fb10f18d9e..2e21e7b02c 100644 --- a/test/stdlib/test_base_context.py +++ b/test/stdlib/test_base_context.py @@ -197,27 +197,27 @@ def test_view_for_generation_model_bound_returns_full_history_when_under_limit() def test_view_for_generation_model_bound_used_as_upper_bound(): # A tiny context_length=10 triggers truncation; most recent items are retained. - # Each "item N" is 6 chars → cost = max(1, 6 // 4) = 1 token each. - # Effective budget = int(10 * 0.55) = 5; fits the 5 most-recent items. + # Each "item N" renders to ~7 chars → cost = max(1, 7 // 4) = 1 token each. + # Effective budget = int(10 * 0.75) = 7; fits the 7 most-recent items. tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=10) ctx = ChatContext(model_id=tiny) for i in range(15): ctx = ctx.add(CBlock(f"item {i}")) result = ctx.view_for_generation() - assert len(result) == 5 + assert len(result) == 7 assert str(result[-1]) == "item 14" def test_view_for_generation_token_budget_drops_oldest(): - # Each item is a 100-char string → cost = 100 // 4 = 25 tokens. - # context_length=130: effective budget = int(130 * 0.55) = 71. - # Fits 2 items (50 tokens); the 3rd would push spent to 75 > 71. + # Each item is a 100-char string → cost = 101 // 4 = 25 tokens. + # context_length=130: effective budget = int(130 * 0.75) = 97. + # Fits 3 items (75 tokens); the 4th would push spent to 100 > 97. tiny = ModelIdentifier(hf_model_name="org/tiny-model", context_length=130) ctx = ChatContext(model_id=tiny) for i in range(7): ctx = ctx.add(CBlock("x" * 100 + str(i))) # 101 chars → cost 25 result = ctx.view_for_generation() - assert len(result) == 2 + assert len(result) == 3 # Newest items are kept. assert str(result[-1]).endswith("6") assert str(result[-2]).endswith("5") @@ -395,15 +395,15 @@ def test_clone_does_not_double_bind_replaced_root_context(): def test_as_list_token_budget_fits_exactly(): # Verify > vs >= boundary: an item whose cost equals the remaining budget # is INCLUDED (condition is `spent + cost > effective_budget`, so equality passes). - # context_length=20 → effective = int(20 * 0.55) = 11. + # context_length=20 → effective = int(20 * 0.75) = 15. # Items of 8 chars → cost = max(1, 8 // 4) = 2 tokens each. - # Five such items = 10 tokens ≤ 11 (all included); 6th = 12 > 11 (excluded). + # Seven items = 14 tokens ≤ 15 (all included); 8th = 16 > 15 (excluded). tiny = ModelIdentifier(hf_model_name="org/exact-fit-model", context_length=20) ctx = ChatContext(model_id=tiny) - for _ in range(6): + for _ in range(8): ctx = ctx.add(CBlock("abcdefgh")) # 8 chars → cost 2 result = ctx.view_for_generation() - assert len(result) == 5 + assert len(result) == 7 def test_build_table_raises_on_collision(monkeypatch): From dd0c8675c649a74413b5874ace1717081f387432 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Tue, 23 Jun 2026 14:28:10 -0400 Subject: [PATCH 10/11] allow arbitrary token context length limit Signed-off-by: AngeloDanducci --- mellea/stdlib/context.py | 100 ++++++++++++++++++------------- test/stdlib/test_base_context.py | 54 ++++++++++++++++- 2 files changed, 112 insertions(+), 42 deletions(-) diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index f8891e8d8e..941169b05a 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -24,14 +24,18 @@ class ChatContext(Context): Args: window_size (int | None): Maximum number of context turns to include when calling `view_for_generation`. `None` (the default) means the full - history is returned, unless a `model_id` is bound and has a known - context length — in that case the model's context length (in tokens) - is used as an upper bound on the number of items returned. + history is returned, unless a `token_context_length_limit` or + `model_id` is set. Explicit `window_size` always takes priority over + token-based limits. + token_context_length_limit (int | None): Explicit token budget cap for + `view_for_generation`. Overrides the model-derived limit when set, + but is itself overridden by `window_size`. Useful for tightening (or + widening) the effective context beyond what the model table reports. model_id (str | ModelIdentifier | None): Optional model identifier used - for automatic context-window sizing. When set and `window_size` is - `None`, `view_for_generation` looks up the model's known context - length and uses it as the window size. Explicit `window_size` always - takes priority. + for automatic context-window sizing. When set and neither + `window_size` nor `token_context_length_limit` are provided, + `view_for_generation` looks up the model's known context length and + uses it as the token budget. Class Attributes: _propagated_fields: Instance-attribute names copied by `add()` and @@ -41,35 +45,42 @@ class ChatContext(Context): Note: `context_length` is measured in tokens; `window_size` counts context - items (CBlocks / Components). When both `window_size` and `model_id` - are set, `window_size` always takes priority — the token-budget path is - skipped entirely. When only `model_id` is bound, `view_for_generation` - estimates per-item token counts (via ``len(rendered) // 4``) and walks - history newest-first, dropping the oldest items until the running sum - fits within `context_length`. Set `window_size` explicitly to enforce - an item-count limit instead of a token budget. - - Per-item token count is estimated as ``len(rendered) // 4`` where - ``rendered`` is the string produced by `TemplateFormatter` (the same - renderer used at generation time), so the estimate reflects actual prompt - content. A 0.75 headroom factor (retaining 75 % of the model's rated - context length) is applied to leave capacity for the system prompt, - injected tool schemas, the current action, and the model's response — - none of which are tracked here. Use ``window_size`` for precise - item-count control. + items (CBlocks / Components). Resolution priority in + `view_for_generation`: + + 1. `window_size` — item-count limit, always wins. + 2. `token_context_length_limit` — explicit token cap, overrides the + model table. + 3. Model-derived limit via `model_id` and `get_context_length`. + 4. No limit — full history. + + Per-item token count is estimated as `len(rendered) // 4` (1 token ≈ + 4 characters) where `rendered` is the string produced by + `TemplateFormatter` (the same renderer used at generation time), so the + estimate reflects actual prompt content. A 0.75 headroom factor + (retaining 75 % of the token budget) is applied to leave capacity for + the system prompt, injected tool schemas, the current action, and the + model's response — none of which are tracked here. Use `window_size` + for precise item-count control. """ - _propagated_fields: tuple[str, ...] = ("_window_size", "_model_id") + _propagated_fields: tuple[str, ...] = ( + "_window_size", + "_token_context_length_limit", + "_model_id", + ) def __init__( self, *, window_size: int | None = None, + token_context_length_limit: int | None = None, model_id: str | ModelIdentifier | None = None, ): - """Initialize ChatContext with an optional sliding-window size and model binding.""" + """Initialize ChatContext with an optional sliding-window size, token budget, and model binding.""" super().__init__() self._window_size = window_size + self._token_context_length_limit = token_context_length_limit self._model_id: str | ModelIdentifier | None = model_id @property @@ -117,15 +128,16 @@ def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext: return self._make_root(model_id) def new_instance(self) -> ChatContext: - """Return a new empty root `ChatContext`, preserving `window_size` and `model_id`. + """Return a new empty root `ChatContext`, preserving `window_size`, `token_context_length_limit`, and `model_id`. Use this instead of `reset_to_new()` when you need to preserve the model binding and window size from an existing instance. `reset_to_new()` is a classmethod that returns a bare `ChatContext()` with no configuration. Returns: - ChatContext: A fresh root context with the same `window_size` - and `model_id` as this instance, but no history. + ChatContext: A fresh root context with the same `window_size`, + `token_context_length_limit`, and `model_id` as this instance, but + no history. """ return self._make_root(self._model_id) @@ -137,7 +149,7 @@ def add(self, c: Component | CBlock) -> ChatContext: Returns: ChatContext: A new `ChatContext` with the added entry, preserving - both `window_size` and any `model_id` binding. + `window_size`, `token_context_length_limit`, and any `model_id` binding. """ new = ChatContext.from_previous(self, c) for field in self._propagated_fields: @@ -149,11 +161,13 @@ def view_for_generation(self) -> list[Component | CBlock] | None: Window size resolution priority: - 1. Explicit `window_size` passed at construction — item-count limit, always wins. - 2. Model-derived context length looked up via `model_id` — token-budget truncation. + 1. Explicit `window_size` — item-count limit, always wins. + 2. Explicit `token_context_length_limit` — token cap, overrides model table. + 3. Model-derived context length looked up via `model_id` — token-budget truncation. Items are added newest-first until adding the next item would exceed the budget. - Token count is estimated as ``len(rendered) // 4`` via `TemplateFormatter`. - 3. No limit — return the full history. + Token count is estimated as `len(rendered) // 4` (1 token ≈ 4 characters) + via `TemplateFormatter`. + 4. No limit — return the full history. Returns: list[Component | CBlock] | None: Ordered list of context entries up to @@ -162,6 +176,9 @@ def view_for_generation(self) -> list[Component | CBlock] | None: if self._window_size is not None: return self.as_list(self._window_size) + if self._token_context_length_limit is not None: + return self._as_list_token_budget(self._token_context_length_limit) + if self._model_id is not None: token_budget = get_context_length(self._model_id) if token_budget is not None: @@ -176,18 +193,20 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: adding the next item would exceed the budget. The returned list is in chronological order (oldest-first), matching `as_list` behaviour. - Per-item token count is estimated as ``len(rendered) // 4`` where - ``rendered`` is the string produced by `TemplateFormatter` for the - bound model — the same renderer used at generation time. A 0.75 + Per-item token count is estimated as `len(rendered) // 4` (1 token ≈ + 4 characters) where `rendered` is the string produced by + `TemplateFormatter` for the bound model — the same renderer used at + generation time. A 0.75 headroom factor (retaining 75 % of the rated context length) reserves capacity for the system prompt, injected tool schemas, the current - action, and the model's response. Use ``window_size`` for precise + action, and the model's response. Use `window_size` for precise item-count control. """ - assert self._model_id is not None from ..formatters import TemplateFormatter # deferred to avoid circular import - formatter = TemplateFormatter(self._model_id) + formatter = ( + TemplateFormatter(self._model_id) if self._model_id is not None else None + ) effective_budget = int(token_budget * 0.75) collected: list[Component | CBlock] = [] spent = 0 @@ -196,7 +215,8 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: while not current.is_root_node: item = current.node_data assert item is not None - cost = max(1, len(formatter.print(item)) // 4) + rendered = formatter.print(item) if formatter is not None else str(item) + cost = max(1, len(rendered) // 4) total += 1 if spent + cost > effective_budget: break diff --git a/test/stdlib/test_base_context.py b/test/stdlib/test_base_context.py index 2e21e7b02c..bf97ec4635 100644 --- a/test/stdlib/test_base_context.py +++ b/test/stdlib/test_base_context.py @@ -150,9 +150,10 @@ def test__bind_model_does_not_mutate_original(): def test__bind_model_preserves_window_size(): - ctx = ChatContext(window_size=5) + ctx = ChatContext(window_size=5, token_context_length_limit=200) bound = ctx._bind_model(IBM_GRANITE_4_1_8B) assert bound._window_size == 5 + assert bound._token_context_length_limit == 200 def test__bind_model_raises_on_non_root_context(): @@ -248,6 +249,52 @@ def test_view_for_generation_unknown_model_returns_full_history(): assert len(result) == 5 +def test_view_for_generation_token_context_length_limit_truncates(): + # Each item is a 100-char string → cost = 101 // 4 = 25 tokens. + # token_context_length_limit=130: effective budget = int(130 * 0.75) = 97. + # Fits 3 items (75 tokens); the 4th would push spent to 100 > 97. + ctx = ChatContext(token_context_length_limit=130) + for i in range(7): + ctx = ctx.add(CBlock("x" * 100 + str(i))) + result = ctx.view_for_generation() + assert len(result) == 3 + assert str(result[-1]).endswith("6") + + +def test_view_for_generation_token_context_length_limit_beats_model(): + # Explicit token limit should override the model-derived one. + # Model has a large context; token_context_length_limit=130 forces truncation. + ctx = ChatContext(token_context_length_limit=130, model_id=IBM_GRANITE_4_1_8B) + for i in range(7): + ctx = ctx.add(CBlock("x" * 100 + str(i))) + result = ctx.view_for_generation() + assert len(result) == 3 + + +def test_view_for_generation_window_size_beats_token_context_length_limit(): + # window_size wins over token_context_length_limit. + ctx = ChatContext(window_size=2, token_context_length_limit=130) + for i in range(7): + ctx = ctx.add(CBlock("x" * 100 + str(i))) + result = ctx.view_for_generation() + assert len(result) == 2 + + +def test_token_context_length_limit_propagates_through_add(): + ctx = ChatContext(token_context_length_limit=130) + ctx = ctx.add(CBlock("hello")) + assert ctx._token_context_length_limit == 130 + + +def test_new_instance_preserves_token_context_length_limit(): + ctx = ChatContext(token_context_length_limit=130) + for i in range(3): + ctx = ctx.add(CBlock(f"msg {i}")) + new_ctx = ctx.new_instance() + assert new_ctx._token_context_length_limit == 130 + assert new_ctx.is_root_node + + # --------------------------------------------------------------------------- # reset_to_new and session.reset() # --------------------------------------------------------------------------- @@ -273,12 +320,15 @@ def test_reset_to_new_classmethod_does_not_preserve_config(): def test_new_instance_preserves_model_id_and_window_size(): - ctx = ChatContext(model_id=IBM_GRANITE_4_1_8B, window_size=5) + ctx = ChatContext( + model_id=IBM_GRANITE_4_1_8B, window_size=5, token_context_length_limit=200 + ) for i in range(3): ctx = ctx.add(CBlock(f"msg {i}")) new_ctx = ctx.new_instance() assert new_ctx.model_id is IBM_GRANITE_4_1_8B assert new_ctx._window_size == 5 + assert new_ctx._token_context_length_limit == 200 def test_session_reset_rebinds_model_id(): From d01be4bccb3db699ad7868ccae00f38aa421882b Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Wed, 24 Jun 2026 11:00:24 -0400 Subject: [PATCH 11/11] review feedback Signed-off-by: AngeloDanducci --- mellea/stdlib/context.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/mellea/stdlib/context.py b/mellea/stdlib/context.py index 941169b05a..4759305b77 100644 --- a/mellea/stdlib/context.py +++ b/mellea/stdlib/context.py @@ -210,22 +210,31 @@ def _as_list_token_budget(self, token_budget: int) -> list[Component | CBlock]: effective_budget = int(token_budget * 0.75) collected: list[Component | CBlock] = [] spent = 0 - total = 0 + chain_length = 0 + node: Context = self + while not node.is_root_node: + chain_length += 1 + node = node.previous_node # type: ignore[assignment] current: Context = self while not current.is_root_node: item = current.node_data - assert item is not None + if item is None: # pragma: no cover + raise RuntimeError( + "Malformed context chain: node_data is None at a non-root node" + ) rendered = formatter.print(item) if formatter is not None else str(item) cost = max(1, len(rendered) // 4) - total += 1 if spent + cost > effective_budget: break collected.append(item) spent += cost prev = current.previous_node - assert prev is not None + if prev is None: # pragma: no cover + raise RuntimeError( + "Malformed context chain: previous_node is None at a non-root node" + ) current = prev - dropped = total - len(collected) + dropped = chain_length - len(collected) if dropped: logger.debug( "Context truncated: dropped %d item(s) to stay within %d-token budget "