From 4ce254193e40a80ebd95c9962dc20fdefff36c00 Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Mon, 12 Oct 2026 19:14:02 -0400 Subject: [PATCH] Treat thinking as progress in stream timeout; trim replayed history The 30s time-to-first-token guard measured from the turn's first model call to the first answer token, so thinking, tool execution and the second call's prefill all counted against it. Long reasoning on a long conversation was cancelled mid-generation (llama-server: stop: cancel task at n_gen=1287). - stream_query: the timeout now fires only when an in-flight model call goes silent for TTFT_TIMEOUT_SECONDS. Any chunk (thinking or empty) resets the clock, each model call starts a fresh one, and time spent in tools between calls doesn't count. - Trim replayed history (blueprints/conversation/context.py): tool results and user messages from earlier turns are capped at 4000 chars (attachments included) and their images replaced by a placeholder. The current turn is untouched. The window no longer starts with orphaned ToolMessages. - Log per-call prefill time, duration and prompt/output tokens (stream_usage=True) plus max_prompt_tokens in the stream_request line. Co-Authored-By: Claude Opus 5.5 --- .env.example | 2 + blueprints/conversation/__init__.py | 76 +++++++++-- blueprints/conversation/agents.py | 9 +- blueprints/conversation/context.py | 67 +++++++++ tests/unit/test_context_trim.py | 78 +++++++++++ tests/unit/test_stream_stall.py | 203 ++++++++++++++++++++++++++++ 6 files changed, 419 insertions(+), 16 deletions(-) create mode 100644 blueprints/conversation/context.py create mode 100644 tests/unit/test_context_trim.py create mode 100644 tests/unit/test_stream_stall.py diff --git a/.env.example b/.env.example index 7084c7d..2fafeb8 100644 --- a/.env.example +++ b/.env.example @@ -24,6 +24,8 @@ BASE_URL=192.168.1.5:8000 # If set, uses llama-server as the primary LLM backend with OpenAI as fallback LLAMA_SERVER_URL=http://192.168.1.213:8080/v1 LLAMA_MODEL_NAME=llama-3.1-8b-instruct +# Abandon a turn when an in-flight model call produces no output (prefill included) +# for this long. Thinking tokens count as output, so long reasoning is fine. TTFT_TIMEOUT_SECONDS=30 # Embedding Server Configuration diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index 33434c7..e2c2752 100644 --- a/blueprints/conversation/__init__.py +++ b/blueprints/conversation/__init__.py @@ -48,7 +48,35 @@ from .models import ( from .tool_events import tool_end_payload, tool_name_from_message, tool_start_payload from .tool_labels import tool_label -TTFT_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30")) +# Longest an in-flight model call may go without producing any output (prompt +# prefill included) before the turn is abandoned as stalled. Thinking tokens and +# empty chunks count as output; time spent running tools does not count at all. +MODEL_STALL_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30")) + + +def stall_timeout(last_model_activity: float | None, now: float) -> float | None: + """Seconds left before the in-flight model call counts as stalled. + + None when no model call is in flight (e.g. while tools run): the httpx read + timeout and the tools' own timeouts cover those phases. + """ + if last_model_activity is None: + return None + return max(MODEL_STALL_TIMEOUT - (now - last_model_activity), 1.0) + + +def _model_call_summary(call: dict) -> dict: + """Per-LLM-call timing and token counts for the stream_request log line.""" + start = call["start"] + first_chunk = call.get("first_chunk") + end = call.get("end") + return { + "prefill_ms": round((first_chunk - start) * 1000) if first_chunk else None, + "ms": round((end - start) * 1000) if end else None, + "prompt_tokens": call.get("prompt_tokens"), + "output_tokens": call.get("output_tokens"), + } + _THINK_OPEN = "" _THINK_CLOSE = "" @@ -381,6 +409,9 @@ async def stream_query(): first_content_at: float | None = None tool_timers: dict[str, float] = {} # run_id -> start time tool_spans: list[dict] = [] + # Last output from the in-flight model call; None between calls. + last_model_activity: float | None = None + model_calls: dict[str, dict] = {} # run_id -> timing / token usage n_tokens = 0 model_name: str | None = None cancelled = False @@ -403,6 +434,8 @@ async def stream_query(): ttft_ms = ( round((first_content_at - t0) * 1000) if first_content_at else None ) + calls = [_model_call_summary(c) for c in model_calls.values()] + prompt_tokens = [c["prompt_tokens"] for c in calls if c["prompt_tokens"]] logging.info( json.dumps( { @@ -414,6 +447,8 @@ async def stream_query(): "ttft_ms": ttft_ms, "tokens": n_tokens, "tools": tool_spans, + "model_calls": calls, + "max_prompt_tokens": max(prompt_tokens, default=None), "model": model_name, "cancelled": cancelled, **({"error": error} if error else {}), @@ -426,15 +461,10 @@ async def stream_query(): payload, version="v2", config=agent_config ).__aiter__() while True: - # Apply a short TTFT timeout only during the window between - # on_chat_model_start and the first content token. Before the - # LLM call starts (tool phase) and after the first token, let - # the httpx read timeout handle stalls instead. - if llm_start_at is not None and first_content_at is None: - ttft_wait = TTFT_TIMEOUT - (time.monotonic() - llm_start_at) - timeout = max(ttft_wait, 1.0) - else: - timeout = None + # Only a silent model call counts as a stall. A model that is + # thinking (or streaming empty chunks) keeps resetting the clock, + # and each new call starts a fresh one. + timeout = stall_timeout(last_model_activity, time.monotonic()) try: if timeout is not None: @@ -446,7 +476,9 @@ async def stream_query(): except StopAsyncIteration: break except asyncio.TimeoutError: - _log_summary(error=f"TTFT timeout after {TTFT_TIMEOUT:.0f}s") + _log_summary( + error=f"model stalled: no output for {MODEL_STALL_TIMEOUT:.0f}s" + ) yield f"data: {json.dumps({'type': 'error', 'message': 'Simba took too long to respond. Please try again.', 'stalled': True})}\n\n" return @@ -498,14 +530,34 @@ async def stream_query(): elif event_type == "on_chat_model_start": turn_has_content = False + now = time.monotonic() + last_model_activity = now + model_calls[event["run_id"]] = {"start": now} if llm_start_at is None: - llm_start_at = time.monotonic() + llm_start_at = now if model_name is None: model_name = event.get("metadata", {}).get( "ls_model_name" ) or event.get("name") + elif event_type == "on_chat_model_end": + last_model_activity = None + call = model_calls.get(event["run_id"]) + if call is not None: + call["end"] = time.monotonic() + usage = getattr( + event.get("data", {}).get("output"), "usage_metadata", None + ) + if usage: + call["prompt_tokens"] = usage.get("input_tokens") + call["output_tokens"] = usage.get("output_tokens") + elif event_type == "on_chat_model_stream": + now = time.monotonic() + last_model_activity = now + call = model_calls.get(event.get("run_id")) + if call is not None and "first_chunk" not in call: + call["first_chunk"] = now chunk = event.get("data", {}).get("chunk") raw = getattr(chunk, "content", None) if isinstance(raw, str) and raw: diff --git a/blueprints/conversation/agents.py b/blueprints/conversation/agents.py index 16a00b3..c9079ac 100644 --- a/blueprints/conversation/agents.py +++ b/blueprints/conversation/agents.py @@ -24,6 +24,7 @@ from blueprints.conversation.artifacts_util import ( truncate, unique, ) +from blueprints.conversation.context import trim_history from blueprints.conversation.memory import save_memory from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT from blueprints.rag.logic import query_vector_store @@ -54,6 +55,8 @@ if llama_url: timeout=httpx.Timeout(LLAMA_READ_TIMEOUT, connect=10.0), max_retries=1, model_kwargs={"reasoning_effort": "low"}, + # Ask llama-server for token usage on the final chunk (prompt size logging). + stream_usage=True, ) else: llama_chat = None @@ -1106,10 +1109,8 @@ def _build_prompt(state, config: RunnableConfig = None) -> list: # LangGraph passes the full state dict {"messages": [...]}, not just the list messages = state["messages"] if isinstance(state, dict) else state non_system = [m for m in messages if not isinstance(m, SystemMessage)] - # Keep only the most recent messages to stay within the 32k context window. - # Tool call/result pairs each count as 2 messages, so 12 ≈ 4-6 conversational turns. - non_system = non_system[-12:] - return [SystemMessage(content=system_prompt)] + non_system + # Window + trim earlier turns to stay within the 32k context window. + return [SystemMessage(content=system_prompt)] + trim_history(non_system) def build_agent(checkpointer=None, extra_tools=()): diff --git a/blueprints/conversation/context.py b/blueprints/conversation/context.py new file mode 100644 index 0000000..1e0afb5 --- /dev/null +++ b/blueprints/conversation/context.py @@ -0,0 +1,67 @@ +"""Shape the message history replayed to the LLM so it fits the context window. + +The LangGraph checkpoint keeps every message at full size: whole search hits, +full notes, attached files, images. Replaying all of that each turn pushes the +prompt toward the 32k-token llama-server slot, which makes every prefill slow. +The current turn keeps everything; earlier turns are trimmed. +""" + +from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage + +# Matches TOOL_RESULT_MAX_CHARS in logic.py (what gets persisted to Postgres). +EARLIER_MESSAGE_MAX_CHARS = 4000 +TRUNCATION_NOTE = "\n[truncated: earlier context]" +IMAGE_PLACEHOLDER = "[image from an earlier message]" +# 12 ≈ 4-6 conversational turns (tool call/result pairs count as 2 messages). +HISTORY_WINDOW = 12 + + +def _cap(text: str) -> str: + if len(text) <= EARLIER_MESSAGE_MAX_CHARS: + return text + return text[:EARLIER_MESSAGE_MAX_CHARS] + TRUNCATION_NOTE + + +def _shrink(message: BaseMessage) -> BaseMessage: + """Cap text and drop images from a message that belongs to an earlier turn.""" + content = message.content + if isinstance(content, str): + new_content = _cap(content) + else: + new_content = [] + for part in content: + if isinstance(part, dict) and part.get("type") in ("image_url", "image"): + new_content.append({"type": "text", "text": IMAGE_PLACEHOLDER}) + elif isinstance(part, dict) and isinstance(part.get("text"), str): + new_content.append({**part, "text": _cap(part["text"])}) + elif isinstance(part, str): + new_content.append(_cap(part)) + else: + new_content.append(part) + if new_content == content: + return message + return message.model_copy(update={"content": new_content}) + + +def trim_history(messages: list[BaseMessage]) -> list[BaseMessage]: + """Window the history and trim tool results/user messages from earlier turns. + + The current turn starts at the last HumanMessage; it and everything after it + (this turn's tool calls and results) are left untouched. + """ + current_turn = max( + (i for i, m in enumerate(messages) if isinstance(m, HumanMessage)), + default=len(messages), + ) + trimmed = [ + _shrink(m) + if i < current_turn and isinstance(m, (ToolMessage, HumanMessage)) + else m + for i, m in enumerate(messages) + ] + window = trimmed[-HISTORY_WINDOW:] + # A window that starts mid tool-exchange would begin with ToolMessages whose + # tool_calls were cut off; OpenAI-style APIs reject those. + while window and isinstance(window[0], ToolMessage): + window = window[1:] + return window diff --git a/tests/unit/test_context_trim.py b/tests/unit/test_context_trim.py new file mode 100644 index 0000000..88fb7f8 --- /dev/null +++ b/tests/unit/test_context_trim.py @@ -0,0 +1,78 @@ +"""Tests for trimming the replayed history (blueprints/conversation/context.py).""" + +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from blueprints.conversation.context import ( + EARLIER_MESSAGE_MAX_CHARS, + HISTORY_WINDOW, + IMAGE_PLACEHOLDER, + TRUNCATION_NOTE, + trim_history, +) + +BIG = "x" * (EARLIER_MESSAGE_MAX_CHARS * 3) + + +def _call(call_id: str, name: str = "simba_search") -> AIMessage: + return AIMessage(content="", tool_calls=[{"id": call_id, "name": name, "args": {}}]) + + +def _result(call_id: str, text: str = BIG) -> ToolMessage: + return ToolMessage(content=text, tool_call_id=call_id, name="simba_search") + + +def test_earlier_tool_results_are_capped_current_turn_untouched(): + messages = [ + HumanMessage(content="first question"), + _call("c1"), + _result("c1"), + AIMessage(content="first answer"), + HumanMessage(content="second question"), + _call("c2"), + _result("c2"), + ] + + out = trim_history(messages) + + earlier, current = out[2], out[6] + assert earlier.content == BIG[:EARLIER_MESSAGE_MAX_CHARS] + TRUNCATION_NOTE + assert current.content == BIG + # The checkpointed originals are never mutated. + assert messages[2].content == BIG + + +def test_earlier_attachments_and_images_are_shrunk(): + image = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}} + messages = [ + HumanMessage(content=[{"type": "text", "text": BIG}, image]), + AIMessage(content="nice photo"), + HumanMessage(content=[{"type": "text", "text": "now this one"}, image]), + ] + + out = trim_history(messages) + + earlier = out[0].content + assert earlier[0]["text"].endswith(TRUNCATION_NOTE) + assert earlier[1] == {"type": "text", "text": IMAGE_PLACEHOLDER} + # The current turn keeps its image. + assert out[2].content[1] == image + + +def test_window_never_starts_with_orphaned_tool_results(): + # Many parallel tool results so the window cut lands between an AIMessage's + # tool_calls and their ToolMessages. + messages = [HumanMessage(content="q"), _call("c0")] + messages += [_result("c0", "r") for _ in range(HISTORY_WINDOW)] + messages += [AIMessage(content="answer"), HumanMessage(content="next")] + + out = trim_history(messages) + + assert len(out) <= HISTORY_WINDOW + assert not isinstance(out[0], ToolMessage) + + +def test_short_history_is_returned_as_is(): + messages = [HumanMessage(content="hi"), AIMessage(content="hello")] + out = trim_history(messages) + assert out == messages + assert out[0] is messages[0] diff --git a/tests/unit/test_stream_stall.py b/tests/unit/test_stream_stall.py new file mode 100644 index 0000000..629cf7f --- /dev/null +++ b/tests/unit/test_stream_stall.py @@ -0,0 +1,203 @@ +"""Tests for the stream-query stall timeout and per-call model metrics.""" + +import asyncio +import json +import logging + +import pytest +from langchain_core.messages import AIMessage, AIMessageChunk +from quart import Quart +from quart_jwt_extended import JWTManager, create_access_token +from tortoise import Tortoise + +import blueprints.conversation as conv +from blueprints.conversation.models import Conversation +from blueprints.users.models import User + +# Shortened stall limit for the tests (seconds). +LIMIT = 0.3 + + +@pytest.fixture +async def owner(): + await Tortoise.init( + db_url="sqlite://:memory:", + modules={ + "models": [ + "blueprints.users.models", + "blueprints.conversation.models", + "blueprints.artifacts.models", + ] + }, + ) + await Tortoise.generate_schemas() + user = await User.create(username="owner", email="owner@example.com") + conversation = await Conversation.create(name="c", user=user) + yield user, conversation + await Tortoise.close_connections() + + +class _TimedAgent: + """Yields (delay_seconds, event) pairs, sleeping before each event.""" + + def __init__(self, script): + self._script = script + + async def astream_events(self, payload, version, config): + for delay, event in self._script: + await asyncio.sleep(delay) + yield event + + +def _start(run_id): + return {"event": "on_chat_model_start", "run_id": run_id, "name": "llm", "data": {}} + + +def _chunk(run_id, text): + return { + "event": "on_chat_model_stream", + "run_id": run_id, + "data": {"chunk": AIMessageChunk(content=text)}, + } + + +def _end(run_id, prompt_tokens=None): + usage = ( + { + "input_tokens": prompt_tokens, + "output_tokens": 5, + "total_tokens": prompt_tokens + 5, + } + if prompt_tokens + else None + ) + return { + "event": "on_chat_model_end", + "run_id": run_id, + "data": {"output": AIMessage(content="", usage_metadata=usage)}, + } + + +def _final(text): + return { + "event": "on_chain_end", + "name": "agent", + "data": {"output": {"messages": [AIMessage(content=text)]}}, + } + + +async def _run(owner, monkeypatch, script): + user, conversation = owner + monkeypatch.setattr(conv, "MODEL_STALL_TIMEOUT", LIMIT) + monkeypatch.setattr(conv, "get_agent_for", lambda u: _TimedAgent(script)) + + async def _prompt(user_id): + return "system" + + monkeypatch.setattr(conv, "build_system_prompt_with_memories", _prompt) + + app = Quart(__name__) + app.config["JWT_SECRET_KEY"] = "test-secret" + JWTManager(app) + app.register_blueprint(conv.conversation_blueprint) + async with app.app_context(): + token = create_access_token(identity=str(user.id)) + resp = await app.test_client().post( + "/api/conversation/stream-query", + json={"query": "hi", "conversation_id": str(conversation.id)}, + headers={"Authorization": f"Bearer {token}"}, + ) + body = (await resp.get_data()).decode() + return [ + json.loads(line[len("data: ") :]) + for line in body.splitlines() + if line.startswith("data: {") + ] + + +def test_stall_timeout_only_while_a_model_call_is_in_flight(monkeypatch): + monkeypatch.setattr(conv, "MODEL_STALL_TIMEOUT", 30.0) + assert conv.stall_timeout(None, 100.0) is None + assert conv.stall_timeout(100.0, 110.0) == 20.0 + assert conv.stall_timeout(100.0, 200.0) == 1.0 # floor + + +async def test_long_thinking_is_not_a_stall(owner, monkeypatch): + # The whole call takes ~4x the limit, but output never pauses for longer + # than half of it; the old time-to-first-content deadline killed this. + script = [(0, _start("m1"))] + script += [(LIMIT / 2, _chunk("m1", "hmm ")) for _ in range(8)] + script += [ + (0, _chunk("m1", "Answer")), + (0, _end("m1")), + (0, _final("Answer")), + ] + + frames = await _run(owner, monkeypatch, script) + + types = [f["type"] for f in frames] + assert "error" not in types + assert frames[-1] == {"type": "response", "message": "Answer"} + + +async def test_tool_time_between_model_calls_does_not_count(owner, monkeypatch): + tool = {"event": "on_tool_start", "name": "web_search", "run_id": "t1", "data": {}} + tool_end = { + "event": "on_tool_end", + "name": "web_search", + "run_id": "t1", + "data": {}, + } + script = [ + (0, _start("m1")), + (0, _chunk("m1", "")), + (0, _end("m1")), + (0, tool), + (LIMIT * 3, tool_end), # slow tool, no model call in flight + (0, _start("m2")), + (LIMIT / 2, _chunk("m2", "Done")), + (0, _end("m2")), + (0, _final("Done")), + ] + + frames = await _run(owner, monkeypatch, script) + + assert "error" not in [f["type"] for f in frames] + assert frames[-1]["message"] == "Done" + + +async def test_silent_model_call_is_a_stall(owner, monkeypatch): + script = [ + (0, _start("m1")), + (0, _chunk("m1", "start")), + (LIMIT * 4, _chunk("m1", "never arrives")), + ] + + frames = await _run(owner, monkeypatch, script) + + assert frames[-1]["type"] == "error" + assert frames[-1]["stalled"] is True + + +async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog): + script = [ + (0, _start("m1")), + (0, _chunk("m1", "Hi")), + (0, _end("m1", prompt_tokens=27262)), + (0, _final("Hi")), + ] + + with caplog.at_level(logging.INFO): + await _run(owner, monkeypatch, script) + + summaries = [ + json.loads(r.getMessage()) + for r in caplog.records + if r.getMessage().startswith('{"event": "stream_request"') + ] + assert summaries, "stream_request summary not logged" + summary = summaries[-1] + assert summary["max_prompt_tokens"] == 27262 + (call,) = summary["model_calls"] + assert call["prompt_tokens"] == 27262 + assert call["prefill_ms"] is not None and call["ms"] is not None -- 2.52.0