.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

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>"
 _THINK_CLOSE = "</think>"
@@ -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:

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

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

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]

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", "<think>hmm ")) for _ in range(8)]
+    script += [
+        (0, _chunk("m1", "</think>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", "<think>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