fix: cleanup models, better retry strat

application/cache.py

@@ -207,8 +207,18 @@ def stream_cache(func):
 
         stream_cache_data = []
         cacheable = True
+        # Skip caching streams that produced no visible content — a
+        # reasoning-only stop (thoughts only, no str deltas) would
+        # otherwise be replayed for the whole TTL on every identical
+        # request, poisoning the cache and denying the reasoning-only
+        # recovery path any chance to run against a fresh provider call
+        # (subsequent identical requests replay the cached empty stream,
+        # trip the recovery guard, and produce another silent-loss).
+        had_content = False
         for chunk in func(self, model, messages, stream, tools, *args, **kwargs):
             yield chunk
+            if isinstance(chunk, str) and chunk:
+                had_content = True
             if isinstance(chunk, (str, dict, list, int, float, bool, type(None))):
                 try:
                     json.dumps(chunk)
@@ -218,7 +228,7 @@ def stream_cache(func):
             else:
                 cacheable = False
 
-        if redis_client and cacheable:
+        if redis_client and cacheable and had_content:
             try:
                 payload = {"version": 1, "chunks": stream_cache_data}
                 redis_client.set(cache_key, json.dumps(payload), ex=1800)

application/core/model_settings.py

@@ -22,8 +22,6 @@ class ModelProvider(str, Enum):
     HUGGINGFACE = "huggingface"
     LLAMA_CPP = "llama.cpp"
     DOCSGPT = "docsgpt"
-    PREMAI = "premai"
-    SAGEMAKER = "sagemaker"
     NOVITA = "novita"
 
 

application/core/settings.py

@@ -168,14 +168,12 @@ class Settings(BaseSettings):
     ELASTIC_URL: Optional[str] = None  # url for elasticsearch
     ELASTIC_INDEX: Optional[str] = "docsgpt"  # index name for elasticsearch
 
-    # SageMaker config
-    SAGEMAKER_ENDPOINT: Optional[str] = None  # SageMaker endpoint name
-    SAGEMAKER_REGION: Optional[str] = None  # SageMaker region name
-    SAGEMAKER_ACCESS_KEY: Optional[str] = None  # SageMaker access key
-    SAGEMAKER_SECRET_KEY: Optional[str] = None  # SageMaker secret key
-
-    # prem ai project id
-    PREMAI_PROJECT_ID: Optional[str] = None
+    # Legacy AWS credentials from the retired SageMaker LLM provider.
+    # Still read as a deprecated fallback by S3 storage (see the S3_*
+    # block below); do not use for new deployments.
+    SAGEMAKER_REGION: Optional[str] = None
+    SAGEMAKER_ACCESS_KEY: Optional[str] = None
+    SAGEMAKER_SECRET_KEY: Optional[str] = None
 
     # Qdrant vectorstore config
     QDRANT_COLLECTION_NAME: Optional[str] = "docsgpt"

application/llm/base.py

@@ -2,6 +2,9 @@ import logging
 from abc import ABC, abstractmethod
 from typing import ClassVar
 
+import httpx
+import openai
+
 from application.cache import gen_cache, stream_cache
 
 from application.core.settings import settings
@@ -10,6 +13,26 @@ from application.usage import gen_token_usage, stream_token_usage
 logger = logging.getLogger(__name__)
 
 
+# Errors safe to retry the primary once on: the request either never
+# reached the server (connect), or the peer closed the connection before
+# a proper end-of-stream marker (read/protocol). None of these leave a
+# side effect on the upstream, so a repeat call is cheap and idempotent.
+# Excludes API-level errors (4xx / 5xx status) which are not transport
+# retries — RateLimitError needs backoff, BadRequestError won't get any
+# better on retry, and the existing fallback handles both.
+_STREAM_RETRYABLE_TRANSPORT_ERRORS = (
+    httpx.RemoteProtocolError,
+    httpx.ReadError,
+    httpx.ReadTimeout,
+    httpx.WriteError,
+    httpx.WriteTimeout,
+    httpx.ConnectError,
+    httpx.ConnectTimeout,
+    httpx.PoolTimeout,
+    openai.APIConnectionError,
+)
+
+
 class BaseLLM(ABC):
     # Stamped onto the ``llm_stream_start`` event so dashboards can group
     # calls by vendor. Subclasses override.
