345 lines
11 KiB
Python
345 lines
11 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.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.agents.middlewares.llm_metrics_middleware import (
|
|
LlmMetricsMiddleware,
|
|
_extract_usage,
|
|
)
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.llm_metrics import SqlLlmMetricsStore
|
|
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery
|
|
|
|
|
|
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:
|
|
# Importing the package registers the LlmCallMetricRow with Base.metadata.
|
|
import deerflow.persistence.llm_metrics.model # noqa: F401
|
|
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
sf = async_sessionmaker(engine, expire_on_commit=False)
|
|
store = SqlLlmMetricsStore(sf)
|
|
try:
|
|
yield store
|
|
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
|
|
|
|
|
|
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
|