Stream the chat route, so the timeout is on silence (#174) #184

Merged
dev-agent merged 1 commits from streaming into main 2026-10-26 22:50:43 +00:00
5 changed files with 374 additions and 30 deletions
Showing only changes of commit 4904fa0b63 - Show all commits

No files matched your search

+15 -10
View File
@@ -135,18 +135,23 @@ a local model can be held to less without affecting the frontier roles.
### `llm.request_timeout_s` — default 600 ### `llm.request_timeout_s` — default 600
How long one model request may take before it is abandoned. It applies to How long a model may go silent before its request is abandoned. On the chat
every request, on both routes. route, which every local model takes, replies are streamed (#174): this is the
longest wait for the next piece of a reply, including the wait for the first
one while the model reads its prompt. A slow model that keeps writing is never
cut off. On the passthrough route, which only hosted Anthropic models take,
it is still the whole reply.
- **Too low:** a slow model's long answer is cut off, and the driver's run - **Too low:** a model reading a long prompt is cut off before it answers, and
fails on work that was going well. At 120 s, the old fixed value, a dense the run fails on work that was going well.
27B model on the RTX 4090 timed out writing one module in one turn. - **Too high:** a hung model holds the run for that long before anything
- **Too high:** a hung connection holds the milestone for that long before notices.
anything notices.
- **Deciding:** this is a safety net, not a budget. `max_output_tokens` and - **Deciding:** this is a safety net, not a budget. `max_output_tokens` and
`max_turns` bound the work; set this above the slowest single answer you `max_turns` bound the work. Set it above the longest the model takes to
expect, which is the model's output ceiling divided by its generation read its largest prompt (its context window divided by its prompt speed),
speed. and, for hosted Anthropic models, above their slowest whole answer.
Abandoning a streamed request also frees the inference server; abandoning
a plain one, measured on the 4090, left it finishing the reply anyway.
### `tuning.driver.max_turns` — default 40 ### `tuning.driver.max_turns` — default 40
@@ -1,7 +1,7 @@
"""A model call that a stop can break off (L2). """A model call that a stop can break off (L2).
A ``Cancel`` takes effect at the next tool boundary, but one model call can A ``Cancel`` takes effect at the next tool boundary, but one model call can
run for many minutes (``request_timeout_s``), longer than a service manager run for many minutes, longer than a service manager
waits before it kills the process. ``Interruptible`` wraps the client and waits before it kills the process. ``Interruptible`` wraps the client and
knows when a call is in flight; ``interrupt``, called from the signal knows when a call is in flight; ``interrupt``, called from the signal
handler, raises ``Stopped`` inside that call so it ends at once. handler, raises ``Stopped`` inside that call so it ends at once.
+135 -17
View File
@@ -42,7 +42,7 @@ RATE_LIMIT_RETRY_S = 20.0
# The proxy, or what's behind it, briefly gone: a 502/503/504 or a dropped # The proxy, or what's behind it, briefly gone: a 502/503/504 or a dropped
# connection. Outages delay, never fail (L5): waited out with backoff before # connection. Outages delay, never fail (L5): waited out with backoff before
# the work is stopped. Not a 500 (it may be the request) and not a timeout # the work is stopped. Not a 500 (it may be the request) and not a timeout
# (a hung model has had its request_timeout_s already). # (a model silent for request_timeout_s is hung, not slow).
OUTAGE_STATUSES = frozenset({502, 503, 504}) OUTAGE_STATUSES = frozenset({502, 503, 504})
OUTAGE_BACKOFF_S = (5.0, 15.0, 45.0) OUTAGE_BACKOFF_S = (5.0, 15.0, 45.0)
@@ -53,10 +53,14 @@ OUTAGE_BACKOFF_S = (5.0, 15.0, 45.0)
# was observed at 4,064 against a 4,096 limit. # was observed at 4,064 against a 4,096 limit.
DEFAULT_MAX_OUTPUT_TOKENS = 16384 DEFAULT_MAX_OUTPUT_TOKENS = 16384
# How long one request may take before it is abandoned (llm.request_timeout_s). # How long a model may go silent before its request is abandoned
# A safety net for a hung connection, not a budget: a dense local model # (llm.request_timeout_s). A safety net for a hung connection, not a budget.
# writing a whole module in one turn was measured taking over 120 s, and # The chat route streams its replies (#174), so this bounds the wait for the
# timing it out failed work that was going well. # next piece of a reply, not the whole reply: a slow model that keeps writing
# is never cut off, while one that stops is. Abandoning a stream also frees
# the inference server, measured; abandoning a plain request left the 4090
# finishing it for whoever came next. The passthrough route, for hosted
# models only, still waits for the whole reply.
DEFAULT_REQUEST_TIMEOUT_S = 600.0 DEFAULT_REQUEST_TIMEOUT_S = 600.0
@@ -212,10 +216,14 @@ class LiteLLMClient:
# nothing to check against. A caller that knows the model's limits # nothing to check against. A caller that knows the model's limits
# (the context budget) passes a tighter one. # (the context budget) passes a tighter one.
body["max_tokens"] = max_tokens or self._max_tokens body["max_tokens"] = max_tokens or self._max_tokens
# Streamed, so the timeout is on silence, not on the whole reply
# (DEFAULT_REQUEST_TIMEOUT_S); usage then comes in the last piece.
body["stream"] = True
body["stream_options"] = {"include_usage": True}
self.last_exchange = exchange = {"route": "chat", "request": body} self.last_exchange = exchange = {"route": "chat", "request": body}
self.last_rate_limit_wait_s = 0.0 # not the last call's, if this one raises self.last_rate_limit_wait_s = 0.0 # not the last call's, if this one raises
response, self.last_rate_limit_wait_s = _post_with_rate_limit_retry( response, self.last_rate_limit_wait_s = _send(
self._http, "/chat/completions", body, self._sleep self._http, "/chat/completions", body, self._sleep, _assemble_chat
) )
_keep_response(exchange, response) _keep_response(exchange, response)
if response.status_code != 200: if response.status_code != 200:
@@ -353,7 +361,7 @@ class AnthropicPassthroughClient:
body["tool_choice"] = {"type": "any"} # Anthropic's name for it body["tool_choice"] = {"type": "any"} # Anthropic's name for it
self.last_exchange = exchange = {"route": "messages", "request": body} self.last_exchange = exchange = {"route": "messages", "request": body}
self.last_rate_limit_wait_s = 0.0 # not the last call's, if this one raises self.last_rate_limit_wait_s = 0.0 # not the last call's, if this one raises
response, self.last_rate_limit_wait_s = _post_with_rate_limit_retry( response, self.last_rate_limit_wait_s = _send(
self._http, "/v1/messages", body, self._sleep self._http, "/v1/messages", body, self._sleep
) )
_keep_response(exchange, response) _keep_response(exchange, response)
@@ -415,22 +423,42 @@ class AnthropicPassthroughClient:
self._http.close() self._http.close()
def _post_with_rate_limit_retry( def _send(
http: httpx.Client, path: str, body: dict[str, Any], sleep: Any http: httpx.Client,
path: str,
body: dict[str, Any],
sleep: Any,
assemble: Any = None,
) -> tuple[httpx.Response, float]: ) -> tuple[httpx.Response, float]:
"""The response, and how long was spent waiting: out a rate limit (once) """The response, and how long was spent waiting: out a rate limit (once)
or an outage (with backoff), not on the model.""" or an outage (with backoff), not on the model.
A streamed reply is put back together by ``assemble`` into the response
the server would have sent unstreamed, so everything after reads one
shape. A server that answers a streamed request with plain JSON is
taken as it is."""
waited = 0.0 waited = 0.0
rate_limited = False rate_limited = False
outages = iter(OUTAGE_BACKOFF_S) outages = iter(OUTAGE_BACKOFF_S)
while True: while True:
try: try:
response = http.post(path, json=body) with http.stream("POST", path, json=body) as response:
streamed = response.headers.get("content-type", "").startswith(
"text/event-stream"
)
if response.status_code == 200 and streamed and assemble:
data = assemble(response.iter_lines())
return _assembled(response, data), waited
response.read()
except httpx.TimeoutException as exc: except httpx.TimeoutException as exc:
raise UpstreamUnavailable(f"LLM request failed: {exc}") from exc raise UpstreamUnavailable(
f"LLM request failed: the model said nothing for "
f"{http.timeout.read:g} s ({exc})"
) from exc
except httpx.HTTPError as exc: except httpx.HTTPError as exc:
# Could not reach the proxy at all: nothing about the request was # Could not reach the proxy, or lost it mid-reply: nothing about
# wrong, so callers holding work in progress should keep it. # the request was wrong, so callers holding work in progress
# should keep it.
pause = next(outages, None) pause = next(outages, None)
if pause is None: if pause is None:
raise UpstreamUnavailable(f"LLM request failed: {exc}") from exc raise UpstreamUnavailable(f"LLM request failed: {exc}") from exc
@@ -453,6 +481,88 @@ def _post_with_rate_limit_retry(
return response, waited return response, waited
def _assembled(streamed: httpx.Response, data: dict[str, Any]) -> httpx.Response:
headers = {}
if cost := streamed.headers.get("x-litellm-response-cost"):
headers["x-litellm-response-cost"] = cost
return httpx.Response(200, json=data, headers=headers, request=streamed.request)
def _assemble_chat(lines: Any) -> dict[str, Any]:
"""A streamed chat completion, put back together as the one response it
stands for: the text, the reasoning, each tool call's arguments joined
from their pieces, the finish reason and the usage."""
model = ""
text: list[str] = []
reasoning: list[str] = []
calls: dict[int, dict[str, Any]] = {}
finish = None
usage: dict[str, Any] = {}
done = False
for line in lines:
if not line.startswith("data:"):
continue # blank lines, comments, keep-alives
payload = line[len("data:") :].strip()
if payload == "[DONE]":
done = True
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError as exc:
raise LLMError(
f"malformed piece of a streamed reply: {payload[:200]!r}"
) from exc
if not isinstance(chunk, dict):
raise LLMError(f"malformed piece of a streamed reply: {payload[:200]!r}")
if "error" in chunk:
# The proxy failing mid-reply: upstream trouble, as a non-200
# would have been before the reply began.
said = json.dumps(chunk["error"])[:300]
if _looks_like_budget_refusal(said):
raise BudgetExhausted(f"the proxy refused on budget grounds: {said}")
raise UpstreamUnavailable(f"the reply broke off with an error: {said}")
model = chunk.get("model") or model
if isinstance(chunk.get("usage"), dict):
usage = chunk["usage"]
for choice in chunk.get("choices") or []:
delta = choice.get("delta") or {}
if isinstance(delta.get("content"), str):
text.append(delta["content"])
if isinstance(delta.get("reasoning_content"), str):
reasoning.append(delta["reasoning_content"])
for piece in delta.get("tool_calls") or []:
call = calls.setdefault(
int(piece.get("index") or 0),
{
"id": "",
"type": "function",
"function": {"name": "", "arguments": ""},
},
)
call["id"] = call["id"] or piece.get("id") or ""
function = piece.get("function") or {}
# The name comes once, the arguments in pieces.
call["function"]["name"] = (
call["function"]["name"] or function.get("name") or ""
)
call["function"]["arguments"] += function.get("arguments") or ""
if choice.get("finish_reason"):
finish = choice["finish_reason"]
if not done and finish is None:
# Cut off without a word: the connection, not the model.
raise UpstreamUnavailable("the streamed reply stopped before it finished")
message: dict[str, Any] = {"role": "assistant", "content": "".join(text) or None}
if reasoning:
message["reasoning_content"] = "".join(reasoning)
if calls:
message["tool_calls"] = [calls[i] for i in sorted(calls)]
return {
"model": model,
"choices": [{"index": 0, "message": message, "finish_reason": finish}],
"usage": usage,
}
# Model families we have a native renderer adapter for. L9: delegations take # Model families we have a native renderer adapter for. L9: delegations take
# the vendor's passthrough route because a normalising middlebox is lossy # the vendor's passthrough route because a normalising middlebox is lossy
# exactly where a long-running delegation lives. # exactly where a long-running delegation lives.
@@ -467,9 +577,17 @@ def _keep_response(exchange: dict[str, Any], response: httpx.Response) -> None:
def _response_cost(response: httpx.Response) -> float: def _response_cost(response: httpx.Response) -> float:
"""LiteLLM's price for the call: a header on a plain reply, the usage's
``cost`` on a streamed one (sent last, once it's known)."""
found: Any = response.headers.get("x-litellm-response-cost")
if not found:
try: try:
return float(response.headers.get("x-litellm-response-cost") or 0.0) found = (response.json().get("usage") or {}).get("cost")
except ValueError: except (ValueError, AttributeError):
found = None
try:
return float(found or 0.0)
except (TypeError, ValueError):
return 0.0 return 0.0
+220
View File
@@ -602,3 +602,223 @@ def test_a_budget_refusal_is_not_waited_on() -> None:
with pytest.raises(BudgetExhausted): with pytest.raises(BudgetExhausted):
make_client(handler).complete(EVENTS, tools=None, model="local/qwen") make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
assert len(calls) == 1 assert len(calls) == 1
# --- streamed replies on the chat route (#174) -----------------------------------
def sse(*chunks: dict, done: bool = True) -> httpx.Response:
lines = "".join(f"data: {json.dumps(c)}\n\n" for c in chunks)
return httpx.Response(
200,
headers={"content-type": "text/event-stream"},
content=lines + ("data: [DONE]\n\n" if done else ""),
)
def piece(delta: dict | None = None, finish: str | None = None, **extra) -> dict:
choice: dict = {"index": 0, "delta": delta or {}}
if finish:
choice["finish_reason"] = finish
return {"model": "qwen-real", "choices": [choice], **extra}
class Failing(httpx.SyncByteStream):
"""A reply that sends its first piece, then fails."""
def __init__(self, first: bytes, failure: Exception) -> None:
self._first, self._failure = first, failure
def __iter__(self):
yield self._first
raise self._failure
def test_the_chat_route_asks_for_a_stream_with_usage() -> None:
seen: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
seen.update(json.loads(request.content))
return sse(piece({"content": "hi"}, "stop"))
make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
assert seen["stream"] is True
assert seen["stream_options"] == {"include_usage": True}
def test_a_streamed_reply_is_put_back_together() -> None:
usage = {"prompt_tokens": 7, "completion_tokens": 3, "cost": 0.25}
usage["prompt_tokens_details"] = {"cached_tokens": 4}
client = make_client(
lambda request: sse(
piece({"role": "assistant", "reasoning_content": "think"}),
piece({"reasoning_content": "ing"}),
piece({"content": "hel"}),
piece({"content": "lo"}),
piece(finish="stop"),
{"model": "qwen-real", "choices": [], "usage": usage},
)
)
[event] = client.complete(EVENTS, tools=None, model="local/qwen")
assert event == AssistantText(text="hello", model="qwen-real")
assert client.last_usage == {"input": 7, "output": 3, "cached": 4}
assert client.last_cost == 0.25
assert client.last_model == "qwen-real"
# The call record keeps the reply as one response, reasoning included.
message = client.last_exchange["response"]["choices"][0]["message"]
assert message["reasoning_content"] == "thinking"
assert message["content"] == "hello"
def test_streamed_tool_calls_are_joined_from_their_pieces() -> None:
def call(index: int, **function) -> dict:
return {"tool_calls": [{"index": index, "function": function}]}
client = make_client(
lambda request: sse(
piece(call(0, name="write_file", arguments='{"path":')),
piece(call(0, arguments=' "a.txt"}')),
piece(call(1, name="read_file", arguments="{")),
piece(call(1, arguments='"path": "b"}')),
piece(finish="tool_calls"),
)
)
first, second = client.complete(EVENTS, tools=None, model="local/qwen")
assert isinstance(first, ToolCall) and isinstance(second, ToolCall)
assert (first.tool, first.arguments) == ("write_file", {"path": "a.txt"})
assert (second.tool, second.arguments) == ("read_file", {"path": "b"})
def test_a_streamed_reply_cut_off_at_the_limit_is_truncated() -> None:
client = make_client(lambda request: sse(piece({"content": "half"}, "length")))
with pytest.raises(TruncatedResponse):
client.complete(EVENTS, tools=None, model="local/qwen")
def test_a_stream_that_stops_before_it_finishes_is_upstream_trouble() -> None:
client = make_client(lambda request: sse(piece({"content": "half"}), done=False))
with pytest.raises(UpstreamUnavailable, match="stopped before it finished"):
client.complete(EVENTS, tools=None, model="local/qwen")
def test_an_error_mid_stream_is_upstream_trouble() -> None:
broke = make_client(lambda request: sse({"error": {"message": "backend died"}}))
with pytest.raises(UpstreamUnavailable, match="backend died"):
broke.complete(EVENTS, tools=None, model="local/qwen")
spent = make_client(
lambda request: sse({"error": {"message": "Budget has been exceeded"}})
)
with pytest.raises(BudgetExhausted):
spent.complete(EVENTS, tools=None, model="local/qwen")
def test_a_model_gone_quiet_mid_reply_times_out_once() -> None:
calls: list[int] = []
def handler(request: httpx.Request) -> httpx.Response:
calls.append(1)
first = f"data: {json.dumps(piece({'content': 'hel'}))}\n\n".encode()
return httpx.Response(
200,
headers={"content-type": "text/event-stream"},
stream=Failing(first, httpx.ReadTimeout("quiet")),
)
with pytest.raises(UpstreamUnavailable, match="said nothing for 600 s"):
make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
assert len(calls) == 1 # a hung model is not asked again
def test_a_connection_lost_mid_reply_is_retried() -> None:
calls: list[int] = []
def handler(request: httpx.Request) -> httpx.Response:
calls.append(1)
if len(calls) == 1:
first = f"data: {json.dumps(piece({'content': 'hel'}))}\n\n".encode()
return httpx.Response(
200,
headers={"content-type": "text/event-stream"},
stream=Failing(first, httpx.RemoteProtocolError("peer closed")),
)
return sse(piece({"content": "hello"}, "stop"))
[event] = make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
assert event == AssistantText(text="hello", model="qwen-real") and len(calls) == 2
def test_the_passthrough_still_asks_for_the_whole_reply() -> None:
seen: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
seen.update(json.loads(request.content))
return httpx.Response(
200,
json={"model": "m", "content": [{"type": "text", "text": "hi"}]},
)
AnthropicPassthroughClient(
"https://litellm.test", "tok", transport=httpx.MockTransport(handler)
).complete(EVENTS, tools=None, model="claude")
assert "stream" not in seen
class _Server:
"""A real HTTP server that streams a reply in pieces, ``gap`` seconds
apart: the timeout is httpx's own, so only a real socket shows what it
measures."""
def __init__(self, pieces: int, gap: float) -> None:
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
class Handler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
self.rfile.read(int(self.headers["content-length"]))
self.send_response(200)
self.send_header("content-type", "text/event-stream")
self.end_headers()
import time
try:
for n in range(pieces):
time.sleep(gap)
chunk = piece({"content": str(n)})
self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
self.wfile.flush()
self.wfile.write(b"data: [DONE]\n\n")
except (BrokenPipeError, ConnectionResetError):
pass # the client gave up, as it should
def log_message(self, *args: Any) -> None:
pass
self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
self.url = f"http://127.0.0.1:{self.server.server_address[1]}"
threading.Thread(target=self.server.serve_forever, daemon=True).start()
def close(self) -> None:
self.server.shutdown()
self.server.server_close()
def test_a_slow_reply_that_keeps_coming_outlasts_the_timeout() -> None:
# Six pieces 0.1 s apart take longer than the 0.35 s timeout, but no gap
# does: a model still writing is never cut off (#174).
server = _Server(pieces=6, gap=0.1)
try:
client = LiteLLMClient(server.url, "tok", timeout_s=0.35)
[event] = client.complete(EVENTS, tools=None, model="local/qwen")
finally:
server.close()
assert isinstance(event, AssistantText) and event.text == "012345"
def test_a_reply_that_goes_quiet_times_out() -> None:
server = _Server(pieces=2, gap=1.0)
try:
client = LiteLLMClient(server.url, "tok", timeout_s=0.2)
with pytest.raises(UpstreamUnavailable, match="said nothing"):
client.complete(EVENTS, tools=None, model="local/qwen")
finally:
server.close()
+3 -2
View File
@@ -40,8 +40,9 @@ llm: # L9: LiteLLM is the only LLM destination
# Ceiling on each response. The model's own max_output_tokens in LiteLLM # Ceiling on each response. The model's own max_output_tokens in LiteLLM
# caps it further, which is where to hold a local model to less. # caps it further, which is where to hold a local model to less.
max_output_tokens: 16384 max_output_tokens: 16384
# How long one request may take before it is abandoned. A safety net for a # How long the model may go silent before its request is abandoned. Local
# hung connection; slow local models can take minutes over a long answer. # models' replies are streamed, so a slow one that keeps writing is never
# cut off; this bounds reading a long prompt, or a hung model.
request_timeout_s: 600 request_timeout_s: 600
roles: # L9: role -> LiteLLM model alias roles: # L9: role -> LiteLLM model alias