@@ -267,10 +290,22 @@ class BaseLLM(ABC):
         streaming generators raise exceptions during iteration. This wrapper
         ensures that if the primary LLM fails at any point during streaming
         (creation or mid-stream), we fall back to the backup model.
+
+        Transport errors that fire before any chunk was yielded (Azure
+        Front Door reset, backend deploy hiccup, TLS blip) get one same-
+        primary retry before we engage the fallback — the request never
+        made it to a state the peer produced output for, so a repeat is
+        idempotent. Once we've yielded anything, retrying would duplicate
+        delivered content; those errors go straight to the fallback path
+        as before.
         """
         self._responding_provider = self.provider_name
+        chunks_yielded = 0
         try:
-            yield from decorated_method()
+            for chunk in decorated_method():
+                chunks_yielded += 1
+                yield chunk
+            return
         except Exception as e:
             if getattr(self, "_stream_reached_finish", False):
                 # The primary already delivered a finish signal — only
@@ -284,6 +319,55 @@ class BaseLLM(ABC):
                     f"not engaging fallback. Error: {str(e)}"
                 )
                 raise
+            # Same-primary retry once for transport errors before any chunk
+            # was yielded. Covers the observed Azure Responses-API pattern
+            # where the peer resets the SSE stream within tens of seconds
+            # with a RemoteProtocolError, well before a legitimate fallback
+            # scenario. A yielded chunk = downstream already saw content,
+            # so replaying would duplicate it — skip retry in that case.
+            if (
+                chunks_yielded == 0
+                and isinstance(e, _STREAM_RETRYABLE_TRANSPORT_ERRORS)
+            ):
+                logger.warning(
+                    f"Primary LLM transport error before any output; "
+                    f"retrying once. Error: {str(e)}"
+                )
+                # Emit a fresh stream-start so dashboards get one start/
+                # finish pair per attempt (stream_token_usage fires a
+                # finish per decorated_method invocation).
+                self._emit_stream_start_log(
+                    kwargs.get("model") or getattr(self, "model_id", None),
+                    kwargs.get("messages"),
+                    kwargs.get("tools"),
+                    bool(
+                        kwargs.get("_usage_attachments")
+                        or kwargs.get("attachments")
+                    ),
+                )
+                try:
+                    for chunk in decorated_method():
+                        chunks_yielded += 1
+                        yield chunk
+                    return
+                except Exception as retry_e:
+                    # Retry delivered a full stream but died on a trailing
+                    # frame (usage chunk, [DONE]) — same guard the outer
+                    # except uses. Without this, the fallback would run
+                    # and the user would receive the whole answer twice
+                    # (and any tool calls would be executed twice).
+                    if getattr(self, "_stream_reached_finish", False):
+                        logger.warning(
+                            f"Primary LLM retry delivered its finish "
+                            f"signal, then failed on a trailing frame; "
+                            f"not engaging fallback. Error: {str(retry_e)}"
+                        )
+                        raise
+                    logger.warning(
+                        f"Primary LLM retry also failed; engaging fallback. "
+                        f"Error: {str(retry_e)}"
+                    )
+                    e = retry_e
             if not self.fallback_llm:
                 logger.error(
                     f"Primary LLM failed and no fallback configured: {str(e)}"

application/llm/handlers/base.py

@@ -1275,6 +1275,7 @@ class LLMHandler(ABC):
         tools_dict: Dict,
         messages: List[Dict],
         _iteration: int = 0,
+        _answer_recovered: bool = False,
     ) -> Generator:
         """
         Handle streaming response flow.
@@ -1284,6 +1285,9 @@ class LLMHandler(ABC):
             response: Current LLM response
             tools_dict: Available tools dictionary
             messages: Conversation history
+            _answer_recovered: Set by the recovery branch to bound the
+                empty-answer rescue to one attempt per stream — a recovery
+                that itself ends reasoning-only must not recurse.
 
         Yields:
             Streaming response chunks
@@ -1292,6 +1296,11 @@ class LLMHandler(ABC):
         tool_calls = {}
         reasoning_buffer = ""
         finish_reason = None
+        # Tracks whether any visible content reached the caller this
+        # round. Distinct from ``buffer``/``finish_reason``: providers
+        # emit content as bare strings (yielded straight through) and
+        # also as ``parsed.content`` on choice objects; either counts.
+        answered = False
 
         # Consume the provider stream to exhaustion before acting on its
         # finish reason. Acting mid-iteration (as this loop used to) left
