diff --git a/skillopt/model/openai_compatible_backend.py b/skillopt/model/openai_compatible_backend.py index e2f3001..7462be4 100644 --- a/skillopt/model/openai_compatible_backend.py +++ b/skillopt/model/openai_compatible_backend.py @@ -33,7 +33,6 @@ from typing import Any from openai import OpenAI from skillopt.model.common import ( - CompatAssistantMessage, TokenTracker, compat_message_from_chat_message, default_model_for_backend, diff --git a/tests/test_openai_compatible_backend.py b/tests/test_openai_compatible_backend.py new file mode 100644 index 0000000..4482e20 --- /dev/null +++ b/tests/test_openai_compatible_backend.py @@ -0,0 +1,146 @@ +"""Tests for the generic OpenAI-compatible model backend.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +import skillopt.model as model +from skillopt.model import backend_config +from skillopt.model import openai_compatible_backend as backend + + +class _CompletionRecorder: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def create(self, **kwargs: Any) -> Any: + self.calls.append(kwargs) + message = SimpleNamespace(content="ok", tool_calls=[]) + usage = SimpleNamespace(prompt_tokens=2, completion_tokens=3, total_tokens=5) + return SimpleNamespace(choices=[SimpleNamespace(message=message)], usage=usage) + + +class _Client: + def __init__(self, recorder: _CompletionRecorder) -> None: + self.chat = SimpleNamespace(completions=recorder) + + +@pytest.fixture(autouse=True) +def isolate_backend_state(monkeypatch: pytest.MonkeyPatch): + optimizer_backend = backend_config.get_optimizer_backend() + target_backend = backend_config.get_target_backend() + optimizer_config = vars(backend.OPTIMIZER_CONFIG).copy() + target_config = vars(backend.TARGET_CONFIG).copy() + backend.reset_token_tracker() + yield + backend.reset_token_tracker() + vars(backend.OPTIMIZER_CONFIG).update(optimizer_config) + vars(backend.TARGET_CONFIG).update(target_config) + backend_config.set_optimizer_backend(optimizer_backend) + backend_config.set_target_backend(target_backend) + backend._reset_clients() + + +def test_configure_preserves_role_specific_values() -> None: + model.configure_openai_compatible( + base_url="https://shared.example/v1", + api_key="shared-key", + model="shared-model", + optimizer_base_url="https://optimizer.example/v1", + optimizer_api_key="optimizer-key", + optimizer_model="optimizer-model", + target_base_url="https://target.example/v1", + target_api_key="target-key", + target_model="target-model", + ) + + assert backend.OPTIMIZER_CONFIG.base_url == "https://optimizer.example/v1" + assert backend.OPTIMIZER_CONFIG.api_key == "optimizer-key" + assert backend.OPTIMIZER_CONFIG.deployment == "optimizer-model" + assert backend.TARGET_CONFIG.base_url == "https://target.example/v1" + assert backend.TARGET_CONFIG.api_key == "target-key" + assert backend.TARGET_CONFIG.deployment == "target-model" + + +def test_optimizer_and_target_route_to_their_own_clients(monkeypatch: pytest.MonkeyPatch) -> None: + optimizer_calls = _CompletionRecorder() + target_calls = _CompletionRecorder() + monkeypatch.setattr( + backend, + "_get_client", + lambda role: _Client(optimizer_calls if role == "optimizer" else target_calls), + ) + model.set_optimizer_backend("openai_compatible") + model.set_target_backend("openai_compatible") + backend.OPTIMIZER_CONFIG.deployment = "optimizer-model" + backend.TARGET_CONFIG.deployment = "target-model" + + model.chat_optimizer("system", "user", retries=1) + model.chat_target_messages([{"role": "user", "content": "question"}], retries=1) + + assert optimizer_calls.calls[0]["model"] == "optimizer-model" + assert target_calls.calls[0]["model"] == "target-model" + + +def test_client_creation_is_lazy(monkeypatch: pytest.MonkeyPatch) -> None: + builds: list[str] = [] + + def build(config: backend.OpenAICompatibleConfig) -> _Client: + builds.append(config.deployment) + return _Client(_CompletionRecorder()) + + backend._reset_clients() + monkeypatch.setattr(backend, "_build_client", build) + assert builds == [] + + backend._get_client("optimizer") + backend._get_client("optimizer") + + assert builds == [backend.OPTIMIZER_CONFIG.deployment] + + +def test_combined_token_summary_counts_each_backend_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def make(stage: str) -> dict[str, dict[str, int]]: + return { + stage: { + "calls": 1, + "prompt_tokens": 2, + "completion_tokens": 3, + "total_tokens": 5, + }, + "_total": { + "calls": 1, + "prompt_tokens": 2, + "completion_tokens": 3, + "total_tokens": 5, + }, + } + + monkeypatch.setattr(model._openai, "get_token_summary", lambda: make("azure")) + monkeypatch.setattr(model._claude, "get_token_summary", lambda: make("claude")) + monkeypatch.setattr(model._qwen, "get_token_summary", lambda: make("qwen")) + monkeypatch.setattr(model._minimax, "get_token_summary", lambda: make("minimax")) + monkeypatch.setattr(model._openai_compat, "get_token_summary", lambda: make("openai_compatible")) + + combined = model.get_token_summary() + + assert set(combined) - {"_total"} == {"azure", "claude", "qwen", "minimax", "openai_compatible"} + expected_stage_total = { + "calls": 1, + "prompt_tokens": 2, + "completion_tokens": 3, + "total_tokens": 5, + } + for stage in {"azure", "claude", "qwen", "minimax", "openai_compatible"}: + assert combined[stage] == expected_stage_total + assert combined["_total"] == { + "calls": 5, + "prompt_tokens": 10, + "completion_tokens": 15, + "total_tokens": 25, + }