diff --git a/application/cache.py b/application/cache.py index 99f738e4..4daee8d4 100644 --- a/application/cache.py +++ b/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) diff --git a/application/core/model_settings.py b/application/core/model_settings.py index 19fe0318..73a39b40 100644 --- a/application/core/model_settings.py +++ b/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" diff --git a/application/core/settings.py b/application/core/settings.py index a10f90df..42cb3918 100644 --- a/application/core/settings.py +++ b/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" diff --git a/application/llm/base.py b/application/llm/base.py index 7399320b..df6ed773 100644 --- a/application/llm/base.py +++ b/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)}" diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 22a9fef5..21d485ee 100644 --- a/application/llm/handlers/base.py +++ b/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, ) diff --git a/application/llm/premai.py b/application/llm/premai.py deleted file mode 100644 index b25b0299..00000000 --- a/application/llm/premai.py +++ /dev/null @@ -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"] diff --git a/application/llm/providers/__init__.py b/application/llm/providers/__init__.py index c712c89e..8c189955 100644 --- a/application/llm/providers/__init__.py +++ b/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} diff --git a/application/llm/providers/premai.py b/application/llm/providers/premai.py deleted file mode 100644 index 6f98ece1..00000000 --- a/application/llm/providers/premai.py +++ /dev/null @@ -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 diff --git a/application/llm/providers/sagemaker.py b/application/llm/providers/sagemaker.py deleted file mode 100644 index b36822be..00000000 --- a/application/llm/providers/sagemaker.py +++ /dev/null @@ -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 diff --git a/application/llm/sagemaker.py b/application/llm/sagemaker.py deleted file mode 100644 index d56d7a9c..00000000 --- a/application/llm/sagemaker.py +++ /dev/null @@ -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": ["", "###"], - }, - } - 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": ["", "###"], - }, - } - 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 ["", "###"]: - print(data["token"]["text"], end="") - yield data["token"]["text"] diff --git a/docs/content/Models/cloud-providers.mdx b/docs/content/Models/cloud-providers.mdx index b11973a8..7adb3710 100644 --- a/docs/content/Models/cloud-providers.mdx +++ b/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. diff --git a/tests/core/test_model_settings.py b/tests/core/test_model_settings.py index f00de9c4..19074ba0 100644 --- a/tests/core/test_model_settings.py +++ b/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" diff --git a/tests/llm/handlers/test_reasoning_only_recovery.py b/tests/llm/handlers/test_reasoning_only_recovery.py new file mode 100644 index 00000000..180f4bbf --- /dev/null +++ b/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] diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index 330a6fca..518dfa67 100644 --- a/tests/llm/test_base.py +++ b/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) diff --git a/tests/llm/test_fallback.py b/tests/llm/test_fallback.py index d5e76652..e06b1edd 100644 --- a/tests/llm/test_fallback.py +++ b/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 ): diff --git a/tests/llm/test_premai.py b/tests/llm/test_premai.py deleted file mode 100644 index 06f57e36..00000000 --- a/tests/llm/test_premai.py +++ /dev/null @@ -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 diff --git a/tests/llm/test_sagemaker.py b/tests/llm/test_sagemaker.py deleted file mode 100644 index 2b893a9a..00000000 --- a/tests/llm/test_sagemaker.py +++ /dev/null @@ -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": ["", "###"] - } - } - self.payload_stream = { - "inputs": self.prompt, - "stream": True, - "parameters": { - "do_sample": True, - "temperature": 0.1, - "max_new_tokens": 512, - "repetition_penalty": 1.03, - "stop": ["", "###"] - } - } - 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() \ No newline at end of file diff --git a/tests/test_cache.py b/tests/test_cache.py index 16118960..0f61eda1 100644 --- a/tests/test_cache.py +++ b/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): diff --git a/tests/test_coverage_gaps.py b/tests/test_coverage_gaps.py index 22683df9..c0d6a307 100644 --- a/tests/test_coverage_gaps.py +++ b/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": ""}}\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) # --------------------------------------------------------------------------- diff --git a/tests/test_remaining_coverage.py b/tests/test_remaining_coverage.py index 6899cb81..425d2406 100644 --- a/tests/test_remaining_coverage.py +++ b/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) # ---------------------------------------------------------------------------