wehub-resource-sync
1955 lines
64 KiB
Python
1955 lines
64 KiB
Python
# 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"<decoded:{len(tokens)} tokens>"
|
|
|
|
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="<think>Need to reason.</think>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="<think>Need to reason.</think>Hello!",
|
|
new_text="<think>Need to reason.</think>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=(
|
|
"<think>Need to inspect first."
|
|
'<tool_call>{"name":"exec_command","arguments":{"cmd":"ls"}}</tool_call>'
|
|
"Then continue.</think>"
|
|
),
|
|
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 "<tool_call>" not in reasoning_items[0]["summary"][0]["text"]
|
|
assert len(message_items) == 1
|
|
assert message_items[0]["content"][0]["text"] == ""
|
|
assert "<tool_call>" 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=(
|
|
"<think>Need to inspect first."
|
|
'<tool_call>{"name":"get_weather","arguments":{"city":"SF"}}</tool_call>'
|
|
"Then continue.</think>"
|
|
),
|
|
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 "<tool_call>" 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=(
|
|
"<think>Need to inspect first."
|
|
'<tool_call>{"name":"get_weather","arguments":{"city":"SF"}}</tool_call>'
|
|
"Then continue.</think>"
|
|
),
|
|
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 "<tool_call>" 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
|