@@ -1348,6 +1357,7 @@ class LLMHandler(ABC):
                 yield chunk
                 continue
             if isinstance(chunk, str):
+                answered = True
                 yield chunk
                 continue
             parsed = self._parse_for_response(agent, chunk)
@@ -1389,12 +1399,90 @@ class LLMHandler(ABC):
                 continue
             if parsed.content:
                 buffer += parsed.content
+                answered = True
                 yield buffer
                 buffer = ""
             if parsed.finish_reason == "stop" and finish_reason is None:
                 finish_reason = "stop"
 
         if finish_reason != "tool_calls":
+            # Silent-loss recovery: the stream ended cleanly (finish=stop
+            # or plain exhaustion) but never yielded a visible answer,
+            # and the model actually did work (thoughts non-empty). One
+            # plain re-send is far cheaper than a message row saved as
+            # status=complete with response="". Deliberately preserves
+            # the model's natural output distribution — no system-message
+            # nudge — since a "produce answer now, no thinking"
+            # instruction measurably shortens the model's reasoning by
+            # ~30% in our A/B, which hurts answer quality more than the
+            # tail failure rate justifies. Bounded to one attempt per
+            # stream by ``_answer_recovered`` so a recovery that itself
+            # reasons-only-stops does not recurse. Skipped when there is
+            # nothing to recover from (no reasoning) since the model
+            # genuinely chose to say nothing.
+            #
+            # Skipped, too, when the agent is running with structured
+            # output (json_schema / json_object). The primary call was
+            # constructed with response_format / response_schema that
+            # this handler cannot safely rebuild here without duplicating
+            # ``BaseAgent._llm_gen``; an unconstrained rescue would
+            # produce non-schema output that downstream consumers reject,
+            # which is worse than the empty-answer failure.
+            structured_output = (
+                getattr(agent, "json_schema", None)
+                or getattr(agent, "json_object", False)
+            )
+            if (
+                not answered
+                and reasoning_buffer
+                and not _answer_recovered
+                and not structured_output
+            ):
+                logger.warning(
+                    "Stream ended reasoning-only (%d chars of thought, "
+                    "0 answer bytes); re-sending once.",
+                    len(reasoning_buffer),
+                )
+                enforce = getattr(agent, "_enforce_context_window", None)
+                if callable(enforce):
+                    messages = enforce(messages)
+                # Mirror the finalize-round contract at the bottom of
+                # this method: when the cap or context limit already
+                # forced ``tools=None`` (with an accompanying "no more
+                # tools" system message), the recovery must not reopen
+                # tools, or the model could run one extra tool call past
+                # the cap. ``_iteration >= MAX_TOOL_ITERATIONS`` catches
+                # the finalize round.
+                recovery_tools = (
+                    None
+                    if _iteration >= MAX_TOOL_ITERATIONS
+                    or getattr(agent, "context_limit_reached", False)
+                    else agent.tools
+                )
+                recovery_kwargs = {
+                    "model": getattr(agent.llm, "model_id", None) or agent.model_id,
+                    "messages": messages,
+                    "tools": recovery_tools,
+                }
+                # Forward the sampling params the primary got
+                # (temperature, max_tokens, top_p, …). ``response_format``
+                # / ``response_schema`` are intentionally NOT forwarded
+                # — the structured_output gate above bails before we get
+                # here. ``previous_response_id`` is Responses-chain state
+                # and belongs to the failed turn, not a fresh recovery.
+                llm_params = getattr(agent, "llm_params", None)
+                if llm_params:
+                    recovery_kwargs.update(llm_params)
+                recovery_response = agent.llm.gen_stream(**recovery_kwargs)
+                self.llm_calls.append(build_stack_data(agent.llm))
+                yield from self.handle_streaming(
+                    agent,
+                    recovery_response,
+                    tools_dict,
+                    messages,
+                    _iteration=_iteration,
+                    _answer_recovered=True,
+                )
             return
 
         tool_handler_gen = self.handle_tool_calls(
@@ -1474,4 +1562,9 @@ class LLMHandler(ABC):
         yield from self.handle_streaming(
             agent, response, tools_dict, messages,
             _iteration=next_iteration,
+            # Carry the recovery guard through tool-loop recursion so a
+            # re-send that emitted tool_calls and then reasons-only-stops
+            # doesn't trigger a second re-send — the bound is per outer
+            # stream, not per tool round.
+            _answer_recovered=_answer_recovered,
         )

application/llm/premai.py

@@ -1,39 +0,0 @@
-from application.llm.base import BaseLLM
-from application.core.settings import settings
-
-
-class PremAILLM(BaseLLM):
-    provider_name = "premai"
-
-    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
-        from premai import Prem
-
-        super().__init__(*args, **kwargs)
-        self.client = Prem(api_key=api_key)
-        self.api_key = api_key
-        self.user_api_key = user_api_key
-        self.project_id = settings.PREMAI_PROJECT_ID
-
-    def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
-        response = self.client.chat.completions.create(
-            model=model,
-            project_id=self.project_id,
-            messages=messages,
-            stream=stream,
-            **kwargs
-        )
-
-        return response.choices[0].message["content"]
-
-    def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):
-        response = self.client.chat.completions.create(
-            model=model,
-            project_id=self.project_id,
-            messages=messages,
-            stream=stream,
-            **kwargs
-        )
-
-        for line in response:
-            if line.choices[0].delta["content"] is not None:
-                yield line.choices[0].delta["content"]

