wehub-resource-sync
810 lines
28 KiB
Python
810 lines
28 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for BatchedEngine and BaseEngine modules.
|
|
|
|
Tests cover:
|
|
- GenerationOutput: dataclass behavior
|
|
- BaseEngine ABC: interface verification
|
|
- BatchedEngine initialization
|
|
- _apply_chat_template(): chat template application
|
|
- _preprocess_messages(): Harmony preprocessing
|
|
- get_stats(), get_cache_stats()
|
|
|
|
Note: mlx_lm.load() is mocked to avoid loading real models.
|
|
"""
|
|
|
|
from abc import ABC
|
|
from types import SimpleNamespace
|
|
from typing import Any, Dict, List, Optional
|
|
from unittest.mock import MagicMock, patch, AsyncMock
|
|
|
|
import pytest
|
|
|
|
from omlx.engine.base import BaseEngine, BaseNonStreamingEngine, GenerationOutput
|
|
|
|
|
|
class FakeStreamingCore:
|
|
"""Minimal async engine core for stream cleanup tests."""
|
|
|
|
def __init__(self):
|
|
self.aborted_request_id = None
|
|
|
|
async def add_request(self, **kwargs):
|
|
return "request-1"
|
|
|
|
async def stream_outputs(self, request_id):
|
|
yield SimpleNamespace(
|
|
output_text="partial",
|
|
new_text="partial",
|
|
prompt_tokens=1,
|
|
completion_tokens=1,
|
|
finished=False,
|
|
finish_reason=None,
|
|
tool_calls=None,
|
|
cached_tokens=0,
|
|
)
|
|
|
|
async def abort_request(self, request_id):
|
|
self.aborted_request_id = request_id
|
|
|
|
|
|
class TestGenerationOutput:
|
|
"""Tests for GenerationOutput dataclass."""
|
|
|
|
def test_default_values(self):
|
|
"""Test GenerationOutput has correct defaults."""
|
|
output = GenerationOutput(text="Hello, world!")
|
|
|
|
assert output.text == "Hello, world!"
|
|
assert output.tokens == []
|
|
assert output.prompt_tokens == 0
|
|
assert output.completion_tokens == 0
|
|
assert output.finish_reason == "stop"
|
|
assert output.new_text == ""
|
|
assert output.finished is True
|
|
assert output.tool_calls is None
|
|
|
|
def test_custom_values(self):
|
|
"""Test GenerationOutput with custom values."""
|
|
output = GenerationOutput(
|
|
text="Generated text",
|
|
tokens=[100, 101, 102],
|
|
prompt_tokens=10,
|
|
completion_tokens=3,
|
|
finish_reason="length",
|
|
new_text="partial",
|
|
finished=False,
|
|
tool_calls=[{"name": "test_tool", "arguments": "{}"}],
|
|
)
|
|
|
|
assert output.text == "Generated text"
|
|
assert output.tokens == [100, 101, 102]
|
|
assert output.prompt_tokens == 10
|
|
assert output.completion_tokens == 3
|
|
assert output.finish_reason == "length"
|
|
assert output.new_text == "partial"
|
|
assert output.finished is False
|
|
assert output.tool_calls == [{"name": "test_tool", "arguments": "{}"}]
|
|
|
|
def test_streaming_output(self):
|
|
"""Test GenerationOutput for streaming use case."""
|
|
output = GenerationOutput(
|
|
text="",
|
|
new_text="Hello",
|
|
prompt_tokens=5,
|
|
completion_tokens=1,
|
|
finished=False,
|
|
finish_reason=None,
|
|
)
|
|
|
|
assert output.text == ""
|
|
assert output.new_text == "Hello"
|
|
assert output.finished is False
|
|
assert output.finish_reason is None
|
|
|
|
|
|
class TestBaseEngine:
|
|
"""Tests for BaseEngine ABC."""
|
|
|
|
def test_is_abstract(self):
|
|
"""Test BaseEngine is abstract and cannot be instantiated."""
|
|
with pytest.raises(TypeError):
|
|
BaseEngine()
|
|
|
|
def test_abstract_methods(self):
|
|
"""Test BaseEngine defines required abstract methods."""
|
|
abstract_methods = {
|
|
"model_name",
|
|
"tokenizer",
|
|
"start",
|
|
"stop",
|
|
"generate",
|
|
"stream_generate",
|
|
"chat",
|
|
"stream_chat",
|
|
"model_type",
|
|
"get_stats",
|
|
"get_cache_stats",
|
|
}
|
|
|
|
# Check abstractmethods
|
|
actual_methods = set()
|
|
for name in dir(BaseEngine):
|
|
if not name.startswith("_"):
|
|
attr = getattr(BaseEngine, name)
|
|
if getattr(attr, "__isabstractmethod__", False):
|
|
actual_methods.add(name)
|
|
|
|
assert actual_methods == abstract_methods
|
|
|
|
def test_concrete_implementation(self):
|
|
"""Test a concrete implementation can be created."""
|
|
|
|
class ConcreteEngine(BaseEngine):
|
|
@property
|
|
def model_name(self) -> str:
|
|
return "test-model"
|
|
|
|
@property
|
|
def tokenizer(self) -> Any:
|
|
return None
|
|
|
|
async def start(self) -> None:
|
|
pass
|
|
|
|
async def stop(self) -> None:
|
|
pass
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
stop: Optional[List[str]] = None,
|
|
**kwargs,
|
|
) -> GenerationOutput:
|
|
return GenerationOutput(text="test")
|
|
|
|
async def stream_generate(
|
|
self,
|
|
prompt: str,
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
stop: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
yield GenerationOutput(text="test")
|
|
|
|
async def chat(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
tools: Optional[List[dict]] = None,
|
|
**kwargs,
|
|
) -> GenerationOutput:
|
|
return GenerationOutput(text="test")
|
|
|
|
async def stream_chat(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
tools: Optional[List[dict]] = None,
|
|
**kwargs,
|
|
):
|
|
yield GenerationOutput(text="test")
|
|
|
|
@property
|
|
def model_type(self) -> Optional[str]:
|
|
return "test"
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
return {}
|
|
|
|
def get_cache_stats(self) -> Optional[Dict[str, Any]]:
|
|
return None
|
|
|
|
engine = ConcreteEngine()
|
|
assert engine.model_name == "test-model"
|
|
|
|
|
|
class TestBaseNonStreamingEngine:
|
|
"""Tests for BaseNonStreamingEngine ABC."""
|
|
|
|
def test_is_abstract(self):
|
|
"""Test BaseNonStreamingEngine is abstract."""
|
|
with pytest.raises(TypeError):
|
|
BaseNonStreamingEngine()
|
|
|
|
def test_abstract_methods(self):
|
|
"""Test BaseNonStreamingEngine defines required abstract methods."""
|
|
abstract_methods = {"model_name", "start", "stop", "get_stats"}
|
|
|
|
actual_methods = set()
|
|
for name in dir(BaseNonStreamingEngine):
|
|
if not name.startswith("_"):
|
|
attr = getattr(BaseNonStreamingEngine, name)
|
|
if getattr(attr, "__isabstractmethod__", False):
|
|
actual_methods.add(name)
|
|
|
|
assert actual_methods == abstract_methods
|
|
|
|
|
|
class TestBatchedEngineInitialization:
|
|
"""Tests for BatchedEngine initialization."""
|
|
|
|
def test_init_stores_parameters(self):
|
|
"""Test BatchedEngine stores initialization parameters."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(
|
|
model_name="test-model",
|
|
trust_remote_code=False,
|
|
stream_interval=5,
|
|
enable_thinking=True,
|
|
)
|
|
|
|
assert engine._model_name == "test-model"
|
|
assert engine._trust_remote_code is False
|
|
assert engine._stream_interval == 5
|
|
assert engine._enable_thinking is True
|
|
assert engine._loaded is False
|
|
assert engine._model is None
|
|
assert engine._tokenizer is None
|
|
assert engine._engine is None
|
|
|
|
def test_init_default_values(self):
|
|
"""Test BatchedEngine default values."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
# Issue #926: default flipped to False so HF repos can't auto-execute custom Python.
|
|
assert engine._trust_remote_code is False
|
|
assert engine._scheduler_config is None
|
|
assert engine._stream_interval == 1
|
|
assert engine._enable_thinking is None
|
|
|
|
def test_model_name_property(self):
|
|
"""Test model_name property."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="my-model")
|
|
|
|
assert engine.model_name == "my-model"
|
|
|
|
def test_tokenizer_property_before_load(self):
|
|
"""Test tokenizer property returns None before loading."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
assert engine.tokenizer is None
|
|
|
|
def test_model_type_property_before_load(self):
|
|
"""Test model_type property returns None before loading."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
assert engine.model_type is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_clears_wrapper_teardown_references(self):
|
|
"""stop() releases wrapper-side native helper references."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
inner_engine = MagicMock()
|
|
engine._engine = MagicMock()
|
|
engine._engine.stop = AsyncMock()
|
|
engine._engine.engine = inner_engine
|
|
engine._model = object()
|
|
engine._tokenizer = object()
|
|
engine._grammar_compiler = object()
|
|
engine._grammar_compiler_init_attempted = True
|
|
engine._loaded = True
|
|
|
|
await engine.stop()
|
|
|
|
assert engine._engine is None
|
|
assert engine._model is None
|
|
assert engine._tokenizer is None
|
|
assert engine._grammar_compiler is None
|
|
assert engine._grammar_compiler_init_attempted is False
|
|
assert engine._loaded is False
|
|
inner_engine.close.assert_called_once()
|
|
|
|
|
|
class TestBatchedEngineStreamingCleanup:
|
|
"""Tests for streaming generator cleanup paths."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_abort_uses_captured_engine_if_engine_cleared(self):
|
|
"""Generator finalization aborts on the original engine reference."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
fake_engine = FakeStreamingCore()
|
|
engine = BatchedEngine(model_name="test-model")
|
|
engine._loaded = True
|
|
engine._engine = fake_engine
|
|
|
|
stream = engine.stream_generate("hello")
|
|
first = await stream.__anext__()
|
|
assert first.text == "partial"
|
|
|
|
engine._engine = None
|
|
await stream.aclose()
|
|
|
|
assert fake_engine.aborted_request_id == "request-1"
|
|
|
|
|
|
class TestBatchedEngineApplyChatTemplate:
|
|
"""Tests for BatchedEngine._apply_chat_template()."""
|
|
|
|
def test_apply_chat_template_with_tokenizer(self):
|
|
"""Test _apply_chat_template when tokenizer has apply_chat_template."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
# Mock tokenizer with apply_chat_template
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "<formatted prompt>"
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
|
|
result = engine._apply_chat_template(messages)
|
|
|
|
assert result == "<formatted prompt>"
|
|
mock_tokenizer.apply_chat_template.assert_called_once()
|
|
|
|
def test_apply_chat_template_with_tools(self):
|
|
"""Test _apply_chat_template passes tools to tokenizer."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "<formatted>"
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
engine._apply_chat_template(messages, tools=tools)
|
|
|
|
# Verify tools were passed
|
|
call_kwargs = mock_tokenizer.apply_chat_template.call_args[1]
|
|
assert "tools" in call_kwargs
|
|
assert call_kwargs["tools"] == tools
|
|
|
|
def test_apply_chat_template_with_enable_thinking(self):
|
|
"""Test _apply_chat_template passes enable_thinking."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model", enable_thinking=True)
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "<formatted>"
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
|
|
engine._apply_chat_template(messages)
|
|
|
|
call_kwargs = mock_tokenizer.apply_chat_template.call_args[1]
|
|
assert call_kwargs.get("enable_thinking") is True
|
|
|
|
def test_apply_chat_template_fallback(self):
|
|
"""Test _apply_chat_template fallback when tokenizer lacks method."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
# Tokenizer without apply_chat_template
|
|
mock_tokenizer = MagicMock(spec=[])
|
|
del mock_tokenizer.apply_chat_template # Explicitly remove
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi"},
|
|
]
|
|
|
|
result = engine._apply_chat_template(messages)
|
|
|
|
assert "user: Hello" in result
|
|
assert "assistant: Hi" in result
|
|
assert result.endswith("assistant:")
|
|
|
|
def test_apply_chat_template_handles_type_error(self):
|
|
"""Test _apply_chat_template handles TypeError from tokenizer."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model", enable_thinking=True)
|
|
|
|
# Tokenizer that raises TypeError for enable_thinking
|
|
mock_tokenizer = MagicMock()
|
|
call_count = 0
|
|
|
|
def side_effect(*args, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 1 and "enable_thinking" in kwargs:
|
|
raise TypeError("enable_thinking not supported")
|
|
return "<formatted>"
|
|
|
|
mock_tokenizer.apply_chat_template.side_effect = side_effect
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
|
|
result = engine._apply_chat_template(messages)
|
|
|
|
assert result == "<formatted>"
|
|
assert call_count == 2 # Called twice (first fails, second succeeds)
|
|
|
|
|
|
class TestBatchedEnginePreprocessMessages:
|
|
"""Tests for BatchedEngine._preprocess_messages()."""
|
|
|
|
def test_preprocess_messages_non_harmony(self):
|
|
"""Test _preprocess_messages returns unchanged for non-Harmony models."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
engine._model = MagicMock()
|
|
engine._model.config = MagicMock()
|
|
engine._model.config.model_type = "llama" # Not gpt_oss
|
|
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
|
|
result = engine._preprocess_messages(messages)
|
|
|
|
assert result == messages
|
|
|
|
def test_preprocess_messages_model_type_none(self):
|
|
"""Test _preprocess_messages when model_type is None."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
# No model loaded
|
|
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
|
|
result = engine._preprocess_messages(messages)
|
|
|
|
assert result == messages
|
|
|
|
|
|
class TestBatchedEngineStats:
|
|
"""Tests for BatchedEngine statistics methods."""
|
|
|
|
def test_get_stats_before_load(self):
|
|
"""Test get_stats() before model is loaded."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model", stream_interval=3)
|
|
|
|
stats = engine.get_stats()
|
|
|
|
assert stats["engine_type"] == "batched"
|
|
assert stats["model_name"] == "test-model"
|
|
assert stats["loaded"] is False
|
|
assert stats["stream_interval"] == 3
|
|
|
|
def test_get_stats_includes_engine_stats(self):
|
|
"""Test get_stats() includes engine stats when loaded."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
# Mock loaded engine
|
|
mock_engine = MagicMock()
|
|
mock_engine.get_stats.return_value = {
|
|
"running": True,
|
|
"steps_executed": 100,
|
|
}
|
|
engine._engine = mock_engine
|
|
engine._loaded = True
|
|
|
|
stats = engine.get_stats()
|
|
|
|
assert stats["loaded"] is True
|
|
assert stats["running"] is True
|
|
assert stats["steps_executed"] == 100
|
|
|
|
def test_get_cache_stats_before_load(self):
|
|
"""Test get_cache_stats() before model is loaded."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
stats = engine.get_cache_stats()
|
|
|
|
assert stats is None
|
|
|
|
def test_get_cache_stats_after_load(self):
|
|
"""Test get_cache_stats() when engine is loaded."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_engine = MagicMock()
|
|
mock_engine.get_cache_stats.return_value = {"entries": 10}
|
|
engine._engine = mock_engine
|
|
|
|
stats = engine.get_cache_stats()
|
|
|
|
assert stats == {"entries": 10}
|
|
|
|
|
|
class TestBatchedEngineModelType:
|
|
"""Tests for BatchedEngine.model_type property."""
|
|
|
|
def test_model_type_from_config(self):
|
|
"""Test model_type from model.config."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.config = MagicMock()
|
|
mock_model.config.model_type = "llama"
|
|
engine._model = mock_model
|
|
|
|
assert engine.model_type == "llama"
|
|
|
|
def test_model_type_from_config_dict(self):
|
|
"""Test model_type from dict-style config."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
# Create a mock model where config is a dict
|
|
mock_model = MagicMock(spec=["config"])
|
|
# Use a real dict for config
|
|
mock_model.config = {"model_type": "qwen2"}
|
|
engine._model = mock_model
|
|
|
|
# The code checks hasattr(config, 'model_type') first,
|
|
# dicts don't have model_type as attribute, so it checks isinstance(config, dict)
|
|
assert engine.model_type == "qwen2"
|
|
|
|
def test_model_type_from_args(self):
|
|
"""Test model_type from model.args."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_model = MagicMock(spec=["args"])
|
|
mock_model.args = MagicMock()
|
|
mock_model.args.model_type = "gpt_oss"
|
|
engine._model = mock_model
|
|
|
|
assert engine.model_type == "gpt_oss"
|
|
|
|
def test_model_type_none_when_not_available(self):
|
|
"""Test model_type returns None when not available."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_model = MagicMock(spec=[])
|
|
engine._model = mock_model
|
|
|
|
assert engine.model_type is None
|
|
|
|
|
|
class TestApplyChatTemplatePartialMode:
|
|
"""Tests for partial mode support in _apply_chat_template()."""
|
|
|
|
def test_partial_mode_sets_continue_final_message(self):
|
|
"""Final assistant message with partial=True sets continue_final_message."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "<formatted>"
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [
|
|
{"role": "user", "content": "Generate JSON"},
|
|
{"role": "assistant", "content": "{", "partial": True},
|
|
]
|
|
|
|
engine._apply_chat_template(messages)
|
|
|
|
call_kwargs = mock_tokenizer.apply_chat_template.call_args[1]
|
|
assert call_kwargs["add_generation_prompt"] is False
|
|
assert call_kwargs["continue_final_message"] is True
|
|
|
|
def test_partial_non_assistant_ignored(self):
|
|
"""partial=True on a non-assistant message does not trigger partial mode."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "<formatted>"
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [
|
|
{"role": "user", "content": "Hello", "partial": True},
|
|
]
|
|
|
|
engine._apply_chat_template(messages)
|
|
|
|
call_kwargs = mock_tokenizer.apply_chat_template.call_args[1]
|
|
assert call_kwargs["add_generation_prompt"] is True
|
|
assert "continue_final_message" not in call_kwargs
|
|
|
|
def test_partial_field_stripped_before_template(self):
|
|
"""partial field is removed from messages before calling apply_chat_template."""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "<formatted>"
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [
|
|
{"role": "user", "content": "Hello", "partial": False},
|
|
{"role": "assistant", "content": "{", "partial": True},
|
|
]
|
|
|
|
engine._apply_chat_template(messages)
|
|
|
|
# Check the messages passed to apply_chat_template
|
|
call_args = mock_tokenizer.apply_chat_template.call_args[0][0]
|
|
for msg in call_args:
|
|
assert "partial" not in msg
|
|
|
|
def test_partial_true_continues_vs_new_turn(self):
|
|
"""Verify the core partial toggle: partial=True continues, absent starts new turn.
|
|
|
|
With partial=True on the final assistant message, the engine must pass
|
|
add_generation_prompt=False and continue_final_message=True so the
|
|
model continues from the assistant's content rather than starting a
|
|
new turn. Without partial, the default add_generation_prompt=True
|
|
appends the generation prompt (e.g. <|im_start|>assistant) for a
|
|
fresh response.
|
|
|
|
NOTE: The `name` field is not tested here because its rendering is
|
|
model-template-specific — many templates silently ignore it, so
|
|
assertions on template output would be fragile and model-dependent.
|
|
"""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_tokenizer = MagicMock()
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
base_messages = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Return JSON with keys: name, age"},
|
|
]
|
|
|
|
# --- partial=True: continue from prefill ---
|
|
mock_tokenizer.apply_chat_template.return_value = "...assistant\n{"
|
|
partial_messages = base_messages + [
|
|
{"role": "assistant", "content": "{", "partial": True},
|
|
]
|
|
engine._apply_chat_template(partial_messages)
|
|
|
|
partial_kwargs = mock_tokenizer.apply_chat_template.call_args[1]
|
|
assert partial_kwargs["add_generation_prompt"] is False
|
|
assert partial_kwargs["continue_final_message"] is True
|
|
|
|
# --- no partial: new turn ---
|
|
mock_tokenizer.apply_chat_template.reset_mock()
|
|
mock_tokenizer.apply_chat_template.return_value = "...<|im_start|>assistant\n"
|
|
normal_messages = list(base_messages) # no trailing assistant
|
|
engine._apply_chat_template(normal_messages)
|
|
|
|
normal_kwargs = mock_tokenizer.apply_chat_template.call_args[1]
|
|
assert normal_kwargs["add_generation_prompt"] is True
|
|
assert "continue_final_message" not in normal_kwargs
|
|
|
|
def test_partial_with_streaming(self):
|
|
"""partial mode kwargs are the same regardless of downstream streaming.
|
|
|
|
The engine's _apply_chat_template is called identically for streaming
|
|
and non-streaming — the partial toggle affects template kwargs only,
|
|
not the generation path. This test confirms the kwargs are set
|
|
correctly when the messages would be used in a streaming context.
|
|
"""
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "...1."
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [
|
|
{"role": "user", "content": "List 3 colors"},
|
|
{"role": "assistant", "content": "1.", "partial": True},
|
|
]
|
|
|
|
engine._apply_chat_template(messages)
|
|
|
|
call_kwargs = mock_tokenizer.apply_chat_template.call_args[1]
|
|
assert call_kwargs["add_generation_prompt"] is False
|
|
assert call_kwargs["continue_final_message"] is True
|
|
# partial stripped from message dicts
|
|
call_msgs = mock_tokenizer.apply_chat_template.call_args[0][0]
|
|
assert "partial" not in call_msgs[-1]
|
|
|
|
def test_count_then_apply_chat_template_idempotent_under_partial_mode(self):
|
|
"""Server flow: count_chat_tokens then _apply_chat_template on the
|
|
same messages list must render with identical partial-mode flags.
|
|
|
|
Mimics the post-fix server contract: detect_and_strip_partial once
|
|
at the API boundary, forward the resolved value to engine methods
|
|
via an explicit is_partial parameter, and assert both phases pass
|
|
the same partial-mode flags to apply_chat_template.
|
|
"""
|
|
from omlx.api.utils import detect_and_strip_partial
|
|
from omlx.engine.batched import BatchedEngine
|
|
|
|
engine = BatchedEngine(model_name="test-model")
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.apply_chat_template.return_value = "<formatted>"
|
|
mock_tokenizer.encode.return_value = [1, 2, 3]
|
|
engine._tokenizer = mock_tokenizer
|
|
|
|
messages = [
|
|
{"role": "user", "content": "Generate JSON"},
|
|
{"role": "assistant", "content": "{", "partial": True},
|
|
]
|
|
|
|
# Server flow: detect_and_strip_partial once at the API boundary,
|
|
# forward the resolved value to all engine methods.
|
|
is_partial = detect_and_strip_partial(messages)
|
|
assert is_partial is True
|
|
|
|
# Phase 1: count.
|
|
engine.count_chat_tokens(messages, is_partial=is_partial)
|
|
count_kwargs = dict(mock_tokenizer.apply_chat_template.call_args.kwargs)
|
|
|
|
# Phase 2: chat. Operates on the same (now-stripped) messages list.
|
|
engine._apply_chat_template(messages, is_partial=is_partial)
|
|
chat_kwargs = dict(mock_tokenizer.apply_chat_template.call_args.kwargs)
|
|
|
|
# Both phases must render with identical partial-mode flags.
|
|
assert count_kwargs.get("continue_final_message") == chat_kwargs.get(
|
|
"continue_final_message"
|
|
), (
|
|
"continue_final_message diverged across phases: "
|
|
f"count={count_kwargs.get('continue_final_message')}, "
|
|
f"chat={chat_kwargs.get('continue_final_message')}"
|
|
)
|
|
assert (
|
|
count_kwargs["add_generation_prompt"]
|
|
== chat_kwargs["add_generation_prompt"]
|
|
), (
|
|
"add_generation_prompt diverged across phases: "
|
|
f"count={count_kwargs['add_generation_prompt']}, "
|
|
f"chat={chat_kwargs['add_generation_prompt']}"
|
|
)
|
|
|
|
# Specific contract: with partial=True forwarded, both phases use
|
|
# continue_final_message=True (not add_generation_prompt=True).
|
|
assert count_kwargs["continue_final_message"] is True
|
|
assert count_kwargs["add_generation_prompt"] is False
|