wehub-resource-sync
285 lines
11 KiB
Python
285 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for context window validation feature."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from omlx.model_settings import ModelSettings
|
|
|
|
|
|
class TestGetMaxContextWindow:
|
|
"""Tests for get_max_context_window() priority logic."""
|
|
|
|
def _make_server_state(
|
|
self, global_max_ctx=32768, policy_cap=None
|
|
):
|
|
"""Create a mock server state with given global
|
|
``max_context_window`` fallback and optional
|
|
``max_context_window_policy`` cap."""
|
|
from omlx.server import SamplingDefaults
|
|
|
|
state = MagicMock()
|
|
state.sampling = SamplingDefaults(
|
|
max_context_window=global_max_ctx,
|
|
max_context_window_policy=policy_cap,
|
|
)
|
|
state.settings_manager = None
|
|
# Discovery-tier (#1308) lookups are exercised in TestGetMaxContextWindow
|
|
# in test_server.py; nulling the pool here keeps these focused on the
|
|
# per-model-setting → global fallback path.
|
|
state.engine_pool = None
|
|
return state
|
|
|
|
def test_returns_global_default(self):
|
|
"""Test returns global default when no model settings."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=32768)
|
|
with patch("omlx.server._server_state", state):
|
|
result = get_max_context_window()
|
|
assert result == 32768
|
|
|
|
def test_model_setting_overrides_global(self):
|
|
"""Test model-specific setting takes priority over global."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=32768)
|
|
mock_manager = MagicMock()
|
|
mock_manager.get_settings_for_request.return_value = ModelSettings(
|
|
max_context_window=4096
|
|
)
|
|
state.settings_manager = mock_manager
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
result = get_max_context_window("test-model")
|
|
assert result == 4096
|
|
|
|
def test_falls_back_to_global_when_model_not_set(self):
|
|
"""Test falls back to global when model has no max_context_window."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=65536)
|
|
mock_manager = MagicMock()
|
|
mock_manager.get_settings_for_request.return_value = ModelSettings(
|
|
max_context_window=None
|
|
)
|
|
state.settings_manager = mock_manager
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
result = get_max_context_window("test-model")
|
|
assert result == 65536
|
|
|
|
def test_no_model_id_returns_global(self):
|
|
"""Test returns global when model_id is None."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=16384)
|
|
with patch("omlx.server._server_state", state):
|
|
result = get_max_context_window(None)
|
|
assert result == 16384
|
|
|
|
def _mount_native_and_policy(
|
|
self, native_ctx: int | None, policy_cap: int | None
|
|
):
|
|
"""Mount a server state with a model that has the given native
|
|
context length, the policy field set to ``policy_cap``, and no
|
|
per-model override."""
|
|
state = self._make_server_state(
|
|
global_max_ctx=32768, policy_cap=policy_cap
|
|
)
|
|
mock_manager = MagicMock()
|
|
mock_manager.get_settings_for_request.return_value = ModelSettings(
|
|
max_context_window=None
|
|
)
|
|
state.settings_manager = mock_manager
|
|
|
|
mock_pool = MagicMock()
|
|
mock_entry = MagicMock()
|
|
mock_entry.model_context_length = native_ctx
|
|
mock_pool.get_entry.return_value = mock_entry
|
|
state.engine_pool = mock_pool
|
|
return state
|
|
|
|
def test_policy_unset_native_wins_unchanged(self):
|
|
"""With ``max_context_window_policy`` unset, the model's
|
|
native context length is returned verbatim — existing
|
|
installs see no behavior change after this PR."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._mount_native_and_policy(
|
|
native_ctx=262_144, policy_cap=None
|
|
)
|
|
with patch("omlx.server._server_state", state):
|
|
assert get_max_context_window("big-model") == 262_144
|
|
|
|
def test_policy_set_clamps_native(self):
|
|
"""With ``max_context_window_policy=128_000`` and a model that
|
|
natively declares 256 K, the effective cap is the policy."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._mount_native_and_policy(
|
|
native_ctx=262_144, policy_cap=128_000
|
|
)
|
|
with patch("omlx.server._server_state", state):
|
|
assert get_max_context_window("big-model") == 128_000, (
|
|
"Policy of 128k must clamp a model that natively declares 256k"
|
|
)
|
|
|
|
def test_policy_set_native_below_policy_wins(self):
|
|
"""When the model's native length is already below the policy,
|
|
the native value wins — policy is a ceiling, not a floor."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._mount_native_and_policy(
|
|
native_ctx=32_768, policy_cap=128_000
|
|
)
|
|
with patch("omlx.server._server_state", state):
|
|
assert get_max_context_window("small-model") == 32_768
|
|
|
|
def test_per_model_override_escapes_policy(self):
|
|
"""A per-model override is the operator's explicit per-model
|
|
choice; the global policy cap does NOT clamp it. This is the
|
|
operator's escape hatch for individual models that should
|
|
exceed the policy."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
state = self._mount_native_and_policy(
|
|
native_ctx=100_000, policy_cap=64_000
|
|
)
|
|
# Add a per-model override above both native and policy
|
|
state.settings_manager = MagicMock()
|
|
state.settings_manager.get_settings_for_request.return_value = ModelSettings(
|
|
max_context_window=200_000
|
|
)
|
|
with patch("omlx.server._server_state", state):
|
|
assert get_max_context_window("override-model") == 200_000, (
|
|
"Per-model override must escape the policy clamp"
|
|
)
|
|
|
|
def test_policy_does_not_apply_to_fallback_path(self):
|
|
"""When the model has no discoverable native context AND no
|
|
per-model override, the fallback default applies — the policy
|
|
is documented as clamping the *native* path only. Existing
|
|
``settings.json`` files with the historical 32768 fallback
|
|
therefore keep working unchanged even when a policy is later
|
|
added to the install."""
|
|
from omlx.server import get_max_context_window
|
|
|
|
# native_ctx=None: model config doesn't expose a context length
|
|
state = self._mount_native_and_policy(
|
|
native_ctx=None, policy_cap=16_000
|
|
)
|
|
with patch("omlx.server._server_state", state):
|
|
# Fallback (32768) returned, not the policy (16_000).
|
|
assert get_max_context_window("no-native-model") == 32_768
|
|
|
|
|
|
class TestValidateContextWindow:
|
|
"""Tests for validate_context_window()."""
|
|
|
|
def _make_server_state(self, global_max_ctx=32768):
|
|
from omlx.server import SamplingDefaults
|
|
|
|
state = MagicMock()
|
|
state.sampling = SamplingDefaults(max_context_window=global_max_ctx)
|
|
state.settings_manager = None
|
|
return state
|
|
|
|
def test_passes_when_under_limit(self):
|
|
"""Test no exception when token count is under limit."""
|
|
from omlx.server import validate_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=1000)
|
|
with patch("omlx.server._server_state", state):
|
|
# Should not raise
|
|
validate_context_window(500)
|
|
|
|
def test_passes_at_exact_limit(self):
|
|
"""Test no exception when token count equals limit."""
|
|
from omlx.server import validate_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=1000)
|
|
with patch("omlx.server._server_state", state):
|
|
# Should not raise (equal is OK)
|
|
validate_context_window(1000)
|
|
|
|
def test_raises_when_over_limit(self):
|
|
"""Test HTTPException raised when token count exceeds limit."""
|
|
from omlx.server import validate_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=1000)
|
|
with patch("omlx.server._server_state", state):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
validate_context_window(1001)
|
|
assert exc_info.value.status_code == 400
|
|
assert "1001 tokens" in exc_info.value.detail
|
|
assert "1000 tokens" in exc_info.value.detail
|
|
|
|
def test_raises_with_model_specific_limit(self):
|
|
"""Test uses model-specific limit when available."""
|
|
from omlx.server import validate_context_window
|
|
|
|
state = self._make_server_state(global_max_ctx=32768)
|
|
mock_manager = MagicMock()
|
|
mock_manager.get_settings_for_request.return_value = ModelSettings(
|
|
max_context_window=100
|
|
)
|
|
state.settings_manager = mock_manager
|
|
|
|
with patch("omlx.server._server_state", state):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
validate_context_window(200, "test-model")
|
|
assert exc_info.value.status_code == 400
|
|
assert "200 tokens" in exc_info.value.detail
|
|
assert "100 tokens" in exc_info.value.detail
|
|
|
|
|
|
class TestCountChatTokens:
|
|
"""Tests for BatchedEngine.count_chat_tokens()."""
|
|
|
|
def test_count_chat_tokens(self):
|
|
"""Test token counting with mocked tokenizer."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine.__new__(BatchedEngine)
|
|
engine._loaded = True
|
|
|
|
# Mock tokenizer
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "formatted prompt"
|
|
mock_tokenizer.encode.return_value = [1, 2, 3, 4, 5]
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
# Mock model (not gpt_oss)
|
|
engine._model = MagicMock(spec=[])
|
|
engine._enable_thinking = None
|
|
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
count = engine.count_chat_tokens(messages)
|
|
|
|
assert count == 5
|
|
mock_tokenizer.apply_chat_template.assert_called_once()
|
|
mock_tokenizer.encode.assert_called_once_with("formatted prompt")
|
|
|
|
def test_count_chat_tokens_with_tools(self):
|
|
"""Test token counting includes tools in template."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine.__new__(BatchedEngine)
|
|
engine._loaded = True
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "prompt with tools"
|
|
mock_tokenizer.encode.return_value = [1, 2, 3, 4, 5, 6, 7]
|
|
engine._tokenizer = mock_tokenizer
|
|
engine._model = MagicMock(spec=[])
|
|
engine._enable_thinking = None
|
|
|
|
messages = [{"role": "user", "content": "Call a tool"}]
|
|
tools = [{"type": "function", "function": {"name": "test"}}]
|
|
count = engine.count_chat_tokens(messages, tools)
|
|
|
|
assert count == 7
|