application/llm/providers/__init__.py

@@ -21,8 +21,6 @@ from application.llm.providers.novita import NovitaProvider
 from application.llm.providers.openai import OpenAIProvider
 from application.llm.providers.openai_compatible import OpenAICompatibleProvider
 from application.llm.providers.openrouter import OpenRouterProvider
-from application.llm.providers.premai import PremAIProvider
-from application.llm.providers.sagemaker import SagemakerProvider
 
 # Order here is the order the registry iterates providers (and therefore
 # the order ``/api/models`` reports them). Match the historical order
@@ -40,8 +38,6 @@ ALL_PROVIDERS: List[Provider] = [
     NovitaProvider(),
     HuggingFaceProvider(),
     LlamaCppProvider(),
-    PremAIProvider(),
-    SagemakerProvider(),
 ]
 
 PROVIDERS_BY_NAME: Dict[str, Provider] = {p.name: p for p in ALL_PROVIDERS}

application/llm/providers/premai.py

@@ -1,19 +0,0 @@
-from __future__ import annotations
-
-from typing import Optional
-
-from application.llm.premai import PremAILLM
-from application.llm.providers.base import Provider
-
-
-class PremAIProvider(Provider):
-    """LLMCreator-only plugin: invocable via LLM_PROVIDER but not in the catalog."""
-
-    name = "premai"
-    llm_class = PremAILLM
-
-    def get_api_key(self, settings) -> Optional[str]:
-        return settings.API_KEY
-
-    def is_enabled(self, settings) -> bool:
-        return False

application/llm/providers/sagemaker.py

@@ -1,24 +0,0 @@
-from __future__ import annotations
-
-from typing import Optional
-
-from application.llm.sagemaker import SagemakerAPILLM
-from application.llm.providers.base import Provider
-
-
-class SagemakerProvider(Provider):
-    """LLMCreator-only plugin: invocable via LLM_PROVIDER but not in the catalog.
-
-    SageMaker reads its credentials from ``SAGEMAKER_*`` settings inside
-    the LLM class itself; this plugin's ``get_api_key`` exists only for
-    LLMCreator's symmetry.
-    """
-
-    name = "sagemaker"
-    llm_class = SagemakerAPILLM
-
-    def get_api_key(self, settings) -> Optional[str]:
-        return settings.API_KEY
-
-    def is_enabled(self, settings) -> bool:
-        return False

application/llm/sagemaker.py

