342 lines
11 KiB
Python
342 lines
11 KiB
Python
"""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, *, response_timeout=60, direct=False):
|
|
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, query=""):
|
|
return "system"
|
|
|
|
monkeypatch.setattr(conv, "build_system_prompt_with_memories", _prompt)
|
|
|
|
app = Quart(__name__)
|
|
app.config["RESPONSE_TIMEOUT"] = response_timeout
|
|
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))
|
|
kwargs = {
|
|
"json": {"query": "hi", "conversation_id": str(conversation.id)},
|
|
"headers": {"Authorization": f"Bearer {token}"},
|
|
}
|
|
if direct:
|
|
async with app.test_request_context(
|
|
"/api/conversation/stream-query", method="POST", **kwargs
|
|
):
|
|
resp = await app.full_dispatch_request()
|
|
body = (await resp.get_data()).decode()
|
|
else:
|
|
resp = await app.test_client().post("/api/conversation/stream-query", **kwargs)
|
|
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_active_stream_outlasts_response_timeout(owner, monkeypatch):
|
|
script = [
|
|
(0, _start("m1")),
|
|
(0, _chunk("m1", "<think>")),
|
|
(0.15, _chunk("m1", "working")),
|
|
(0.15, _chunk("m1", "</think>Answer")),
|
|
(0, _end("m1")),
|
|
(0, _final("Answer")),
|
|
]
|
|
|
|
frames = await _run(owner, monkeypatch, script, response_timeout=0.1)
|
|
|
|
assert "error" not in [f["type"] for f in frames]
|
|
assert frames[-1] == {"type": "response", "message": "Answer"}
|
|
|
|
|
|
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
|
|
|
|
|
|
async def test_token_counts_use_usage_not_chunk_count(owner, monkeypatch, caplog):
|
|
script = [
|
|
(0, _start("m1")),
|
|
(0, _chunk("m1", "Hello there")),
|
|
(0, _chunk("m1", " world again")),
|
|
(0, _end("m1", prompt_tokens=100)),
|
|
(0, _final("Hello there world again")),
|
|
]
|
|
with caplog.at_level(logging.INFO):
|
|
await _run(owner, monkeypatch, script)
|
|
summary = [
|
|
json.loads(r.getMessage())
|
|
for r in caplog.records
|
|
if r.getMessage().startswith('{"event": "stream_request"')
|
|
][-1]
|
|
assert summary["tokens"] == 5
|
|
assert summary["content_chunks"] == 2
|
|
assert summary["usage_complete"] is True
|
|
assert "thinking_ms" not in summary
|
|
assert summary["model_to_first_content_ms"] is not None
|
|
|
|
|
|
async def test_stalled_call_reports_partial_usage_and_outcome(
|
|
owner, monkeypatch, caplog
|
|
):
|
|
results = []
|
|
durations = []
|
|
monkeypatch.setattr(
|
|
conv, "count_result", lambda **kwargs: results.append(kwargs["result"])
|
|
)
|
|
monkeypatch.setattr(
|
|
conv, "observe_response", lambda **kwargs: durations.append(kwargs["seconds"])
|
|
)
|
|
script = [
|
|
(0, _start("m1")),
|
|
(0, _chunk("m1", "Hi")),
|
|
(0, _end("m1", prompt_tokens=100)),
|
|
(0, _start("m2")),
|
|
(LIMIT * 4, _chunk("m2", "late")),
|
|
]
|
|
with caplog.at_level(logging.INFO):
|
|
await _run(owner, monkeypatch, script)
|
|
summary = [
|
|
json.loads(r.getMessage())
|
|
for r in caplog.records
|
|
if r.getMessage().startswith('{"event": "stream_request"')
|
|
][-1]
|
|
assert summary["tokens"] is None
|
|
assert summary["completed_output_tokens"] == 5
|
|
assert summary["usage_complete"] is False
|
|
assert results == ["stalled"]
|
|
assert len(durations) == 1
|
|
|
|
|
|
async def test_no_response_is_counted_as_error(owner, monkeypatch):
|
|
results = []
|
|
monkeypatch.setattr(
|
|
conv, "count_result", lambda **kwargs: results.append(kwargs["result"])
|
|
)
|
|
frames = await _run(
|
|
owner, monkeypatch, [(0, _start("m1")), (0, _end("m1", prompt_tokens=100))]
|
|
)
|
|
assert frames[-1]["type"] == "error"
|
|
assert results == ["error"]
|
|
|
|
|
|
@pytest.mark.parametrize("cancelled", [False, True])
|
|
async def test_final_persistence_exit_has_exactly_one_outcome(
|
|
owner, monkeypatch, caplog, cancelled
|
|
):
|
|
results, durations, writes = [], [], []
|
|
original = conv.add_message_to_conversation
|
|
|
|
async def persist(**kwargs):
|
|
if kwargs["speaker"] == "simba":
|
|
writes.append(kwargs)
|
|
if cancelled:
|
|
raise asyncio.CancelledError()
|
|
raise RuntimeError("persistence failed")
|
|
return await original(**kwargs)
|
|
|
|
monkeypatch.setattr(conv, "add_message_to_conversation", persist)
|
|
monkeypatch.setattr(
|
|
conv, "count_result", lambda **kwargs: results.append(kwargs["result"])
|
|
)
|
|
monkeypatch.setattr(
|
|
conv, "observe_response", lambda **kwargs: durations.append(kwargs["seconds"])
|
|
)
|
|
script = [
|
|
(0, _start("m1")),
|
|
(0, _chunk("m1", "Answer")),
|
|
(0, _end("m1", prompt_tokens=100)),
|
|
(0, _final("Answer")),
|
|
]
|
|
with caplog.at_level(logging.INFO):
|
|
if cancelled:
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await _run(owner, monkeypatch, script, direct=True)
|
|
else:
|
|
frames = await _run(owner, monkeypatch, script, direct=True)
|
|
assert frames[-1]["type"] == "error"
|
|
assert "persistence failed" in frames[-1]["message"]
|
|
assert len(writes) == 1
|
|
assert results == ["cancelled" if cancelled else "error"]
|
|
assert len(durations) == 1
|
|
summaries = [
|
|
json.loads(r.getMessage())
|
|
for r in caplog.records
|
|
if r.getMessage().startswith('{"event": "stream_request"')
|
|
]
|
|
assert len(summaries) == 1
|
|
assert summaries[0]["cancelled"] is cancelled
|