# SPDX-License-Identifier: Apache-2.0 """ Integration tests for oMLX server endpoints. Tests the FastAPI endpoints using TestClient with mocked EnginePool and Engine to verify request/response formats without loading actual models. """ import json from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any, Dict, List, Optional from unittest.mock import AsyncMock import pytest from fastapi.testclient import TestClient from omlx.api.responses_utils import ResponseStore from omlx.engine.base import BaseEngine from omlx.engine.embedding import EmbeddingEngine from omlx.engine.reranker import RerankerEngine from omlx.mcp.types import MCPToolResult @dataclass class MockEmbeddingOutput: """Mock embedding output for testing.""" embeddings: List[List[float]] = field( default_factory=lambda: [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] ) total_tokens: int = 10 dimensions: int = 3 @dataclass class MockRerankOutput: """Mock rerank output for testing.""" scores: List[float] = field(default_factory=lambda: [0.9, 0.5, 0.3]) indices: List[int] = field(default_factory=lambda: [0, 1, 2]) total_tokens: int = 50 @dataclass class MockGenerationOutput: """Mock generation output for testing.""" text: str = "Hello, I am a helpful assistant." tokens: List[int] = field(default_factory=lambda: [1, 2, 3, 4, 5]) prompt_tokens: int = 10 completion_tokens: int = 5 finish_reason: str = "stop" new_text: str = "" finished: bool = True tool_calls: Optional[List[Dict[str, Any]]] = None cached_tokens: int = 0 class MockEmbeddingEngineImpl(EmbeddingEngine): """Mock embedding engine for testing that inherits from EmbeddingEngine.""" def __init__(self, model_name: str = "test-embedding-model"): # Don't call super().__init__ to avoid loading real model self._model_name = model_name self._model = None # Set as None but present self.calls: List[Dict[str, Any]] = [] @property def model_name(self) -> str: return self._model_name async def start(self) -> None: pass async def stop(self) -> None: pass async def embed(self, texts, **kwargs) -> MockEmbeddingOutput: self.calls.append({"texts": list(texts), "kwargs": dict(kwargs)}) return MockEmbeddingOutput( embeddings=[[0.1, 0.2, 0.3] for _ in texts], total_tokens=len(texts) * 5, dimensions=3, ) def get_stats(self) -> Dict[str, Any]: return {"model_name": self._model_name, "loaded": True} class MockRerankerEngineImpl(RerankerEngine): """Mock reranker engine for testing that inherits from RerankerEngine.""" def __init__(self, model_name: str = "test-reranker-model"): # Don't call super().__init__ to avoid loading real model self._model_name = model_name self._model = None # Set as None but present @property def model_name(self) -> str: return self._model_name async def start(self) -> None: pass async def stop(self) -> None: pass async def rerank( self, query: str, documents: List[str], top_n: Optional[int] = None, **kwargs ) -> MockRerankOutput: n_docs = len(documents) scores = [0.9 - i * 0.2 for i in range(n_docs)] indices = list(range(n_docs)) if top_n: indices = indices[:top_n] return MockRerankOutput( scores=scores, indices=indices, total_tokens=n_docs * 20, ) def get_stats(self) -> Dict[str, Any]: return {"model_name": self._model_name, "loaded": True} class MockTokenizer: """Mock tokenizer for testing.""" def __init__(self): self.eos_token_id = 2 def encode(self, text: str) -> List[int]: # Simple simulation: split by words 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 for testing.""" def __init__(self, model_name: str = "test-llm-model"): self._model_name = model_name self._tokenizer = MockTokenizer() self._model_type = "llama" @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 async def generate(self, prompt: str, **kwargs) -> MockGenerationOutput: return MockGenerationOutput(text="Generated response.") async def stream_generate(self, prompt: str, **kwargs): yield MockGenerationOutput( text="Hello", new_text="Hello", finished=False, ) yield MockGenerationOutput( text="Hello world", new_text=" world", 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.") async def stream_chat(self, messages: List[Dict], **kwargs): yield MockGenerationOutput( text="Hello", new_text="Hello", finished=False, ) yield MockGenerationOutput( text="Hello from chat", new_text=" from chat", finished=True, finish_reason="stop", ) def get_stats(self) -> Dict[str, Any]: return {} def get_cache_stats(self): return None class RecordingResponsesEngine(MockBaseEngine): """Mock engine that records request messages across /v1/responses calls.""" def __init__(self, outputs: Optional[List[MockGenerationOutput]] = None): super().__init__() self._outputs = list(outputs or []) self.recorded_messages: List[List[Dict[str, Any]]] = [] self._model_type = "gpt_oss" async def chat(self, messages: List[Dict], **kwargs) -> MockGenerationOutput: self.recorded_messages.append(messages) if self._outputs: return self._outputs.pop(0) return MockGenerationOutput(text="Chat response.") class MockEnginePool: """Mock engine pool for testing.""" def __init__( self, llm_engine: Optional[MockBaseEngine] = None, embedding_engine: Optional[MockEmbeddingEngineImpl] = None, reranker_engine: Optional[MockRerankerEngineImpl] = None, ): self._llm_engine = llm_engine or MockBaseEngine() self._embedding_engine = embedding_engine self._reranker_engine = reranker_engine self._models = [ {"id": "test-model", "loaded": True, "pinned": False, "size": 1000000} ] self._entries: Dict[str, Any] = {} self.get_engine_calls: List[Dict[str, Any]] = [] self.release_calls: List[str] = [] self.abort_requested_models: set[str] = set() @property def model_count(self) -> int: return len(self._models) @property def loaded_model_count(self) -> int: return sum(1 for m in self._models if m["loaded"]) @property def max_model_memory(self) -> int: return 32 * 1024 * 1024 * 1024 # 32GB @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, "loaded_count": self.loaded_model_count, "max_model_memory": self.max_model_memory, } async def get_engine( self, model_id: str, _lease: bool = False, runtime_settings=None, ): # _lease mirrors the real EnginePool's acquire-vs-use lease (#1667); # the mock has no eviction so it just accepts the flag. # runtime_settings mirrors exposed-profile variant loads. self.get_engine_calls.append( { "model_id": model_id, "_lease": _lease, "runtime_settings": runtime_settings, } ) # Return appropriate engine based on model name pattern if "embed" in model_id.lower(): if self._embedding_engine: return self._embedding_engine raise ValueError(f"No embedding engine for {model_id}") elif "rerank" in model_id.lower(): if self._reranker_engine: return self._reranker_engine raise ValueError(f"No reranker engine for {model_id}") return self._llm_engine async def release_engine(self, model_id: str) -> None: # No-op release counterpart of the in-use lease (#1667). self.release_calls.append(model_id) return None def is_abort_requested(self, model_id: str) -> bool: return model_id in self.abort_requested_models @pytest.fixture def mock_llm_engine(): """Create a mock LLM engine.""" return MockBaseEngine() @pytest.fixture def mock_embedding_engine(): """Create a mock embedding engine.""" return MockEmbeddingEngineImpl() @pytest.fixture def mock_reranker_engine(): """Create a mock reranker engine.""" return MockRerankerEngineImpl() @pytest.fixture def mock_engine_pool(mock_llm_engine, mock_embedding_engine, mock_reranker_engine): """Create a mock engine pool.""" return MockEnginePool( llm_engine=mock_llm_engine, embedding_engine=mock_embedding_engine, reranker_engine=mock_reranker_engine, ) @pytest.fixture def client(mock_engine_pool): """Create a test client with mocked server state.""" from omlx.server import app, _server_state # Store original state original_pool = _server_state.engine_pool original_default = _server_state.default_model # Set mock state _server_state.engine_pool = mock_engine_pool _server_state.default_model = "test-model" yield TestClient(app) # Restore original state _server_state.engine_pool = original_pool _server_state.default_model = original_default class TestHealthEndpoint: """Tests for the /health endpoint.""" def test_health_returns_healthy_status(self, client): """Test that health endpoint returns healthy status.""" response = client.get("/health") assert response.status_code == 200 data = response.json() assert data["status"] == "healthy" def test_health_contains_required_fields(self, client): """Test that health response contains required fields.""" response = client.get("/health") assert response.status_code == 200 data = response.json() assert "status" in data assert "default_model" in data assert "engine_pool" in data def test_health_engine_pool_info(self, client): """Test that health response contains engine pool info.""" response = client.get("/health") assert response.status_code == 200 data = response.json() pool_info = data["engine_pool"] assert "model_count" in pool_info assert "loaded_count" in pool_info assert "final_ceiling" in pool_info assert "current_model_memory" in pool_info class TestModelsEndpoint: """Tests for the /v1/models endpoint.""" def test_models_returns_list(self, client): """Test that models endpoint returns a list.""" response = client.get("/v1/models") assert response.status_code == 200 data = response.json() assert data["object"] == "list" assert "data" in data def test_models_format(self, client): """Test that model entries have correct format.""" response = client.get("/v1/models") assert response.status_code == 200 data = response.json() if data["data"]: model = data["data"][0] assert "id" in model assert "object" in model class TestResponsesEndpoint: def test_responses_uses_llm_lease(self, client, mock_engine_pool): response = client.post( "/v1/responses", json={"model": "test-model", "input": "Hello"}, ) assert response.status_code == 200 assert mock_engine_pool.get_engine_calls[-1]["_lease"] is True assert mock_engine_pool.release_calls == ["test-model"] def test_response_endpoint_includes_reasoning_item_for_think_blocks( self, client, mock_llm_engine ): mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text="Need to reason.Hello!", prompt_tokens=3, completion_tokens=6, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/responses", json={"model": "test-model", "input": "Hello"}, ) assert response.status_code == 200 data = response.json() assert [item["type"] for item in data["output"]] == ["reasoning", "message"] assert data["output"][0]["summary"][0]["text"] == "Need to reason." assert data["output"][1]["content"][0]["text"] == "Hello!" assert data["usage"]["output_tokens_details"]["reasoning_tokens"] == 3 def test_response_stream_includes_reasoning_item_for_think_blocks( self, client, mock_llm_engine ): async def stream_chat(messages, **kwargs): yield MockGenerationOutput( text="Need to reason.Hello!", new_text="Need to reason.Hello!", prompt_tokens=3, completion_tokens=6, finish_reason="stop", finished=True, ) mock_llm_engine.stream_chat = stream_chat response = client.post( "/v1/responses", json={"model": "test-model", "input": "Hello", "stream": True}, ) assert response.status_code == 200 events = [ json.loads(line.removeprefix("data: ")) for line in response.text.splitlines() if line.startswith("data: ") ] reasoning_deltas = [ event["delta"] for event in events if event.get("type") == "response.reasoning_summary_text.delta" ] assert "".join(reasoning_deltas) == "Need to reason." added_items = [ event for event in events if event.get("type") == "response.output_item.added" ] assert added_items[0]["item"]["type"] == "reasoning" assert added_items[0]["output_index"] == 0 assert added_items[1]["item"]["type"] == "message" assert added_items[1]["output_index"] == 1 completed = next( event for event in events if event.get("type") == "response.completed" ) output = completed["response"]["output"] assert [item["type"] for item in output] == ["reasoning", "message"] assert output[0]["summary"][0]["text"] == "Need to reason." assert output[1]["content"][0]["text"] == "Hello!" usage = completed["response"]["usage"] assert usage["output_tokens_details"]["reasoning_tokens"] == 3 def test_response_endpoint_recovers_tool_call_from_thinking(self, tmp_path): from omlx.server import app, _server_state state_dir = tmp_path / "response-state" engine = RecordingResponsesEngine( outputs=[ MockGenerationOutput( text=( "Need to inspect first." '{"name":"exec_command","arguments":{"cmd":"ls"}}' "Then continue." ), finish_reason="stop", ), ] ) pool = MockEnginePool(llm_engine=engine) original_pool = _server_state.engine_pool original_default = _server_state.default_model original_store = _server_state.responses_store try: _server_state.engine_pool = pool _server_state.default_model = "test-model" _server_state.responses_store = ResponseStore(state_dir=state_dir) client = TestClient(app) response = client.post( "/v1/responses", json={ "model": "test-model", "input": "Explore the code", "tools": [ { "type": "function", "name": "exec_command", "description": "Run a shell command", "parameters": { "type": "object", "properties": {"cmd": {"type": "string"}}, "required": ["cmd"], }, } ], }, ) assert response.status_code == 200 output_items = response.json()["output"] message_items = [item for item in output_items if item["type"] == "message"] reasoning_items = [ item for item in output_items if item["type"] == "reasoning" ] function_items = [ item for item in output_items if item["type"] == "function_call" ] assert len(reasoning_items) == 1 assert reasoning_items[0]["summary"][0]["text"] == ( "Need to inspect first.Then continue." ) assert "" not in reasoning_items[0]["summary"][0]["text"] assert len(message_items) == 1 assert message_items[0]["content"][0]["text"] == "" assert "" not in message_items[0]["content"][0]["text"] assert len(function_items) == 1 assert function_items[0]["name"] == "exec_command" assert function_items[0]["arguments"] == '{"cmd": "ls"}' finally: _server_state.engine_pool = original_pool _server_state.default_model = original_default _server_state.responses_store = original_store def test_previous_response_id_persists_across_store_restart(self, tmp_path): from omlx.server import app, _server_state state_dir = tmp_path / "response-state" engine = RecordingResponsesEngine( outputs=[ MockGenerationOutput( text="", finish_reason="tool_calls", tool_calls=[ { "id": "call_123", "name": "exec_command", "arguments": '{"cmd":"ls"}', } ], ), MockGenerationOutput(text="Done.", finish_reason="stop"), ] ) pool = MockEnginePool(llm_engine=engine) original_pool = _server_state.engine_pool original_default = _server_state.default_model original_store = _server_state.responses_store try: _server_state.engine_pool = pool _server_state.default_model = "test-model" _server_state.responses_store = ResponseStore(state_dir=state_dir) client = TestClient(app) first = client.post( "/v1/responses", json={"model": "test-model", "input": "Explore the code"}, ) assert first.status_code == 200 first_id = first.json()["id"] # Simulate a restart by rebuilding the store from disk. _server_state.responses_store = ResponseStore(state_dir=state_dir) second = client.post( "/v1/responses", json={ "model": "test-model", "previous_response_id": first_id, "input": [ { "type": "function_call_output", "call_id": "call_123", "output": "file1.txt\nfile2.txt", } ], }, ) assert second.status_code == 200 replayed = engine.recorded_messages[1] assert replayed[0] == {"role": "user", "content": "Explore the code"} assert replayed[1]["role"] == "assistant" assert replayed[1]["tool_calls"][0]["id"] == "call_123" assert replayed[2] == { "role": "tool", "tool_call_id": "call_123", "content": "file1.txt\nfile2.txt", } finally: _server_state.engine_pool = original_pool _server_state.default_model = original_default _server_state.responses_store = original_store def test_missing_previous_response_id_returns_404(self, tmp_path): from omlx.server import app, _server_state engine = RecordingResponsesEngine(outputs=[MockGenerationOutput(text="Done.")]) pool = MockEnginePool(llm_engine=engine) original_pool = _server_state.engine_pool original_default = _server_state.default_model original_store = _server_state.responses_store try: _server_state.engine_pool = pool _server_state.default_model = "test-model" _server_state.responses_store = ResponseStore( state_dir=tmp_path / "response-state" ) client = TestClient(app) response = client.post( "/v1/responses", json={ "model": "test-model", "previous_response_id": "resp_missing", "input": "Continue", }, ) assert response.status_code == 404 finally: _server_state.engine_pool = original_pool _server_state.default_model = original_default _server_state.responses_store = original_store class TestModelsStatusEndpoint: """Tests for the /v1/models/status endpoint.""" def test_models_status_returns_details(self, client): """Test that models status returns detailed info.""" response = client.get("/v1/models/status") assert response.status_code == 200 data = response.json() assert "models" in data def test_models_status_includes_model_alias(self, client): """Model aliases should be available to clients that join status metadata.""" from omlx.server import _server_state class Settings: model_alias = "gpt-4o" max_context_window = 32768 max_tokens = 8192 class SettingsManager: def get_settings(self, model_id): return Settings() def get_settings_for_request(self, model_id, resolved_model_id=None): return Settings() original_settings_manager = _server_state.settings_manager try: _server_state.settings_manager = SettingsManager() response = client.get("/v1/models/status") finally: _server_state.settings_manager = original_settings_manager assert response.status_code == 200 model = response.json()["models"][0] assert model["id"] == "test-model" assert model["model_alias"] == "gpt-4o" assert model["max_context_window"] == 32768 assert model["max_tokens"] == 8192 class TestCompletionEndpoint: """Tests for the /v1/completions endpoint.""" def test_completion_uses_llm_lease(self, client, mock_engine_pool): """LLM completion keeps a pool lease until the response body finishes.""" response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Hello, world!", }, ) assert response.status_code == 200 assert mock_engine_pool.get_engine_calls[-1]["_lease"] is True assert mock_engine_pool.release_calls == ["test-model"] def test_completion_basic_request(self, client): """Test basic completion request.""" response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Hello, world!", }, ) assert response.status_code == 200 data = response.json() assert "choices" in data assert len(data["choices"]) > 0 assert "text" in data["choices"][0] def test_completion_response_format(self, client): """Test completion response has correct format.""" response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Test prompt", "max_tokens": 100, }, ) assert response.status_code == 200 data = response.json() assert data["object"] == "text_completion" assert "model" in data assert "choices" in data assert "usage" in data def test_completion_with_list_prompt(self, client): """Test completion with list of prompts.""" response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": ["First prompt", "Second prompt"], }, ) assert response.status_code == 200 data = response.json() assert "choices" in data def test_completion_includes_cached_tokens_on_cache_hit( self, client, mock_llm_engine ): """Non-streaming completion responses should expose cached token counts.""" mock_llm_engine.generate = AsyncMock( return_value=MockGenerationOutput( text="Generated response.", prompt_tokens=2215, completion_tokens=5, cached_tokens=2048, ) ) response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Cache hit prompt", }, ) assert response.status_code == 200 data = response.json() assert data["usage"]["prompt_tokens_details"]["cached_tokens"] == 2048 def test_completion_forwards_thinking_budget(self, client, mock_llm_engine): """Non-streaming /v1/completions must forward thinking_budget to the engine. Regression for #1825: the field was absent from CompletionRequest, so Pydantic dropped it and it never reached generate().""" mock_llm_engine.generate = AsyncMock( return_value=MockGenerationOutput(text="Generated response.") ) response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Explain why the sky is blue.", "thinking_budget": 300, }, ) assert response.status_code == 200 assert mock_llm_engine.generate.call_args.kwargs["thinking_budget"] == 300 def test_completion_streaming_forwards_thinking_budget( self, client, mock_llm_engine ): """Streaming /v1/completions must forward thinking_budget to the engine (companion to the non-streaming path; see #1825).""" captured = {} async def recording_stream_generate(prompt, **kwargs): captured.update(kwargs) yield MockGenerationOutput(text="Hi", new_text="Hi", finished=False) yield MockGenerationOutput( text="Hi there", new_text=" there", finished=True, finish_reason="stop", ) mock_llm_engine.stream_generate = recording_stream_generate response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Explain why the sky is blue.", "stream": True, "thinking_budget": 300, }, ) assert response.status_code == 200 assert captured.get("thinking_budget") == 300 def test_completion_thinking_budget_from_model_settings( self, client, mock_llm_engine ): """A model-level thinking budget (admin settings) applies to /v1/completions even when the request omits the parameter, matching /v1/chat/completions.""" from omlx.model_settings import ModelSettings from omlx.server import _server_state class StubSettingsManager: def get_settings(self, model_id): return ModelSettings( thinking_budget_enabled=True, thinking_budget_tokens=256, ) mock_llm_engine.generate = AsyncMock( return_value=MockGenerationOutput(text="Generated response.") ) original_settings_manager = _server_state.settings_manager _server_state.settings_manager = StubSettingsManager() try: response = client.post( "/v1/completions", json={ "model": "test-model", "prompt": "Explain why the sky is blue.", }, ) finally: _server_state.settings_manager = original_settings_manager assert response.status_code == 200 assert mock_llm_engine.generate.call_args.kwargs["thinking_budget"] == 256 class TestChatCompletionEndpoint: """Tests for the /v1/chat/completions endpoint.""" def test_chat_completion_uses_llm_lease(self, client, mock_engine_pool): """Chat completion keeps a pool lease until the response body finishes.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], }, ) assert response.status_code == 200 assert mock_engine_pool.get_engine_calls[-1]["_lease"] is True assert mock_engine_pool.release_calls == ["test-model"] def test_chat_completion_basic(self, client): """Test basic chat completion request.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], }, ) assert response.status_code == 200 data = response.json() assert "choices" in data assert len(data["choices"]) > 0 def test_chat_completion_response_format(self, client): """Test chat completion response format.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hi!"}, ], }, ) assert response.status_code == 200 data = response.json() assert data["object"] == "chat.completion" assert "model" in data assert "choices" in data assert data["choices"][0]["message"]["role"] == "assistant" assert "usage" in data def test_chat_completion_with_parameters(self, client): """Test chat completion with sampling parameters.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Test"}], "temperature": 0.7, "top_p": 0.9, "max_tokens": 256, }, ) assert response.status_code == 200 def test_chat_completion_includes_cached_tokens_on_cache_hit( self, client, mock_llm_engine ): """Non-streaming chat responses should expose cached token counts.""" mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text="Chat response.", prompt_tokens=2215, completion_tokens=5, cached_tokens=2048, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Cache hit prompt"}], }, ) assert response.status_code == 200 data = response.json() assert data["usage"]["prompt_tokens_details"]["cached_tokens"] == 2048 def test_chat_completion_sanitizes_reasoning_tool_call_markup( self, client, mock_llm_engine ): """Thinking-only tool calls should become structured tool_calls without leaked markup.""" mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text=( "Need to inspect first." '{"name":"get_weather","arguments":{"city":"SF"}}' "Then continue." ), prompt_tokens=10, completion_tokens=5, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Hi"}], "tools": [ { "type": "function", "function": { "name": "get_weather", "description": "Get weather", "parameters": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, }, } ], }, ) assert response.status_code == 200 data = response.json() message = data["choices"][0]["message"] assert message["reasoning_content"] == "Need to inspect first.Then continue." assert "" not in message["reasoning_content"] assert len(message["tool_calls"]) == 1 assert message["tool_calls"][0]["function"]["name"] == "get_weather" assert message["tool_calls"][0]["function"]["arguments"] == '{"city": "SF"}' assert data["choices"][0]["finish_reason"] == "tool_calls" @pytest.mark.parametrize( ("tool_choice", "request_tools", "expect_tools"), [ (None, None, True), ("auto", None, True), ("none", None, False), ( "none", [ { "type": "function", "function": { "name": "user_search", "description": "Search from request tools", "parameters": { "type": "object", "properties": { "query": {"type": "string"}, }, }, }, } ], False, ), ], ) def test_chat_completion_tool_choice_controls_mcp_tools( self, client, mock_llm_engine, tool_choice, request_tools, expect_tools, ): """tool_choice='none' should suppress request and globally configured tools.""" from omlx.server import _server_state class RecordingMCPManager: def __init__(self): self.calls = [] def get_merged_tools(self, user_tools=None): self.calls.append(user_tools) return [ { "type": "function", "function": { "name": "mcp_search", "description": "Search via MCP", "parameters": { "type": "object", "properties": { "query": {"type": "string"}, }, }, }, } ] recorded_count_tools = [] recorded_chat_kwargs = [] def count_chat_tokens(messages, tools=None, **kwargs): recorded_count_tools.append(tools) return 1 async def chat(messages, **kwargs): recorded_chat_kwargs.append(kwargs) return MockGenerationOutput( text="Plain response.", prompt_tokens=1, completion_tokens=1, finish_reason="stop", ) original_mcp_manager = _server_state.mcp_manager manager = RecordingMCPManager() mock_llm_engine.count_chat_tokens = count_chat_tokens mock_llm_engine.chat = chat payload = { "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], } if tool_choice is not None: payload["tool_choice"] = tool_choice if request_tools is not None: payload["tools"] = request_tools try: _server_state.mcp_manager = manager response = client.post("/v1/chat/completions", json=payload) finally: _server_state.mcp_manager = original_mcp_manager assert response.status_code == 200 assert recorded_chat_kwargs if expect_tools: assert manager.calls == [request_tools] assert recorded_count_tools[0] is not None assert "tools" in recorded_chat_kwargs[0] else: assert manager.calls == [] assert recorded_count_tools == [None] assert "tools" not in recorded_chat_kwargs[0] class TestAnthropicMessagesEndpoint: """Tests for the /v1/messages endpoint (Anthropic format).""" def test_anthropic_messages_uses_llm_lease(self, client, mock_engine_pool): response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}], }, ) assert response.status_code == 200 assert mock_engine_pool.get_engine_calls[-1]["_lease"] is True assert mock_engine_pool.release_calls == ["test-model"] def test_anthropic_messages_basic(self, client): """Test basic Anthropic messages request.""" response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}], }, ) assert response.status_code == 200 data = response.json() assert data["type"] == "message" assert data["role"] == "assistant" def test_anthropic_messages_response_format(self, client): """Test Anthropic messages response format.""" response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hi there!"}], }, ) assert response.status_code == 200 data = response.json() assert "id" in data assert "content" in data assert "usage" in data assert "input_tokens" in data["usage"] assert "output_tokens" in data["usage"] def test_anthropic_messages_with_system(self, client): """Test Anthropic messages with system prompt.""" response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "system": "You are a helpful assistant.", "messages": [{"role": "user", "content": "Hello!"}], }, ) assert response.status_code == 200 def test_anthropic_messages_sanitize_thinking_tool_call_markup( self, client, mock_llm_engine ): """Anthropic thinking blocks should not expose raw tool-call markup.""" mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text=( "Need to inspect first." '{"name":"get_weather","arguments":{"city":"SF"}}' "Then continue." ), prompt_tokens=10, completion_tokens=5, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/messages", json={ "model": "test-model", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hi"}], "tools": [ { "name": "get_weather", "description": "Get weather", "input_schema": { "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], }, } ], }, ) assert response.status_code == 200 data = response.json() thinking_blocks = [ block for block in data["content"] if block["type"] == "thinking" ] tool_use_blocks = [ block for block in data["content"] if block["type"] == "tool_use" ] assert len(thinking_blocks) == 1 assert thinking_blocks[0]["thinking"] == "Need to inspect first.Then continue." assert "" not in thinking_blocks[0]["thinking"] assert len(tool_use_blocks) == 1 assert tool_use_blocks[0]["name"] == "get_weather" assert tool_use_blocks[0]["input"] == {"city": "SF"} assert data["stop_reason"] == "tool_use" def test_anthropic_messages_with_exposed_profile_model( self, client, mock_llm_engine, mock_engine_pool, tmp_path ): """A model:profile id resolves and overlays profile settings on /v1/messages.""" from omlx.model_settings import ModelSettings, ModelSettingsManager from omlx.server import _server_state manager = ModelSettingsManager(tmp_path) manager.set_settings("test-model", ModelSettings(temperature=0.1)) manager.save_profile( model_id="test-model", name="thinking", display_name="Thinking", description=None, settings={"temperature": 0.9, "enable_thinking": True}, expose_as_model=True, ) def resolve_model_id(model_id, settings_manager=None): if settings_manager is not None: source = settings_manager.get_exposed_profile_source_model_id(model_id) if source: return source return model_id recorded_chat_kwargs = [] async def chat(messages, **kwargs): recorded_chat_kwargs.append(kwargs) return MockGenerationOutput( text="Hello from the profile.", prompt_tokens=1, completion_tokens=1, finish_reason="stop", ) mock_llm_engine.chat = chat original_resolve = mock_engine_pool.resolve_model_id original_settings_manager = _server_state.settings_manager mock_engine_pool.resolve_model_id = resolve_model_id try: _server_state.settings_manager = manager response = client.post( "/v1/messages", json={ "model": "test-model:thinking", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}], }, ) finally: _server_state.settings_manager = original_settings_manager mock_engine_pool.resolve_model_id = original_resolve assert response.status_code == 200 data = response.json() assert data["type"] == "message" assert data["role"] == "assistant" # The profile's settings — not the base model's — reached the engine. assert recorded_chat_kwargs assert recorded_chat_kwargs[0]["temperature"] == 0.9 ct_kwargs = recorded_chat_kwargs[0].get("chat_template_kwargs") or {} assert ct_kwargs.get("enable_thinking") is True class TestEmbeddingsEndpoint: """Tests for the /v1/embeddings endpoint.""" def test_embeddings_single_input(self, client, mock_engine_pool): """Test embeddings with single input.""" mock_engine_pool._models.append( {"id": "test-embed-model", "loaded": True, "pinned": False, "size": 500000} ) response = client.post( "/v1/embeddings", json={ "model": "test-embed-model", "input": "Hello, world!", }, ) assert response.status_code == 200 data = response.json() assert data["object"] == "list" assert "data" in data assert len(data["data"]) == 1 assert data["data"][0]["object"] == "embedding" def test_embeddings_multiple_inputs(self, client, mock_engine_pool): """Test embeddings with multiple inputs.""" mock_engine_pool._models.append( {"id": "test-embed-model", "loaded": True, "pinned": False, "size": 500000} ) response = client.post( "/v1/embeddings", json={ "model": "test-embed-model", "input": ["First text", "Second text"], }, ) assert response.status_code == 200 data = response.json() assert len(data["data"]) == 2 def test_embeddings_use_discovered_context_length(self, client, mock_engine_pool): """Embedding requests should not fall back to mlx-embeddings' 512 default.""" mock_engine_pool._models.append( {"id": "test-embed-model", "loaded": True, "pinned": False, "size": 500000} ) mock_engine_pool._entries["test-embed-model"] = SimpleNamespace( model_context_length=40960 ) response = client.post( "/v1/embeddings", json={ "model": "test-embed-model", "input": "hello", }, ) assert response.status_code == 200 kwargs = mock_engine_pool._embedding_engine.calls[-1]["kwargs"] assert kwargs["max_length"] == 40960 assert kwargs["truncation"] is True def test_embeddings_request_max_length_overrides_default( self, client, mock_engine_pool ): """Explicit max_length should be forwarded to the embedding engine.""" mock_engine_pool._models.append( {"id": "test-embed-model", "loaded": True, "pinned": False, "size": 500000} ) mock_engine_pool._entries["test-embed-model"] = SimpleNamespace( model_context_length=40960 ) response = client.post( "/v1/embeddings", json={ "model": "test-embed-model", "input": "hello", "max_length": 1024, "truncation": False, }, ) assert response.status_code == 200 kwargs = mock_engine_pool._embedding_engine.calls[-1]["kwargs"] assert kwargs["max_length"] == 1024 assert kwargs["truncation"] is False def test_embeddings_response_format(self, client, mock_engine_pool): """Test embeddings response format.""" mock_engine_pool._models.append( {"id": "test-embed-model", "loaded": True, "pinned": False, "size": 500000} ) response = client.post( "/v1/embeddings", json={ "model": "test-embed-model", "input": "Test text", }, ) assert response.status_code == 200 data = response.json() assert "model" in data assert "usage" in data assert "prompt_tokens" in data["usage"] assert "total_tokens" in data["usage"] assert "embedding" in data["data"][0] assert isinstance(data["data"][0]["embedding"], list) def test_embeddings_structured_items_input(self, client, mock_engine_pool): """Test embeddings with structured multimodal items.""" mock_engine_pool._models.append( {"id": "test-embed-model", "loaded": True, "pinned": False, "size": 500000} ) image_data_uri = ( "data:image/png;base64," "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/" "x8AAwMCAO+/p9sAAAAASUVORK5CYII=" ) response = client.post( "/v1/embeddings", json={ "model": "test-embed-model", "items": [ {"text": "hello"}, {"image": image_data_uri}, { "text": "hello", "image": image_data_uri, }, ], }, ) assert response.status_code == 200 data = response.json() assert len(data["data"]) == 3 def test_embeddings_rejects_mixed_input_sources(self, client, mock_engine_pool): """Test embeddings rejects input and items together.""" mock_engine_pool._models.append( {"id": "test-embed-model", "loaded": True, "pinned": False, "size": 500000} ) response = client.post( "/v1/embeddings", json={ "model": "test-embed-model", "input": "hello", "items": [{"text": "hello"}], }, ) assert response.status_code == 422 class TestRerankEndpoint: """Tests for the /v1/rerank endpoint.""" def test_rerank_basic(self, client, mock_engine_pool): """Test basic rerank request.""" mock_engine_pool._models.append( { "id": "test-rerank-model", "loaded": True, "pinned": False, "size": 500000, } ) response = client.post( "/v1/rerank", json={ "model": "test-rerank-model", "query": "What is machine learning?", "documents": [ "ML is a subset of AI.", "The weather is nice today.", ], }, ) assert response.status_code == 200 data = response.json() assert "results" in data assert len(data["results"]) == 2 def test_rerank_with_top_n(self, client, mock_engine_pool): """Test rerank with top_n parameter.""" mock_engine_pool._models.append( { "id": "test-rerank-model", "loaded": True, "pinned": False, "size": 500000, } ) response = client.post( "/v1/rerank", json={ "model": "test-rerank-model", "query": "Test query", "documents": ["Doc 1", "Doc 2", "Doc 3"], "top_n": 2, }, ) assert response.status_code == 200 data = response.json() assert len(data["results"]) == 2 def test_rerank_response_format(self, client, mock_engine_pool): """Test rerank response format.""" mock_engine_pool._models.append( { "id": "test-rerank-model", "loaded": True, "pinned": False, "size": 500000, } ) response = client.post( "/v1/rerank", json={ "model": "test-rerank-model", "query": "Test", "documents": ["Document 1"], "return_documents": True, }, ) assert response.status_code == 200 data = response.json() assert "id" in data assert "model" in data assert "results" in data result = data["results"][0] assert "index" in result assert "relevance_score" in result assert "document" in result class TestTokenCountEndpoint: """Tests for the /v1/messages/count_tokens endpoint.""" def test_token_count_uses_llm_lease(self, client, mock_engine_pool): response = client.post( "/v1/messages/count_tokens", json={ "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], }, ) assert response.status_code == 200 assert mock_engine_pool.get_engine_calls[-1]["_lease"] is True assert mock_engine_pool.release_calls == ["test-model"] def test_token_count_basic(self, client): """Test basic token counting.""" response = client.post( "/v1/messages/count_tokens", json={ "model": "test-model", "messages": [{"role": "user", "content": "Hello world"}], }, ) assert response.status_code == 200 data = response.json() assert "input_tokens" in data assert isinstance(data["input_tokens"], int) def test_token_count_with_system(self, client): """Test token counting with system prompt.""" response = client.post( "/v1/messages/count_tokens", json={ "model": "test-model", "system": "You are helpful.", "messages": [{"role": "user", "content": "Hi!"}], }, ) assert response.status_code == 200 data = response.json() assert "input_tokens" in data class TestMCPEndpoints: """Tests for MCP-related endpoints.""" def test_mcp_tools_empty(self, client): """Test MCP tools endpoint when no MCP configured.""" response = client.get("/v1/mcp/tools") assert response.status_code == 200 data = response.json() assert "tools" in data assert "count" in data assert data["count"] == 0 def test_mcp_servers_empty(self, client): """Test MCP servers endpoint when no MCP configured.""" response = client.get("/v1/mcp/servers") assert response.status_code == 200 data = response.json() assert "servers" in data def test_mcp_execute_no_config(self, client): """Test MCP execute fails when not configured.""" response = client.post( "/v1/mcp/execute", json={ "tool_name": "test_tool", "arguments": {}, }, ) # Should return 503 when MCP not configured assert response.status_code == 503 def test_mcp_execute_accepts_tool_alias(self, client): """Test MCP execute accepts tool as an alias for tool_name.""" from omlx.server import _server_state original_mcp_manager = _server_state.mcp_manager manager = AsyncMock() manager.execute_tool.return_value = MCPToolResult( tool_name="test_tool", content={"ok": True}, ) try: _server_state.mcp_manager = manager response = client.post( "/v1/mcp/execute", json={ "tool": "test_tool", "arguments": {"query": "hello"}, }, ) finally: _server_state.mcp_manager = original_mcp_manager assert response.status_code == 200 assert response.json() == { "tool_name": "test_tool", "content": {"ok": True}, "is_error": False, "error_message": None, } manager.execute_tool.assert_awaited_once_with( "test_tool", {"query": "hello"}, ) def test_mcp_execute_tool_name_field(self, client): """Test MCP execute happy path with tool_name field.""" from omlx.server import _server_state original_mcp_manager = _server_state.mcp_manager manager = AsyncMock() manager.execute_tool.return_value = MCPToolResult( tool_name="my_tool", content="ok", ) try: _server_state.mcp_manager = manager response = client.post( "/v1/mcp/execute", json={ "tool_name": "my_tool", "arguments": {"q": "x"}, }, ) finally: _server_state.mcp_manager = original_mcp_manager assert response.status_code == 200 manager.execute_tool.assert_awaited_once_with("my_tool", {"q": "x"}) def test_mcp_execute_tool_name_wins_over_tool(self, client): """Test tool_name takes precedence when both fields are present.""" from omlx.server import _server_state original_mcp_manager = _server_state.mcp_manager manager = AsyncMock() manager.execute_tool.return_value = MCPToolResult( tool_name="canonical", content="ok", ) try: _server_state.mcp_manager = manager response = client.post( "/v1/mcp/execute", json={ "tool_name": "canonical", "tool": "alias_should_lose", "arguments": {}, }, ) finally: _server_state.mcp_manager = original_mcp_manager assert response.status_code == 200 manager.execute_tool.assert_awaited_once_with("canonical", {}) def test_mcp_execute_rejects_missing_tool(self, client): """Test MCP execute returns 422 when neither tool nor tool_name is present.""" response = client.post( "/v1/mcp/execute", json={"arguments": {"q": "x"}}, ) assert response.status_code == 422 class TestErrorHandling: """Tests for error handling in endpoints.""" def test_missing_model(self, client): """Test error when model is not specified.""" # For Anthropic endpoint, missing model should raise validation error response = client.post( "/v1/messages", json={ "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}], }, ) assert response.status_code == 422 # Validation error def test_empty_messages(self, client): """Test error when messages is empty.""" response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [], }, ) # Empty messages may be allowed or raise error depending on implementation # Just verify we get a response assert response.status_code in [200, 400, 422] def test_invalid_request_format(self, client): """Test error for invalid request format.""" response = client.post( "/v1/chat/completions", json={ "invalid_field": "test", }, ) assert response.status_code == 422 class TestJsonOutputParsing: """Tests for parse_json_output in non-streaming endpoints.""" def test_chat_completion_parses_markdown_json(self, client, mock_llm_engine): """Markdown-wrapped JSON should be parsed when response_format=json_object.""" import json mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text='```json\n{"name": "test", "age": 25}\n```', prompt_tokens=10, completion_tokens=8, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Return JSON"}], "response_format": {"type": "json_object"}, }, ) assert response.status_code == 200 data = response.json() content = data["choices"][0]["message"]["content"] parsed = json.loads(content) assert parsed == {"name": "test", "age": 25} def test_chat_completion_clean_json_unchanged(self, client, mock_llm_engine): """Already-clean JSON should pass through without corruption.""" import json mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text='{"key": "value"}', prompt_tokens=10, completion_tokens=5, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/chat/completions", json={ "model": "test-model", "messages": [{"role": "user", "content": "Return JSON"}], "response_format": {"type": "json_object"}, }, ) assert response.status_code == 200 data = response.json() content = data["choices"][0]["message"]["content"] parsed = json.loads(content) assert parsed == {"key": "value"} def test_responses_parses_markdown_json(self, client, mock_llm_engine): """Responses API should parse markdown-wrapped JSON with text.format.""" import json mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text='```json\n{"city": "Seoul", "temp": 15}\n```', prompt_tokens=10, completion_tokens=8, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/responses", json={ "model": "test-model", "input": "Return weather JSON", "text": { "format": {"type": "json_object"}, }, }, ) assert response.status_code == 200 data = response.json() output_text = data["output"][0]["content"][0]["text"] parsed = json.loads(output_text) assert parsed == {"city": "Seoul", "temp": 15} def test_responses_without_format_unchanged(self, client, mock_llm_engine): """Responses API without text.format should return raw text.""" mock_llm_engine.chat = AsyncMock( return_value=MockGenerationOutput( text="Hello, how can I help?", prompt_tokens=10, completion_tokens=5, finish_reason="stop", finished=True, ) ) response = client.post( "/v1/responses", json={ "model": "test-model", "input": "Hi", }, ) assert response.status_code == 200 data = response.json() output_text = data["output"][0]["content"][0]["text"] assert "Hello" in output_text