wehub-resource-sync
451 lines
15 KiB
Python
451 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for omlx.request module."""
|
|
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from omlx.request import (
|
|
RequestStatus,
|
|
SamplingParams,
|
|
Request,
|
|
RequestOutput,
|
|
)
|
|
|
|
|
|
class TestRequestStatus:
|
|
"""Test cases for RequestStatus enum."""
|
|
|
|
def test_status_values(self):
|
|
"""Test that status enum has expected values."""
|
|
assert RequestStatus.WAITING is not None
|
|
assert RequestStatus.RUNNING is not None
|
|
assert RequestStatus.PREEMPTED is not None
|
|
assert RequestStatus.FINISHED_STOPPED is not None
|
|
assert RequestStatus.FINISHED_LENGTH_CAPPED is not None
|
|
assert RequestStatus.FINISHED_ABORTED is not None
|
|
|
|
def test_status_ordering(self):
|
|
"""Test that finished statuses are greater than active statuses."""
|
|
assert RequestStatus.WAITING < RequestStatus.FINISHED_STOPPED
|
|
assert RequestStatus.RUNNING < RequestStatus.FINISHED_STOPPED
|
|
assert RequestStatus.PREEMPTED < RequestStatus.FINISHED_STOPPED
|
|
|
|
def test_is_finished_active_states(self):
|
|
"""Test is_finished returns False for active states."""
|
|
assert RequestStatus.is_finished(RequestStatus.WAITING) is False
|
|
assert RequestStatus.is_finished(RequestStatus.RUNNING) is False
|
|
assert RequestStatus.is_finished(RequestStatus.PREEMPTED) is False
|
|
|
|
def test_is_finished_finished_states(self):
|
|
"""Test is_finished returns True for finished states."""
|
|
assert RequestStatus.is_finished(RequestStatus.FINISHED_STOPPED) is True
|
|
assert RequestStatus.is_finished(RequestStatus.FINISHED_LENGTH_CAPPED) is True
|
|
assert RequestStatus.is_finished(RequestStatus.FINISHED_ABORTED) is True
|
|
|
|
def test_get_finish_reason_stopped(self):
|
|
"""Test get_finish_reason for FINISHED_STOPPED."""
|
|
reason = RequestStatus.get_finish_reason(RequestStatus.FINISHED_STOPPED)
|
|
assert reason == "stop"
|
|
|
|
def test_get_finish_reason_length_capped(self):
|
|
"""Test get_finish_reason for FINISHED_LENGTH_CAPPED."""
|
|
reason = RequestStatus.get_finish_reason(RequestStatus.FINISHED_LENGTH_CAPPED)
|
|
assert reason == "length"
|
|
|
|
def test_get_finish_reason_aborted(self):
|
|
"""Test get_finish_reason for FINISHED_ABORTED."""
|
|
reason = RequestStatus.get_finish_reason(RequestStatus.FINISHED_ABORTED)
|
|
assert reason == "abort"
|
|
|
|
def test_get_finish_reason_active_states(self):
|
|
"""Test get_finish_reason returns None for active states."""
|
|
assert RequestStatus.get_finish_reason(RequestStatus.WAITING) is None
|
|
assert RequestStatus.get_finish_reason(RequestStatus.RUNNING) is None
|
|
assert RequestStatus.get_finish_reason(RequestStatus.PREEMPTED) is None
|
|
|
|
|
|
class TestSamplingParams:
|
|
"""Test cases for SamplingParams dataclass."""
|
|
|
|
def test_default_values(self):
|
|
"""Test default sampling parameter values."""
|
|
params = SamplingParams()
|
|
assert params.max_tokens == 256
|
|
assert params.temperature == 0.7
|
|
assert params.top_p == 0.9
|
|
assert params.top_k == 0
|
|
assert params.min_p == 0.0
|
|
assert params.xtc_probability == 0.0
|
|
assert params.xtc_threshold == 0.1
|
|
assert params.repetition_penalty == 1.0
|
|
assert params.presence_penalty == 0.0
|
|
assert params.stop == []
|
|
assert params.stop_token_ids == []
|
|
assert params.logprobs is False
|
|
assert params.top_logprobs is None
|
|
|
|
def test_custom_values(self):
|
|
"""Test custom sampling parameter values."""
|
|
params = SamplingParams(
|
|
max_tokens=1024,
|
|
temperature=0.5,
|
|
top_p=0.95,
|
|
top_k=40,
|
|
min_p=0.05,
|
|
xtc_probability=0.5,
|
|
xtc_threshold=0.1,
|
|
repetition_penalty=1.1,
|
|
presence_penalty=0.5,
|
|
stop=["###", "END"],
|
|
stop_token_ids=[2, 100],
|
|
logprobs=True,
|
|
top_logprobs=5,
|
|
)
|
|
assert params.max_tokens == 1024
|
|
assert params.temperature == 0.5
|
|
assert params.top_p == 0.95
|
|
assert params.top_k == 40
|
|
assert params.min_p == 0.05
|
|
assert params.xtc_probability == 0.5
|
|
assert params.xtc_threshold == 0.1
|
|
assert params.repetition_penalty == 1.1
|
|
assert params.presence_penalty == 0.5
|
|
assert params.stop == ["###", "END"]
|
|
assert params.stop_token_ids == [2, 100]
|
|
assert params.logprobs is True
|
|
assert params.top_logprobs == 5
|
|
|
|
def test_post_init_none_stop(self):
|
|
"""Test that None stop sequences are converted to empty lists."""
|
|
params = SamplingParams(stop=None, stop_token_ids=None)
|
|
assert params.stop == []
|
|
assert params.stop_token_ids == []
|
|
|
|
def test_greedy_sampling(self):
|
|
"""Test parameters for greedy sampling (temperature=0)."""
|
|
params = SamplingParams(temperature=0.0, top_k=1)
|
|
assert params.temperature == 0.0
|
|
assert params.top_k == 1
|
|
|
|
|
|
class TestRequest:
|
|
"""Test cases for Request dataclass."""
|
|
|
|
def test_basic_creation(self):
|
|
"""Test basic request creation."""
|
|
request = Request(
|
|
request_id="test-001",
|
|
prompt="Hello, world!",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
assert request.request_id == "test-001"
|
|
assert request.prompt == "Hello, world!"
|
|
assert request.status == RequestStatus.WAITING
|
|
assert request.output_token_ids == []
|
|
assert request.output_text == ""
|
|
|
|
def test_creation_with_token_ids(self):
|
|
"""Test request creation with token IDs as prompt."""
|
|
request = Request(
|
|
request_id="test-002",
|
|
prompt=[1, 2, 3, 4, 5],
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
assert request.prompt == [1, 2, 3, 4, 5]
|
|
|
|
def test_arrival_time_auto_set(self):
|
|
"""Test that arrival_time is automatically set."""
|
|
before = time.monotonic()
|
|
request = Request(
|
|
request_id="test-003",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
after = time.monotonic()
|
|
assert before <= request.arrival_time <= after
|
|
|
|
def test_num_output_tokens_property(self):
|
|
"""Test num_output_tokens property."""
|
|
request = Request(
|
|
request_id="test-004",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
assert request.num_output_tokens == 0
|
|
|
|
request.output_token_ids = [100, 200, 300]
|
|
assert request.num_output_tokens == 3
|
|
|
|
def test_num_tokens_property(self):
|
|
"""Test num_tokens property (prompt + output)."""
|
|
request = Request(
|
|
request_id="test-005",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request.num_prompt_tokens = 10
|
|
request.output_token_ids = [100, 200, 300]
|
|
assert request.num_tokens == 13
|
|
|
|
def test_max_tokens_property(self):
|
|
"""Test max_tokens property from sampling_params."""
|
|
request = Request(
|
|
request_id="test-006",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(max_tokens=512),
|
|
)
|
|
assert request.max_tokens == 512
|
|
|
|
def test_is_finished_method(self):
|
|
"""Test is_finished method."""
|
|
request = Request(
|
|
request_id="test-007",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
assert request.is_finished() is False
|
|
|
|
request.status = RequestStatus.FINISHED_STOPPED
|
|
assert request.is_finished() is True
|
|
|
|
def test_get_finish_reason_method(self):
|
|
"""Test get_finish_reason method."""
|
|
request = Request(
|
|
request_id="test-008",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
|
|
# No finish reason while active
|
|
assert request.get_finish_reason() is None
|
|
|
|
# From status
|
|
request.status = RequestStatus.FINISHED_STOPPED
|
|
assert request.get_finish_reason() == "stop"
|
|
|
|
# Override with explicit finish_reason
|
|
request.finish_reason = "custom_reason"
|
|
assert request.get_finish_reason() == "custom_reason"
|
|
|
|
def test_append_output_token(self):
|
|
"""Test append_output_token method."""
|
|
request = Request(
|
|
request_id="test-009",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request.append_output_token(100)
|
|
request.append_output_token(200)
|
|
|
|
assert request.output_token_ids == [100, 200]
|
|
assert request.num_computed_tokens == 2
|
|
|
|
def test_set_finished(self):
|
|
"""Test set_finished method."""
|
|
request = Request(
|
|
request_id="test-010",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request.set_finished(RequestStatus.FINISHED_STOPPED)
|
|
|
|
assert request.status == RequestStatus.FINISHED_STOPPED
|
|
assert request.finish_reason == "stop"
|
|
|
|
def test_set_finished_with_reason(self):
|
|
"""Test set_finished with custom reason."""
|
|
request = Request(
|
|
request_id="test-011",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request.set_finished(RequestStatus.FINISHED_ABORTED, reason="user_cancelled")
|
|
|
|
assert request.status == RequestStatus.FINISHED_ABORTED
|
|
assert request.finish_reason == "user_cancelled"
|
|
|
|
def test_comparison_by_priority(self):
|
|
"""Test request comparison by priority."""
|
|
request1 = Request(
|
|
request_id="test-012",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
priority=1,
|
|
)
|
|
request2 = Request(
|
|
request_id="test-013",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
priority=2,
|
|
)
|
|
# Lower priority value = higher priority
|
|
assert request1 < request2
|
|
|
|
def test_comparison_by_arrival_time(self):
|
|
"""Test request comparison by arrival time (same priority)."""
|
|
request1 = Request(
|
|
request_id="test-014",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
arrival_time=100.0,
|
|
)
|
|
request2 = Request(
|
|
request_id="test-015",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
arrival_time=200.0,
|
|
)
|
|
# Earlier arrival time = higher priority
|
|
assert request1 < request2
|
|
|
|
def test_hash(self):
|
|
"""Test request hash is based on request_id."""
|
|
request1 = Request(
|
|
request_id="test-016",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request2 = Request(
|
|
request_id="test-016",
|
|
prompt="Different prompt",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
assert hash(request1) == hash(request2)
|
|
|
|
def test_equality(self):
|
|
"""Test request equality is based on request_id."""
|
|
request1 = Request(
|
|
request_id="test-017",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request2 = Request(
|
|
request_id="test-017",
|
|
prompt="Different",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request3 = Request(
|
|
request_id="test-018",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
assert request1 == request2
|
|
assert request1 != request3
|
|
|
|
def test_equality_with_non_request(self):
|
|
"""Test equality with non-Request objects."""
|
|
request = Request(
|
|
request_id="test-019",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
assert request != "test-019"
|
|
assert request != 123
|
|
assert request != None
|
|
|
|
def test_reasoning_model_fields(self):
|
|
"""Test reasoning model support fields."""
|
|
request = Request(
|
|
request_id="test-020",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
needs_think_prefix=True,
|
|
)
|
|
assert request.needs_think_prefix is True
|
|
assert request.think_prefix_sent is False
|
|
|
|
def test_harmony_model_field(self):
|
|
"""Test Harmony model field."""
|
|
request = Request(
|
|
request_id="test-021",
|
|
prompt="Test",
|
|
sampling_params=SamplingParams(),
|
|
is_harmony_model=True,
|
|
)
|
|
assert request.is_harmony_model is True
|
|
|
|
def test_multimodal_fields(self):
|
|
"""Test multimodal content fields."""
|
|
request = Request(
|
|
request_id="test-022",
|
|
prompt="Describe this image",
|
|
sampling_params=SamplingParams(),
|
|
images=["image_data_1", "image_data_2"],
|
|
videos=["video_data_1"],
|
|
)
|
|
assert request.images == ["image_data_1", "image_data_2"]
|
|
assert request.videos == ["video_data_1"]
|
|
|
|
|
|
class TestRequestOutput:
|
|
"""Test cases for RequestOutput dataclass."""
|
|
|
|
def test_basic_creation(self):
|
|
"""Test basic RequestOutput creation."""
|
|
output = RequestOutput(request_id="test-001")
|
|
assert output.request_id == "test-001"
|
|
assert output.new_token_ids == []
|
|
assert output.new_text == ""
|
|
assert output.output_token_ids == []
|
|
assert output.output_text == ""
|
|
assert output.finished is False
|
|
assert output.finish_reason is None
|
|
|
|
def test_with_tokens(self):
|
|
"""Test RequestOutput with tokens."""
|
|
output = RequestOutput(
|
|
request_id="test-002",
|
|
new_token_ids=[100, 200],
|
|
new_text="Hello",
|
|
output_token_ids=[100, 200, 300, 400],
|
|
output_text="Hello world",
|
|
)
|
|
assert output.new_token_ids == [100, 200]
|
|
assert output.new_text == "Hello"
|
|
assert output.output_token_ids == [100, 200, 300, 400]
|
|
assert output.output_text == "Hello world"
|
|
|
|
def test_finished_output(self):
|
|
"""Test finished RequestOutput."""
|
|
output = RequestOutput(
|
|
request_id="test-003",
|
|
finished=True,
|
|
finish_reason="stop",
|
|
)
|
|
assert output.finished is True
|
|
assert output.finish_reason == "stop"
|
|
|
|
def test_usage_property(self):
|
|
"""Test usage property."""
|
|
output = RequestOutput(
|
|
request_id="test-004",
|
|
prompt_tokens=10,
|
|
completion_tokens=20,
|
|
)
|
|
usage = output.usage
|
|
assert usage["prompt_tokens"] == 10
|
|
assert usage["completion_tokens"] == 20
|
|
assert usage["total_tokens"] == 30
|
|
|
|
def test_usage_property_zero(self):
|
|
"""Test usage property with zero tokens."""
|
|
output = RequestOutput(request_id="test-005")
|
|
usage = output.usage
|
|
assert usage["prompt_tokens"] == 0
|
|
assert usage["completion_tokens"] == 0
|
|
assert usage["total_tokens"] == 0
|
|
|
|
def test_tool_calls(self):
|
|
"""Test RequestOutput with tool calls."""
|
|
tool_calls = [
|
|
{"id": "call_1", "function": {"name": "get_weather", "arguments": "{}"}},
|
|
{"id": "call_2", "function": {"name": "search", "arguments": "{}"}},
|
|
]
|
|
output = RequestOutput(
|
|
request_id="test-006",
|
|
tool_calls=tool_calls,
|
|
)
|
|
assert output.tool_calls == tool_calls
|
|
assert len(output.tool_calls) == 2
|