732 lines
25 KiB
Python
732 lines
25 KiB
Python
"""Tests for the LLM call metrics middleware + repository.
|
|
|
|
Covers:
|
|
- Successful calls emit a row with duration, token counts and tokens_per_sec.
|
|
- Failed calls record the exception class as ``error_type`` plus a trimmed
|
|
``error_message``.
|
|
- The middleware never blocks on a backend exception (best-effort writes).
|
|
- The SQL store applies filters and pagination for the admin list endpoint.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from langchain_core.messages import AIMessage
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway.admin_audit_retention import sweep_admin_audit_retention
|
|
from deerflow.agents.middlewares.llm_metrics_middleware import (
|
|
LlmMetricsMiddleware,
|
|
_extract_usage,
|
|
)
|
|
from deerflow.persistence.admin_stats.model import AdminLeaderboardDailyStatRow
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.llm_metrics import SqlLlmMetricsStore
|
|
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery
|
|
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
|
|
from deerflow.persistence.run import RunRepository
|
|
from deerflow.persistence.run.model import RunRow
|
|
from deerflow.persistence.tool_metrics import SqlToolMetricsStore
|
|
from deerflow.persistence.tool_metrics.base import ToolMetricRecord, ToolMetricsQuery
|
|
from deerflow.persistence.tool_metrics.model import ToolCallMetricRow
|
|
from deerflow.persistence.types import BEIJING_TZ
|
|
|
|
|
|
class _RecordingStore:
|
|
"""Minimal in-memory store used to assert what the middleware persists."""
|
|
|
|
def __init__(self) -> None:
|
|
self.records: list[LlmMetricRecord] = []
|
|
|
|
async def record(self, metric: LlmMetricRecord) -> None:
|
|
self.records.append(metric)
|
|
|
|
async def list(self, query: LlmMetricsQuery):
|
|
return [r.__dict__ for r in self.records]
|
|
|
|
async def count(self, query: LlmMetricsQuery) -> int:
|
|
return len(self.records)
|
|
|
|
async def iter_all(self, query: LlmMetricsQuery):
|
|
for record in self.records:
|
|
yield record.__dict__
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def sql_store(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'metrics.db'}")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(
|
|
lambda sync_conn: Base.metadata.create_all(
|
|
sync_conn,
|
|
tables=[LlmCallMetricRow.__table__],
|
|
)
|
|
)
|
|
sf = async_sessionmaker(engine, expire_on_commit=False)
|
|
store = SqlLlmMetricsStore(sf)
|
|
try:
|
|
yield store
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def tool_sql_store(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'tool-metrics.db'}")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(
|
|
lambda sync_conn: Base.metadata.create_all(
|
|
sync_conn,
|
|
tables=[ToolCallMetricRow.__table__],
|
|
)
|
|
)
|
|
store = SqlToolMetricsStore(async_sessionmaker(engine, expire_on_commit=False))
|
|
try:
|
|
yield store
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def run_sql_store(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(
|
|
lambda sync_conn: Base.metadata.create_all(
|
|
sync_conn,
|
|
tables=[RunRow.__table__],
|
|
)
|
|
)
|
|
store = RunRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
try:
|
|
yield store
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def snapshot_session_factory(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'snapshots.db'}")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(
|
|
lambda sync_conn: Base.metadata.create_all(
|
|
sync_conn,
|
|
tables=[AdminLeaderboardDailyStatRow.__table__],
|
|
)
|
|
)
|
|
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
try:
|
|
yield session_factory
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
def _make_request():
|
|
return type(
|
|
"FakeRequest",
|
|
(),
|
|
{
|
|
"model": type("FakeModel", (), {"model_name": "gpt-test"}),
|
|
"runtime": type(
|
|
"FakeRuntime",
|
|
(),
|
|
{
|
|
"config": {"configurable": {"thread_id": "t1"}},
|
|
"context": {"agent_name": "lead_agent", "run_id": "r1"},
|
|
},
|
|
),
|
|
"state": {"thread_id": "t1"},
|
|
},
|
|
)()
|
|
|
|
|
|
def test_middleware_records_success_with_speed():
|
|
import time as _time
|
|
|
|
store = _RecordingStore()
|
|
mw = LlmMetricsMiddleware(store=store)
|
|
response = AIMessage(
|
|
content="hi",
|
|
usage_metadata={"input_tokens": 5, "output_tokens": 20, "total_tokens": 25},
|
|
)
|
|
|
|
def slow_handler(_req):
|
|
# Force a measurable elapsed window so tokens_per_sec is populated.
|
|
# 50ms is comfortably above Windows monotonic-clock granularity.
|
|
_time.sleep(0.05)
|
|
return response
|
|
|
|
result = mw.wrap_model_call(_make_request(), slow_handler)
|
|
|
|
assert result is response
|
|
# awrap_model_call uses a fire-and-forget task; here the sync path runs
|
|
# asyncio.run synchronously, so the record landed already.
|
|
assert len(store.records) == 1
|
|
record = store.records[0]
|
|
assert record.status == "success"
|
|
assert record.model_name == "gpt-test"
|
|
assert record.thread_id == "t1"
|
|
assert record.run_id == "r1"
|
|
assert record.input_tokens == 5
|
|
assert record.output_tokens == 20
|
|
assert record.total_tokens == 25
|
|
# tokens_per_sec may be very high in tests because the handler is
|
|
# ~instant; we only assert that the field is populated.
|
|
assert record.tokens_per_sec is not None
|
|
assert record.duration_ms >= 0
|
|
|
|
|
|
def test_middleware_records_error_with_classification():
|
|
store = _RecordingStore()
|
|
mw = LlmMetricsMiddleware(store=store)
|
|
|
|
class BoomError(RuntimeError):
|
|
pass
|
|
|
|
def explode(_req):
|
|
raise BoomError("provider exploded")
|
|
|
|
with pytest.raises(BoomError):
|
|
mw.wrap_model_call(_make_request(), explode)
|
|
|
|
assert len(store.records) == 1
|
|
record = store.records[0]
|
|
assert record.status == "error"
|
|
assert record.error_type == "BoomError"
|
|
assert record.error_message == "provider exploded"
|
|
assert record.tokens_per_sec is None
|
|
|
|
|
|
def test_middleware_swallows_store_errors():
|
|
"""A persistence outage must not break the agent runtime."""
|
|
|
|
class FaultyStore:
|
|
async def record(self, metric):
|
|
raise RuntimeError("disk gone")
|
|
|
|
async def list(self, q):
|
|
return []
|
|
|
|
async def count(self, q):
|
|
return 0
|
|
|
|
async def iter_all(self, q):
|
|
if False:
|
|
yield {}
|
|
|
|
mw = LlmMetricsMiddleware(store=FaultyStore())
|
|
# The wrap returns normally even though _flush would raise; the call below
|
|
# must not raise.
|
|
response = AIMessage(
|
|
content="ok",
|
|
usage_metadata={"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
|
)
|
|
out = mw.wrap_model_call(_make_request(), lambda req: response)
|
|
assert out is response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_store_round_trip(sql_store):
|
|
base = datetime.now(UTC)
|
|
for i in range(3):
|
|
await sql_store.record(
|
|
LlmMetricRecord(
|
|
id=uuid.uuid4().hex,
|
|
created_at=base - timedelta(minutes=i),
|
|
user_id="alice",
|
|
model_name="gpt-test" if i % 2 == 0 else "gpt-other",
|
|
duration_ms=100 * (i + 1),
|
|
output_tokens=10 * (i + 1),
|
|
tokens_per_sec=10.0 * (i + 1),
|
|
status="success",
|
|
)
|
|
)
|
|
rows = await sql_store.list(LlmMetricsQuery(limit=10))
|
|
assert len(rows) == 3
|
|
# Newest first.
|
|
assert rows[0]["duration_ms"] == 100
|
|
|
|
filtered = await sql_store.list(LlmMetricsQuery(model_name="gpt-test", limit=10))
|
|
assert all(r["model_name"] == "gpt-test" for r in filtered)
|
|
assert len(filtered) == 2
|
|
|
|
assert await sql_store.count(LlmMetricsQuery(model_name="gpt-test")) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_store_filters_errors_only(sql_store):
|
|
base = datetime.now(UTC)
|
|
await sql_store.record(
|
|
LlmMetricRecord(id=uuid.uuid4().hex, created_at=base, status="success", duration_ms=10, output_tokens=1)
|
|
)
|
|
await sql_store.record(
|
|
LlmMetricRecord(
|
|
id=uuid.uuid4().hex,
|
|
created_at=base,
|
|
status="error",
|
|
duration_ms=5,
|
|
error_type="APITimeoutError",
|
|
error_message="timed out",
|
|
)
|
|
)
|
|
only_errors = await sql_store.list(LlmMetricsQuery(status="error", limit=10))
|
|
assert len(only_errors) == 1
|
|
assert only_errors[0]["error_type"] == "APITimeoutError"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_store_filters_by_user_ids(sql_store):
|
|
"""``user_ids`` (set by a fuzzy username search) matches several accounts."""
|
|
base = datetime.now(UTC)
|
|
for uid in ("alice", "alicia", "bob"):
|
|
await sql_store.record(
|
|
LlmMetricRecord(
|
|
id=uuid.uuid4().hex,
|
|
created_at=base,
|
|
user_id=uid,
|
|
duration_ms=1,
|
|
status="success",
|
|
)
|
|
)
|
|
# Fuzzy "ali" → {alice, alicia}; bob must be excluded.
|
|
rows = await sql_store.list(LlmMetricsQuery(user_ids=["alice", "alicia"], limit=10))
|
|
assert {r["user_id"] for r in rows} == {"alice", "alicia"}
|
|
assert await sql_store.count(LlmMetricsQuery(user_ids=["alice", "alicia"])) == 2
|
|
# user_ids takes precedence over a single user_id.
|
|
rows = await sql_store.list(LlmMetricsQuery(user_id="bob", user_ids=["alice"], limit=10))
|
|
assert {r["user_id"] for r in rows} == {"alice"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_store_iter_all_paginates(sql_store):
|
|
base = datetime.now(UTC)
|
|
for i in range(1200):
|
|
await sql_store.record(
|
|
LlmMetricRecord(
|
|
id=uuid.uuid4().hex,
|
|
created_at=base - timedelta(seconds=i),
|
|
duration_ms=1,
|
|
status="success",
|
|
)
|
|
)
|
|
seen = 0
|
|
async for _ in sql_store.iter_all(LlmMetricsQuery()):
|
|
seen += 1
|
|
assert seen == 1200
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_metric_purge_before_keeps_rows_at_or_after_cutoff(sql_store):
|
|
"""The retention boundary is exclusive: a record exactly at cutoff stays."""
|
|
|
|
cutoff = datetime.now(UTC) - timedelta(days=90)
|
|
for record_id, created_at in (
|
|
("old", cutoff - timedelta(days=1)),
|
|
("boundary", cutoff),
|
|
("fresh", cutoff + timedelta(days=1)),
|
|
):
|
|
await sql_store.record(
|
|
LlmMetricRecord(id=record_id, created_at=created_at, duration_ms=1)
|
|
)
|
|
|
|
assert await sql_store.purge_before(cutoff) == 1
|
|
# A subsequent cleanup pass is idempotent.
|
|
assert await sql_store.purge_before(cutoff) == 0
|
|
|
|
rows = await sql_store.list(LlmMetricsQuery(limit=10))
|
|
assert {row["id"] for row in rows} == {"boundary", "fresh"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_metric_purge_before_keeps_rows_at_or_after_cutoff(tool_sql_store):
|
|
"""Tool-call audit retention has the same strict-before-cutoff behavior."""
|
|
|
|
cutoff = datetime.now(UTC) - timedelta(days=90)
|
|
for record_id, created_at in (
|
|
("old", cutoff - timedelta(days=1)),
|
|
("boundary", cutoff),
|
|
("fresh", cutoff + timedelta(days=1)),
|
|
):
|
|
await tool_sql_store.record(
|
|
ToolMetricRecord(
|
|
id=record_id,
|
|
created_at=created_at,
|
|
tool_name="web_search",
|
|
duration_ms=1,
|
|
)
|
|
)
|
|
|
|
assert await tool_sql_store.purge_before(cutoff) == 1
|
|
assert await tool_sql_store.purge_before(cutoff) == 0
|
|
|
|
rows = await tool_sql_store.list(ToolMetricsQuery(limit=10))
|
|
assert {row["id"] for row in rows} == {"boundary", "fresh"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminal_run_summary_purge_keeps_boundary_and_live_runs(run_sql_store):
|
|
"""Audit-summary retention never deletes a current, boundary, or late-completed run."""
|
|
|
|
cutoff = datetime.now(UTC) - timedelta(days=90)
|
|
async with run_sql_store._sf() as session:
|
|
session.add_all(
|
|
[
|
|
RunRow(
|
|
run_id="old-terminal",
|
|
thread_id="thread-old-terminal",
|
|
status="success",
|
|
created_at=cutoff - timedelta(seconds=1),
|
|
updated_at=cutoff - timedelta(seconds=1),
|
|
),
|
|
RunRow(
|
|
run_id="boundary-terminal",
|
|
thread_id="thread-boundary-terminal",
|
|
status="error",
|
|
created_at=cutoff,
|
|
updated_at=cutoff,
|
|
),
|
|
RunRow(
|
|
run_id="old-live",
|
|
thread_id="thread-old-live",
|
|
status="running",
|
|
created_at=cutoff - timedelta(days=1),
|
|
updated_at=cutoff - timedelta(days=1),
|
|
),
|
|
RunRow(
|
|
run_id="old-created-recently-completed",
|
|
thread_id="thread-old-created-recently-completed",
|
|
status="success",
|
|
created_at=cutoff - timedelta(days=1),
|
|
updated_at=cutoff + timedelta(seconds=1),
|
|
),
|
|
]
|
|
)
|
|
await session.commit()
|
|
|
|
assert await run_sql_store.purge_terminal_before(cutoff) == 1
|
|
assert await run_sql_store.get("old-terminal", user_id=None) is None
|
|
assert await run_sql_store.get("boundary-terminal", user_id=None) is not None
|
|
assert await run_sql_store.get("old-live", user_id=None) is not None
|
|
assert await run_sql_store.get("old-created-recently-completed", user_id=None) is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_metric_purge_caps_one_large_batch_at_500(sql_store):
|
|
"""A single metrics delete never expands into an unbounded SQL statement."""
|
|
|
|
cutoff = datetime.now(UTC) - timedelta(days=90)
|
|
async with sql_store._sf() as session:
|
|
session.add_all(
|
|
[
|
|
LlmCallMetricRow(
|
|
id=f"expired-{index}",
|
|
created_at=cutoff - timedelta(days=1),
|
|
duration_ms=1,
|
|
)
|
|
for index in range(501)
|
|
]
|
|
)
|
|
await session.commit()
|
|
|
|
# Even a much larger requested size must respect the portable 500-row cap.
|
|
assert await sql_store.purge_before(cutoff, batch_size=10_000) == 500
|
|
assert await sql_store.count(LlmMetricsQuery()) == 1
|
|
assert await sql_store.purge_before(cutoff, batch_size=10_000) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminal_run_summary_purge_caps_one_large_batch_at_500(run_sql_store):
|
|
"""Run-summary retention uses the same bounded primary-key deletion rule."""
|
|
|
|
cutoff = datetime.now(UTC) - timedelta(days=90)
|
|
async with run_sql_store._sf() as session:
|
|
session.add_all(
|
|
[
|
|
RunRow(
|
|
run_id=f"expired-run-{index}",
|
|
thread_id=f"thread-{index}",
|
|
status="success",
|
|
created_at=cutoff - timedelta(days=1),
|
|
updated_at=cutoff - timedelta(days=1),
|
|
)
|
|
for index in range(501)
|
|
]
|
|
)
|
|
await session.commit()
|
|
|
|
assert await run_sql_store.purge_terminal_before(cutoff, batch_size=10_000) == 500
|
|
async with run_sql_store._sf() as session:
|
|
remaining = list((await session.scalars(select(RunRow.run_id))).all())
|
|
assert len(remaining) == 1
|
|
assert await run_sql_store.purge_terminal_before(cutoff, batch_size=10_000) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admin_retention_sweep_deletes_only_one_configured_batch_per_store(
|
|
sql_store,
|
|
run_sql_store,
|
|
):
|
|
"""A small sweep budget leaves the remaining historic backlog for later ticks."""
|
|
|
|
now = datetime(2026, 10, 26, tzinfo=UTC)
|
|
cutoff = now - timedelta(days=90)
|
|
async with sql_store._sf() as session:
|
|
session.add_all(
|
|
[
|
|
LlmCallMetricRow(
|
|
id=f"metric-{index}",
|
|
created_at=cutoff - timedelta(days=1),
|
|
duration_ms=1,
|
|
)
|
|
for index in range(11)
|
|
]
|
|
)
|
|
await session.commit()
|
|
async with run_sql_store._sf() as session:
|
|
session.add_all(
|
|
[
|
|
RunRow(
|
|
run_id=f"run-{index}",
|
|
thread_id=f"thread-{index}",
|
|
status="success",
|
|
created_at=cutoff - timedelta(days=1),
|
|
updated_at=cutoff - timedelta(days=1),
|
|
)
|
|
for index in range(11)
|
|
]
|
|
)
|
|
await session.commit()
|
|
|
|
result = await sweep_admin_audit_retention(
|
|
session_factory=None,
|
|
llm_metrics_store=sql_store,
|
|
run_store=run_sql_store,
|
|
now=now,
|
|
batch_size=10,
|
|
max_batches_per_store=1,
|
|
inter_batch_pause_seconds=0,
|
|
max_sweep_seconds=60,
|
|
)
|
|
|
|
assert result["llm_metrics"] == 10
|
|
assert result["run_summaries"] == 10
|
|
assert await sql_store.count(LlmMetricsQuery()) == 1
|
|
async with run_sql_store._sf() as session:
|
|
remaining_runs = list((await session.scalars(select(RunRow.run_id))).all())
|
|
assert len(remaining_runs) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admin_retention_snapshot_batch_does_not_delete_an_entire_old_date(
|
|
snapshot_session_factory,
|
|
):
|
|
"""Snapshot cleanup selects composite keys, not every settings hash for a date."""
|
|
|
|
now = datetime(2026, 10, 26, tzinfo=UTC)
|
|
cutoff = now - timedelta(days=90)
|
|
cutoff_date = cutoff.astimezone(BEIJING_TZ).date().isoformat()
|
|
old_date = (cutoff - timedelta(days=1)).astimezone(BEIJING_TZ).date().isoformat()
|
|
async with snapshot_session_factory() as session:
|
|
session.add_all(
|
|
[
|
|
AdminLeaderboardDailyStatRow(
|
|
stat_date=old_date,
|
|
settings_hash=f"old-hash-{index:02d}",
|
|
)
|
|
for index in range(12)
|
|
]
|
|
+ [
|
|
AdminLeaderboardDailyStatRow(
|
|
stat_date=cutoff_date,
|
|
settings_hash="boundary-hash",
|
|
)
|
|
]
|
|
)
|
|
await session.commit()
|
|
|
|
result = await sweep_admin_audit_retention(
|
|
session_factory=snapshot_session_factory,
|
|
now=now,
|
|
batch_size=10,
|
|
max_batches_per_store=1,
|
|
inter_batch_pause_seconds=0,
|
|
max_sweep_seconds=60,
|
|
)
|
|
|
|
assert result["leaderboard_snapshots"] == 10
|
|
async with snapshot_session_factory() as session:
|
|
rows = list((await session.scalars(select(AdminLeaderboardDailyStatRow))).all())
|
|
assert sum(row.stat_date == old_date for row in rows) == 2
|
|
assert [(row.stat_date, row.settings_hash) for row in rows if row.stat_date == cutoff_date] == [
|
|
(cutoff_date, "boundary-hash")
|
|
]
|
|
|
|
|
|
class _AlwaysFullRetentionStore:
|
|
"""A store that models an arbitrarily large historical backlog."""
|
|
|
|
def __init__(self, maximum_calls: int) -> None:
|
|
self.calls: list[tuple[datetime, int]] = []
|
|
self.maximum_calls = maximum_calls
|
|
|
|
async def purge_before(self, cutoff: datetime, *, batch_size: int) -> int:
|
|
self.calls.append((cutoff, batch_size))
|
|
# Fail fast if a future implementation loses the sweep budget instead
|
|
# of letting this regression test loop forever.
|
|
if len(self.calls) > self.maximum_calls:
|
|
raise AssertionError("retention sweep exceeded its per-store work budget")
|
|
return batch_size
|
|
|
|
|
|
class _BusyRetentionStore:
|
|
"""Models a table temporarily locked by another database writer."""
|
|
|
|
async def purge_before(self, cutoff: datetime, *, batch_size: int) -> int:
|
|
raise RuntimeError("database is locked")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admin_retention_sweep_has_fixed_per_store_work_budget():
|
|
"""One scheduler tick stops after a bounded amount of historic work."""
|
|
|
|
batch_size = 37
|
|
max_batches = 3
|
|
store = _AlwaysFullRetentionStore(maximum_calls=max_batches)
|
|
now = datetime(2026, 10, 26, tzinfo=UTC)
|
|
|
|
result = await sweep_admin_audit_retention(
|
|
session_factory=None,
|
|
llm_metrics_store=store,
|
|
now=now,
|
|
batch_size=batch_size,
|
|
max_batches_per_store=max_batches,
|
|
inter_batch_pause_seconds=0,
|
|
max_sweep_seconds=60,
|
|
)
|
|
|
|
expected_cutoff = now - timedelta(days=90)
|
|
assert store.calls == [(expected_cutoff, batch_size)] * max_batches
|
|
assert result == {
|
|
"llm_metrics": batch_size * max_batches,
|
|
"tool_metrics": 0,
|
|
"run_summaries": 0,
|
|
"leaderboard_snapshots": 0,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admin_retention_busy_store_does_not_block_other_sources():
|
|
"""A lock on one audit table defers only that source to a later tick."""
|
|
|
|
available_store = _AlwaysFullRetentionStore(maximum_calls=1)
|
|
result = await sweep_admin_audit_retention(
|
|
session_factory=None,
|
|
llm_metrics_store=_BusyRetentionStore(),
|
|
tool_metrics_store=available_store,
|
|
now=datetime(2026, 10, 26, tzinfo=UTC),
|
|
batch_size=10,
|
|
max_batches_per_store=1,
|
|
inter_batch_pause_seconds=0,
|
|
max_sweep_seconds=60,
|
|
)
|
|
|
|
assert available_store.calls
|
|
assert result == {
|
|
"llm_metrics": 0,
|
|
"tool_metrics": 10,
|
|
"run_summaries": 0,
|
|
"leaderboard_snapshots": 0,
|
|
}
|
|
|
|
|
|
def test_extract_usage_prefers_usage_metadata():
|
|
response = AIMessage(
|
|
content="x",
|
|
usage_metadata={"input_tokens": 3, "output_tokens": 7, "total_tokens": 10},
|
|
)
|
|
out = _extract_usage(response)
|
|
assert out == {"input_tokens": 3, "output_tokens": 7, "total_tokens": 10}
|
|
|
|
|
|
def test_extract_usage_falls_back_to_response_metadata_token_usage():
|
|
"""DeepSeek / vLLM-style: only total in usage_metadata, breakdown elsewhere."""
|
|
response = AIMessage(
|
|
content="x",
|
|
usage_metadata={"input_tokens": 0, "output_tokens": 0, "total_tokens": 11093},
|
|
response_metadata={
|
|
"token_usage": {
|
|
"prompt_tokens": 9000,
|
|
"completion_tokens": 2093,
|
|
"total_tokens": 11093,
|
|
},
|
|
},
|
|
)
|
|
out = _extract_usage(response)
|
|
assert out["input_tokens"] == 9000
|
|
assert out["output_tokens"] == 2093
|
|
assert out["total_tokens"] == 11093
|
|
|
|
|
|
def test_extract_usage_unwraps_model_response():
|
|
"""The handler hands us a ModelResponse dataclass, not a bare AIMessage."""
|
|
from langchain.agents.middleware.types import ModelResponse
|
|
|
|
inner = AIMessage(
|
|
content="hi",
|
|
usage_metadata={"input_tokens": 7, "output_tokens": 9, "total_tokens": 16},
|
|
)
|
|
response = ModelResponse(result=[inner])
|
|
out = _extract_usage(response)
|
|
assert out == {"input_tokens": 7, "output_tokens": 9, "total_tokens": 16}
|
|
|
|
|
|
def test_extract_usage_unwraps_extended_model_response():
|
|
"""ExtendedModelResponse wraps a ModelResponse in a ``model_response`` field."""
|
|
from langchain.agents.middleware.types import ExtendedModelResponse, ModelResponse
|
|
|
|
inner = AIMessage(
|
|
content="hi",
|
|
usage_metadata={"input_tokens": 1, "output_tokens": 2, "total_tokens": 3},
|
|
)
|
|
extended = ExtendedModelResponse(model_response=ModelResponse(result=[inner]))
|
|
out = _extract_usage(extended)
|
|
assert out["output_tokens"] == 2
|
|
|
|
|
|
def test_extract_usage_derives_missing_field_when_two_present():
|
|
response = AIMessage(
|
|
content="x",
|
|
usage_metadata={"input_tokens": 9000, "output_tokens": 0, "total_tokens": 11093},
|
|
)
|
|
out = _extract_usage(response)
|
|
# output should be derived as total - input = 2093.
|
|
assert out["output_tokens"] == 2093
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_middleware_records_success():
|
|
store = _RecordingStore()
|
|
mw = LlmMetricsMiddleware(store=store)
|
|
|
|
async def handler(_req):
|
|
return AIMessage(
|
|
content="hi",
|
|
usage_metadata={"input_tokens": 1, "output_tokens": 4, "total_tokens": 5},
|
|
)
|
|
|
|
response = await mw.awrap_model_call(_make_request(), handler)
|
|
assert isinstance(response, AIMessage)
|
|
|
|
# awrap_model_call schedules the write as an asyncio.task; give the loop a
|
|
# chance to drain.
|
|
await asyncio.sleep(0)
|
|
assert len(store.records) == 1
|
|
assert store.records[0].status == "success"
|
|
assert store.records[0].output_tokens == 4
|