90 lines
4.3 KiB
Python
90 lines
4.3 KiB
Python
from unittest.mock import MagicMock
|
|
import pytest
|
|
from src.openai_compatible import (
|
|
NormalizedResponse,
|
|
OpenAICompatibleRequest,
|
|
send_openai_compatible,
|
|
)
|
|
from src.vendor_capabilities import VendorCapabilities, register
|
|
|
|
@pytest.fixture
|
|
def caps() -> VendorCapabilities:
|
|
return VendorCapabilities(vendor="test", model="test-model", context_window=8192, cost_input_per_mtok=1.0, cost_output_per_mtok=2.0)
|
|
|
|
def _mock_completion(text: str = "hello", tool_calls=None, usage_input: int = 10, usage_output: int = 5):
|
|
m = MagicMock()
|
|
m.choices = [MagicMock()]
|
|
m.choices[0].message.content = text
|
|
m.choices[0].message.tool_calls = tool_calls or []
|
|
m.usage.prompt_tokens = usage_input
|
|
m.usage.completion_tokens = usage_output
|
|
m.usage.prompt_tokens_details = None
|
|
m.usage.completion_tokens_details = None
|
|
return m
|
|
|
|
def test_send_non_streaming_returns_text_in_result(caps: VendorCapabilities) -> None:
|
|
client = MagicMock()
|
|
client.chat.completions.create.return_value = _mock_completion("hi", usage_input=20, usage_output=10)
|
|
request = OpenAICompatibleRequest(messages=[{"role": "user", "content": "ping"}], model="m", max_tokens=100)
|
|
result = send_openai_compatible(client, request, capabilities=caps)
|
|
assert result.ok
|
|
assert result.data == "hi"
|
|
assert result.errors == []
|
|
|
|
def test_send_streaming_aggregates_chunks(caps: VendorCapabilities) -> None:
|
|
client = MagicMock()
|
|
chunks = [
|
|
MagicMock(choices=[MagicMock(delta=MagicMock(content="hel", tool_calls=None))]),
|
|
MagicMock(choices=[MagicMock(delta=MagicMock(content="lo", tool_calls=None))]),
|
|
MagicMock(choices=[MagicMock(delta=MagicMock(content="", tool_calls=None))], usage=MagicMock(prompt_tokens=15, completion_tokens=5)),
|
|
]
|
|
client.chat.completions.create.return_value = iter(chunks)
|
|
received: list = []
|
|
request = OpenAICompatibleRequest(messages=[{"role": "user", "content": "ping"}], model="m", stream=True, stream_callback=received.append)
|
|
result = send_openai_compatible(client, request, capabilities=caps)
|
|
assert result.ok
|
|
assert result.data == "hello"
|
|
assert received == ["hel", "lo"]
|
|
|
|
def test_tool_call_detection_in_blocking_response(caps: VendorCapabilities) -> None:
|
|
from src.openai_compatible import _send_blocking
|
|
tool_call = MagicMock()
|
|
tool_call.id = "call_1"
|
|
tool_call.function.name = "read_file"
|
|
tool_call.function.arguments = '{"path": "/tmp/x"}'
|
|
completion = _mock_completion(text="", tool_calls=[tool_call])
|
|
client = MagicMock()
|
|
client.chat.completions.create.return_value = completion
|
|
kwargs = {"model": "m", "messages": [{"role": "user", "content": "ping"}], "temperature": 0.0, "top_p": 1.0, "max_tokens": 8192, "stream": False}
|
|
response = _send_blocking(client, kwargs)
|
|
assert len(response.tool_calls) == 1
|
|
assert response.tool_calls[0]["function"]["name"] == "read_file"
|
|
assert response.tool_calls[0]["id"] == "call_1"
|
|
|
|
def test_vision_multimodal_message(caps: VendorCapabilities) -> None:
|
|
client = MagicMock()
|
|
client.chat.completions.create.return_value = _mock_completion("looks like a cat")
|
|
messages = [{"role": "user", "content": [{"type": "text", "text": "what is this?"}, {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}}]}]
|
|
request = OpenAICompatibleRequest(messages=messages, model="m")
|
|
result = send_openai_compatible(client, request, capabilities=caps)
|
|
sent_messages = client.chat.completions.create.call_args.kwargs["messages"]
|
|
assert sent_messages[0]["content"] == messages[0]["content"]
|
|
assert result.data == "looks like a cat"
|
|
|
|
def test_error_classification_429_to_rate_limit(caps: VendorCapabilities) -> None:
|
|
from openai import RateLimitError
|
|
from src.result_types import Result, ErrorKind
|
|
client = MagicMock()
|
|
client.chat.completions.create.side_effect = RateLimitError("rate limited", response=MagicMock(status_code=429), body=None)
|
|
request = OpenAICompatibleRequest(messages=[{"role": "user", "content": "ping"}], model="m")
|
|
result = send_openai_compatible(client, request, capabilities=caps)
|
|
assert isinstance(result, Result)
|
|
assert not result.ok
|
|
assert result.errors[0].kind == ErrorKind.RATE_LIMIT
|
|
|
|
def test_normalized_response_is_frozen_dataclass() -> None:
|
|
from dataclasses import FrozenInstanceError
|
|
r = NormalizedResponse(text="x", tool_calls=[], usage_input_tokens=0, usage_output_tokens=0, usage_cache_read_tokens=0, usage_cache_creation_tokens=0, raw_response=None)
|
|
with pytest.raises(FrozenInstanceError):
|
|
r.text = "y"
|