@@ -1,141 +0,0 @@
-from application.llm.base import BaseLLM
-from application.core.settings import settings
-import json
-import io
-
-
-class LineIterator:
-    """
-    A helper class for parsing the byte stream input.
-
-    The output of the model will be in the following format:
-    ```
-    b'{"outputs": [" a"]}\n'
-    b'{"outputs": [" challenging"]}\n'
-    b'{"outputs": [" problem"]}\n'
-    ...
-    ```
-
-    While usually each PayloadPart event from the event stream will contain a byte array
-    with a full json, this is not guaranteed and some of the json objects may be split across
-    PayloadPart events. For example:
-    ```
-    {'PayloadPart': {'Bytes': b'{"outputs": '}}
-    {'PayloadPart': {'Bytes': b'[" problem"]}\n'}}
-    ```
-
-    This class accounts for this by concatenating bytes written via the 'write' function
-    and then exposing a method which will return lines (ending with a '\n' character) within
-    the buffer via the 'scan_lines' function. It maintains the position of the last read
-    position to ensure that previous bytes are not exposed again.
-    """
-
-    def __init__(self, stream):
-        self.byte_iterator = iter(stream)
-        self.buffer = io.BytesIO()
-        self.read_pos = 0
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        while True:
-            self.buffer.seek(self.read_pos)
-            line = self.buffer.readline()
-            if line and line[-1] == ord("\n"):
-                self.read_pos += len(line)
-                return line[:-1]
-            try:
-                chunk = next(self.byte_iterator)
-            except StopIteration:
-                if self.read_pos < self.buffer.getbuffer().nbytes:
-                    continue
-                raise
-            if "PayloadPart" not in chunk:
-                print("Unknown event type:" + chunk)
-                continue
-            self.buffer.seek(0, io.SEEK_END)
-            self.buffer.write(chunk["PayloadPart"]["Bytes"])
-
-
-class SagemakerAPILLM(BaseLLM):
-    provider_name = "sagemaker"
-
-    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
-        import boto3
-
-        runtime = boto3.client(
-            "runtime.sagemaker",
-            aws_access_key_id="xxx",
-            aws_secret_access_key="xxx",
-            region_name="us-west-2",
-        )
-
-        super().__init__(*args, **kwargs)
-        self.api_key = api_key
-        self.user_api_key = user_api_key
-        self.endpoint = settings.SAGEMAKER_ENDPOINT
-        self.runtime = runtime
-
-    def _raw_gen(self, baseself, model, messages, stream=False, tools=None, **kwargs):
-        context = messages[0]["content"]
-        user_question = messages[-1]["content"]
-        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
-
-        # Construct payload for endpoint
-        payload = {
-            "inputs": prompt,
-            "stream": False,
-            "parameters": {
-                "do_sample": True,
-                "temperature": 0.1,
-                "max_new_tokens": 30,
-                "repetition_penalty": 1.03,
-                "stop": ["</s>", "###"],
-            },
-        }
-        body_bytes = json.dumps(payload).encode("utf-8")
-
-        # Invoke the endpoint
-        response = self.runtime.invoke_endpoint(
-            EndpointName=self.endpoint, ContentType="application/json", Body=body_bytes
-        )
-        result = json.loads(response["Body"].read().decode())
-        import sys
-
-        print(result[0]["generated_text"], file=sys.stderr)
-        return result[0]["generated_text"][len(prompt) :]
-
-    def _raw_gen_stream(self, baseself, model, messages, stream=True, tools=None, **kwargs):
-        context = messages[0]["content"]
-        user_question = messages[-1]["content"]
-        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
-
-        # Construct payload for endpoint
-        payload = {
-            "inputs": prompt,
-            "stream": True,
-            "parameters": {
-                "do_sample": True,
-                "temperature": 0.1,
-                "max_new_tokens": 512,
-                "repetition_penalty": 1.03,
-                "stop": ["</s>", "###"],
-            },
-        }
-        body_bytes = json.dumps(payload).encode("utf-8")
-
-        # Invoke the endpoint
-        response = self.runtime.invoke_endpoint_with_response_stream(
-            EndpointName=self.endpoint, ContentType="application/json", Body=body_bytes
-        )
-        # result = json.loads(response['Body'].read().decode())
-        event_stream = response["Body"]
-        start_json = b"{"
-        for line in LineIterator(event_stream):
-            if line != b"" and start_json in line:
-                # print(line)
-                data = json.loads(line[line.find(start_json) :].decode("utf-8"))
-                if data["token"]["text"] not in ["</s>", "###"]:
-                    print(data["token"]["text"], end="")
-                    yield data["token"]["text"]

docs/content/Models/cloud-providers.mdx

@@ -32,8 +32,6 @@ DocsGPT offers direct, streamlined support for the following cloud LLM providers
 | OpenRouter                   | `openrouter`   | (See OpenRouter docs)       |
 | Novita AI                    | `novita`       | (See Novita docs)           |
 | HuggingFace Inference API    | `huggingface`  | `meta-llama/Llama-3.1-8B-Instruct` |
