90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
import os
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
import openrouter_judge as oj
|
|
|
|
|
|
def _model(model_id, prompt, completion, *, image=True, response_format=True, batch=False):
|
|
return {
|
|
"id": model_id,
|
|
"architecture": {
|
|
"input_modalities": ["text", "image"] if image else ["text"],
|
|
"output_modalities": ["text"],
|
|
},
|
|
"supported_parameters": ["response_format"] if response_format else ["temperature"],
|
|
"pricing": {"prompt": str(prompt), "completion": str(completion)},
|
|
}
|
|
|
|
|
|
class ModelSelectionTests(unittest.TestCase):
|
|
def test_candidates_are_image_json_models_sorted_by_price(self):
|
|
catalog = [
|
|
_model("expensive/model", 0.10, 0.40),
|
|
_model("cheap/model", 0.01, 0.02),
|
|
_model("text-only", 0.001, 0.001, image=False),
|
|
_model("no-json", 0.001, 0.001, response_format=False),
|
|
_model("batch/model:batch", 0.0001, 0.0001),
|
|
]
|
|
self.assertEqual(oj.vision_candidates(catalog), ["cheap/model", "expensive/model"])
|
|
|
|
def test_explicit_model_wins_over_auto_and_env(self):
|
|
with mock.patch.dict(os.environ, {"GIG_LINE_JUDGE_MODEL": "env/model"}):
|
|
self.assertEqual(oj.resolve_models("explicit/model", catalog=[]),
|
|
["explicit/model"])
|
|
|
|
def test_auto_uses_env_override(self):
|
|
with mock.patch.dict(os.environ, {"GIG_LINE_JUDGE_MODEL": "env/model"}):
|
|
self.assertEqual(oj.resolve_models("auto", catalog=[]), ["env/model"])
|
|
|
|
def test_auto_uses_catalog_and_falls_back_offline(self):
|
|
catalog = [_model("cheap/model", 0.01, 0.02)]
|
|
self.assertEqual(oj.resolve_models("auto", catalog=catalog), ["cheap/model"])
|
|
self.assertEqual(oj.resolve_models("auto", catalog=[]), oj.FALLBACK_MODELS[:6])
|
|
|
|
|
|
class FallbackTests(unittest.TestCase):
|
|
def test_fallback_switches_model_after_provider_error(self):
|
|
seen = []
|
|
|
|
def call(body):
|
|
seen.append(body["model"])
|
|
if body["model"] == "bad/model":
|
|
raise ValueError("OpenRouter HTTP 400: provider rejected image")
|
|
return {"ok": True}
|
|
|
|
model, payload = oj.call_with_fallback(
|
|
["bad/model", "good/model"], lambda m: {"model": m}, call, attempts=1)
|
|
self.assertEqual(model, "good/model")
|
|
self.assertEqual(payload, {"ok": True})
|
|
self.assertEqual(seen, ["bad/model", "good/model"])
|
|
|
|
def test_auth_error_does_not_try_another_model(self):
|
|
seen = []
|
|
|
|
def call(body):
|
|
seen.append(body["model"])
|
|
raise ValueError("OpenRouter HTTP 401: API key not valid")
|
|
|
|
with self.assertRaises(ValueError):
|
|
oj.call_with_fallback(["one", "two"], lambda m: {"model": m}, call)
|
|
self.assertEqual(seen, ["one"])
|
|
|
|
def test_transient_error_retries_same_model(self):
|
|
seen = []
|
|
|
|
def call(body):
|
|
seen.append(body["model"])
|
|
if len(seen) == 1:
|
|
raise ValueError("OpenRouter HTTP 500")
|
|
return {"ok": True}
|
|
|
|
model, _ = oj.call_with_fallback(
|
|
["one"], lambda m: {"model": m}, call, attempts=2)
|
|
self.assertEqual(model, "one")
|
|
self.assertEqual(seen, ["one", "one"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|