# SPDX-License-Identifier: Apache-2.0 """ End-to-end streaming tests for oMLX server. Tests streaming response formats for OpenAI and Anthropic APIs using mock AsyncIterator without loading actual models. """ import json import pytest from dataclasses import dataclass, field from typing import Any, AsyncIterator, Dict, List, Optional from unittest.mock import AsyncMock, MagicMock, patch from omlx.engine.base import BaseEngine @dataclass class MockGenerationOutput: """Mock generation output for streaming tests.""" text: str = "" tokens: List[int] = field(default_factory=list) prompt_tokens: int = 10 completion_tokens: int = 0 finish_reason: Optional[str] = None new_text: str = "" finished: bool = False tool_calls: Optional[List[Dict[str, Any]]] = None cached_tokens: int = 0 class MockTokenizer: """Mock tokenizer for testing.""" def __init__(self): self.eos_token_id = 2 def encode(self, text: str) -> List[int]: return [100 + i for i, _ in enumerate(text.split())] def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str: return f"" def apply_chat_template( self, messages: List[Dict], tokenize: bool = False, **kwargs ) -> str: parts = [] for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") parts.append(f"{role}: {content}") return "\n".join(parts) class MockBaseEngine(BaseEngine): """Mock LLM engine with streaming support for testing.""" def __init__(self, model_name: str = "test-model"): self._model_name = model_name self._tokenizer = MockTokenizer() self._model_type = "llama" # Configurable streaming responses self._stream_outputs: List[MockGenerationOutput] = [] @property def model_name(self) -> str: return self._model_name @property def tokenizer(self): return self._tokenizer @property def model_type(self) -> Optional[str]: return self._model_type @property def prefix_cache_enabled(self) -> bool: return False async def start(self) -> None: pass async def stop(self) -> None: pass def set_stream_outputs(self, outputs: List[MockGenerationOutput]): """Set custom streaming outputs for testing.""" self._stream_outputs = outputs async def generate(self, prompt: str, **kwargs) -> MockGenerationOutput: return MockGenerationOutput( text="Generated response.", completion_tokens=5, finish_reason="stop", finished=True, ) async def stream_generate(self, prompt: str, **kwargs) -> AsyncIterator[MockGenerationOutput]: if self._stream_outputs: for output in self._stream_outputs: yield output else: yield MockGenerationOutput( text="Hello", new_text="Hello", completion_tokens=1, finished=False, ) yield MockGenerationOutput( text="Hello world", new_text=" world", completion_tokens=2, finished=True, finish_reason="stop", ) def count_chat_tokens( self, messages: List[Dict], tools=None, chat_template_kwargs=None, **kwargs ) -> int: prompt = self._tokenizer.apply_chat_template(messages, tokenize=False) return len(self._tokenizer.encode(prompt)) async def chat(self, messages: List[Dict], **kwargs) -> MockGenerationOutput: return MockGenerationOutput( text="Chat response.", completion_tokens=5, finish_reason="stop", finished=True, ) async def stream_chat(self, messages: List[Dict], **kwargs) -> AsyncIterator[MockGenerationOutput]: if self._stream_outputs: for output in self._stream_outputs: yield output else: yield MockGenerationOutput( text="Hi", new_text="Hi", completion_tokens=1, finished=False, ) yield MockGenerationOutput( text="Hi there", new_text=" there", completion_tokens=2, finished=True, finish_reason="stop", ) def get_stats(self) -> Dict[str, Any]: return {} def get_cache_stats(self) -> Optional[Dict[str, Any]]: return None class MockEnginePool: """Mock engine pool for testing.""" def __init__(self, engine: Optional[MockBaseEngine] = None): self._engine = engine or MockBaseEngine() self._models = [ {"id": "test-model", "loaded": True, "pinned": False, "size": 1000000} ] self._entries: Dict[str, Any] = {} @property def model_count(self) -> int: return len(self._models) @property def loaded_model_count(self) -> int: return 1 @property def max_model_memory(self) -> int: return 32 * 1024 * 1024 * 1024 @property def current_model_memory(self) -> int: return 1000000 def get_entry(self, model_id: str): return self._entries.get(model_id) def resolve_model_id(self, model_id_or_alias, settings_manager=None): return model_id_or_alias def get_model_ids(self) -> List[str]: return [m["id"] for m in self._models] def get_status(self) -> Dict[str, Any]: return {"models": self._models} async def get_engine(self, model_id: str, _lease: bool = False): return self._engine async def release_engine(self, model_id: str) -> None: return None def parse_sse_events(response_text: str) -> List[Dict]: """Parse SSE events from response text.""" events = [] for line in response_text.strip().split("\n"): if line.startswith("data: "): data = line[6:] # Remove "data: " prefix if data == "[DONE]": events.append({"done": True}) else: try: events.append(json.loads(data)) except json.JSONDecodeError: pass return events class TestOpenAIStreamingFormat: """Tests for OpenAI streaming response format (SSE).""" @pytest.fixture def mock_engine(self): """Create mock engine.""" return MockBaseEngine() @pytest.fixture def mock_engine_pool(self, mock_engine): """Create mock engine pool.""" return MockEnginePool(mock_engine) @pytest.fixture def client(self, mock_engine_pool): """Create test client with mocked state.""" from fastapi.testclient import TestClient from omlx.server import app, _server_state original_pool = _server_state.engine_pool original_default = _server_state.default_model _server_state.engine_pool = mock_engine_pool _server_state.default_model = "test-model" yield TestClient(app) _server_state.engine_pool = original_pool _server_state.default_model = original_default @pytest.mark.slow @pytest.mark.integration def test_chat_completion_streaming_format(self, client): """Test that streaming chat completion returns SSE format.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], "stream": True, }, ) assert response.status_code == 200 assert "text/event-stream" in response.headers["content-type"] @pytest.mark.slow @pytest.mark.integration def test_chat_completion_streaming_events(self, client): """Test streaming events structure.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], "stream": True, }, ) assert response.status_code == 200 events = parse_sse_events(response.text) # Should have at least one content event and a [DONE] event assert len(events) >= 2 # Check that last event is DONE assert events[-1].get("done") is True # Check structure of first chunk first_chunk = events[0] assert "id" in first_chunk assert first_chunk["object"] == "chat.completion.chunk" assert "choices" in first_chunk @pytest.mark.slow @pytest.mark.integration def test_chat_completion_streaming_role_in_first_chunk(self, client): """Test that first chunk contains role.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Test"}], "stream": True, }, ) events = parse_sse_events(response.text) non_done_events = [e for e in events if not e.get("done")] if non_done_events: first_chunk = non_done_events[0] # First chunk should have role in delta assert "choices" in first_chunk delta = first_chunk["choices"][0].get("delta", {}) assert delta.get("role") == "assistant" @pytest.mark.slow @pytest.mark.integration def test_completion_streaming_format(self, client): """Test that streaming completion returns SSE format.""" response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Hello", "stream": True, }, ) assert response.status_code == 200 assert "text/event-stream" in response.headers["content-type"] @pytest.mark.slow @pytest.mark.integration def test_completion_streaming_events(self, client): """Test completion streaming events structure.""" response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Once upon a time", "stream": True, }, ) events = parse_sse_events(response.text) # Should have events and end with DONE assert len(events) >= 2 assert events[-1].get("done") is True # Check first content event structure first_chunk = events[0] assert first_chunk["object"] == "text_completion" assert "choices" in first_chunk assert "text" in first_chunk["choices"][0] class TestAnthropicStreamingFormat: """Tests for Anthropic streaming response format (SSE events).""" @pytest.fixture def mock_engine(self): """Create mock engine.""" return MockBaseEngine() @pytest.fixture def mock_engine_pool(self, mock_engine): """Create mock engine pool.""" return MockEnginePool(mock_engine) @pytest.fixture def client(self, mock_engine_pool): """Create test client with mocked state.""" from fastapi.testclient import TestClient from omlx.server import app, _server_state original_pool = _server_state.engine_pool original_default = _server_state.default_model _server_state.engine_pool = mock_engine_pool _server_state.default_model = "test-model" yield TestClient(app) _server_state.engine_pool = original_pool _server_state.default_model = original_default @pytest.mark.slow @pytest.mark.integration def test_anthropic_streaming_format(self, client): """Test that Anthropic streaming returns SSE format.""" response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}], "stream": True, }, ) assert response.status_code == 200 assert "text/event-stream" in response.headers["content-type"] @pytest.mark.slow @pytest.mark.integration def test_anthropic_streaming_event_order(self, client): """Test Anthropic streaming event order.""" response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hi"}], "stream": True, }, ) events = parse_sse_events(response.text) non_done_events = [e for e in events if not e.get("done")] if len(non_done_events) >= 3: # Check event type order event_types = [e.get("type") for e in non_done_events] # Should start with message_start assert "message_start" in event_types # Should have content_block_start assert "content_block_start" in event_types # Should end with message_stop assert "message_stop" in event_types @pytest.mark.slow @pytest.mark.integration def test_anthropic_message_start_event(self, client): """Test Anthropic message_start event structure.""" response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Test"}], "stream": True, }, ) events = parse_sse_events(response.text) message_start = None for event in events: if event.get("type") == "message_start": message_start = event break assert message_start is not None assert "message" in message_start @pytest.mark.slow @pytest.mark.integration def test_anthropic_message_delta_event(self, client): """Test Anthropic message_delta event has stop_reason.""" response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Quick test"}], "stream": True, }, ) events = parse_sse_events(response.text) message_delta = None for event in events: if event.get("type") == "message_delta": message_delta = event break if message_delta: assert "delta" in message_delta assert "usage" in message_delta class TestStreamingHelperFunctions: """Tests for streaming helper functions in server module.""" @pytest.mark.asyncio async def test_stream_completion_yields_sse(self): """Test stream_completion yields SSE formatted strings.""" from omlx.server import stream_completion from omlx.api.openai_models import CompletionRequest engine = MockBaseEngine() request = CompletionRequest(model="test-model", prompt="Hello", stream=True) events = [] async for event in stream_completion(engine, "Hello", request): events.append(event) # Should have content events and DONE assert len(events) >= 2 assert events[-1] == "data: [DONE]\n\n" # Check SSE format for event in events[:-1]: assert event.startswith("data: ") assert event.endswith("\n\n") @pytest.mark.asyncio async def test_stream_completion_json_content(self): """Test stream_completion events contain valid JSON.""" from omlx.server import stream_completion from omlx.api.openai_models import CompletionRequest engine = MockBaseEngine() request = CompletionRequest(model="test-model", prompt="Test", stream=True) async for event in stream_completion(engine, "Test", request): if event != "data: [DONE]\n\n": json_str = event[6:-2] # Remove "data: " and "\n\n" data = json.loads(json_str) assert "id" in data assert "model" in data assert "choices" in data @pytest.mark.asyncio async def test_stream_chat_completion_yields_sse(self): """Test stream_chat_completion yields SSE formatted strings.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40 ): events.append(event) assert len(events) >= 2 assert events[-1] == "data: [DONE]\n\n" @pytest.mark.asyncio async def test_stream_chat_completion_first_chunk_has_role(self): """Test first streaming chunk has assistant role.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hello")], stream=True, ) first_event = None messages = [{"role": "user", "content": "Hello"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40 ): if event != "data: [DONE]\n\n": first_event = event break assert first_event is not None json_str = first_event[6:-2] data = json.loads(json_str) assert data["choices"][0]["delta"].get("role") == "assistant" @pytest.mark.asyncio async def test_stream_chat_completion_prompt_opened_thinking_streams_as_reasoning(self): """If the rendered prompt opens , initial deltas are reasoning.""" from omlx.api.openai_models import ChatCompletionRequest, Message from omlx.server import stream_chat_completion class PromptOpenedThinkingTokenizer(MockTokenizer): think_start = "" think_end = "" think_start_id = 999 think_end_id = 998 unk_token_id = -1 def convert_tokens_to_ids(self, token: str): if token == self.think_start: return self.think_start_id if token == self.think_end: return self.think_end_id return self.unk_token_id def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: ids = [101] if text.rstrip().endswith(self.think_start): ids.append(self.think_start_id) return ids def apply_chat_template( self, messages: list[dict], tokenize: bool = False, **kwargs ) -> str: return "user: Hi\nassistant:" engine = MockBaseEngine() engine._tokenizer = PromptOpenedThinkingTokenizer() engine.set_stream_outputs([ MockGenerationOutput( text="Need to inspect first.", new_text="Need to inspect first.", completion_tokens=1, finished=False, ), MockGenerationOutput( text="Need to inspect first.Done.", new_text="Done.", completion_tokens=2, finished=True, finish_reason="stop", ), ]) request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, ) payloads = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, ): if event.startswith("data: {"): payloads.append(json.loads(event[6:-2])) reasoning_deltas = [] content_deltas = [] for payload in payloads: choices = payload.get("choices", []) if not choices: continue delta = choices[0].get("delta", {}) if delta.get("reasoning_content"): reasoning_deltas.append(delta["reasoning_content"]) if delta.get("content"): content_deltas.append(delta["content"]) assert "".join(reasoning_deltas) == "Need to inspect first." assert content_deltas == ["Done."] @pytest.mark.asyncio async def test_stream_chat_completion_partial_prompt_opened_thinking(self): """Partial-mode prompt detection must mirror continue_final_message.""" from omlx.api.openai_models import ChatCompletionRequest, Message from omlx.server import stream_chat_completion class PartialThinkingTokenizer(MockTokenizer): think_start = "" think_end = "" think_start_id = 999 think_end_id = 998 unk_token_id = -1 def __init__(self): super().__init__() self.template_kwargs = None def convert_tokens_to_ids(self, token: str): if token == self.think_start: return self.think_start_id if token == self.think_end: return self.think_end_id return self.unk_token_id def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: ids = [101] if text.rstrip().endswith(self.think_start): ids.append(self.think_start_id) return ids def apply_chat_template( self, messages: list[dict], tokenize: bool = False, **kwargs ) -> str: self.template_kwargs = dict(kwargs) assert all("partial" not in message for message in messages) if ( kwargs.get("add_generation_prompt") is False and kwargs.get("continue_final_message") is True ): return "user: Hi\nassistant:" return "user: Hi\nassistant:" engine = MockBaseEngine() engine._tokenizer = PartialThinkingTokenizer() engine.set_stream_outputs([ MockGenerationOutput( text="Hidden reasoning.", new_text="Hidden reasoning.", completion_tokens=1, finished=False, ), MockGenerationOutput( text="Hidden reasoning.Visible.", new_text="Visible.", completion_tokens=2, finished=True, finish_reason="stop", ), ]) request = ChatCompletionRequest( model="test-model", messages=[ Message(role="user", content="Hi"), Message(role="assistant", content="", partial=True), ], stream=True, ) payloads = [] messages = [ {"role": "user", "content": "Hi"}, {"role": "assistant", "content": "", "partial": True}, ] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, is_partial=True, ): if event.startswith("data: {"): payloads.append(json.loads(event[6:-2])) reasoning_deltas = [] content_deltas = [] for payload in payloads: choices = payload.get("choices", []) if not choices: continue delta = choices[0].get("delta", {}) if delta.get("reasoning_content"): reasoning_deltas.append(delta["reasoning_content"]) if delta.get("content"): content_deltas.append(delta["content"]) assert engine.tokenizer.template_kwargs["add_generation_prompt"] is False assert engine.tokenizer.template_kwargs["continue_final_message"] is True assert messages[-1]["partial"] is True assert "".join(reasoning_deltas) == "Hidden reasoning." assert content_deltas == ["Visible."] @pytest.mark.asyncio async def test_stream_chat_completion_with_tools_streams_content_incrementally(self): """Tool availability must not force full buffering of normal text deltas.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Hello", new_text="Hello", completion_tokens=1, finished=False, ), MockGenerationOutput( text="Hello world", new_text=" world", completion_tokens=2, finished=True, finish_reason="stop", ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] content_deltas = [ payload["choices"][0].get("delta", {}).get("content") for payload in payloads if payload.get("choices") ] content_deltas = [delta for delta in content_deltas if delta] # With two streamed model outputs, we expect two incremental content chunks, # not one buffered chunk emitted only at completion. assert content_deltas == ["Hello", " world"] @pytest.mark.asyncio async def test_stream_chat_completion_sanitizes_tool_call_markup_inside_reasoning(self): """Reasoning deltas should keep prose while suppressing tool-call markup.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Need to inspect first.", new_text="Need to inspect first.", completion_tokens=1, finished=False, ), MockGenerationOutput( text=( "Need to inspect first." '{"name":"get_weather","arguments":{"city":"SF"}}' "Then continue." ), new_text=( '{"name":"get_weather","arguments":{"city":"SF"}}' "Then continue." ), completion_tokens=2, finished=True, finish_reason="stop", tool_calls=None, ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] reasoning_deltas = [] content_deltas = [] tool_call_deltas = [] finish_reasons = [] for payload in payloads: choices = payload.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) reasoning = delta.get("reasoning_content") if reasoning: reasoning_deltas.append(reasoning) content = delta.get("content") if content: content_deltas.append(content) if delta.get("tool_calls"): tool_call_deltas.extend(delta["tool_calls"]) finish_reason = choice.get("finish_reason") if finish_reason: finish_reasons.append(finish_reason) streamed_reasoning = "".join(reasoning_deltas) assert "" not in streamed_reasoning assert "" not in streamed_reasoning assert "Need to inspect first." in streamed_reasoning assert "Then continue." in streamed_reasoning assert content_deltas == [] assert len(tool_call_deltas) == 1 assert tool_call_deltas[0]["function"]["name"] == "get_weather" assert json.loads(tool_call_deltas[0]["function"]["arguments"]) == {"city": "SF"} assert "tool_calls" in finish_reasons @pytest.mark.asyncio async def test_stream_chat_completion_sanitizes_fragmented_reasoning_tool_call_markup(self): """Fragmented tool-call tags inside reasoning should never leak to the client.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Need to inspect ", new_text="Need to inspect ", completion_tokens=1, finished=False, ), MockGenerationOutput( text="" not in streamed_reasoning assert "" not in streamed_reasoning assert streamed_reasoning.strip() == "Need to inspect" assert len(tool_call_deltas) == 1 assert tool_call_deltas[0]["function"]["name"] == "get_weather" assert json.loads(tool_call_deltas[0]["function"]["arguments"]) == {"city": "SF"} assert "tool_calls" in finish_reasons @pytest.mark.asyncio async def test_stream_anthropic_messages_sanitizes_tool_call_markup_inside_thinking(self): """Anthropic thinking blocks should hide raw tool-call markup and emit tool_use.""" from omlx.server import stream_anthropic_messages from omlx.api.anthropic_models import MessagesRequest engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Need to inspect first.", new_text="Need to inspect first.", completion_tokens=1, finished=False, ), MockGenerationOutput( text=( "Need to inspect first." '{"name":"get_weather","arguments":{"city":"SF"}}' "Then continue." ), new_text=( '{"name":"get_weather","arguments":{"city":"SF"}}' "Then continue." ), completion_tokens=2, finished=True, finish_reason="stop", tool_calls=None, ), ]) anthropic_tools = [{ "name": "get_weather", "description": "Get weather", "input_schema": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }] internal_tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = MessagesRequest( model="test-model", max_tokens=256, messages=[{"role": "user", "content": "Hi"}], stream=True, tools=anthropic_tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_anthropic_messages( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=internal_tools, ): events.append(event) parsed_events = [] for event in events: for line in event.split("\n"): if line.startswith("data: "): try: parsed_events.append(json.loads(line[6:])) except json.JSONDecodeError: pass thinking_deltas = [] tool_use_blocks = [] block_start_indices = [] stop_reasons = [] for event in parsed_events: if event.get("type") == "content_block_delta": delta = event.get("delta", {}) if delta.get("type") == "thinking_delta": thinking_deltas.append(delta["thinking"]) elif event.get("type") == "content_block_start": block_start_indices.append(event["index"]) content_block = event.get("content_block", {}) if content_block.get("type") == "tool_use": tool_use_blocks.append(content_block) elif event.get("type") == "message_delta": stop_reason = event.get("delta", {}).get("stop_reason") if stop_reason: stop_reasons.append(stop_reason) streamed_thinking = "".join(thinking_deltas) assert "" not in streamed_thinking assert "" not in streamed_thinking assert "Need to inspect first." in streamed_thinking assert "Then continue." in streamed_thinking assert block_start_indices == list(range(len(block_start_indices))) assert len(tool_use_blocks) == 1 assert tool_use_blocks[0]["name"] == "get_weather" assert "tool_use" in stop_reasons @pytest.mark.asyncio async def test_stream_anthropic_messages_starts_parser_when_prompt_opens_thinking(self): """Prompt-opened thinking must stream as thinking, not public text.""" from omlx.server import stream_anthropic_messages from omlx.api.anthropic_models import MessagesRequest class PromptOpenedThinkingTokenizer(MockTokenizer): think_start = "" think_end = "" think_start_id = 9001 think_end_id = 9002 def encode(self, text: str, add_special_tokens: bool = False) -> List[int]: if text.rstrip().endswith(self.think_start): return [101, self.think_start_id] return [101] def apply_chat_template( self, messages: List[Dict], tokenize: bool = False, **kwargs ) -> str: return "system: hidden\nuser: hi\nassistant:" engine = MockBaseEngine() engine._tokenizer = PromptOpenedThinkingTokenizer() engine.set_stream_outputs([ MockGenerationOutput( text="The user asked a simple informational question.", new_text="The user asked a simple informational question.", completion_tokens=1, finished=False, ), MockGenerationOutput( text=( "The user asked a simple informational question." "Visible answer." ), new_text="Visible answer.", completion_tokens=2, finished=True, finish_reason="stop", ), ]) request = MessagesRequest( model="test-model", max_tokens=256, messages=[{"role": "user", "content": "Hi"}], stream=True, ) events = [] async for event in stream_anthropic_messages( engine, [{"role": "user", "content": "Hi"}], request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, ): events.append(event) parsed_events = [] for event in events: for line in event.split("\n"): if line.startswith("data: "): parsed_events.append(json.loads(line[6:])) thinking_deltas = [] text_deltas = [] for event in parsed_events: if event.get("type") != "content_block_delta": continue delta = event.get("delta", {}) if delta.get("type") == "thinking_delta": thinking_deltas.append(delta["thinking"]) elif delta.get("type") == "text_delta": text_deltas.append(delta["text"]) streamed_thinking = "".join(thinking_deltas) streamed_text = "".join(text_deltas) assert "The user asked a simple informational question." in streamed_thinking assert "The user asked a simple informational question." not in streamed_text assert streamed_text == "Visible answer." @pytest.mark.asyncio async def test_anthropic_tool_only_stream_starts_with_tool_use_block(self): """A tool-only response should not emit an empty text block before tool_use.""" from omlx.server import stream_anthropic_messages from omlx.api.anthropic_models import MessagesRequest engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="", new_text="", completion_tokens=1, finished=True, finish_reason="tool_calls", tool_calls=[{ "name": "get_weather", "arguments": "{\"city\":\"SF\"}", }], ), ]) anthropic_tools = [{ "name": "get_weather", "description": "Get weather", "input_schema": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }] internal_tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": anthropic_tools[0]["input_schema"], }, }] request = MessagesRequest( model="test-model", max_tokens=256, messages=[{"role": "user", "content": "Hi"}], stream=True, tools=anthropic_tools, ) events = [] async for event in stream_anthropic_messages( engine, [{"role": "user", "content": "Hi"}], request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=internal_tools, ): events.append(event) parsed_events = [] for event in events: for line in event.split("\n"): if line.startswith("data: "): parsed_events.append(json.loads(line[6:])) block_starts = [ event for event in parsed_events if event.get("type") == "content_block_start" ] assert block_starts[0]["index"] == 0 assert block_starts[0]["content_block"]["type"] == "tool_use" assert all( event["content_block"]["type"] != "text" for event in block_starts ) @pytest.mark.asyncio async def test_stream_chat_completion_with_tools_and_tool_calls_keeps_prior_content(self): """A tool_call finish should end the turn, not suppress already-generated text.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Let me check that for you.", new_text="Let me check that for you.", completion_tokens=1, finished=False, ), MockGenerationOutput( text="Let me check that for you.", new_text="", completion_tokens=1, finished=True, finish_reason="tool_calls", tool_calls=[{ "name": "get_weather", "arguments": "{\"city\":\"SF\"}", }], ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] content_deltas = [] saw_tool_call_delta = False finish_reasons = [] for payload in payloads: choices = payload.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) content = delta.get("content") if content: content_deltas.append(content) if delta.get("tool_calls"): saw_tool_call_delta = True finish_reason = choice.get("finish_reason") if finish_reason: finish_reasons.append(finish_reason) assert content_deltas == ["Let me check that for you."] assert saw_tool_call_delta is True assert "tool_calls" in finish_reasons @pytest.mark.asyncio async def test_stream_chat_completion_with_tools_parsed_from_text_does_not_leak_markup(self): """Parsed-from-text tool calls must not appear in streamed content deltas.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() # Non-Harmony path: no output.tool_calls; tool calls parsed from inline text. engine.set_stream_outputs([ MockGenerationOutput( text="Let me check", new_text="Let me check", completion_tokens=1, finished=False, finish_reason=None, tool_calls=None, ), MockGenerationOutput( text=( "Let me check that for you." "{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}" ), new_text=( " that for you." "{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}" ), completion_tokens=2, finished=True, finish_reason="stop", tool_calls=None, ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] content_deltas = [] tool_call_deltas = [] finish_reasons = [] content_event_indexes = [] tool_call_event_indexes = [] for idx, payload in enumerate(payloads): choices = payload.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) content = delta.get("content") if content: content_deltas.append(content) content_event_indexes.append(idx) if delta.get("tool_calls"): tool_call_deltas.extend(delta["tool_calls"]) tool_call_event_indexes.append(idx) finish_reason = choice.get("finish_reason") if finish_reason: finish_reasons.append(finish_reason) streamed_content = "".join(content_deltas) assert "" not in streamed_content assert "" not in streamed_content assert content_deltas == ["Let me check", " that for you."] assert streamed_content == "Let me check that for you." assert streamed_content.count("Let me check that for you.") == 1 assert len(tool_call_deltas) == 1 assert tool_call_deltas[0]["function"]["name"] == "get_weather" assert json.loads(tool_call_deltas[0]["function"]["arguments"]) == {"city": "SF"} assert content_event_indexes assert tool_call_event_indexes assert max(content_event_indexes) < min(tool_call_event_indexes) assert "tool_calls" in finish_reasons @pytest.mark.asyncio async def test_stream_chat_completion_with_tools_parsed_from_split_text_does_not_leak_partial_markup(self): """Split markup across chunks must not leak raw fragments in content deltas.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() # Non-Harmony path with fragmented tool-call markup across chunks. engine.set_stream_outputs([ MockGenerationOutput( text="Let me check that for you.{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}" ), new_text="ol_call>", completion_tokens=3, finished=True, finish_reason="stop", tool_calls=None, ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] content_deltas = [] tool_call_deltas = [] finish_reasons = [] content_event_indexes = [] tool_call_event_indexes = [] for idx, payload in enumerate(payloads): choices = payload.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) content = delta.get("content") if content: content_deltas.append(content) content_event_indexes.append(idx) if delta.get("tool_calls"): tool_call_deltas.extend(delta["tool_calls"]) tool_call_event_indexes.append(idx) finish_reason = choice.get("finish_reason") if finish_reason: finish_reasons.append(finish_reason) streamed_content = "".join(content_deltas) assert "" not in streamed_content assert "" not in streamed_content assert content_deltas == ["Let me check that for you."] assert streamed_content == "Let me check that for you." assert len(tool_call_deltas) == 1 assert tool_call_deltas[0]["function"]["name"] == "get_weather" assert json.loads(tool_call_deltas[0]["function"]["arguments"]) == {"city": "SF"} assert content_event_indexes assert tool_call_event_indexes assert max(content_event_indexes) < min(tool_call_event_indexes) assert "tool_calls" in finish_reasons @pytest.mark.asyncio async def test_stream_chat_completion_with_tools_parsed_from_namespaced_text_does_not_leak_markup(self): """Supported namespaced tags must not leak into streamed content deltas.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Let me check ", new_text="Let me check ", completion_tokens=1, finished=False, finish_reason=None, tool_calls=None, ), MockGenerationOutput( text=( "Let me check " "that for you." "" "" "\"SF\"" "" "" ), new_text=( "that for you." "" "" "\"SF\"" "" "" ), completion_tokens=2, finished=True, finish_reason="stop", tool_calls=None, ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] content_deltas = [] tool_call_deltas = [] finish_reasons = [] content_event_indexes = [] tool_call_event_indexes = [] for idx, payload in enumerate(payloads): choices = payload.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) content = delta.get("content") if content: content_deltas.append(content) content_event_indexes.append(idx) if delta.get("tool_calls"): tool_call_deltas.extend(delta["tool_calls"]) tool_call_event_indexes.append(idx) finish_reason = choice.get("finish_reason") if finish_reason: finish_reasons.append(finish_reason) streamed_content = "".join(content_deltas) assert "" not in streamed_content assert "" not in streamed_content assert "{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}<|/to" ), new_text=( "ol|>{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}<|/to" ), completion_tokens=2, finished=False, finish_reason=None, tool_calls=None, ), MockGenerationOutput( text=( "Let me check that for you." "<|tool|>{\"name\":\"get_weather\",\"arguments\":{\"city\":\"SF\"}}<|/tool|>" ), new_text="ol|>", completion_tokens=3, finished=True, finish_reason="stop", tool_calls=None, ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] content_deltas = [] tool_call_deltas = [] finish_reasons = [] content_event_indexes = [] tool_call_event_indexes = [] for idx, payload in enumerate(payloads): choices = payload.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) content = delta.get("content") if content: content_deltas.append(content) content_event_indexes.append(idx) if delta.get("tool_calls"): tool_call_deltas.extend(delta["tool_calls"]) tool_call_event_indexes.append(idx) finish_reason = choice.get("finish_reason") if finish_reason: finish_reasons.append(finish_reason) streamed_content = "".join(content_deltas) assert "<|to" not in streamed_content assert "<|tool|>" not in streamed_content assert "<|/tool|>" not in streamed_content assert len(content_deltas) >= 1 assert streamed_content == "Let me check that for you." assert len(tool_call_deltas) == 1 assert tool_call_deltas[0]["function"]["name"] == "get_weather" assert json.loads(tool_call_deltas[0]["function"]["arguments"]) == {"city": "SF"} assert content_event_indexes assert tool_call_event_indexes assert max(content_event_indexes) < min(tool_call_event_indexes) assert "tool_calls" in finish_reasons @pytest.mark.asyncio async def test_stream_chat_completion_suppresses_unmatched_tool_like_literal_suffix_under_clean_output_strict(self): """Clean-output strict contract suppresses unmatched tool-like suffixes.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Use literal marker " "" "\"SF\"" "" "" ), new_text=( "Let me check." "" "" "\"SF\"" "" "" ), completion_tokens=1, finished=True, finish_reason="stop", tool_calls=None, ), ]) tools = [{ "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, }] request = ChatCompletionRequest( model="test-model", messages=[Message(role="user", content="Hi")], stream=True, tools=tools, ) events = [] messages = [{"role": "user", "content": "Hi"}] async for event in stream_chat_completion( engine, messages, request, max_tokens=256, temperature=0.7, top_p=0.9, top_k=40, tools=tools, ): events.append(event) payloads = [ json.loads(event[6:-2]) for event in events if event.startswith("data: {") ] content_deltas = [] tool_call_deltas = [] finish_reasons = [] for payload in payloads: choices = payload.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) content = delta.get("content") if content: content_deltas.append(content) if delta.get("tool_calls"): tool_call_deltas.extend(delta["tool_calls"]) finish_reason = choice.get("finish_reason") if finish_reason: finish_reasons.append(finish_reason) streamed_content = "".join(content_deltas) assert "" not in streamed_content assert "" not in streamed_content assert streamed_content == "Let me check." assert len(tool_call_deltas) == 1 assert tool_call_deltas[0]["function"]["name"] == "get_weather" assert json.loads(tool_call_deltas[0]["function"]["arguments"]) == {"city": "SF"} assert "tool_calls" in finish_reasons @pytest.mark.asyncio async def test_stream_chat_completion_preserves_non_tool_namespaced_like_suffix_literal(self): """Trailing namespaced-looking literals that are not tool_call tags should be preserved.""" from omlx.server import stream_chat_completion from omlx.api.openai_models import ChatCompletionRequest, Message engine = MockBaseEngine() engine.set_stream_outputs([ MockGenerationOutput( text="Keep literal suffix