-| Prem AI                      | `premai`       | (See Prem AI docs)          |
-| AWS SageMaker                | `sagemaker`    | (See SageMaker docs)        |
 
 DocsGPT also ships a **model catalog** (`application/core/models/*.yaml`) that the in-app model picker reads, so common models from these providers — including DeepSeek — appear ready to select once the matching API key is set.
 

tests/core/test_model_settings.py

@@ -34,8 +34,6 @@ class TestModelProvider:
         assert ModelProvider.HUGGINGFACE == "huggingface"
         assert ModelProvider.NOVITA == "novita"
         assert ModelProvider.OPENROUTER == "openrouter"
-        assert ModelProvider.SAGEMAKER == "sagemaker"
-        assert ModelProvider.PREMAI == "premai"
         assert ModelProvider.LLAMA_CPP == "llama.cpp"
 
 

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()

tests/test_cache.py

@@ -446,6 +446,55 @@ def test_stream_cache_bypasses_when_tools_provided(mock_make_redis):
     mock_redis_instance.get.assert_not_called()
 
 
+@pytest.mark.unit
+@patch("application.cache.get_redis_instance")
+def test_stream_cache_skips_write_when_no_content_deltas(mock_make_redis):
+    """A stream that emits only reasoning ("thought") dicts and a
+    finish chunk — i.e. reasoning-only-stop, the silent-loss bug's
+    signature — must NOT be cached. Otherwise the empty stream is
+    replayed for the TTL on every identical request, poisoning the
+    cache and denying the reasoning-only recovery any chance to hit
+    a fresh provider call.
+    """
+    mock_redis_instance = MagicMock()
+    mock_make_redis.return_value = mock_redis_instance
+    mock_redis_instance.get.return_value = None
+
+    @stream_cache
+    def mock_function(self, model, messages, stream, tools):
+        yield {"type": "thought", "thought": "thinking hard"}
+        yield {"type": "thought", "thought": " and harder"}
+        yield {"type": "stop"}
+
+    messages = [{"role": "user", "content": "test"}]
+    result = list(mock_function(None, "model", messages, stream=True, tools=None))
+
+    assert len(result) == 3
+    mock_redis_instance.set.assert_not_called()
+
+
+@pytest.mark.unit
+@patch("application.cache.get_redis_instance")
+def test_stream_cache_writes_when_any_content_chunk_seen(mock_make_redis):
+    """The mirror case: a stream with even one str content delta is
+    cached normally (the poison guard is minimal — only reasoning-only
+    streams are dropped)."""
+    mock_redis_instance = MagicMock()
+    mock_make_redis.return_value = mock_redis_instance
+    mock_redis_instance.get.return_value = None
+
+    @stream_cache
+    def mock_function(self, model, messages, stream, tools):
+        yield {"type": "thought", "thought": "brief thought"}
+        yield "the answer"
+
+    messages = [{"role": "user", "content": "test"}]
+    result = list(mock_function(None, "model", messages, stream=True, tools=None))
+
+    assert result == [{"type": "thought", "thought": "brief thought"}, "the answer"]
+    mock_redis_instance.set.assert_called_once()
+
+
 @pytest.mark.unit
 @patch("application.cache.get_redis_instance")
 def test_stream_cache_no_redis(mock_make_redis):

tests/test_coverage_gaps.py

@@ -124,114 +124,6 @@ class TestBaseConnectorAbstract:
         assert "secret" not in result
 
 
