202 lines
6.8 KiB
Python
202 lines
6.8 KiB
Python
"""Stream truncation event persistence + model-layer classification tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
import deerflow.persistence.stream_truncation as stream_truncation_module
|
|
from deerflow.models.patched_openai import _ToolCallStreamDiagnostics
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.stream_truncation import (
|
|
SqlStreamTruncationStore,
|
|
StreamTruncationEvent,
|
|
StreamTruncationQuery,
|
|
)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def session_factory(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'stream-truncation.db'}")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
yield factory
|
|
await engine.dispose()
|
|
|
|
|
|
def _event(**overrides) -> StreamTruncationEvent:
|
|
payload = dict(
|
|
id="evt-1",
|
|
created_at=datetime.now(UTC),
|
|
model_name="relay-model",
|
|
run_id="run-1",
|
|
event_type="silent_no_terminal",
|
|
finish_reason="none",
|
|
status="success",
|
|
chunk_count=12,
|
|
duration_ms=3_000,
|
|
first_chunk_ms=120,
|
|
)
|
|
payload.update(overrides)
|
|
return StreamTruncationEvent(**payload)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_record_list_count_and_breakdown(session_factory) -> None:
|
|
store = SqlStreamTruncationStore(session_factory)
|
|
await store.record(_event())
|
|
await store.record(
|
|
_event(
|
|
id="evt-2",
|
|
model_name="relay-model",
|
|
event_type="length_truncated",
|
|
finish_reason="length",
|
|
)
|
|
)
|
|
await store.record(_event(id="evt-3", model_name="other-model"))
|
|
|
|
items = await store.list(StreamTruncationQuery(limit=10))
|
|
assert len(items) == 3
|
|
total = await store.count(StreamTruncationQuery(model_name="relay-model"))
|
|
assert total == 2
|
|
|
|
filtered = await store.list(
|
|
StreamTruncationQuery(model_name="relay-model", event_type="length_truncated")
|
|
)
|
|
assert [item["id"] for item in filtered] == ["evt-2"]
|
|
|
|
breakdown = await store.model_breakdown(StreamTruncationQuery())
|
|
by_model = {item["model_name"]: item for item in breakdown}
|
|
assert by_model["relay-model"]["total"] == 2
|
|
assert by_model["relay-model"]["length_truncated"] == 1
|
|
assert by_model["relay-model"]["silent_no_terminal"] == 1
|
|
assert by_model["other-model"]["total"] == 1
|
|
assert by_model["relay-model"]["last_at"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retention_cap_deletes_rows_older_than_one_month(session_factory) -> None:
|
|
store = SqlStreamTruncationStore(session_factory)
|
|
old = _event(
|
|
id="evt-old",
|
|
created_at=datetime.now(UTC) - timedelta(days=40),
|
|
)
|
|
await store.record(old)
|
|
# The sweep after every insert must already have removed the 40-day-old row.
|
|
assert await store.count(StreamTruncationQuery()) == 0
|
|
|
|
fresh = _event(id="evt-fresh", created_at=datetime.now(UTC) - timedelta(days=3))
|
|
await store.record(fresh)
|
|
assert await store.count(StreamTruncationQuery()) == 1
|
|
|
|
purged = await store.purge_before(datetime.now(UTC) - timedelta(days=30))
|
|
assert purged == 0
|
|
assert await store.count(StreamTruncationQuery()) == 1
|
|
|
|
|
|
def _diagnostics(**overrides) -> _ToolCallStreamDiagnostics:
|
|
diagnostics = _ToolCallStreamDiagnostics(
|
|
model_name="relay-model",
|
|
run_id="run-abc",
|
|
strict_stream_integrity=False,
|
|
)
|
|
for key, value in overrides.items():
|
|
setattr(diagnostics, key, value)
|
|
return diagnostics
|
|
|
|
|
|
def _capture_recorded(monkeypatch: pytest.MonkeyPatch):
|
|
captured: list[dict] = []
|
|
|
|
def fake_record(**kwargs):
|
|
captured.append(kwargs)
|
|
|
|
monkeypatch.setattr(
|
|
stream_truncation_module,
|
|
"record_stream_truncation_event",
|
|
fake_record,
|
|
)
|
|
return captured
|
|
|
|
|
|
def _flush(diagnostics: _ToolCallStreamDiagnostics, *, status: str = "success", error=None, error_type: str = "none") -> None:
|
|
diagnostics._record_truncation_event(
|
|
status=status,
|
|
error=error,
|
|
error_type=error_type,
|
|
total_ms=1_500,
|
|
first_chunk_ms=100,
|
|
)
|
|
|
|
|
|
def test_classification_silent_no_terminal(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured = _capture_recorded(monkeypatch)
|
|
diagnostics = _diagnostics(chunk_count=8)
|
|
_flush(diagnostics)
|
|
assert len(captured) == 1
|
|
assert captured[0]["event_type"] == "silent_no_terminal"
|
|
assert captured[0]["finish_reason"] == "none"
|
|
assert captured[0]["model_name"] == "relay-model"
|
|
assert captured[0]["strict_mode"] is False
|
|
assert captured[0]["run_id"] == "run-abc"
|
|
|
|
|
|
def test_classification_skips_clean_and_empty_streams(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured = _capture_recorded(monkeypatch)
|
|
# Normal stop with terminal frame -> nothing to record.
|
|
_flush(_diagnostics(chunk_count=8, finish_reason="stop"))
|
|
# No terminal frame but zero chunks -> empty response, not a half answer.
|
|
_flush(_diagnostics(chunk_count=0))
|
|
# Cancelled by the user -> not a truncation.
|
|
_flush(_diagnostics(chunk_count=8), status="cancelled", error=asyncio.CancelledError())
|
|
assert captured == []
|
|
|
|
|
|
def test_classification_intentional_bridge_stop_is_skipped(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured = _capture_recorded(monkeypatch)
|
|
_flush(_diagnostics(chunk_count=8, intentional_early_finish=True))
|
|
assert captured == []
|
|
|
|
|
|
def test_classification_strict_incomplete_and_midstream_errors(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured = _capture_recorded(monkeypatch)
|
|
_flush(
|
|
_diagnostics(chunk_count=8, strict_stream_integrity=True),
|
|
status="error",
|
|
error=RuntimeError("incomplete"),
|
|
error_type="IncompleteChatStreamError",
|
|
)
|
|
_flush(
|
|
_diagnostics(chunk_count=30),
|
|
status="error",
|
|
error=TimeoutError(),
|
|
error_type="ReadTimeout",
|
|
)
|
|
# Errors before the first chunk are connection failures, not truncations.
|
|
_flush(
|
|
_diagnostics(chunk_count=0),
|
|
status="error",
|
|
error=TimeoutError(),
|
|
error_type="ReadTimeout",
|
|
)
|
|
assert [item["event_type"] for item in captured] == [
|
|
"incomplete_stream_error",
|
|
"midstream_error",
|
|
]
|
|
assert captured[0]["strict_mode"] is True
|
|
|
|
|
|
def test_classification_length_and_content_filter(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured = _capture_recorded(monkeypatch)
|
|
_flush(_diagnostics(chunk_count=99, finish_reason="length"))
|
|
_flush(_diagnostics(chunk_count=99, finish_reason="content_filter"))
|
|
assert [item["event_type"] for item in captured] == [
|
|
"length_truncated",
|
|
"content_filter",
|
|
]
|