fix: cleanup models, better retry strat

tests/llm/handlers/test_reasoning_only_recovery.py

@@ -0,0 +1,389 @@
+"""Regression tests for the silent-loss recovery in ``handle_streaming``.
+
+Bug the recovery guards against: a stream ends cleanly (``finish_reason=stop``
+or plain exhaustion) after emitting only reasoning ("thought") chunks and
+zero visible content — observed in prod when cross-provider fallback lands
+on a model that returns ``reasoning_content`` deltas without any ``content``
+deltas. Downstream then saves ``response=""`` and marks the message
+``status=complete``: the user sees nothing.
+
+The recovery is a plain, bounded re-send of the same request (same
+messages, same tools, no system nudge). Behavior-preserving because a
+system message like "produce final answer now, no more thinking or tools"
+measurably shortens the model's reasoning by ~30% in A/B tests, which
+hurts answer quality on the common path more than the tail failure rate
+justifies. Critically: does not fire when the model genuinely chose to
+say nothing (no thoughts either), and does not recurse if the recovery
+itself also lands reasoning-only.
+"""
+
+from types import SimpleNamespace
+from typing import Any, Dict, Generator, List
+
+import pytest
+
+from application.llm.handlers.base import LLMHandler, LLMResponse, ToolCall
+
+
+# Scripted streams. Each element is either:
+#  * a str                                    -> content delta (yielded)
+#  * {"type": "thought", "thought": "..."}    -> reasoning delta
+#  * "STOP"                                   -> synthesises a stop choice
+# The scripts mirror what OpenAILLM._raw_gen_stream / _responses_gen_stream
+# actually yield to the handler.
+
+STOP_SENTINEL = "__STOP__"
+
+
+def _stop_choice():
+    """Provider-agnostic stop-finish chunk (matches OpenAI chat/completions
+    delta with content=None and finish_reason='stop')."""
+    delta = SimpleNamespace(content=None, tool_calls=None)
+    return SimpleNamespace(delta=delta, finish_reason="stop", message=None)
+
+
+class ScriptedHandler(LLMHandler):
+    """Concrete LLMHandler that iterates a supplied script.
+
+    Note: because ``FakeLLM._responding_provider = "openai"`` (set below),
+    ``LLMHandler._parse_for_response`` routes choice-object parsing to
+    ``OpenAILLMHandler.parse_response`` rather than to this class's
+    ``parse_response``. The scripted stop chunks are shape-compatible
+    with an OpenAI chat/completions delta (``SimpleNamespace(delta=…,
+    finish_reason="stop")``) so OpenAI's parser handles them cleanly.
+    ``parse_response`` here exists only for the abstract-class contract
+    and as a fallback if a test flips the provider off.
+    """
+
+    def parse_response(self, response: Any) -> LLMResponse:
+        # Only reached when ``agent.llm._responding_provider`` is None /
+        # non-string; otherwise ``_parse_for_response`` routes to the
+        # per-provider handler.
+        finish = getattr(response, "finish_reason", "") or ""
+        delta = getattr(response, "delta", None)
+        content = getattr(delta, "content", None) if delta else None
+        return LLMResponse(
+            content=content or "",
+            tool_calls=[],
+            finish_reason=finish,
+            raw_response=response,
+        )
+
+    def create_tool_message(self, tool_call: ToolCall, result: Any) -> Dict:
+        return {"role": "tool", "tool_call_id": tool_call.id, "content": str(result)}
+
+    def _iterate_stream(self, response: Any) -> Generator:
+        for item in response:
+            if item == STOP_SENTINEL:
+                yield _stop_choice()
+            else:
+                yield item
+
+
+def _script_generator(script):
+    """Return a fresh generator each time (agent.llm.gen_stream is called
+    for the recovery re-stream)."""
+
+    def factory():
+        for item in script:
+            yield item
+
+    return factory
+
+
+class FakeLLM:
+    """Minimal agent.llm stand-in — recovery reads ``model_id``, calls
+    ``gen_stream`` for the re-stream, and _handle_response consults it.
+    Tracks how many times ``gen_stream`` was invoked and what ``tools``
+    argument it received (so tests can pin that the recovery preserves
+    the original tools list — no ``tools=None`` behavior change)."""
+
+    def __init__(self, recovery_script=None):
+        self.model_id = "test-model"
+        self._responding_provider = "openai"
+        self._fallback_llm = None
+        self._stream_reached_finish = False
+        self._recovery_script = recovery_script
+        self.gen_stream_calls = 0
+        self.gen_stream_tools_seen = []
+
+    def gen_stream(self, model, messages, tools=None, **kwargs):
+        self.gen_stream_calls += 1
+        self.gen_stream_tools_seen.append(tools)
+        # First recovery call returns the scripted stream. Subsequent
+        # calls (there shouldn't be any) return empty — the recovery is
+        # supposed to be one-shot.
+        if self._recovery_script is None:
+            return iter([])
+        script = self._recovery_script
+        self._recovery_script = None  # one-shot
+        return iter(script)
+
+
+def _fake_agent(recovery_script=None):
+    llm = FakeLLM(recovery_script=recovery_script)
+    return SimpleNamespace(llm=llm, model_id="test-model", tools=None), llm
+
+
+def _drain(handler, agent, primary_script) -> List[Any]:
+    """Feed ``primary_script`` into handle_streaming and return every event
+    it yielded downstream. The response arg is any iterable — the handler
+    calls ``_iterate_stream(response)`` on it."""
+    events = []
+    for event in handler.handle_streaming(agent, iter(primary_script), {}, []):
+        events.append(event)
+    return events
+
+
+class TestReasoningOnlyRecovery:
+    def test_recovers_when_stream_ends_reasoning_only(self):
+        """Primary emits thoughts + finish=stop with no content →
+        recovery call fires and its content lands as visible chunks."""
+        primary = [
+            {"type": "thought", "thought": "Composing analysis..."},
+            {"type": "thought", "thought": " almost done."},
+            STOP_SENTINEL,
+        ]
+        recovery = ["Final answer: ", "hello world.", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        handler = ScriptedHandler()
+
+        events = _drain(handler, agent, primary)
+
+        thoughts = [e for e in events if isinstance(e, dict) and e.get("type") == "thought"]
+        strs = [e for e in events if isinstance(e, str)]
+
+        assert len(thoughts) == 2  # primary thoughts flow through
+        assert "".join(strs) == "Final answer: hello world."
+        assert llm.gen_stream_calls == 1  # exactly one recovery call
+
+    def test_no_recovery_when_answer_already_yielded(self):
+        """Primary emits some content before stop → no recovery, no extra
+        provider call."""
+        primary = [
+            {"type": "thought", "thought": "Thinking briefly..."},
+            "The answer is 42.",
+            STOP_SENTINEL,
+        ]
+        agent, llm = _fake_agent(recovery_script=["SHOULD NOT FIRE", STOP_SENTINEL])
+        handler = ScriptedHandler()
+
+        events = _drain(handler, agent, primary)
+
+        strs = [e for e in events if isinstance(e, str)]
+        assert "".join(strs) == "The answer is 42."
+        assert llm.gen_stream_calls == 0  # no recovery
+
+    def test_no_recovery_when_reasoning_buffer_empty(self):
+        """Primary emits only finish=stop (model genuinely said nothing)
+        → no recovery; a rescue call would just repeat the same failure
+        and burn budget for nothing."""
+        primary = [STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=["SHOULD NOT FIRE", STOP_SENTINEL])
+        handler = ScriptedHandler()
+
+        events = _drain(handler, agent, primary)
+
+        assert events == []  # nothing to yield, nothing recovered
+        assert llm.gen_stream_calls == 0
+
+    def test_recovery_that_also_reasons_only_does_not_recurse(self):
+        """If the recovery stream itself ends reasoning-only, the second
+        pass must NOT trigger another recovery (bounded to one attempt
+        per outer stream via ``_answer_recovered``)."""
+        primary = [
+            {"type": "thought", "thought": "primary reasoning"},
+            STOP_SENTINEL,
+        ]
+        # Recovery emits more thoughts and stops — still no content. The
+        # outer recovery guard must catch that we already tried, so we
+        # exit with 0 answers instead of looping.
+        recovery = [
+            {"type": "thought", "thought": "recovery reasoning"},
+            STOP_SENTINEL,
+        ]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        handler = ScriptedHandler()
+
+        events = _drain(handler, agent, primary)
+
+        strs = [e for e in events if isinstance(e, str)]
+        assert strs == []  # no answer surfaced (bug still visible)
+        assert llm.gen_stream_calls == 1  # but exactly one recovery, not many
+
+    def test_recovery_is_a_plain_resend_no_nudge_added(self):
+        """The recovery must NOT add a system nudge to messages — a
+        behavior-modifying instruction ("no thinking, write the answer")
+        measurably shortens the model's reasoning, hurting answer
+        quality on the common path. The re-send is deliberately
+        behavior-preserving: same messages, same tools."""
+        primary = [
+            {"type": "thought", "thought": "reasoning..."},
+            STOP_SENTINEL,
+        ]
+        recovery = ["recovered", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+
+        seen_messages: List[List[Dict]] = []
+        orig_gen_stream = llm.gen_stream
+
+        def capturing_gen_stream(model, messages, tools=None, **kwargs):
+            seen_messages.append(list(messages))
+            return orig_gen_stream(model, messages, tools=tools, **kwargs)
+
+        llm.gen_stream = capturing_gen_stream
+        starting = [{"role": "user", "content": "hi"}]
+        handler = ScriptedHandler()
+
+        list(handler.handle_streaming(agent, iter(primary), {}, starting))
+
+        assert len(seen_messages) == 1, "recovery must fire exactly once"
+        assert seen_messages[0] == starting, (
+            "recovery must forward the original messages verbatim — no "
+            "system nudge, no other mutation"
+        )
+
+    def test_recovery_preserves_agent_tools(self):
+        """The recovery must forward ``agent.tools`` unchanged — passing
+        ``tools=None`` would prevent the model from calling further tools
+        even when they might legitimately be part of finishing the
+        answer, and is another behavior change we want to avoid."""
+        primary = [
+            {"type": "thought", "thought": "reasoning..."},
+            STOP_SENTINEL,
+        ]
+        recovery = ["recovered", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        agent.tools = [{"type": "function", "function": {"name": "search"}}]
+        handler = ScriptedHandler()
+
+        list(handler.handle_streaming(agent, iter(primary), {}, []))
+
+        assert llm.gen_stream_calls == 1
+        assert llm.gen_stream_tools_seen == [agent.tools]
+
+    @pytest.mark.parametrize(
+        "primary_last_chunk", [STOP_SENTINEL, None],
+        ids=["explicit-stop", "no-stop-chunk-just-exhaustion"],
+    )
+    def test_recovery_fires_on_both_explicit_stop_and_bare_exhaustion(
+        self, primary_last_chunk,
+    ):
+        """Providers that emit finish_reason=stop AND providers whose
+        stream just exhausts (no parseable stop chunk) must both trigger
+        recovery when reasoning-only."""
+        primary = [{"type": "thought", "thought": "hmm"}]
+        if primary_last_chunk:
+            primary.append(primary_last_chunk)
+        recovery = ["ok done", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        handler = ScriptedHandler()
+
+        events = _drain(handler, agent, primary)
+        strs = [e for e in events if isinstance(e, str)]
+        assert "".join(strs) == "ok done"
+        assert llm.gen_stream_calls == 1
+
+    @pytest.mark.parametrize(
+        "attr,value",
+        [("json_schema", {"type": "object"}), ("json_object", True)],
+        ids=["json_schema", "json_object"],
+    )
+    def test_no_recovery_when_structured_output_configured(self, attr, value):
+        """Structured-output agents must NOT trigger recovery. The
+        primary call was built with response_format/response_schema; the
+        recovery here would drop those (it doesn't rebuild _llm_gen's
+        kwargs), so a rescue would emit unconstrained content that
+        downstream schema validation rejects. Better to surface the
+        empty answer than silently violate the contract."""
+        primary = [
+            {"type": "thought", "thought": "reasoning..."},
+            STOP_SENTINEL,
+        ]
+        recovery = ["SHOULD NOT FIRE", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        setattr(agent, attr, value)
+        handler = ScriptedHandler()
+
+        events = _drain(handler, agent, primary)
+
+        assert llm.gen_stream_calls == 0
+        assert not [e for e in events if isinstance(e, str)]
+
+    def test_recovery_forwards_agent_llm_params(self):
+        """Sampling params (temperature, max_tokens, top_p, …) live on
+        the agent, and the primary call gets them via ``_llm_gen``. The
+        recovery must forward them too so the rescued answer matches
+        the intended output distribution — otherwise you'd get a
+        default-temperature response for an agent configured with,
+        say, temperature=0.1."""
+        primary = [
+            {"type": "thought", "thought": "reasoning..."},
+            STOP_SENTINEL,
+        ]
+        recovery = ["ok", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        agent.llm_params = {"temperature": 0.1, "max_tokens": 256}
+
+        seen_kwargs: List[Dict] = []
+        orig = llm.gen_stream
+
+        def capturing(model, messages, tools=None, **kwargs):
+            seen_kwargs.append(dict(kwargs))
+            return orig(model, messages, tools=tools, **kwargs)
+
+        llm.gen_stream = capturing
+        handler = ScriptedHandler()
+        list(handler.handle_streaming(agent, iter(primary), {}, []))
+
+        assert len(seen_kwargs) == 1
+        assert seen_kwargs[0]["temperature"] == 0.1
+        assert seen_kwargs[0]["max_tokens"] == 256
+
+    def test_recovery_forces_tools_none_at_cap(self):
+        """When ``_iteration >= MAX_TOOL_ITERATIONS`` the current call
+        is the cap-triggered finalize round: it was sent with
+        ``tools=None`` plus a "no more tools" system message. If it
+        reasons-only-stops, the recovery MUST NOT reopen tools — the
+        finalize contract exists precisely so no more tool calls run."""
+        from application.llm.handlers.base import MAX_TOOL_ITERATIONS
+
+        primary = [
+            {"type": "thought", "thought": "reasoning..."},
+            STOP_SENTINEL,
+        ]
+        recovery = ["ok", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        agent.tools = [{"type": "function", "function": {"name": "search"}}]
+        handler = ScriptedHandler()
+
+        list(handler.handle_streaming(
+            agent, iter(primary), {}, [],
+            _iteration=MAX_TOOL_ITERATIONS,
+        ))
+
+        assert llm.gen_stream_calls == 1
+        assert llm.gen_stream_tools_seen == [None], (
+            "recovery at MAX_TOOL_ITERATIONS must pass tools=None to "
+            "honour the finalize round's contract"
+        )
+
+    def test_recovery_forces_tools_none_when_context_limit_reached(self):
+        """Same contract as the cap case, different trigger: when the
+        agent flagged ``context_limit_reached=True`` mid-loop, the
+        preceding round was already sent with ``tools=None`` + a wrap-
+        up system message. Recovery must not reopen tools there either."""
+        primary = [
+            {"type": "thought", "thought": "reasoning..."},
+            STOP_SENTINEL,
+        ]
+        recovery = ["ok", STOP_SENTINEL]
+        agent, llm = _fake_agent(recovery_script=recovery)
+        agent.tools = [{"type": "function", "function": {"name": "search"}}]
+        agent.context_limit_reached = True
+        handler = ScriptedHandler()
+
+        list(handler.handle_streaming(agent, iter(primary), {}, []))
+
+        assert llm.gen_stream_calls == 1
+        assert llm.gen_stream_tools_seen == [None]

tests/llm/test_base.py

@@ -405,8 +405,6 @@ class TestProviderNameRegistry:
         from application.llm.novita import NovitaLLM
         from application.llm.open_router import OpenRouterLLM
         from application.llm.openai import OpenAILLM
-        from application.llm.premai import PremAILLM
-        from application.llm.sagemaker import SagemakerAPILLM
 
         assert OpenAILLM.provider_name == "openai"
         assert GoogleLLM.provider_name == "google"
@@ -415,9 +413,7 @@ class TestProviderNameRegistry:
         assert NovitaLLM.provider_name == "novita"
         assert OpenRouterLLM.provider_name == "openrouter"
         assert DocsGPTAPILLM.provider_name == "docsgpt"
-        assert PremAILLM.provider_name == "premai"
         assert LlamaCpp.provider_name == "llama_cpp"
-        assert SagemakerAPILLM.provider_name == "sagemaker"
 
     @patch("application.llm.base.gen_cache", lambda f: f)
     @patch("application.llm.base.gen_token_usage", lambda f: f)

tests/llm/test_fallback.py

@@ -6,6 +6,7 @@ per-agent backup model is used before the global FALLBACK_* settings.
 
 from unittest.mock import MagicMock
 
+import httpx
 import pytest
 
 from application.llm.base import BaseLLM
@@ -22,10 +23,18 @@ class FakeLLM(BaseLLM):
         # 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
@@ -44,11 +53,30 @@ class FakeLLM(BaseLLM):
 
     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
         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 self.fail_at is not None and yielded >= self.fail_at:
-                raise RuntimeError("mid-stream failure")
+            if local_fail_at is not None and yielded >= local_fail_at:
+                raise local_error_class("mid-stream failure")
             yield chunk
             yielded += 1
 
@@ -290,6 +318,192 @@ class TestStreamingFallback:
         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(
+            "application.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
     ):

tests/llm/test_premai.py

@@ -1,190 +0,0 @@
-"""Unit tests for application/llm/premai.py — PremAILLM.
-
-Covers:
-  - Constructor
-  - _raw_gen: API call and return value
-  - _raw_gen_stream: streaming with delta content filtering
-"""
-
-import sys
-import types
-
-import pytest
-
-
-# ---------------------------------------------------------------------------
-# Fake premai module
-# ---------------------------------------------------------------------------
-
-
-class _FakeMessage:
-    def __init__(self, content):
-        self.message = {"content": content}
-
-
-class _FakeDelta:
-    def __init__(self, content):
-        self.delta = {"content": content}
-
-
-class _FakeChoice:
-    def __init__(self, content):
-        self.message = {"content": content}
-
-
-class _FakeStreamChoice:
-    def __init__(self, content):
-        self.delta = {"content": content}
-
-
-class _FakeResponse:
-    def __init__(self, content="result_text"):
-        self.choices = [_FakeChoice(content)]
-
-
-class _FakeStreamLine:
-    def __init__(self, content):
-        self.choices = [_FakeStreamChoice(content)]
-
-
-class _FakeChatCompletions:
-    def __init__(self):
-        self.last_kwargs = None
-
-    def create(self, **kwargs):
-        self.last_kwargs = kwargs
-        if kwargs.get("stream"):
-            return [
-                _FakeStreamLine("chunk1"),
-                _FakeStreamLine("chunk2"),
-                _FakeStreamLine(None),  # None content should be filtered
-            ]
-        return _FakeResponse()
-
-
-class _FakeChat:
-    def __init__(self):
-        self.completions = _FakeChatCompletions()
-
-
-class _FakePrem:
-    def __init__(self, api_key=None):
-        self.api_key = api_key
-        self.chat = _FakeChat()
-
-
-@pytest.fixture(autouse=True)
-def patch_premai(monkeypatch):
-    fake_mod = types.ModuleType("premai")
-    fake_mod.Prem = _FakePrem
-    sys.modules["premai"] = fake_mod
-
-    if "application.llm.premai" in sys.modules:
-        del sys.modules["application.llm.premai"]
-    yield
-    sys.modules.pop("premai", None)
-    if "application.llm.premai" in sys.modules:
-        del sys.modules["application.llm.premai"]
-
-
-@pytest.fixture
-def llm():
-    from application.llm.premai import PremAILLM
-
-    return PremAILLM(api_key="test-key")
-
-
-# ---------------------------------------------------------------------------
-# Constructor
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.unit
-class TestPremAIConstructor:
-
-    def test_sets_api_key(self, llm):
-        assert llm.api_key == "test-key"
-
-    def test_sets_user_api_key_none(self, llm):
-        assert llm.user_api_key is None
-
-    def test_client_created(self, llm):
-        assert isinstance(llm.client, _FakePrem)
-
-    def test_project_id_from_settings(self, llm):
-        from application.core.settings import settings
-
-        assert llm.project_id == settings.PREMAI_PROJECT_ID
-
-
-# ---------------------------------------------------------------------------
-# _raw_gen
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.unit
-class TestRawGen:
-
-    def test_returns_content(self, llm):
-        msgs = [{"role": "user", "content": "hi"}]
-        result = llm._raw_gen(llm, model="model-1", messages=msgs)
-        assert result == "result_text"
-
-    def test_passes_model_and_project_id(self, llm):
-        msgs = [{"role": "user", "content": "hi"}]
-        llm._raw_gen(llm, model="my-model", messages=msgs)
-        kwargs = llm.client.chat.completions.last_kwargs
-        assert kwargs["model"] == "my-model"
-        assert kwargs["project_id"] == llm.project_id
-        assert kwargs["stream"] is False
-
-    def test_passes_messages(self, llm):
-        msgs = [{"role": "user", "content": "hello"}]
-        llm._raw_gen(llm, model="m", messages=msgs)
-        kwargs = llm.client.chat.completions.last_kwargs
-        assert kwargs["messages"] == msgs
-
-    def test_extra_kwargs_forwarded(self, llm):
-        msgs = [{"role": "user", "content": "hi"}]
-        llm._raw_gen(llm, model="m", messages=msgs, temperature=0.5)
-        kwargs = llm.client.chat.completions.last_kwargs
-        assert kwargs["temperature"] == 0.5
-
-
-# ---------------------------------------------------------------------------
-# _raw_gen_stream
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.unit
-class TestRawGenStream:
-
-    def test_yields_non_none_content(self, llm):
-        msgs = [{"role": "user", "content": "hi"}]
-        chunks = list(
-            llm._raw_gen_stream(llm, model="m", messages=msgs, stream=True)
-        )
-        assert chunks == ["chunk1", "chunk2"]
-
-    def test_filters_none_content(self, llm):
-        msgs = [{"role": "user", "content": "hi"}]
-        chunks = list(
-            llm._raw_gen_stream(llm, model="m", messages=msgs, stream=True)
-        )
-        assert None not in chunks
-
-    def test_passes_stream_true(self, llm):
-        msgs = [{"role": "user", "content": "hi"}]
-        list(llm._raw_gen_stream(llm, model="m", messages=msgs))
-        kwargs = llm.client.chat.completions.last_kwargs
-        assert kwargs["stream"] is True
-
-    def test_passes_extra_kwargs(self, llm):
-        msgs = [{"role": "user", "content": "hi"}]
-        list(
-            llm._raw_gen_stream(
-                llm, model="m", messages=msgs, max_tokens=100
-            )
-        )
-        kwargs = llm.client.chat.completions.last_kwargs
-        assert kwargs["max_tokens"] == 100

tests/llm/test_sagemaker.py

@@ -1,106 +0,0 @@
-# FILEPATH: /path/to/test_sagemaker.py
-
-import json
-import unittest
-from unittest.mock import MagicMock, patch
-from application.llm.sagemaker import SagemakerAPILLM, LineIterator
-
-class TestSagemakerAPILLM(unittest.TestCase):
-    
-    def setUp(self):
-        self.sagemaker = SagemakerAPILLM()
-        self.context = "This is the context"
-        self.user_question = "What is the answer?"
-        self.messages = [
-            {"content": self.context},
-            {"content": "Some other message"},
-            {"content": self.user_question}
-        ]
-        self.prompt = f"### Instruction \n {self.user_question} \n ### Context \n {self.context} \n ### Answer \n"
-        self.payload = {
-            "inputs": self.prompt,
-            "stream": False,
-            "parameters": {
-                "do_sample": True,
-                "temperature": 0.1,
-                "max_new_tokens": 30,
-                "repetition_penalty": 1.03,
-                "stop": ["</s>", "###"]
-            }
-        }
-        self.payload_stream = {
-            "inputs": self.prompt,
-            "stream": True,
-            "parameters": {
-                "do_sample": True,
-                "temperature": 0.1,
-                "max_new_tokens": 512,
-                "repetition_penalty": 1.03,
-                "stop": ["</s>", "###"]
-            }
-        }
-        self.body_bytes = json.dumps(self.payload).encode('utf-8')
-        self.body_bytes_stream = json.dumps(self.payload_stream).encode('utf-8')
-        self.response = {
-            "Body": MagicMock()
-        }
-        self.result = [
-            {
-                "generated_text": "This is the generated text"
-            }
-        ]
-        self.response['Body'].read.return_value.decode.return_value = json.dumps(self.result)
-        
-    def test_gen(self):
-        with patch('application.cache.get_redis_instance') as mock_make_redis:
-            mock_redis_instance = mock_make_redis.return_value
-            mock_redis_instance.get.return_value = None
-
-            with patch.object(self.sagemaker.runtime, 'invoke_endpoint', 
-                            return_value=self.response) as mock_invoke_endpoint:
-                output = self.sagemaker.gen(None, self.messages)
-                mock_invoke_endpoint.assert_called_once_with(
-                    EndpointName=self.sagemaker.endpoint,
-                    ContentType='application/json',
-                    Body=self.body_bytes
-                )
-                self.assertEqual(output, 
-                                self.result[0]['generated_text'][len(self.prompt):])
-            mock_make_redis.assert_called_once()
-            mock_redis_instance.set.assert_called_once()
-    
-    def test_gen_stream(self):
-        with patch('application.cache.get_redis_instance') as mock_make_redis:
-            mock_redis_instance = mock_make_redis.return_value
-            mock_redis_instance.get.return_value = None
-
-            with patch.object(self.sagemaker.runtime, 'invoke_endpoint_with_response_stream', 
-                            return_value=self.response) as mock_invoke_endpoint:
-                output = list(self.sagemaker.gen_stream(None, self.messages, tools=None))
-                mock_invoke_endpoint.assert_called_once_with(
-                    EndpointName=self.sagemaker.endpoint,
-                    ContentType='application/json',
-                    Body=self.body_bytes_stream
-                )
-                self.assertEqual(output, [])
-            mock_redis_instance.set.assert_called_once()
-class TestLineIterator(unittest.TestCase):
-    
-    def setUp(self):
-        self.stream = [
-            {'PayloadPart': {'Bytes': b'{"outputs": [" a"]}\n'}},
-            {'PayloadPart': {'Bytes': b'{"outputs": [" challenging"]}\n'}},
-            {'PayloadPart': {'Bytes': b'{"outputs": [" problem"]}\n'}}
-        ]
-        self.line_iterator = LineIterator(self.stream)
-        
-    def test_iter(self):
-        self.assertEqual(iter(self.line_iterator), self.line_iterator)
-        
-    def test_next(self):
-        self.assertEqual(next(self.line_iterator), b'{"outputs": [" a"]}')
-        self.assertEqual(next(self.line_iterator), b'{"outputs": [" challenging"]}')
-        self.assertEqual(next(self.line_iterator), b'{"outputs": [" problem"]}')
-        
-if __name__ == '__main__':
-    unittest.main()