Spaces:
Runtime error
Runtime error
File size: 2,248 Bytes
290ff9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | from __future__ import annotations
import pytest
from ergo_agentic.models.routing import (
ainvoke_structured_with_fallback,
get_fallback_model_id,
)
class _ServerError(Exception):
def __init__(self, status_code: int):
super().__init__(f"status={status_code}")
self.status_code = status_code
class _StubRunnable:
def __init__(self, result=None, exc: Exception | None = None):
self._result = result
self._exc = exc
async def ainvoke(self, _messages):
if self._exc is not None:
raise self._exc
return self._result
class _StubModel:
def __init__(self, runnable: _StubRunnable):
self._runnable = runnable
def with_structured_output(self, _schema):
return self._runnable
@pytest.mark.asyncio
async def test_ainvoke_structured_falls_back_on_transient_error(monkeypatch):
models = {
"google_genai:gemini-2.5-pro": _StubModel(_StubRunnable(exc=_ServerError(503))),
"google_genai:gemini-2.5-flash": _StubModel(_StubRunnable(result={"ok": True})),
}
monkeypatch.setattr(
"ergo_agentic.models.routing.get_chat_model",
lambda model_id, temperature=0.0: models[model_id],
)
result = await ainvoke_structured_with_fallback(
model_id="google_genai:gemini-2.5-pro",
schema=dict,
messages=[],
)
assert result == {"ok": True}
@pytest.mark.asyncio
async def test_ainvoke_structured_does_not_fallback_on_non_transient_error(monkeypatch):
models = {
"google_genai:gemini-2.5-pro": _StubModel(_StubRunnable(exc=ValueError("bad prompt"))),
}
monkeypatch.setattr(
"ergo_agentic.models.routing.get_chat_model",
lambda model_id, temperature=0.0: models[model_id],
)
with pytest.raises(ValueError, match="bad prompt"):
await ainvoke_structured_with_fallback(
model_id="google_genai:gemini-2.5-pro",
schema=dict,
messages=[],
)
def test_get_fallback_model_id_for_gemini_pro():
assert get_fallback_model_id("google_genai:gemini-2.5-pro") == (
"google_genai:gemini-2.5-flash"
)
|