-# ---------------------------------------------------------------------------
-# 20. application/llm/sagemaker.py  (lines 52,60,64,67,74,88,106,140)
-# ---------------------------------------------------------------------------
-@pytest.mark.unit
-class TestSagemakerLineIterator:
-    def test_line_iterator_basic(self):
-        from application.llm.sagemaker import LineIterator
-
-        chunks = [
-            {"PayloadPart": {"Bytes": b'{"outputs": [" hello"]}\n'}},
-            {"PayloadPart": {"Bytes": b'{"outputs": [" world"]}\n'}},
-        ]
-        it = LineIterator(iter(chunks))
-        lines = list(it)
-        assert len(lines) == 2
-        assert b"hello" in lines[0]
-
-    def test_line_iterator_split_json(self):
-        from application.llm.sagemaker import LineIterator
-
-        chunks = [
-            {"PayloadPart": {"Bytes": b'{"outputs": '}},
-            {"PayloadPart": {"Bytes": b'[" split"]}\n'}},
-        ]
-        it = LineIterator(iter(chunks))
-        lines = list(it)
-        assert len(lines) == 1
-
-    def test_line_iterator_unknown_event(self):
-        from application.llm.sagemaker import LineIterator
-
-        # The source code on line 55 does `print("Unknown event type:" + chunk)`
-        # which will TypeError when chunk is a dict. We verify that line 54
-        # is covered by catching the error.
-        chunks = [
-            {"InternalServerException": {"Message": "oops"}},
-            {"PayloadPart": {"Bytes": b'{"outputs": ["ok"]}\n'}},
-        ]
-        it = LineIterator(iter(chunks))
-        # The first chunk triggers line 54 branch, but line 55 raises
-        # TypeError due to str + dict concat bug in source.
-        # We just confirm the branch is reached.
-        with pytest.raises(TypeError):
-            list(it)
-
-    def test_sagemaker_llm_init(self):
-        with patch("boto3.client") as mock_boto:
-            mock_boto.return_value = MagicMock()
-            from application.llm.sagemaker import SagemakerAPILLM
-
-            llm = SagemakerAPILLM(api_key="k", user_api_key="uk")
-            assert llm.api_key == "k"
-            assert llm.user_api_key == "uk"
-            assert llm.runtime is not None
-
-    def test_sagemaker_raw_gen(self):
-        with patch("boto3.client") as mock_boto:
-            mock_runtime = MagicMock()
-            body_content = json.dumps(
-                [{"generated_text": "PREFIX ANSWER"}]
-            ).encode("utf-8")
-            mock_body = MagicMock()
-            mock_body.read.return_value = body_content
-            mock_runtime.invoke_endpoint.return_value = {"Body": mock_body}
-            mock_boto.return_value = mock_runtime
-
-            from application.llm.sagemaker import SagemakerAPILLM
-
-            llm = SagemakerAPILLM()
-            messages = [
-                {"content": "context", "role": "system"},
-                {"content": "question", "role": "user"},
-            ]
-            result = llm._raw_gen(None, "model", messages)
-            assert isinstance(result, str)
-
-    def test_sagemaker_raw_gen_stream(self):
-        with patch("boto3.client") as mock_boto:
-            mock_runtime = MagicMock()
-
-            event_stream = [
-                {
-                    "PayloadPart": {
-                        "Bytes": b'{"token": {"text": "hello"}}\n'
-                    }
-                },
-                {
-                    "PayloadPart": {
-                        "Bytes": b'{"token": {"text": "</s>"}}\n'
-                    }
-                },
-            ]
-            mock_runtime.invoke_endpoint_with_response_stream.return_value = {
-                "Body": iter(event_stream)
-            }
-            mock_boto.return_value = mock_runtime
-
-            from application.llm.sagemaker import SagemakerAPILLM
-
-            llm = SagemakerAPILLM()
-            messages = [
-                {"content": "context", "role": "system"},
-                {"content": "question", "role": "user"},
-            ]
-            chunks = list(llm._raw_gen_stream(None, "model", messages))
-            assert "hello" in chunks
-
-
 # ---------------------------------------------------------------------------
 # 9. application/agents/tools/spec_parser.py  (lines 58-59, 71-82, 173-176, 179-180)
 # ---------------------------------------------------------------------------

tests/test_remaining_coverage.py

@@ -497,24 +497,6 @@ class TestMarkdownParserEmptyHeader:
         assert len(tups) >= 1
 
 
-# ---------------------------------------------------------------------------
-# application/llm/sagemaker.py  (line 52)
-# ---------------------------------------------------------------------------
-@pytest.mark.unit
-class TestSagemakerLineIteratorStopIteration:
-    def test_stop_iteration_with_newline_data(self):
-        """Cover line 52: chunk with newline yields a line."""
-        from application.llm.sagemaker import LineIterator
-
-        # Chunk with newline so it yields
-        chunks = [
-            {"PayloadPart": {"Bytes": b'{"outputs": [" partial"]}\n'}},
-        ]
-        it = LineIterator(iter(chunks))
-        lines = list(it)
-        assert len(lines) == 1
-
-
 # ---------------------------------------------------------------------------
 # application/core/url_validation.py  (lines 89-90)
 # ---------------------------------------------------------------------------