mirror of
https://github.com/usestrix/strix.git
synced 2026-10-12 03:47:56 +00:00
Merge remote-tracking branch 'origin/main' into fix/str-1192-mcp-tool-counts
This commit is contained in:
Alex Schapiro
commit
157a9e25b6
8 changed files with 307 additions and 41 deletions
|
|
@ -5,6 +5,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
|
@ -262,6 +263,8 @@ class _TurnGuardModel(Model):
|
|||
not covered by the request timeout, which resets on any byte (keepalives
|
||||
included). ``LLM_STREAM_IDLE_TIMEOUT`` bounds the gap between events so the
|
||||
turn fails instead of hanging, and the existing retry path replays it.
|
||||
``LLM_STREAM_FIRST_EVENT_TIMEOUT`` and ``LLM_STREAM_TOTAL_TIMEOUT`` bound
|
||||
the first event and the whole stream.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -269,11 +272,15 @@ class _TurnGuardModel(Model):
|
|||
inner: Model,
|
||||
*,
|
||||
max_tool_calls_per_turn: int = 0,
|
||||
stream_idle_timeout: float = 0.0,
|
||||
stream_idle_timeout: float | None = None,
|
||||
first_event_timeout: float | None = None,
|
||||
total_timeout: float | None = None,
|
||||
) -> None:
|
||||
self._inner = inner
|
||||
self._max_tool_calls_per_turn = max_tool_calls_per_turn
|
||||
self._stream_idle_timeout = stream_idle_timeout
|
||||
self._first_event_timeout = first_event_timeout
|
||||
self._total_timeout = total_timeout
|
||||
|
||||
def _limiter(self) -> TurnToolCallLimiter:
|
||||
return TurnToolCallLimiter(self._max_tool_calls_per_turn)
|
||||
|
|
@ -360,7 +367,12 @@ class _TurnGuardModel(Model):
|
|||
conversation_id=conversation_id,
|
||||
prompt=prompt,
|
||||
)
|
||||
async for event in _with_idle_timeout(stream, self._stream_idle_timeout):
|
||||
async for event in _with_timeouts(
|
||||
stream,
|
||||
idle=self._stream_idle_timeout,
|
||||
first_event=self._first_event_timeout,
|
||||
total=self._total_timeout,
|
||||
):
|
||||
guarded = _guard_event(event, rewriter, limiter)
|
||||
if guarded is not None:
|
||||
yield guarded
|
||||
|
|
@ -394,28 +406,68 @@ async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None:
|
|||
await stream.aclose()
|
||||
|
||||
|
||||
async def _with_idle_timeout(
|
||||
stream: AsyncIterator[TResponseStreamEvent], timeout: float
|
||||
async def _with_timeouts(
|
||||
stream: AsyncIterator[TResponseStreamEvent],
|
||||
*,
|
||||
idle: float | None = None,
|
||||
first_event: float | None = None,
|
||||
total: float | None = None,
|
||||
) -> AsyncIterator[TResponseStreamEvent]:
|
||||
if timeout <= 0:
|
||||
async for event in stream:
|
||||
yield event
|
||||
return
|
||||
|
||||
"""Bound the first event, the gap between events and the whole stream; None is no bound."""
|
||||
iterator = stream.__aiter__()
|
||||
started = time.monotonic()
|
||||
while True:
|
||||
limits: list[tuple[float, float, str]] = []
|
||||
# Wait until the soonest limit expires: first-event (else idle) until the
|
||||
# first event arrives, then idle; plus whatever is left of the total.
|
||||
if first_event is not None:
|
||||
limits.append((first_event, first_event, "stream_first_event_timeout"))
|
||||
elif idle is not None:
|
||||
limits.append((idle, idle, "stream_idle_timeout"))
|
||||
if total is not None:
|
||||
left = max(0.0, started + total - time.monotonic())
|
||||
limits.append((left, total, "stream_total_timeout"))
|
||||
try:
|
||||
event = await asyncio.wait_for(iterator.__anext__(), timeout)
|
||||
if limits:
|
||||
event = await _next_event(iterator, *min(limits))
|
||||
else:
|
||||
event = await iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
except TimeoutError:
|
||||
except TimeoutError as exc:
|
||||
await _aclose(stream)
|
||||
message = f"model stream produced no event for {timeout:.0f}s"
|
||||
logger.warning("%s; abandoning the turn", message)
|
||||
raise TimeoutError(message) from None
|
||||
logger.warning("%s; abandoning the turn", exc)
|
||||
raise
|
||||
first_event = None
|
||||
yield event
|
||||
|
||||
|
||||
async def _next_event(
|
||||
iterator: AsyncIterator[TResponseStreamEvent], wait: float, limit: float, reason: str
|
||||
) -> TResponseStreamEvent:
|
||||
# asyncio.timeout() cancels without a message, so the request log could not
|
||||
# tell this from a shutdown; this is asyncio.timeout() with a message.
|
||||
task = asyncio.current_task()
|
||||
assert task is not None
|
||||
cancelling = task.cancelling()
|
||||
expired = False
|
||||
|
||||
def expire() -> None:
|
||||
nonlocal expired
|
||||
expired = True
|
||||
task.cancel(msg=f"{request_log.CANCEL_REASON_PREFIX}{reason}")
|
||||
|
||||
handle = asyncio.get_running_loop().call_later(wait, expire)
|
||||
try:
|
||||
return await iterator.__anext__()
|
||||
except BaseException as exc:
|
||||
if expired and task.uncancel() <= cancelling and isinstance(exc, asyncio.CancelledError):
|
||||
raise TimeoutError(f"model stream hit {reason} ({limit:.0f}s)") from None
|
||||
raise
|
||||
finally:
|
||||
handle.cancel()
|
||||
|
||||
|
||||
def _guard_event(
|
||||
event: TResponseStreamEvent, rewriter: TurnCallIdRewriter, limiter: TurnToolCallLimiter
|
||||
) -> TResponseStreamEvent | None:
|
||||
|
|
@ -550,7 +602,10 @@ class StrixProvider(MultiProvider):
|
|||
def get_model(self, model_name: str | None) -> Model:
|
||||
llm = load_settings().llm
|
||||
slug = codex.subscription_model(model_name)
|
||||
idle_timeout = float(llm.stream_idle_timeout)
|
||||
# The settings use 0 for no bound.
|
||||
idle_timeout = float(llm.stream_idle_timeout) or None
|
||||
first_event_timeout = float(llm.stream_first_event_timeout) or None
|
||||
total_timeout = float(llm.stream_total_timeout) or None
|
||||
if slug:
|
||||
# The ChatGPT subscription backend is always streamed; it has no
|
||||
# non-streaming mode to fall back to, so LLM_DISABLE_STREAMING
|
||||
|
|
@ -591,11 +646,13 @@ class StrixProvider(MultiProvider):
|
|||
# The wrapper emits its single event only once the whole request
|
||||
# is done, so an idle gap is meaningless here; the request
|
||||
# timeout bounds it instead.
|
||||
idle_timeout = 0.0
|
||||
idle_timeout = first_event_timeout = total_timeout = None
|
||||
return _TurnGuardModel(
|
||||
model,
|
||||
max_tool_calls_per_turn=llm.max_tool_calls_per_turn,
|
||||
stream_idle_timeout=idle_timeout,
|
||||
first_event_timeout=first_event_timeout,
|
||||
total_timeout=total_timeout,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -745,6 +802,11 @@ def _install_openrouter_stream_cost_capture() -> None:
|
|||
|
||||
class _StrixOpenRouterStreamingHandler(OpenRouterChatCompletionStreamingHandler):
|
||||
def chunk_parser(self, chunk: dict[str, Any]) -> Any:
|
||||
# Before parsing: LiteLLM raises on an error chunk and drops its
|
||||
# top-level ``provider``.
|
||||
request_log.record_upstream_provider(
|
||||
chunk.get("provider"), _openrouter_error_type(chunk.get("error"))
|
||||
)
|
||||
stream = super().chunk_parser(chunk)
|
||||
usage = chunk.get("usage")
|
||||
response_id = chunk.get("id") or getattr(stream, "id", None)
|
||||
|
|
@ -763,12 +825,22 @@ def _install_openrouter_stream_cost_capture() -> None:
|
|||
json_mode=json_mode,
|
||||
)
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any:
|
||||
# A non-2xx reply names the provider in ``error.metadata``.
|
||||
with contextlib.suppress(Exception):
|
||||
error = json.loads(error_message)["error"]
|
||||
request_log.record_upstream_provider(
|
||||
error["metadata"].get("provider_name"), _openrouter_error_type(error)
|
||||
)
|
||||
return super().get_error_class(error_message, status_code, headers)
|
||||
|
||||
def transform_response(self, *args: Any, **kwargs: Any) -> Any:
|
||||
# Non-streamed replies (LLM_DISABLE_STREAMING) skip the chunk parser.
|
||||
response = super().transform_response(*args, **kwargs)
|
||||
raw_response = kwargs.get("raw_response", args[1] if len(args) > 1 else None)
|
||||
with contextlib.suppress(Exception):
|
||||
body = raw_response.json() # type: ignore[union-attr]
|
||||
request_log.record_upstream_provider(body.get("provider"))
|
||||
if body.get("usage"):
|
||||
record_openrouter_provider(body.get("provider"), body["usage"])
|
||||
return response
|
||||
|
|
@ -789,6 +861,14 @@ def _install_openrouter_stream_cost_capture() -> None:
|
|||
litellm.OpenrouterConfig = _StrixOpenrouterConfig # type: ignore[misc]
|
||||
|
||||
|
||||
def _openrouter_error_type(error: object) -> object:
|
||||
if isinstance(error, dict):
|
||||
metadata = error.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
return metadata.get("error_type")
|
||||
return None
|
||||
|
||||
|
||||
OPENROUTER_ATTRIBUTION_HEADERS = {
|
||||
"HTTP-Referer": "https://strix.ai",
|
||||
"X-Title": "Strix",
|
||||
|
|
|
|||
|
|
@ -79,6 +79,10 @@ class LlmSettings(BaseSettings):
|
|||
timeout: int = Field(default=300, alias="LLM_TIMEOUT")
|
||||
preflight_timeout: int = Field(default=30, ge=1, alias="LLM_PREFLIGHT_TIMEOUT")
|
||||
stream_idle_timeout: int = Field(default=300, ge=0, alias="LLM_STREAM_IDLE_TIMEOUT")
|
||||
stream_first_event_timeout: int = Field(
|
||||
default=120, ge=0, alias="LLM_STREAM_FIRST_EVENT_TIMEOUT"
|
||||
)
|
||||
stream_total_timeout: int = Field(default=600, ge=0, alias="LLM_STREAM_TOTAL_TIMEOUT")
|
||||
max_tool_calls_per_turn: int = Field(
|
||||
default=32,
|
||||
ge=0,
|
||||
|
|
|
|||
|
|
@ -184,9 +184,7 @@ def _is_transient_model_error(exc: BaseException) -> bool:
|
|||
return True
|
||||
code = _model_error_status_code(exc)
|
||||
if code is not None:
|
||||
import litellm
|
||||
|
||||
return bool(litellm._should_retry(code))
|
||||
return code >= 400 and code not in (401, 402, 403, 404)
|
||||
return isinstance(exc, APIError)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -159,11 +159,33 @@ class HttpReply:
|
|||
|
||||
status_code: int | None = None
|
||||
headers: dict[str, str] | None = None
|
||||
upstream_provider: str | None = None
|
||||
upstream_error_type: str | None = None
|
||||
|
||||
|
||||
_http_reply: ContextVar[HttpReply | None] = ContextVar("strix_llm_http_reply", default=None)
|
||||
|
||||
|
||||
def record_upstream_provider(provider: object, error_type: object = None) -> None:
|
||||
"""Remember which provider behind a gateway served the attempt in flight."""
|
||||
reply = _http_reply.get()
|
||||
if reply is None:
|
||||
return
|
||||
if isinstance(provider, str) and provider:
|
||||
reply.upstream_provider = provider
|
||||
if isinstance(error_type, str) and error_type:
|
||||
reply.upstream_error_type = error_type
|
||||
|
||||
|
||||
def _upstream_details(reply: HttpReply | None) -> list[tuple[str, object]]:
|
||||
if reply is None:
|
||||
return []
|
||||
return [
|
||||
("upstream_provider", reply.upstream_provider),
|
||||
("upstream_error_type", reply.upstream_error_type),
|
||||
]
|
||||
|
||||
|
||||
async def record_http_reply(response: Response) -> None:
|
||||
"""httpx ``response`` event hook: remember the reply for the attempt in flight.
|
||||
|
||||
|
|
@ -772,6 +794,7 @@ def _litellm_success(
|
|||
("request", _litellm_request_details(kwargs)),
|
||||
("response", _litellm_response_details(response)),
|
||||
("litellm", _litellm_hidden_details(hidden, slo)),
|
||||
*_upstream_details(_http_reply.get()),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -821,6 +844,7 @@ def _litellm_failure(
|
|||
("request", _litellm_request_details(kwargs)),
|
||||
("error", error or None),
|
||||
("litellm", _litellm_hidden_details(hidden, slo)),
|
||||
*_upstream_details(_http_reply.get()),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -894,15 +918,18 @@ def _observe_sdk_shared_http_client() -> None:
|
|||
|
||||
|
||||
def _log_line_sink(event: LlmRequestEvent) -> None:
|
||||
details = event.details or {}
|
||||
level = logging.DEBUG if event.outcome == "success" else logging.WARNING
|
||||
logger.log(
|
||||
level,
|
||||
"llm_request route=%s provider=%s model=%s host=%s outcome=%s status=%s "
|
||||
"request_id=%s response_id=%s stream=%s duration_ms=%d ttft_ms=%s "
|
||||
"llm_request route=%s provider=%s upstream=%s upstream_error=%s model=%s host=%s "
|
||||
"outcome=%s status=%s request_id=%s response_id=%s stream=%s duration_ms=%d ttft_ms=%s "
|
||||
"req_bytes=%s res_bytes=%s finish=%s "
|
||||
"in=%s out=%s cached=%s cost=%s agent=%s attempt=%d%s",
|
||||
event.route,
|
||||
event.provider or "-",
|
||||
details.get("upstream_provider") or "-",
|
||||
details.get("upstream_error_type") or "-",
|
||||
event.model,
|
||||
event.api_host or "-",
|
||||
event.outcome,
|
||||
|
|
@ -1070,6 +1097,7 @@ class RequestLoggingModel(Model):
|
|||
details=merge_details(
|
||||
("request", request.details),
|
||||
("response", _openai_response_details(response, raw_response)),
|
||||
*_upstream_details(reply),
|
||||
),
|
||||
)
|
||||
status, request_id = _openai_error_fields(exc)
|
||||
|
|
@ -1082,13 +1110,14 @@ class RequestLoggingModel(Model):
|
|||
outcome="error",
|
||||
status_code=status,
|
||||
provider_request_id=request_id or base.provider_request_id,
|
||||
error_type=type(exc).__name__,
|
||||
error_type=_cancel_reason(exc) or type(exc).__name__,
|
||||
error_message=_abandonment_message(exc) or clean_error_message(exc),
|
||||
response_bytes=_exception_body_size(exc, status),
|
||||
response_headers=headers_from_response(error_headers) or base.response_headers,
|
||||
details=merge_details(
|
||||
("request", request.details),
|
||||
("error", _exception_details(exc)),
|
||||
*_upstream_details(reply),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -1318,9 +1347,23 @@ def _is_abandonment(exc: BaseException | None) -> bool:
|
|||
return isinstance(exc, asyncio.CancelledError | GeneratorExit)
|
||||
|
||||
|
||||
# Prefixes the reason a stream timeout gives when it cancels an attempt.
|
||||
CANCEL_REASON_PREFIX = "strix:"
|
||||
|
||||
|
||||
def _cancel_reason(exc: BaseException) -> str | None:
|
||||
"""The stream timeout that cancelled the attempt, e.g. ``stream_idle_timeout``."""
|
||||
if isinstance(exc, asyncio.CancelledError) and exc.args:
|
||||
message = str(exc.args[0])
|
||||
if message.startswith(CANCEL_REASON_PREFIX):
|
||||
return message.removeprefix(CANCEL_REASON_PREFIX)
|
||||
return None
|
||||
|
||||
|
||||
def _abandonment_message(exc: BaseException) -> str | None:
|
||||
if isinstance(exc, asyncio.CancelledError):
|
||||
return "attempt cancelled before the reply was consumed (stream idle timeout or shutdown)"
|
||||
reason = f": {exc.args[0]}" if exc.args else ""
|
||||
return f"attempt cancelled before the reply was consumed{reason}"
|
||||
if isinstance(exc, GeneratorExit):
|
||||
return "stream closed by the caller before it finished"
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
|
@ -389,3 +390,67 @@ def test_openrouter_request_carries_agent_session_id() -> None:
|
|||
assert body()["session_id"] == session_id
|
||||
finally:
|
||||
request_log.reset_call_context(token)
|
||||
|
||||
|
||||
def _openrouter_config() -> Any:
|
||||
_install_openrouter_stream_cost_capture()
|
||||
config = ProviderConfigManager.get_provider_chat_config(
|
||||
model="z-ai/glm-5.3", provider=LlmProviders.OPENROUTER
|
||||
)
|
||||
assert config is not None
|
||||
return config
|
||||
|
||||
|
||||
def test_openrouter_records_upstream_of_stream_and_mid_stream_error() -> None:
|
||||
handler = _openrouter_config().get_model_response_iterator(
|
||||
streaming_response=iter([]), sync_stream=True
|
||||
)
|
||||
reply = request_log.HttpReply()
|
||||
token = request_log._http_reply.set(reply)
|
||||
try:
|
||||
handler.chunk_parser(
|
||||
{
|
||||
"id": "gen-a",
|
||||
"created": 1,
|
||||
"model": "z-ai/glm-5.3",
|
||||
"provider": "Relace",
|
||||
"choices": [{"index": 0, "delta": {"content": "x"}}],
|
||||
}
|
||||
)
|
||||
assert reply.upstream_provider == "Relace"
|
||||
assert reply.upstream_error_type is None
|
||||
with pytest.raises(Exception, match="incomplete tool call"):
|
||||
handler.chunk_parser(
|
||||
{
|
||||
"id": "gen-a",
|
||||
"provider": "Relace",
|
||||
"error": {
|
||||
"code": 400,
|
||||
"message": "Generation stopped with an incomplete tool call.",
|
||||
"metadata": {"error_type": "invalid_request"},
|
||||
},
|
||||
}
|
||||
)
|
||||
finally:
|
||||
request_log._http_reply.reset(token)
|
||||
assert reply.upstream_provider == "Relace"
|
||||
assert reply.upstream_error_type == "invalid_request"
|
||||
|
||||
|
||||
def test_openrouter_records_upstream_of_pre_stream_error() -> None:
|
||||
body = {
|
||||
"error": {
|
||||
"message": "Provider returned error",
|
||||
"code": 400,
|
||||
"metadata": {"raw": "...", "provider_name": "InferenceNet"},
|
||||
}
|
||||
}
|
||||
reply = request_log.HttpReply()
|
||||
token = request_log._http_reply.set(reply)
|
||||
try:
|
||||
_openrouter_config().get_error_class(json.dumps(body), 400, {})
|
||||
_openrouter_config().get_error_class("<html>502 Bad Gateway</html>", 502, {})
|
||||
finally:
|
||||
request_log._http_reply.reset(token)
|
||||
assert reply.upstream_provider == "InferenceNet"
|
||||
assert reply.upstream_error_type is None
|
||||
|
|
|
|||
|
|
@ -79,12 +79,13 @@ def test_content_guardrail_is_not_retried() -> None:
|
|||
assert execution._is_transient_model_error(guardrail) is False
|
||||
|
||||
|
||||
def test_client_errors_are_not_transient() -> None:
|
||||
def test_client_errors_are_transient() -> None:
|
||||
bad_request = BadRequestError(
|
||||
"bad", response=httpx.Response(400, request=_request()), body=None
|
||||
)
|
||||
assert execution._is_transient_model_error(bad_request) is False
|
||||
assert execution._is_transient_model_error(_status_error(404)) is False
|
||||
assert execution._is_transient_model_error(bad_request) is True
|
||||
for status in (401, 402, 403, 404):
|
||||
assert execution._is_transient_model_error(_status_error(status)) is False
|
||||
assert execution._is_transient_model_error(ValueError("nope")) is False
|
||||
|
||||
|
||||
|
|
@ -167,11 +168,8 @@ async def test_run_cycle_gives_up_after_max_retries(
|
|||
async def test_run_cycle_does_not_retry_permanent_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
bad_request = BadRequestError(
|
||||
"bad", response=httpx.Response(400, request=_request()), body=None
|
||||
)
|
||||
streams = [_FakeStream(exc=bad_request), _FakeStream()]
|
||||
with pytest.raises(BadRequestError):
|
||||
streams = [_FakeStream(exc=ValueError("nope")), _FakeStream()]
|
||||
with pytest.raises(ValueError, match="nope"):
|
||||
await _run_once(monkeypatch, streams)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -394,9 +394,30 @@ def test_log_line_sink_formats_without_content(caplog: pytest.LogCaptureFixture)
|
|||
assert "request_id=req_line01" in line
|
||||
assert "status=400" in line
|
||||
assert "provider=anthropic" in line
|
||||
assert "upstream=- upstream_error=-" in line
|
||||
assert "SECRET PROMPT" not in line
|
||||
|
||||
|
||||
def test_log_line_names_upstream_recorded_during_attempt(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
exc = _anthropic_error(400, ANTHROPIC_BLOCK_BODY, {})
|
||||
token = request_log._http_reply.set(request_log.HttpReply())
|
||||
try:
|
||||
request_log.record_upstream_provider("Together", "provider_unavailable")
|
||||
event = request_log.event_from_litellm(
|
||||
_anthropic_kwargs(exc), None, None, None, outcome="error"
|
||||
)
|
||||
finally:
|
||||
request_log._http_reply.reset(token)
|
||||
assert event.details is not None
|
||||
assert event.details["upstream_provider"] == "Together"
|
||||
with caplog.at_level(logging.DEBUG, logger="strix.llm.request_log"):
|
||||
request_log._log_line_sink(event)
|
||||
line = caplog.records[-1].getMessage()
|
||||
assert "upstream=Together upstream_error=provider_unavailable" in line
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# sizes, timing, finish reason, free-form headers and details #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
|
|
|||
|
|
@ -23,7 +23,8 @@ from openai import AsyncOpenAI
|
|||
|
||||
from strix.config import loader
|
||||
from strix.config.loader import load_settings
|
||||
from strix.config.models import StrixProvider, _TurnGuardModel, _with_idle_timeout
|
||||
from strix.config.models import StrixProvider, _TurnGuardModel, _with_timeouts
|
||||
from strix.llm import request_log
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -78,9 +79,21 @@ def stalling_gateway() -> Iterator[str]:
|
|||
server.server_close()
|
||||
|
||||
|
||||
def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]:
|
||||
class _ServedByRelace(OpenAIChatCompletionsModel):
|
||||
async def stream_response(self, *args: Any, **kwargs: Any) -> AsyncIterator[Any]:
|
||||
request_log.record_upstream_provider("Relace")
|
||||
async for event in super().stream_response(*args, **kwargs):
|
||||
yield event
|
||||
|
||||
|
||||
def _stream(base_url: str, *, idle_timeout: float | None) -> AsyncIterator[Any]:
|
||||
client = AsyncOpenAI(api_key="tok", base_url=base_url, max_retries=0, timeout=_STALL_SECONDS)
|
||||
inner: Model = OpenAIChatCompletionsModel(model="gw-model", openai_client=client)
|
||||
inner: Model = request_log.RequestLoggingModel(
|
||||
_ServedByRelace(model="gw-model", openai_client=client),
|
||||
model_name="gw-model",
|
||||
provider="openai",
|
||||
base_url=base_url,
|
||||
)
|
||||
guarded = _TurnGuardModel(inner, stream_idle_timeout=idle_timeout)
|
||||
return guarded.stream_response(
|
||||
None,
|
||||
|
|
@ -96,7 +109,7 @@ def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]:
|
|||
)
|
||||
|
||||
|
||||
async def _drain(base_url: str, *, idle_timeout: float) -> list[Any]:
|
||||
async def _drain(base_url: str, *, idle_timeout: float | None) -> list[Any]:
|
||||
return [event async for event in _stream(base_url, idle_timeout=idle_timeout)]
|
||||
|
||||
|
||||
|
|
@ -105,16 +118,31 @@ async def test_stalled_stream_hangs_without_the_watchdog(stalling_gateway: str)
|
|||
# Repro: tokens arrive, then nothing. Un-watched, the turn just sits there;
|
||||
# the request timeout is far away and would reset on any keepalive byte.
|
||||
with pytest.raises(TimeoutError):
|
||||
await asyncio.wait_for(_drain(stalling_gateway, idle_timeout=0), timeout=2)
|
||||
await asyncio.wait_for(_drain(stalling_gateway, idle_timeout=None), timeout=2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stalled_stream_is_abandoned_by_the_watchdog(stalling_gateway: str) -> None:
|
||||
logged: list[request_log.LlmRequestEvent] = []
|
||||
request_log.register_sink(logged.append)
|
||||
started = time.monotonic()
|
||||
with pytest.raises(TimeoutError, match="produced no event"):
|
||||
await _drain(stalling_gateway, idle_timeout=1)
|
||||
try:
|
||||
with pytest.raises(TimeoutError, match="stream_idle_timeout"):
|
||||
await _drain(stalling_gateway, idle_timeout=1)
|
||||
finally:
|
||||
request_log.unregister_sink(logged.append)
|
||||
|
||||
assert time.monotonic() - started < _STALL_SECONDS
|
||||
# The request log can tell our timeout from a shutdown.
|
||||
assert [
|
||||
(e.error_type, (e.details or {}).get("upstream_provider"), e.error_message) for e in logged
|
||||
] == [
|
||||
(
|
||||
"stream_idle_timeout",
|
||||
"Relace",
|
||||
"attempt cancelled before the reply was consumed: strix:stream_idle_timeout",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -124,14 +152,39 @@ async def test_events_keep_flowing_while_the_stream_is_alive() -> None:
|
|||
await asyncio.sleep(0.05)
|
||||
yield f"event-{i}"
|
||||
|
||||
seen: list[Any] = [event async for event in _with_idle_timeout(_live(), 1.0)]
|
||||
seen: list[Any] = [event async for event in _with_timeouts(_live(), idle=1.0)]
|
||||
|
||||
assert seen == [f"event-{i}" for i in range(5)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_event_and_whole_stream_are_bounded() -> None:
|
||||
async def _slow_start() -> AsyncIterator[Any]:
|
||||
await asyncio.sleep(_STALL_SECONDS)
|
||||
yield "late"
|
||||
|
||||
async def _endless() -> AsyncIterator[Any]:
|
||||
while True:
|
||||
await asyncio.sleep(0.05)
|
||||
yield "event"
|
||||
|
||||
with pytest.raises(TimeoutError, match="stream_first_event_timeout"):
|
||||
[event async for event in _with_timeouts(_slow_start(), idle=5, first_event=0.2)]
|
||||
with pytest.raises(TimeoutError, match="stream_idle_timeout"):
|
||||
[event async for event in _with_timeouts(_slow_start(), idle=0.2)]
|
||||
with pytest.raises(TimeoutError, match="stream_total_timeout"):
|
||||
[event async for event in _with_timeouts(_endless(), idle=5, total=0.3)]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _reset_settings(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
||||
for key in ("STRIX_LLM", "LLM_DISABLE_STREAMING", "LLM_STREAM_IDLE_TIMEOUT"):
|
||||
for key in (
|
||||
"STRIX_LLM",
|
||||
"LLM_DISABLE_STREAMING",
|
||||
"LLM_STREAM_IDLE_TIMEOUT",
|
||||
"LLM_STREAM_FIRST_EVENT_TIMEOUT",
|
||||
"LLM_STREAM_TOTAL_TIMEOUT",
|
||||
):
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
monkeypatch.setattr(loader, "_override", None)
|
||||
|
|
@ -151,11 +204,15 @@ def test_idle_timeout_is_configurable(
|
|||
) -> None:
|
||||
monkeypatch.setattr("strix.config.models.MultiProvider.get_model", lambda *_: _DummyModel())
|
||||
monkeypatch.setenv("LLM_STREAM_IDLE_TIMEOUT", "45")
|
||||
monkeypatch.setenv("LLM_STREAM_FIRST_EVENT_TIMEOUT", "30")
|
||||
monkeypatch.setenv("LLM_STREAM_TOTAL_TIMEOUT", "400")
|
||||
load_settings()
|
||||
|
||||
model = StrixProvider().get_model("openai/gpt-4o-mini")
|
||||
assert isinstance(model, _TurnGuardModel)
|
||||
assert model._stream_idle_timeout == 45
|
||||
assert model._first_event_timeout == 30
|
||||
assert model._total_timeout == 400
|
||||
|
||||
|
||||
def test_idle_timeout_is_off_without_streaming(
|
||||
|
|
@ -170,4 +227,4 @@ def test_idle_timeout_is_off_without_streaming(
|
|||
|
||||
model = StrixProvider().get_model("openai/gpt-4o-mini")
|
||||
assert isinstance(model, _TurnGuardModel)
|
||||
assert model._stream_idle_timeout == 0
|
||||
assert model._stream_idle_timeout is model._first_event_timeout is model._total_timeout is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue