mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 14:12:58 +00:00
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
1794 lines
66 KiB
Python
1794 lines
66 KiB
Python
"""Integration tests for LLM fallback behaviour.
|
|
|
|
Verifies that when a primary model fails (immediately or mid-stream), the
|
|
per-agent backup model is used before the global FALLBACK_* settings.
|
|
"""
|
|
|
|
import copy
|
|
import logging
|
|
import types
|
|
from unittest.mock import MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from docsgpt.llm.anthropic import AnthropicLLM
|
|
from docsgpt.llm.base import BaseLLM
|
|
from docsgpt.llm.google_ai import GoogleLLM
|
|
from docsgpt.llm.groq import GroqLLM
|
|
from docsgpt.llm.openai import OpenAILLM
|
|
|
|
|
|
# Concrete LLM stubs
|
|
|
|
|
|
class FakeLLM(BaseLLM):
|
|
"""Minimal concrete BaseLLM for testing."""
|
|
|
|
def __init__(self, responses=None, stream_chunks=None, fail_at=None, **kwargs):
|
|
# Accept and discard api_key / user_api_key so LLMCreator.create_llm
|
|
# signatures work without errors.
|
|
kwargs.pop("api_key", None)
|
|
kwargs.pop("user_api_key", None)
|
|
# Retry-once tests supply these; strip before BaseLLM.__init__.
|
|
error_class = kwargs.pop("error_class", RuntimeError)
|
|
fail_schedule = kwargs.pop("fail_schedule", None)
|
|
error_schedule = kwargs.pop("error_schedule", None)
|
|
super().__init__(**kwargs)
|
|
self.responses = responses or ["fake response"]
|
|
self.stream_chunks = stream_chunks or ["chunk1", "chunk2"]
|
|
self.fail_at = fail_at # None = no failure, 0 = immediate, N = after N chunks
|
|
self.error_class = error_class
|
|
self.fail_schedule = fail_schedule
|
|
self.error_schedule = error_schedule
|
|
self.stream_calls = 0
|
|
self.user_api_key = None
|
|
self.gen_called = False
|
|
self.gen_stream_called = False
|
|
self.last_model_received = None # tracks the model kwarg passed to gen/gen_stream
|
|
self.last_messages_received = None # tracks the messages kwarg at the raw level
|
|
self.last_kwargs_received = None # tracks the extra gen kwargs at the raw level
|
|
|
|
# Track at the raw-method level. _execute_with_fallback applies
|
|
# decorators to the fallback's raw method directly and
|
|
# never calls .gen() / .gen_stream() on it, so a public-method
|
|
# override would not register fallback hops.
|
|
def _raw_gen(self, baseself, model, messages, stream, tools=None, **kwargs):
|
|
self.gen_called = True
|
|
self.last_model_received = model
|
|
self.last_messages_received = messages
|
|
self.last_kwargs_received = dict(kwargs)
|
|
if self.fail_at is not None:
|
|
raise RuntimeError("primary model unavailable")
|
|
return self.responses[0]
|
|
|
|
def _raw_gen_stream(self, baseself, model, messages, stream, tools=None, **kwargs):
|
|
self.gen_stream_called = True
|
|
self.stream_calls = getattr(self, "stream_calls", 0) + 1
|
|
self.last_model_received = model
|
|
self.last_messages_received = messages
|
|
self.last_kwargs_received = dict(kwargs)
|
|
yielded = 0
|
|
# Per-attempt failure schedule: `fail_schedule[n]` = fail_at value
|
|
# for the n-th call (0-indexed). Falls back to the constant
|
|
# ``fail_at`` when the schedule is exhausted or unset. Lets a test
|
|
# simulate "fail then succeed" without a custom subclass.
|
|
schedule = getattr(self, "fail_schedule", None)
|
|
if schedule is not None and self.stream_calls - 1 < len(schedule):
|
|
local_fail_at = schedule[self.stream_calls - 1]
|
|
else:
|
|
local_fail_at = self.fail_at
|
|
# Per-attempt exception class: same rationale as fail_schedule.
|
|
error_schedule = getattr(self, "error_schedule", None)
|
|
if (
|
|
error_schedule is not None
|
|
and self.stream_calls - 1 < len(error_schedule)
|
|
):
|
|
local_error_class = error_schedule[self.stream_calls - 1]
|
|
else:
|
|
local_error_class = getattr(self, "error_class", RuntimeError)
|
|
for chunk in self.stream_chunks:
|
|
if local_fail_at is not None and yielded >= local_fail_at:
|
|
raise local_error_class("mid-stream failure")
|
|
yield chunk
|
|
yielded += 1
|
|
|
|
|
|
# Helpers
|
|
|
|
|
|
def _noop_decorator(func):
|
|
"""Pass-through decorator that replaces cache / token-usage wrappers."""
|
|
|
|
def wrapper(self_llm, model, messages, stream, tools=None, **kwargs):
|
|
return func(self_llm, model, messages, stream, tools, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
def _noop_stream_decorator(func):
|
|
"""Pass-through generator decorator for streaming wrappers."""
|
|
|
|
def wrapper(self_llm, model, messages, stream, tools=None, **kwargs):
|
|
yield from func(self_llm, model, messages, stream, tools, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_decorators(monkeypatch):
|
|
"""Replace cache & token-usage decorators with no-ops so tests focus on
|
|
fallback logic without needing Redis or token-counting infra."""
|
|
monkeypatch.setattr("docsgpt.llm.base.gen_cache", _noop_decorator)
|
|
monkeypatch.setattr("docsgpt.llm.base.gen_token_usage", _noop_decorator)
|
|
monkeypatch.setattr("docsgpt.llm.base.stream_cache", _noop_stream_decorator)
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.stream_token_usage", _noop_stream_decorator
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def patch_model_utils(monkeypatch):
|
|
"""Patch model_utils functions used by fallback_llm property."""
|
|
|
|
def _apply(get_provider=None, get_api_key=None, create_llm=None):
|
|
if get_provider:
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_provider_from_model_id",
|
|
get_provider,
|
|
)
|
|
if get_api_key:
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_api_key_for_provider",
|
|
get_api_key,
|
|
)
|
|
if create_llm:
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.llm_creator.LLMCreator.create_llm",
|
|
create_llm,
|
|
)
|
|
|
|
return _apply
|
|
|
|
|
|
CALL_ARGS = dict(model="test-model", messages=[{"role": "user", "content": "hi"}])
|
|
|
|
|
|
# Tests — fallback_llm property resolution
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestFallbackLLMResolution:
|
|
|
|
def test_backup_model_preferred_over_global_fallback(self, patch_model_utils):
|
|
"""When agent has backup models configured, the first valid one is used
|
|
as fallback — not the global FALLBACK_* settings."""
|
|
backup_llm = FakeLLM(responses=["backup response"])
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda mid, **_kwargs: "openai",
|
|
get_api_key=lambda prov: "fake-key",
|
|
create_llm=lambda type, **kw: backup_llm,
|
|
)
|
|
|
|
primary = FakeLLM(backup_models=["backup-model-id"])
|
|
fallback = primary.fallback_llm
|
|
|
|
assert fallback is backup_llm
|
|
|
|
def test_global_fallback_used_when_no_backup_models(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
"""When no per-agent backup models exist, global FALLBACK_* is used."""
|
|
global_fallback = FakeLLM(responses=["global fallback"])
|
|
|
|
patch_model_utils(
|
|
create_llm=lambda type, **kw: global_fallback,
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.settings",
|
|
MagicMock(
|
|
FALLBACK_LLM_PROVIDER="openai",
|
|
FALLBACK_LLM_NAME="gpt-4o",
|
|
FALLBACK_LLM_API_KEY="key",
|
|
API_KEY="key",
|
|
),
|
|
)
|
|
|
|
primary = FakeLLM(backup_models=[])
|
|
fallback = primary.fallback_llm
|
|
|
|
assert fallback is global_fallback
|
|
|
|
def test_skips_unresolvable_backup_model_tries_next(self, patch_model_utils):
|
|
"""If the first backup model can't be resolved, skip it and try the next."""
|
|
good_backup = FakeLLM(responses=["good backup"])
|
|
call_count = {"n": 0}
|
|
|
|
def fake_get_provider(model_id, **_kwargs):
|
|
call_count["n"] += 1
|
|
if model_id == "bad-model":
|
|
return None # unresolvable
|
|
return "openai"
|
|
|
|
patch_model_utils(
|
|
get_provider=fake_get_provider,
|
|
get_api_key=lambda prov: "key",
|
|
create_llm=lambda type, **kw: good_backup,
|
|
)
|
|
|
|
primary = FakeLLM(backup_models=["bad-model", "good-model"])
|
|
fallback = primary.fallback_llm
|
|
|
|
assert fallback is good_backup
|
|
assert call_count["n"] == 2 # tried both
|
|
|
|
def test_no_fallback_when_nothing_configured(self, monkeypatch):
|
|
"""No backup models + no global FALLBACK_* → fallback_llm is None."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.settings",
|
|
MagicMock(FALLBACK_LLM_PROVIDER=None),
|
|
)
|
|
primary = FakeLLM(backup_models=[])
|
|
assert primary.fallback_llm is None
|
|
|
|
|
|
# Tests — non-streaming fallback (gen)
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestNonStreamingFallback:
|
|
|
|
def test_primary_success_no_fallback(self):
|
|
"""When primary succeeds, fallback is never touched."""
|
|
primary = FakeLLM(responses=["primary ok"])
|
|
result = primary.gen(**CALL_ARGS)
|
|
assert result == "primary ok"
|
|
|
|
def test_primary_fails_uses_backup_model(self, patch_model_utils):
|
|
"""Primary fails immediately → backup model from agent config is used."""
|
|
backup = FakeLLM(responses=["backup ok"])
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda mid, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(fail_at=0, backup_models=["backup-model"])
|
|
result = primary.gen(**CALL_ARGS)
|
|
assert result == "backup ok"
|
|
assert backup.gen_called
|
|
|
|
def test_no_fallback_raises(self, monkeypatch):
|
|
"""Primary fails and no fallback configured → exception propagates."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.settings",
|
|
MagicMock(FALLBACK_LLM_PROVIDER=None),
|
|
)
|
|
primary = FakeLLM(fail_at=0, backup_models=[])
|
|
with pytest.raises(RuntimeError, match="primary model unavailable"):
|
|
primary.gen(**CALL_ARGS)
|
|
|
|
|
|
# Tests — streaming fallback (gen_stream)
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestStreamingFallback:
|
|
|
|
def test_stream_primary_success(self):
|
|
"""Full stream completes without triggering fallback."""
|
|
primary = FakeLLM(stream_chunks=["a", "b", "c"])
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
assert chunks == ["a", "b", "c"]
|
|
|
|
def test_stream_immediate_failure_uses_backup(self, patch_model_utils):
|
|
"""Primary fails before yielding anything → entire backup stream returned."""
|
|
backup = FakeLLM(stream_chunks=["fallback1", "fallback2"])
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
stream_chunks=["x", "y"],
|
|
fail_at=0, # fail before first chunk
|
|
backup_models=["backup-model"],
|
|
)
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
assert chunks == ["fallback1", "fallback2"]
|
|
assert backup.gen_stream_called
|
|
|
|
def test_stream_mid_stream_failure_uses_backup(self, patch_model_utils):
|
|
"""Primary yields some chunks then fails → backup stream follows partial output."""
|
|
backup = FakeLLM(stream_chunks=["recovery1", "recovery2"])
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
stream_chunks=["ok1", "ok2", "ok3"],
|
|
fail_at=2, # yields ok1, ok2, then fails before ok3
|
|
backup_models=["backup-model"],
|
|
)
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
# First two from primary, then full backup stream
|
|
assert chunks == ["ok1", "ok2", "recovery1", "recovery2"]
|
|
|
|
def test_stream_no_fallback_raises(self, monkeypatch):
|
|
"""Primary stream fails and no fallback → exception propagates."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.settings",
|
|
MagicMock(FALLBACK_LLM_PROVIDER=None),
|
|
)
|
|
primary = FakeLLM(stream_chunks=["x"], fail_at=0, backup_models=[])
|
|
with pytest.raises(RuntimeError, match="mid-stream failure"):
|
|
list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
def test_stream_transport_error_retries_primary_once(
|
|
self, patch_model_utils
|
|
):
|
|
"""Transport error before any yield → retry same primary once,
|
|
succeed, and skip the fallback entirely.
|
|
|
|
Covers the Azure Responses-API pattern (Front Door reset within
|
|
seconds, no output produced): the request never reached a
|
|
content-producing state, so the same primary is safe to replay.
|
|
"""
|
|
backup = FakeLLM(stream_chunks=["fallback"])
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = FakeLLM(
|
|
stream_chunks=["ok1", "ok2"],
|
|
backup_models=["backup-model"],
|
|
)
|
|
# First attempt: transport error before yielding. Second attempt:
|
|
# full stream. The fallback should NEVER be called.
|
|
primary.fail_schedule = [0, None]
|
|
primary.error_schedule = [httpx.RemoteProtocolError, RuntimeError]
|
|
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert chunks == ["ok1", "ok2"]
|
|
assert primary.stream_calls == 2
|
|
assert not backup.gen_stream_called
|
|
|
|
def test_stream_transport_error_retry_then_fails_uses_fallback(
|
|
self, patch_model_utils
|
|
):
|
|
"""Transport error both times → after retry exhausted, fall back."""
|
|
backup = FakeLLM(stream_chunks=["b1", "b2"])
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = FakeLLM(
|
|
stream_chunks=["ok1"],
|
|
backup_models=["backup-model"],
|
|
)
|
|
primary.fail_schedule = [0, 0]
|
|
primary.error_schedule = [
|
|
httpx.RemoteProtocolError,
|
|
httpx.RemoteProtocolError,
|
|
]
|
|
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert chunks == ["b1", "b2"]
|
|
assert primary.stream_calls == 2
|
|
assert backup.gen_stream_called
|
|
|
|
def test_stream_transport_error_after_yield_skips_retry(
|
|
self, patch_model_utils
|
|
):
|
|
"""Transport error AFTER a chunk was yielded → don't retry the
|
|
primary (would duplicate delivered content), go straight to
|
|
fallback.
|
|
"""
|
|
backup = FakeLLM(stream_chunks=["b1"])
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = FakeLLM(
|
|
stream_chunks=["ok1", "ok2", "ok3"],
|
|
fail_at=2, # yields ok1, ok2, then RemoteProtocolError
|
|
error_class=httpx.RemoteProtocolError,
|
|
backup_models=["backup-model"],
|
|
)
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
# Primary emitted two chunks, then straight to fallback (no retry).
|
|
assert chunks == ["ok1", "ok2", "b1"]
|
|
assert primary.stream_calls == 1
|
|
assert backup.gen_stream_called
|
|
|
|
def test_stream_non_transport_error_skips_retry(self, patch_model_utils):
|
|
"""Non-transport errors (e.g. app-level RuntimeError, 4xx) don't
|
|
get the retry — repeating won't help — but fallback still runs."""
|
|
backup = FakeLLM(stream_chunks=["b1"])
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = FakeLLM(
|
|
stream_chunks=["x"],
|
|
fail_at=0,
|
|
error_class=RuntimeError,
|
|
backup_models=["backup-model"],
|
|
)
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert chunks == ["b1"]
|
|
assert primary.stream_calls == 1 # NOT retried
|
|
assert backup.gen_stream_called
|
|
|
|
def test_stream_transport_error_no_fallback_still_retries(
|
|
self, monkeypatch
|
|
):
|
|
"""Retry logic runs even when no fallback is configured — a
|
|
retryable transport blip should not require a backup to recover.
|
|
"""
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.settings",
|
|
MagicMock(FALLBACK_LLM_PROVIDER=None),
|
|
)
|
|
primary = FakeLLM(
|
|
stream_chunks=["ok1"],
|
|
backup_models=[],
|
|
)
|
|
primary.fail_schedule = [0, None]
|
|
primary.error_schedule = [httpx.RemoteProtocolError, RuntimeError]
|
|
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert chunks == ["ok1"]
|
|
assert primary.stream_calls == 2
|
|
|
|
def test_retry_that_reaches_finish_then_trailing_frame_fails_skips_fallback(
|
|
self, patch_model_utils
|
|
):
|
|
"""The retry may deliver the entire answer, set
|
|
``_stream_reached_finish``, and then die on a trailing frame
|
|
(usage-only chunk, [DONE]). Without a re-check of the flag in
|
|
the retry's except, we'd fall through to fallback and the user
|
|
would receive the whole answer twice. This test pins the guard.
|
|
|
|
Setup: primary attempt 1 fails immediately with a retryable
|
|
transport error → retry runs, streams two chunks AND sets
|
|
``_stream_reached_finish=True``, then raises on the last frame.
|
|
Expected: the two retry chunks are yielded, the fallback is
|
|
NOT engaged, and the trailing-frame exception is re-raised
|
|
(which the streaming handler treats as non-fatal via its own
|
|
``_stream_reached_finish`` guard).
|
|
"""
|
|
backup = FakeLLM(stream_chunks=["should-not-appear"])
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
# A primary whose retry attempt yields chunks AND flips the
|
|
# finish flag before raising on the trailing frame. Subclassing
|
|
# keeps the fixture stubs untouched for the other tests.
|
|
class FinishThenFailPrimary(FakeLLM):
|
|
def _raw_gen_stream(
|
|
self, baseself, model, messages, stream, tools=None, **kwargs
|
|
):
|
|
self.stream_calls = getattr(self, "stream_calls", 0) + 1
|
|
if self.stream_calls == 1:
|
|
# Immediate transport failure to trigger the retry.
|
|
raise httpx.RemoteProtocolError("initial reset")
|
|
# Retry attempt: yield the full stream, then die on a
|
|
# trailing frame after the finish signal was delivered.
|
|
for chunk in ("a", "b"):
|
|
yield chunk
|
|
self._stream_reached_finish = True
|
|
raise httpx.RemoteProtocolError("trailing-frame drop")
|
|
|
|
primary = FinishThenFailPrimary(
|
|
stream_chunks=[],
|
|
backup_models=["backup-model"],
|
|
)
|
|
|
|
# The trailing-frame RemoteProtocolError propagates because the
|
|
# guard re-raises. The handler layer swallows it via its own
|
|
# ``_stream_reached_finish`` check; at the base-LLM layer it's
|
|
# the correct behaviour.
|
|
chunks = []
|
|
with pytest.raises(httpx.RemoteProtocolError, match="trailing-frame"):
|
|
for chunk in primary.gen_stream(**CALL_ARGS):
|
|
chunks.append(chunk)
|
|
|
|
assert chunks == ["a", "b"]
|
|
assert primary.stream_calls == 2
|
|
assert not backup.gen_stream_called
|
|
|
|
def test_fallback_emits_stream_start_with_fallback_provider(
|
|
self, patch_model_utils, caplog
|
|
):
|
|
# The fallback raw-stream path bypasses ``gen_stream``, so it must
|
|
# emit its own ``llm_stream_start`` event tagged with the fallback
|
|
# vendor — otherwise dashboards record only the failed primary
|
|
# even when the response came from the backup.
|
|
import logging as _logging
|
|
|
|
class FallbackProvider(FakeLLM):
|
|
provider_name = "fallback-vendor"
|
|
|
|
backup = FallbackProvider(
|
|
stream_chunks=["b1"], model_id="backup-model-id"
|
|
)
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
class PrimaryProvider(FakeLLM):
|
|
provider_name = "primary-vendor"
|
|
|
|
primary = PrimaryProvider(
|
|
stream_chunks=["x"],
|
|
fail_at=0,
|
|
backup_models=["backup-model-id"],
|
|
)
|
|
|
|
with caplog.at_level(_logging.INFO, logger="root"):
|
|
list(
|
|
primary.gen_stream(
|
|
model="primary-model",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
)
|
|
)
|
|
|
|
starts = [r for r in caplog.records if r.message == "llm_stream_start"]
|
|
assert len(starts) == 2
|
|
assert starts[0].provider == "primary-vendor"
|
|
assert starts[0].model == "primary-model"
|
|
assert starts[1].provider == "fallback-vendor"
|
|
assert starts[1].model == "backup-model-id"
|
|
|
|
|
|
# Tests — fallback never re-enters the orchestrator (Option B regression)
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestFallbackNoRecursion:
|
|
"""When the primary fails, _execute_with_fallback applies decorators to
|
|
the fallback's raw method directly. The fallback's own ``fallback_llm``
|
|
property must never be accessed — otherwise a fallback failure would
|
|
re-enter the orchestrator and walk the global FALLBACK_LLM_* chain
|
|
unboundedly."""
|
|
|
|
def test_backup_fallback_llm_property_never_accessed_on_gen_failure(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
backup = FakeLLM(fail_at=0) # backup also fails
|
|
|
|
accessed_on = []
|
|
original_property = BaseLLM.fallback_llm
|
|
|
|
def tracked_fallback_llm(self_llm):
|
|
accessed_on.append(self_llm)
|
|
return original_property.fget(self_llm)
|
|
|
|
monkeypatch.setattr(
|
|
BaseLLM, "fallback_llm", property(tracked_fallback_llm)
|
|
)
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(fail_at=0, backup_models=["backup-model"])
|
|
with pytest.raises(RuntimeError, match="primary model unavailable"):
|
|
primary.gen(**CALL_ARGS)
|
|
|
|
assert primary in accessed_on # primary lazy-loaded its fallback
|
|
assert backup not in accessed_on # backup's chain was never walked
|
|
|
|
def test_backup_fallback_llm_property_never_accessed_on_stream_failure(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
backup = FakeLLM(stream_chunks=["x"], fail_at=0)
|
|
|
|
accessed_on = []
|
|
original_property = BaseLLM.fallback_llm
|
|
|
|
def tracked_fallback_llm(self_llm):
|
|
accessed_on.append(self_llm)
|
|
return original_property.fget(self_llm)
|
|
|
|
monkeypatch.setattr(
|
|
BaseLLM, "fallback_llm", property(tracked_fallback_llm)
|
|
)
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
stream_chunks=["y"], fail_at=0, backup_models=["backup-model"]
|
|
)
|
|
with pytest.raises(RuntimeError, match="mid-stream failure"):
|
|
list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert primary in accessed_on
|
|
assert backup not in accessed_on
|
|
|
|
def test_fallback_failure_propagates_without_chain(self, patch_model_utils):
|
|
"""When both primary and fallback fail, the fallback's exception
|
|
propagates cleanly — no third hop, no extra retries."""
|
|
backup = FakeLLM(fail_at=0)
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(fail_at=0, backup_models=["backup-model"])
|
|
with pytest.raises(RuntimeError, match="primary model unavailable"):
|
|
primary.gen(**CALL_ARGS)
|
|
|
|
assert backup.gen_called # confirms fallback raw method WAS invoked
|
|
|
|
|
|
# Tests — backup model priority over global fallback
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestBackupModelPriority:
|
|
|
|
def test_agent_backup_tried_before_global_on_gen_failure(self, patch_model_utils):
|
|
"""On gen() failure, agent's backup model is used — not the global fallback."""
|
|
backup = FakeLLM(responses=["agent backup"])
|
|
created_models = []
|
|
|
|
def fake_create_llm(type, **kw):
|
|
created_models.append(kw.get("model_id"))
|
|
return backup
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=fake_create_llm,
|
|
)
|
|
|
|
primary = FakeLLM(fail_at=0, backup_models=["agent-backup-model"])
|
|
result = primary.gen(**CALL_ARGS)
|
|
|
|
assert result == "agent backup"
|
|
assert "agent-backup-model" in created_models
|
|
|
|
def test_agent_backup_tried_before_global_on_stream_failure(
|
|
self, patch_model_utils
|
|
):
|
|
"""On gen_stream() failure, agent's backup model is used — not the global."""
|
|
backup = FakeLLM(stream_chunks=["agent-stream"])
|
|
created_models = []
|
|
|
|
def fake_create_llm(type, **kw):
|
|
created_models.append(kw.get("model_id"))
|
|
return backup
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=fake_create_llm,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
stream_chunks=["x"], fail_at=0, backup_models=["agent-backup-model"]
|
|
)
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert chunks == ["agent-stream"]
|
|
assert "agent-backup-model" in created_models
|
|
|
|
def test_global_fallback_used_when_all_backup_models_fail(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
"""If every agent backup model fails to initialize, fall through to global."""
|
|
global_fallback = FakeLLM(responses=["global ok"])
|
|
call_order = []
|
|
|
|
def fake_get_provider(mid, **_kwargs):
|
|
if mid == "broken-backup":
|
|
return "nonexistent_provider"
|
|
return "openai"
|
|
|
|
def fake_create_llm(type, **kw):
|
|
model_id = kw.get("model_id")
|
|
call_order.append(model_id)
|
|
if model_id == "broken-backup":
|
|
raise ValueError("provider init failed")
|
|
return global_fallback
|
|
|
|
patch_model_utils(
|
|
get_provider=fake_get_provider,
|
|
get_api_key=lambda p: "k",
|
|
create_llm=fake_create_llm,
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.settings",
|
|
MagicMock(
|
|
FALLBACK_LLM_PROVIDER="openai",
|
|
FALLBACK_LLM_NAME="global-model",
|
|
FALLBACK_LLM_API_KEY="gk",
|
|
API_KEY="gk",
|
|
),
|
|
)
|
|
|
|
primary = FakeLLM(fail_at=0, backup_models=["broken-backup"])
|
|
result = primary.gen(**CALL_ARGS)
|
|
|
|
assert result == "global ok"
|
|
# Tried broken-backup first, then fell through to global-model
|
|
assert call_order == ["broken-backup", "global-model"]
|
|
|
|
|
|
# Tests — fallback uses its own model_id, not the primary's
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestFallbackModelIdOverride:
|
|
"""The fallback LLM must be called with its own model_id — not the
|
|
primary's. Otherwise providers like Groq receive an unknown model name
|
|
(e.g. a Qwen model_id) and return 404."""
|
|
|
|
def test_gen_fallback_receives_own_model_id(self, patch_model_utils):
|
|
"""Non-streaming: fallback.gen() is called with fallback.model_id."""
|
|
backup = FakeLLM(
|
|
responses=["backup ok"], model_id="groq-gpt-oss-120b"
|
|
)
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "groq",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
fail_at=0,
|
|
model_id="qwen/qwen3-4b-2507",
|
|
backup_models=["groq-gpt-oss-120b"],
|
|
)
|
|
result = primary.gen(**CALL_ARGS)
|
|
|
|
assert result == "backup ok"
|
|
assert backup.last_model_received == "groq-gpt-oss-120b"
|
|
|
|
def test_gen_stream_fallback_receives_own_model_id(self, patch_model_utils):
|
|
"""Streaming: fallback.gen_stream() is called with fallback.model_id."""
|
|
backup = FakeLLM(
|
|
stream_chunks=["ok"], model_id="groq-gpt-oss-120b"
|
|
)
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "groq",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
stream_chunks=["x"],
|
|
fail_at=0,
|
|
model_id="qwen/qwen3-4b-2507",
|
|
backup_models=["groq-gpt-oss-120b"],
|
|
)
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert chunks == ["ok"]
|
|
assert backup.last_model_received == "groq-gpt-oss-120b"
|
|
|
|
def test_mid_stream_fallback_receives_own_model_id(self, patch_model_utils):
|
|
"""Mid-stream failure: fallback still gets its own model_id, not the
|
|
primary's that was already partially streaming."""
|
|
backup = FakeLLM(
|
|
stream_chunks=["recovered"], model_id="groq-gpt-oss-120b"
|
|
)
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "groq",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
stream_chunks=["partial1", "partial2", "boom"],
|
|
fail_at=2,
|
|
model_id="qwen/qwen3-4b-2507",
|
|
backup_models=["groq-gpt-oss-120b"],
|
|
)
|
|
chunks = list(primary.gen_stream(**CALL_ARGS))
|
|
|
|
assert chunks == ["partial1", "partial2", "recovered"]
|
|
assert backup.last_model_received == "groq-gpt-oss-120b"
|
|
|
|
|
|
# Tests — model_user_id (BYOM owner scope) propagates into fallback resolution
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestFallbackModelUserIdScope:
|
|
"""A shared agent dispatched by user B but owned by user A stores
|
|
A's BYOM UUIDs as backup_models. Without the P2 fix the fallback
|
|
property looks up those UUIDs against ``decoded_token['sub']`` (B,
|
|
the caller), which can't see A's per-user layer — backups are
|
|
silently skipped and the global FALLBACK_* settings are used
|
|
instead. These tests pin down that ``model_user_id`` (the owner)
|
|
is used both for the registry lookup and for the recursive
|
|
``LLMCreator.create_llm`` call."""
|
|
|
|
def test_backup_lookup_uses_model_user_id_not_caller(
|
|
self, patch_model_utils
|
|
):
|
|
captured = {"user_id": None}
|
|
|
|
def fake_get_provider(model_id, **kwargs):
|
|
captured["user_id"] = kwargs.get("user_id")
|
|
return "openai"
|
|
|
|
backup = FakeLLM(responses=["ok"])
|
|
patch_model_utils(
|
|
get_provider=fake_get_provider,
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
decoded_token={"sub": "caller-bob"},
|
|
model_user_id="owner-alice",
|
|
backup_models=["alice-byom-uuid"],
|
|
)
|
|
_ = primary.fallback_llm
|
|
assert captured["user_id"] == "owner-alice"
|
|
|
|
def test_backup_create_llm_receives_model_user_id(self, patch_model_utils):
|
|
backup = FakeLLM(responses=["ok"])
|
|
captured = {}
|
|
|
|
def fake_create_llm(type, **kw):
|
|
captured["model_user_id"] = kw.get("model_user_id")
|
|
captured["model_id"] = kw.get("model_id")
|
|
return backup
|
|
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=fake_create_llm,
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
decoded_token={"sub": "caller-bob"},
|
|
model_user_id="owner-alice",
|
|
backup_models=["alice-byom-uuid"],
|
|
)
|
|
_ = primary.fallback_llm
|
|
assert captured["model_user_id"] == "owner-alice"
|
|
assert captured["model_id"] == "alice-byom-uuid"
|
|
|
|
def test_global_fallback_create_llm_receives_model_user_id(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
"""The global FALLBACK_LLM_NAME path must also forward
|
|
``model_user_id`` — operators can configure it to a BYOM UUID
|
|
that's owned by the same user as the primary model."""
|
|
backup = FakeLLM(responses=["ok"])
|
|
captured = {}
|
|
|
|
def fake_create_llm(type, **kw):
|
|
captured["model_user_id"] = kw.get("model_user_id")
|
|
return backup
|
|
|
|
patch_model_utils(create_llm=fake_create_llm)
|
|
monkeypatch.setattr(
|
|
"docsgpt.llm.base.settings",
|
|
MagicMock(
|
|
FALLBACK_LLM_PROVIDER="openai",
|
|
FALLBACK_LLM_NAME="some-uuid",
|
|
FALLBACK_LLM_API_KEY="k",
|
|
API_KEY="k",
|
|
),
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
decoded_token={"sub": "caller-bob"},
|
|
model_user_id="owner-alice",
|
|
backup_models=[],
|
|
)
|
|
_ = primary.fallback_llm
|
|
assert captured["model_user_id"] == "owner-alice"
|
|
|
|
def test_falls_back_to_caller_when_model_user_id_unset(
|
|
self, patch_model_utils
|
|
):
|
|
"""Built-in models / pre-P2 callers don't pass model_user_id.
|
|
In that case the caller's sub is still used — preserving
|
|
existing behaviour."""
|
|
captured = {}
|
|
|
|
def fake_get_provider(model_id, **kwargs):
|
|
captured["user_id"] = kwargs.get("user_id")
|
|
return "openai"
|
|
|
|
patch_model_utils(
|
|
get_provider=fake_get_provider,
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: FakeLLM(responses=["ok"]),
|
|
)
|
|
|
|
primary = FakeLLM(
|
|
decoded_token={"sub": "caller-bob"},
|
|
model_user_id=None,
|
|
backup_models=["some-builtin-id"],
|
|
)
|
|
_ = primary.fallback_llm
|
|
assert captured["user_id"] == "caller-bob"
|
|
|
|
|
|
# Tests — LLMCreator wires model_user_id through to BaseLLM
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLLMCreatorPassesModelUserId:
|
|
"""End-to-end through ``LLMCreator.create_llm``: the constructed
|
|
LLM must store ``model_user_id`` so its fallback property can
|
|
resolve under the right scope."""
|
|
|
|
def test_model_user_id_set_on_constructed_llm(self, monkeypatch):
|
|
from docsgpt.llm.llm_creator import LLMCreator
|
|
from docsgpt.llm.providers import PROVIDERS_BY_NAME
|
|
|
|
captured = {}
|
|
|
|
class _CapturingLLM:
|
|
def __init__(self, api_key, user_api_key, *args, **kwargs):
|
|
captured["model_user_id"] = kwargs.get("model_user_id")
|
|
|
|
# Pick any registered provider — we only need the constructor
|
|
# call to land in our fake.
|
|
monkeypatch.setattr(
|
|
PROVIDERS_BY_NAME["openai"], "llm_class", _CapturingLLM
|
|
)
|
|
|
|
LLMCreator.create_llm(
|
|
type="openai",
|
|
api_key="k",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "caller-bob"},
|
|
model_id=None,
|
|
model_user_id="owner-alice",
|
|
)
|
|
|
|
assert captured["model_user_id"] == "owner-alice"
|
|
|
|
|
|
# Tests — responding-provider tracking (cross-provider fallback handler fix)
|
|
|
|
|
|
class _Google(FakeLLM):
|
|
provider_name = "google"
|
|
|
|
|
|
class _OpenAI(FakeLLM):
|
|
provider_name = "openai"
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestRespondingProviderTracking:
|
|
"""The handler that parses a response must follow the model that
|
|
actually produced it. ``BaseLLM`` exposes ``_responding_provider`` so
|
|
the handler layer can re-route ``parse_response`` after a fallback to a
|
|
different-provider model (Google primary -> OpenAI backup), instead of
|
|
silently dropping the backup's tool calls."""
|
|
|
|
def test_defaults_to_own_provider_before_any_call(self):
|
|
assert _Google()._responding_provider == "google"
|
|
|
|
def test_stream_success_keeps_primary_provider(self):
|
|
primary = _Google(stream_chunks=["a", "b"])
|
|
list(primary.gen_stream(**CALL_ARGS))
|
|
assert primary._responding_provider == "google"
|
|
|
|
def test_stream_fallback_records_backup_provider(self, patch_model_utils):
|
|
backup = _OpenAI(stream_chunks=["x"])
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = _Google(
|
|
stream_chunks=["a"], fail_at=0, backup_models=["backup-model"]
|
|
)
|
|
list(primary.gen_stream(**CALL_ARGS))
|
|
assert primary._responding_provider == "openai"
|
|
|
|
def test_gen_success_keeps_primary_provider(self):
|
|
primary = _Google(responses=["ok"])
|
|
primary.gen(**CALL_ARGS)
|
|
assert primary._responding_provider == "google"
|
|
|
|
def test_gen_fallback_records_backup_provider(self, patch_model_utils):
|
|
backup = _OpenAI(responses=["backup ok"])
|
|
patch_model_utils(
|
|
get_provider=lambda m, **_kwargs: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = _Google(fail_at=0, backup_models=["backup-model"])
|
|
primary.gen(**CALL_ARGS)
|
|
assert primary._responding_provider == "openai"
|
|
|
|
|
|
# Tests — fallback payload size gate
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestFallbackPayloadSizeGate:
|
|
"""A payload that cannot fit the fallback's context window must skip the
|
|
fallback attempt (it would be a guaranteed second rejection) and
|
|
propagate the primary's error instead."""
|
|
|
|
def _primary_with_backup(self, patch_model_utils, backup):
|
|
patch_model_utils(
|
|
get_provider=lambda mid, **_kw: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
return FakeLLM(fail_at=0, backup_models=["backup-model"])
|
|
|
|
BIG_ARGS = dict(
|
|
model="test-model",
|
|
messages=[{"role": "user", "content": "word " * 300}],
|
|
)
|
|
|
|
def test_gen_skips_fallback_when_payload_cannot_fit(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
backup = FakeLLM(responses=["backup ok"])
|
|
primary = self._primary_with_backup(patch_model_utils, backup)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_token_limit",
|
|
lambda mid, user_id=None: 10,
|
|
)
|
|
with pytest.raises(RuntimeError, match="primary model unavailable"):
|
|
primary.gen(**self.BIG_ARGS)
|
|
assert backup.gen_called is False
|
|
|
|
def test_stream_skips_fallback_when_payload_cannot_fit(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
backup = FakeLLM(stream_chunks=["backup chunk"])
|
|
primary = self._primary_with_backup(patch_model_utils, backup)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_token_limit",
|
|
lambda mid, user_id=None: 10,
|
|
)
|
|
with pytest.raises(RuntimeError, match="mid-stream failure"):
|
|
list(primary.gen_stream(**self.BIG_ARGS))
|
|
assert backup.gen_stream_called is False
|
|
|
|
def test_fallback_proceeds_when_payload_fits(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
backup = FakeLLM(responses=["backup ok"])
|
|
primary = self._primary_with_backup(patch_model_utils, backup)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_token_limit",
|
|
lambda mid, user_id=None: 100000,
|
|
)
|
|
assert primary.gen(**self.BIG_ARGS) == "backup ok"
|
|
assert backup.gen_called is True
|
|
|
|
def test_estimation_failure_never_blocks_fallback(
|
|
self, monkeypatch, patch_model_utils
|
|
):
|
|
backup = FakeLLM(responses=["backup ok"])
|
|
primary = self._primary_with_backup(patch_model_utils, backup)
|
|
|
|
def boom(*a, **kw):
|
|
raise ValueError("estimator broken")
|
|
|
|
monkeypatch.setattr("docsgpt.usage._count_prompt_tokens", boom)
|
|
assert primary.gen(**self.BIG_ARGS) == "backup ok"
|
|
assert backup.gen_called is True
|
|
|
|
|
|
# Tests — no fallback restream after the primary delivered its finish signal
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestNoRestreamAfterFinish:
|
|
"""A trailing-frame failure (between the finish chunk and stream end)
|
|
must NOT restream the already-delivered answer from the fallback."""
|
|
|
|
class FinishThenFailLLM(FakeLLM):
|
|
def _raw_gen_stream(self, baseself, model, messages, stream, tools=None, **kwargs):
|
|
self.gen_stream_called = True
|
|
yield "the full answer"
|
|
# Provider marked the stream finished (finish_reason arrived)...
|
|
self._stream_reached_finish = True
|
|
# ...then the trailing usage/[DONE] frame dies.
|
|
raise RuntimeError("connection reset in trailing frame")
|
|
|
|
def test_trailing_frame_failure_skips_fallback(self, patch_model_utils):
|
|
backup = FakeLLM(stream_chunks=["fallback chunk"])
|
|
patch_model_utils(
|
|
get_provider=lambda mid, **_kw: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = self.FinishThenFailLLM(backup_models=["backup-model"])
|
|
|
|
received = []
|
|
with pytest.raises(RuntimeError, match="trailing frame"):
|
|
for chunk in primary.gen_stream(**CALL_ARGS):
|
|
received.append(chunk)
|
|
|
|
assert received == ["the full answer"]
|
|
assert backup.gen_stream_called is False
|
|
|
|
def test_pre_finish_failure_still_falls_back(self, patch_model_utils):
|
|
backup = FakeLLM(stream_chunks=["fallback chunk"])
|
|
patch_model_utils(
|
|
get_provider=lambda mid, **_kw: "openai",
|
|
get_api_key=lambda p: "k",
|
|
create_llm=lambda type, **kw: backup,
|
|
)
|
|
primary = FakeLLM(fail_at=0, backup_models=["backup-model"])
|
|
out = list(primary.gen_stream(**CALL_ARGS))
|
|
assert out == ["fallback chunk"]
|
|
assert backup.gen_stream_called is True
|
|
|
|
|
|
# Tests — fallback message reshaping (parts arrays prepared for the primary)
|
|
#
|
|
# ``prepare_messages_with_attachments`` runs against the *primary* model, so
|
|
# by the time a fallback engages the messages can carry ``file`` parts (whose
|
|
# Files-API ids only the primary's endpoint+credential can resolve) and
|
|
# ``image_url`` parts a non-vision fallback 4xxes on. These tests pin the
|
|
# handoff contract: the fallback must receive content it can accept.
|
|
|
|
|
|
def _parts_messages():
|
|
return [
|
|
{"role": "system", "content": "You are helpful."},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "summarize the attached report"},
|
|
{"type": "file", "file": {"file_id": "assistant-abc123"}},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "data:image/png;base64,AAAA"},
|
|
},
|
|
],
|
|
},
|
|
]
|
|
|
|
|
|
ATTACHMENTS = [
|
|
{"id": "att-1", "filename": "report.pdf", "content": "EXTRACTED REPORT TEXT"}
|
|
]
|
|
|
|
|
|
class VisionFakeLLM(FakeLLM):
|
|
def get_supported_attachment_types(self):
|
|
return ["image/png", "image/jpeg"]
|
|
|
|
|
|
class SharedEndpointPdfFakeLLM(FakeLLM):
|
|
def get_supported_attachment_types(self):
|
|
return ["application/pdf", "image/png"]
|
|
|
|
def _endpoint_scope(self):
|
|
return "scope-shared"
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestFallbackMessageReshaping:
|
|
|
|
def _run_stream(self, primary, messages, attachments):
|
|
return list(
|
|
primary.gen_stream(
|
|
model="test-model",
|
|
messages=messages,
|
|
_usage_attachments=attachments,
|
|
)
|
|
)
|
|
|
|
def test_stream_fallback_gets_flattened_string_content(self):
|
|
fallback = FakeLLM(stream_chunks=["fb"])
|
|
primary = FakeLLM(fail_at=0)
|
|
primary._fallback_llm = fallback
|
|
original = _parts_messages()
|
|
snapshot = copy.deepcopy(original)
|
|
|
|
chunks = self._run_stream(primary, original, ATTACHMENTS)
|
|
|
|
assert chunks == ["fb"]
|
|
received = fallback.last_messages_received
|
|
assert received[0] == {"role": "system", "content": "You are helpful."}
|
|
user_content = received[1]["content"]
|
|
assert isinstance(user_content, str)
|
|
assert "summarize the attached report" in user_content
|
|
assert "EXTRACTED REPORT TEXT" in user_content
|
|
assert "Image attachment omitted" in user_content
|
|
# The primary-scoped Files-API id must never reach another endpoint.
|
|
assert "assistant-abc123" not in user_content
|
|
# The primary's own message array is not mutated.
|
|
assert original == snapshot
|
|
|
|
def test_stream_vision_fallback_keeps_image_parts(self):
|
|
fallback = VisionFakeLLM(stream_chunks=["fb"])
|
|
primary = FakeLLM(fail_at=0)
|
|
primary._fallback_llm = fallback
|
|
|
|
self._run_stream(primary, _parts_messages(), ATTACHMENTS)
|
|
|
|
user_content = fallback.last_messages_received[1]["content"]
|
|
assert isinstance(user_content, list)
|
|
types = [part["type"] for part in user_content]
|
|
assert "image_url" in types
|
|
assert "file" not in types
|
|
joined = " ".join(
|
|
part.get("text", "") for part in user_content if part["type"] == "text"
|
|
)
|
|
assert "EXTRACTED REPORT TEXT" in joined
|
|
|
|
def test_stream_same_endpoint_pdf_fallback_keeps_file_parts(self):
|
|
fallback = SharedEndpointPdfFakeLLM(stream_chunks=["fb"])
|
|
primary = SharedEndpointPdfFakeLLM(fail_at=0)
|
|
primary._fallback_llm = fallback
|
|
|
|
self._run_stream(primary, _parts_messages(), ATTACHMENTS)
|
|
|
|
user_content = fallback.last_messages_received[1]["content"]
|
|
assert isinstance(user_content, list)
|
|
assert {"type": "file", "file": {"file_id": "assistant-abc123"}} in user_content
|
|
|
|
def test_stream_file_part_without_extracted_content_becomes_note(self):
|
|
fallback = FakeLLM(stream_chunks=["fb"])
|
|
primary = FakeLLM(fail_at=0)
|
|
primary._fallback_llm = fallback
|
|
|
|
self._run_stream(primary, _parts_messages(), attachments=None)
|
|
|
|
user_content = fallback.last_messages_received[1]["content"]
|
|
assert isinstance(user_content, str)
|
|
assert "could not be included" in user_content
|
|
assert "assistant-abc123" not in user_content
|
|
|
|
def test_stream_string_messages_pass_through_unchanged(self):
|
|
fallback = FakeLLM(stream_chunks=["fb"])
|
|
primary = FakeLLM(fail_at=0)
|
|
primary._fallback_llm = fallback
|
|
messages = [{"role": "user", "content": "plain text"}]
|
|
|
|
self._run_stream(primary, messages, ATTACHMENTS)
|
|
|
|
assert fallback.last_messages_received == messages
|
|
|
|
def test_gen_fallback_gets_flattened_string_content(self):
|
|
fallback = FakeLLM(responses=["fb answer"])
|
|
primary = FakeLLM(fail_at=0)
|
|
primary._fallback_llm = fallback
|
|
|
|
result = primary.gen(
|
|
model="test-model",
|
|
messages=_parts_messages(),
|
|
_usage_attachments=ATTACHMENTS,
|
|
)
|
|
|
|
assert result == "fb answer"
|
|
user_content = fallback.last_messages_received[1]["content"]
|
|
assert isinstance(user_content, str)
|
|
assert "EXTRACTED REPORT TEXT" in user_content
|
|
assert "assistant-abc123" not in user_content
|
|
|
|
|
|
# Tests — cross-provider structured-output adaptation
|
|
#
|
|
# Structured output is provider-specific: OpenAI-wire classes take
|
|
# ``response_format``, Google takes ``response_schema``. Forwarding the
|
|
# primary's kwarg verbatim either loses enforcement silently (Google swallows
|
|
# ``response_format`` in ``**kwargs``) or raises TypeError inside the OpenAI
|
|
# SDK (``response_schema`` is not a Chat-Completions param) — which turned the
|
|
# fallback into no fallback at all for every structured node.
|
|
|
|
|
|
SCHEMA = {
|
|
"type": "object",
|
|
"properties": {
|
|
"answer": {"type": "string"},
|
|
"score": {"type": "integer"},
|
|
},
|
|
"required": ["answer"],
|
|
}
|
|
|
|
|
|
class _OpenAIWireFake(FakeLLM):
|
|
"""OpenAI-wire double: real declaration + real preparer, fake transport."""
|
|
|
|
provider_name = "openai"
|
|
structured_output_kwarg = "response_format"
|
|
prepare_structured_output_format = OpenAILLM.prepare_structured_output_format
|
|
|
|
def _supports_structured_output(self):
|
|
return True
|
|
|
|
|
|
class _GoogleFake(FakeLLM):
|
|
"""Google double: real declaration + real preparer, fake transport."""
|
|
|
|
provider_name = "google"
|
|
structured_output_kwarg = "response_schema"
|
|
prepare_structured_output_format = GoogleLLM.prepare_structured_output_format
|
|
|
|
def _supports_structured_output(self):
|
|
return True
|
|
|
|
|
|
class _AnthropicFake(FakeLLM):
|
|
"""Provider with no structured-output kwarg at all."""
|
|
|
|
provider_name = "anthropic"
|
|
|
|
|
|
def _openai_envelope(schema=SCHEMA, strict=True):
|
|
"""An OpenAI ``response_format`` built the way a provider would."""
|
|
return OpenAILLM.prepare_structured_output_format(
|
|
_OpenAIWireFake(), schema, strict=strict
|
|
)
|
|
|
|
|
|
def _google_schema(schema=SCHEMA):
|
|
"""The Google ``response_schema`` conversion of ``schema``."""
|
|
return GoogleLLM.prepare_structured_output_format(_GoogleFake(), schema)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestStructuredOutputDeclarations:
|
|
"""One source of truth: the kwarg name lives on the LLM class."""
|
|
|
|
def test_openai_declares_response_format(self):
|
|
assert OpenAILLM.structured_output_kwarg == "response_format"
|
|
|
|
def test_openai_subclasses_inherit_the_declaration(self):
|
|
assert GroqLLM.structured_output_kwarg == "response_format"
|
|
|
|
def test_google_declares_response_schema(self):
|
|
assert GoogleLLM.structured_output_kwarg == "response_schema"
|
|
|
|
def test_anthropic_declares_nothing(self):
|
|
assert AnthropicLLM.structured_output_kwarg is None
|
|
|
|
def test_base_declares_nothing(self):
|
|
assert BaseLLM.structured_output_kwarg is None
|
|
|
|
def test_openai_prepare_records_source(self):
|
|
llm = _OpenAIWireFake()
|
|
llm.prepare_structured_output_format(SCHEMA, strict=False)
|
|
assert llm._structured_output_source == (SCHEMA, False)
|
|
|
|
def test_google_prepare_records_source(self):
|
|
llm = _GoogleFake()
|
|
llm.prepare_structured_output_format(SCHEMA)
|
|
assert llm._structured_output_source == (SCHEMA, True)
|
|
|
|
def test_empty_schema_clears_recorded_source(self):
|
|
llm = _OpenAIWireFake()
|
|
llm.prepare_structured_output_format(SCHEMA)
|
|
llm.prepare_structured_output_format(None)
|
|
assert llm._structured_output_source is None
|
|
|
|
def test_base_records_nothing_by_default(self):
|
|
assert FakeLLM()._structured_output_source is None
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestAdaptStructuredOutputKwargs:
|
|
"""Unit-level contract of the adapter itself."""
|
|
|
|
def test_no_structured_kwargs_returns_equal_copy(self):
|
|
primary = _OpenAIWireFake()
|
|
kwargs = {"model": "m", "messages": [], "temperature": 0.2}
|
|
|
|
adapted = primary._adapt_structured_output_kwargs(_GoogleFake(), kwargs)
|
|
|
|
assert adapted == kwargs
|
|
assert adapted is not kwargs
|
|
|
|
def test_does_not_mutate_the_callers_kwargs(self):
|
|
primary = _OpenAIWireFake()
|
|
primary.prepare_structured_output_format(SCHEMA)
|
|
kwargs = {"model": "m", "response_format": _openai_envelope()}
|
|
snapshot = copy.deepcopy(kwargs)
|
|
|
|
primary._adapt_structured_output_kwargs(_GoogleFake(), kwargs)
|
|
|
|
assert kwargs == snapshot
|
|
|
|
def test_recovers_schema_from_envelope_when_no_source_recorded(self):
|
|
"""A hand-built ``response_format`` (research_agent) never went through
|
|
``prepare_structured_output_format``, so nothing was recorded — the raw
|
|
schema is still readable out of the OpenAI envelope."""
|
|
primary = _OpenAIWireFake()
|
|
assert primary._structured_output_source is None
|
|
|
|
adapted = primary._adapt_structured_output_kwargs(
|
|
_GoogleFake(model_id="gemini-2.5-flash"),
|
|
{"model": "m", "response_format": _openai_envelope()},
|
|
)
|
|
|
|
assert "response_format" not in adapted
|
|
assert adapted["response_schema"]["type"] == "OBJECT"
|
|
assert set(adapted["response_schema"]["properties"]) == {"answer", "score"}
|
|
assert adapted["model"] == "m"
|
|
|
|
def test_google_schema_without_source_is_dropped(self, caplog):
|
|
"""Google's conversion is lossy/type-mapped — not reversible."""
|
|
primary = _GoogleFake()
|
|
assert primary._structured_output_source is None
|
|
|
|
with caplog.at_level(logging.WARNING, logger="docsgpt.llm.base"):
|
|
adapted = primary._adapt_structured_output_kwargs(
|
|
_OpenAIWireFake(model_id="gpt-4o-mini"),
|
|
{"model": "m", "response_schema": _google_schema()},
|
|
)
|
|
|
|
assert "response_schema" not in adapted
|
|
assert "response_format" not in adapted
|
|
assert "gpt-4o-mini" in caplog.text
|
|
|
|
def test_fallback_without_structured_support_drops_and_warns(self, caplog):
|
|
primary = _OpenAIWireFake()
|
|
primary.prepare_structured_output_format(SCHEMA)
|
|
fallback = _GoogleFake(model_id="gemini-2.5-flash")
|
|
fallback._supports_structured_output = lambda: False
|
|
|
|
with caplog.at_level(logging.WARNING, logger="docsgpt.llm.base"):
|
|
adapted = primary._adapt_structured_output_kwargs(
|
|
fallback, {"model": "m", "response_format": _openai_envelope()}
|
|
)
|
|
|
|
assert "response_format" not in adapted
|
|
assert "response_schema" not in adapted
|
|
assert "cannot enforce structured output" in caplog.text
|
|
|
|
def test_non_callable_support_flag_is_honored(self):
|
|
"""Test doubles sometimes set the capability as a plain bool."""
|
|
primary = _OpenAIWireFake()
|
|
primary.prepare_structured_output_format(SCHEMA)
|
|
fallback = _GoogleFake(model_id="gemini-2.5-flash")
|
|
fallback._supports_structured_output = False
|
|
|
|
adapted = primary._adapt_structured_output_kwargs(
|
|
fallback, {"response_format": _openai_envelope()}
|
|
)
|
|
|
|
assert adapted == {}
|
|
|
|
def test_preparer_returning_none_drops_the_kwarg(self, caplog):
|
|
class _NullPreparer(_GoogleFake):
|
|
def prepare_structured_output_format(self, json_schema, strict=True):
|
|
return None
|
|
|
|
primary = _OpenAIWireFake()
|
|
primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
with caplog.at_level(logging.WARNING, logger="docsgpt.llm.base"):
|
|
adapted = primary._adapt_structured_output_kwargs(
|
|
_NullPreparer(model_id="gemini-2.5-flash"),
|
|
{"response_format": _openai_envelope()},
|
|
)
|
|
|
|
assert adapted == {}
|
|
assert "cannot enforce structured output" in caplog.text
|
|
|
|
def test_preparer_raising_never_breaks_the_fallback(self, caplog):
|
|
class _ExplodingPreparer(_GoogleFake):
|
|
def prepare_structured_output_format(self, json_schema, strict=True):
|
|
raise ValueError("boom")
|
|
|
|
primary = _OpenAIWireFake()
|
|
primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
with caplog.at_level(logging.WARNING, logger="docsgpt.llm.base"):
|
|
adapted = primary._adapt_structured_output_kwargs(
|
|
_ExplodingPreparer(model_id="gemini-2.5-flash"),
|
|
{"response_format": _openai_envelope()},
|
|
)
|
|
|
|
assert adapted == {}
|
|
assert "Failed to prepare structured output" in caplog.text
|
|
|
|
def test_strict_flag_survives_the_translation(self):
|
|
primary = _GoogleFake()
|
|
primary.prepare_structured_output_format(SCHEMA)
|
|
primary._structured_output_source = (SCHEMA, False)
|
|
|
|
adapted = primary._adapt_structured_output_kwargs(
|
|
_OpenAIWireFake(model_id="gpt-4o-mini"),
|
|
{"response_schema": _google_schema()},
|
|
)
|
|
|
|
assert adapted["response_format"]["json_schema"]["strict"] is False
|
|
# strict=False leaves the schema untouched (no additionalProperties).
|
|
assert "additionalProperties" not in (
|
|
adapted["response_format"]["json_schema"]["schema"]
|
|
)
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestCrossProviderStructuredOutputFallback:
|
|
"""End-to-end through ``gen`` / ``gen_stream``: the backup must receive
|
|
the schema in *its own* provider's kwarg."""
|
|
|
|
def test_gen_openai_primary_google_fallback_gets_response_schema(self):
|
|
primary = _OpenAIWireFake(fail_at=0)
|
|
fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash")
|
|
primary._fallback_llm = fallback
|
|
response_format = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
result = primary.gen(**CALL_ARGS, response_format=response_format)
|
|
|
|
assert result == "fb"
|
|
assert fallback.last_kwargs_received["response_schema"] == _google_schema()
|
|
assert "response_format" not in fallback.last_kwargs_received
|
|
|
|
def test_stream_openai_primary_google_fallback_gets_response_schema(self):
|
|
primary = _OpenAIWireFake(stream_chunks=["x"], fail_at=0)
|
|
fallback = _GoogleFake(stream_chunks=["fb"], model_id="gemini-2.5-flash")
|
|
primary._fallback_llm = fallback
|
|
response_format = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
chunks = list(primary.gen_stream(**CALL_ARGS, response_format=response_format))
|
|
|
|
assert chunks == ["fb"]
|
|
assert fallback.last_kwargs_received["response_schema"] == _google_schema()
|
|
assert "response_format" not in fallback.last_kwargs_received
|
|
|
|
def test_gen_google_primary_openai_fallback_gets_response_format(self):
|
|
primary = _GoogleFake(fail_at=0)
|
|
fallback = _OpenAIWireFake(responses=["fb"], model_id="gpt-4o-mini")
|
|
primary._fallback_llm = fallback
|
|
response_schema = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
result = primary.gen(**CALL_ARGS, response_schema=response_schema)
|
|
|
|
assert result == "fb"
|
|
received = fallback.last_kwargs_received
|
|
assert "response_schema" not in received
|
|
assert received["response_format"]["type"] == "json_schema"
|
|
assert set(received["response_format"]["json_schema"]["schema"]["properties"]) == {
|
|
"answer",
|
|
"score",
|
|
}
|
|
|
|
def test_stream_google_primary_openai_fallback_gets_response_format(self):
|
|
primary = _GoogleFake(stream_chunks=["x"], fail_at=0)
|
|
fallback = _OpenAIWireFake(stream_chunks=["fb"], model_id="gpt-4o-mini")
|
|
primary._fallback_llm = fallback
|
|
response_schema = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema))
|
|
|
|
assert chunks == ["fb"]
|
|
received = fallback.last_kwargs_received
|
|
assert "response_schema" not in received
|
|
assert received["response_format"]["type"] == "json_schema"
|
|
|
|
def test_same_wire_family_passes_response_format_verbatim(self):
|
|
"""OpenAI -> openai_compatible: no re-preparation, byte-identical."""
|
|
primary = _OpenAIWireFake(fail_at=0)
|
|
fallback = _OpenAIWireFake(responses=["fb"], model_id="qwen3-4b")
|
|
primary._fallback_llm = fallback
|
|
response_format = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
primary.gen(**CALL_ARGS, response_format=response_format)
|
|
|
|
assert fallback.last_kwargs_received["response_format"] is response_format
|
|
assert "response_schema" not in fallback.last_kwargs_received
|
|
|
|
def test_same_wire_family_passes_response_schema_verbatim(self):
|
|
primary = _GoogleFake(fail_at=0)
|
|
fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-pro")
|
|
primary._fallback_llm = fallback
|
|
response_schema = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
primary.gen(**CALL_ARGS, response_schema=response_schema)
|
|
|
|
assert fallback.last_kwargs_received["response_schema"] is response_schema
|
|
assert "response_format" not in fallback.last_kwargs_received
|
|
|
|
def test_json_object_mode_kept_within_the_openai_family(self):
|
|
primary = _OpenAIWireFake(fail_at=0)
|
|
fallback = _OpenAIWireFake(responses=["fb"], model_id="qwen3-4b")
|
|
primary._fallback_llm = fallback
|
|
|
|
primary.gen(**CALL_ARGS, response_format={"type": "json_object"})
|
|
|
|
assert fallback.last_kwargs_received["response_format"] == {
|
|
"type": "json_object"
|
|
}
|
|
|
|
def test_json_object_mode_dropped_for_google_fallback(self):
|
|
"""Google has no json_object equivalent wired — drop, don't crash."""
|
|
primary = _OpenAIWireFake(fail_at=0)
|
|
fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash")
|
|
primary._fallback_llm = fallback
|
|
|
|
result = primary.gen(**CALL_ARGS, response_format={"type": "json_object"})
|
|
|
|
assert result == "fb"
|
|
assert "response_format" not in fallback.last_kwargs_received
|
|
assert "response_schema" not in fallback.last_kwargs_received
|
|
|
|
def test_json_object_mode_dropped_for_google_fallback_streaming(self):
|
|
primary = _OpenAIWireFake(stream_chunks=["x"], fail_at=0)
|
|
fallback = _GoogleFake(stream_chunks=["fb"], model_id="gemini-2.5-flash")
|
|
primary._fallback_llm = fallback
|
|
|
|
chunks = list(
|
|
primary.gen_stream(**CALL_ARGS, response_format={"type": "json_object"})
|
|
)
|
|
|
|
assert chunks == ["fb"]
|
|
assert "response_format" not in fallback.last_kwargs_received
|
|
assert "response_schema" not in fallback.last_kwargs_received
|
|
|
|
def test_anthropic_fallback_gets_neither_kwarg(self):
|
|
"""Anthropic has no structured-output kwarg: unstructured, not broken."""
|
|
primary = _OpenAIWireFake(fail_at=0)
|
|
fallback = _AnthropicFake(responses=["fb"], model_id="claude-sonnet-4")
|
|
primary._fallback_llm = fallback
|
|
response_format = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
result = primary.gen(**CALL_ARGS, response_format=response_format)
|
|
|
|
assert result == "fb"
|
|
assert fallback.last_kwargs_received == {}
|
|
|
|
def test_anthropic_fallback_gets_neither_kwarg_streaming(self):
|
|
primary = _GoogleFake(stream_chunks=["x"], fail_at=0)
|
|
fallback = _AnthropicFake(stream_chunks=["fb"], model_id="claude-sonnet-4")
|
|
primary._fallback_llm = fallback
|
|
response_schema = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema))
|
|
|
|
assert chunks == ["fb"]
|
|
assert fallback.last_kwargs_received == {}
|
|
|
|
def test_unrelated_gen_kwargs_are_forwarded_untouched(self):
|
|
primary = _OpenAIWireFake(fail_at=0)
|
|
fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash")
|
|
primary._fallback_llm = fallback
|
|
response_format = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
primary.gen(
|
|
**CALL_ARGS, response_format=response_format, temperature=0.3
|
|
)
|
|
|
|
assert fallback.last_kwargs_received["temperature"] == 0.3
|
|
|
|
|
|
# The OpenAI SDK's ``chat.completions.create`` has an explicit keyword
|
|
# signature: a forwarded ``response_schema`` raises TypeError before the
|
|
# request is even built, so the Google -> OpenAI hop died with "Fallback LLM
|
|
# also failed". These tests drive the *real* ``OpenAILLM`` raw methods against
|
|
# a client double with the same strictness.
|
|
|
|
|
|
class _StrictChatCompletions:
|
|
"""Chat-Completions double that rejects kwargs the real SDK rejects."""
|
|
|
|
_ACCEPTED = {
|
|
"model",
|
|
"messages",
|
|
"stream",
|
|
"stream_options",
|
|
"tools",
|
|
"tool_choice",
|
|
"parallel_tool_calls",
|
|
"response_format",
|
|
"temperature",
|
|
"top_p",
|
|
"max_completion_tokens",
|
|
"reasoning_effort",
|
|
"presence_penalty",
|
|
"frequency_penalty",
|
|
"seed",
|
|
"stop",
|
|
"n",
|
|
"user",
|
|
}
|
|
|
|
def __init__(self):
|
|
self.last_kwargs = None
|
|
|
|
def create(self, **kwargs):
|
|
unexpected = sorted(set(kwargs) - self._ACCEPTED)
|
|
if unexpected:
|
|
raise TypeError(
|
|
f"Completions.create() got an unexpected keyword argument "
|
|
f"'{unexpected[0]}'"
|
|
)
|
|
self.last_kwargs = kwargs
|
|
if kwargs.get("stream"):
|
|
return [
|
|
_stream_line(content="fb"),
|
|
_stream_line(finish_reason="stop"),
|
|
]
|
|
message = types.SimpleNamespace(content="fb answer", tool_calls=None)
|
|
return types.SimpleNamespace(
|
|
choices=[types.SimpleNamespace(message=message)], usage=None
|
|
)
|
|
|
|
|
|
def _stream_line(content=None, finish_reason=None):
|
|
delta = types.SimpleNamespace(
|
|
content=content, reasoning_content=None, tool_calls=None
|
|
)
|
|
choice = types.SimpleNamespace(delta=delta, finish_reason=finish_reason)
|
|
return types.SimpleNamespace(choices=[choice], usage=None)
|
|
|
|
|
|
def _strict_openai_llm():
|
|
llm = OpenAILLM(api_key="sk-test", user_api_key=None, model_id="gpt-4o-mini")
|
|
llm.client = types.SimpleNamespace(
|
|
chat=types.SimpleNamespace(completions=_StrictChatCompletions())
|
|
)
|
|
return llm
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestRealOpenAIFallbackRejectsForeignKwargs:
|
|
|
|
def test_gen_google_primary_real_openai_fallback_does_not_typeerror(self):
|
|
primary = _GoogleFake(fail_at=0)
|
|
fallback = _strict_openai_llm()
|
|
primary._fallback_llm = fallback
|
|
response_schema = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
result = primary.gen(**CALL_ARGS, response_schema=response_schema)
|
|
|
|
assert result == "fb answer"
|
|
sent = fallback.client.chat.completions.last_kwargs
|
|
assert "response_schema" not in sent
|
|
assert sent["response_format"]["type"] == "json_schema"
|
|
|
|
def test_stream_google_primary_real_openai_fallback_does_not_typeerror(self):
|
|
primary = _GoogleFake(stream_chunks=["x"], fail_at=0)
|
|
fallback = _strict_openai_llm()
|
|
primary._fallback_llm = fallback
|
|
response_schema = primary.prepare_structured_output_format(SCHEMA)
|
|
|
|
chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema))
|
|
|
|
assert chunks == ["fb"]
|
|
sent = fallback.client.chat.completions.last_kwargs
|
|
assert "response_schema" not in sent
|
|
assert sent["response_format"]["type"] == "json_schema"
|