mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 08:16:38 +00:00
arc53-machine
25c82003d7
Correctness - /api/remote never recorded source.created, so URL, GitHub and connector sources had a source.deleted with no matching creation. All three creation paths now go through one _audit_source_created helper. - The prompt-cache rate divided cached tokens by a whole bucket's prompt tokens. A bucket is a day and mixes calls whose provider reports a cache breakdown with calls whose provider does not, so filtering buckets in the client could not separate them and the rate was understated by however much traffic ran on a non-reporting provider. The denominator is now computed in SQL over the reporting rows. - The outcome pill matched values nothing writes. Guardrails emit triggered / not_evaluated and the device feed emits dispatched; the map had blocked / denied / allowed, so a guardrail that fired rendered neutral grey -- the one signal the merged feed exists to surface. Fixtures were seeding the fictional values, so the tests passed on it too. - Stream duration_ms timed the consumer. stream_token_usage is a generator, so start-to-exhaustion includes the agent loop's tool handling and the SSE client's pace; a slow browser recorded ~30s for a sub-second call. It now accumulates only the time spent inside next(). Safety - Activity filters failed open: an unknown facet or unparseable timestamp was dropped, and no filter means every row, so a typo widened an audit view and on the export streamed the full history. Both are now a 400. - The search term was interpolated into an ILIKE pattern, so "100%" matched everything and "q1_report" matched more than it should. Escaped. - 0034 set actor_id NOT NULL with no default. A previous-release process inserting mid-rollout would raise, and in admin/routes.py that insert shares the request transaction, so a role grant beside it would roll back too. Noise and dead code - The per-user panel is a security panel: data-plane events file under the actor, so an active account's routine deletes pushed a denied login out of the 20-row window. It now excludes them; the Activity tab shows everything. - device_audit_log had no created_at-leading index, so the merged feed sequentially scanned that branch every page (migration 0036). - conversation.deleted_all no longer records when nothing was deleted, and agent.updated no longer records an empty field list. - Dropped by_model from /admin/usage (no consumer; an extra aggregate per page load), the duplicate filter surface on AuthEventsRepository that nothing called, and the unreachable FLOW_LABELS.schedule entry. - Type hints on record_event's conn and the remaining unannotated helpers.
399 lines
16 KiB
Python
399 lines
16 KiB
Python
"""Tests for TokenUsageRepository against a real Postgres instance."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
|
|
|
|
|
|
def _repo(conn) -> TokenUsageRepository:
|
|
return TokenUsageRepository(conn)
|
|
|
|
|
|
def _now():
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
class TestInsert:
|
|
def test_inserts_row(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(user_id="u1", prompt_tokens=10, generated_tokens=5)
|
|
total = repo.sum_tokens_in_range(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u1"
|
|
)
|
|
assert total == 15
|
|
|
|
def test_insert_with_api_key(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(api_key="key-1", prompt_tokens=20, generated_tokens=10)
|
|
total = repo.sum_tokens_in_range(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), api_key="key-1"
|
|
)
|
|
assert total == 30
|
|
|
|
def test_cost_defaults_to_zero_and_round_trips(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(user_id="u-cost", prompt_tokens=1, generated_tokens=1)
|
|
repo.insert(user_id="u-cost", prompt_tokens=1, generated_tokens=1, cost=0.00012345)
|
|
costs = pg_conn.execute(
|
|
text("SELECT cost FROM token_usage WHERE user_id = 'u-cost' ORDER BY id")
|
|
).scalars().all()
|
|
assert [float(c) for c in costs] == [0.0, 0.00012345]
|
|
|
|
|
|
class TestRollupExclusion:
|
|
def test_sum_tokens_ignores_schedule_rollups(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(user_id="u-roll", prompt_tokens=10, generated_tokens=5, source="agent_stream")
|
|
repo.insert(user_id="u-roll", prompt_tokens=10, generated_tokens=5, source="schedule")
|
|
total = repo.sum_tokens_in_range(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u-roll"
|
|
)
|
|
assert total == 15
|
|
|
|
|
|
class TestUsageTotals:
|
|
def _seed(self, repo):
|
|
repo.insert(user_id="u-tot", prompt_tokens=100, generated_tokens=10, cost=0.5)
|
|
repo.insert(user_id="u-tot", api_key="k", prompt_tokens=20, generated_tokens=2, cost=0.25)
|
|
repo.insert(user_id="u-tot", prompt_tokens=7, generated_tokens=0, cost=0.125, source="title")
|
|
# A keyless agent (or workflow node): an agent id without a key.
|
|
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
|
|
|
agent = AgentsRepository(repo._conn).create("u-tot", "keyless", "draft")
|
|
repo.insert(
|
|
user_id="u-tot", agent_id=str(agent["id"]),
|
|
prompt_tokens=4, generated_tokens=0, cost=0.0625,
|
|
)
|
|
repo.insert(user_id="u-tot", prompt_tokens=999, generated_tokens=0, source="schedule")
|
|
repo.insert(user_id="u-other", prompt_tokens=999, generated_tokens=0, cost=9)
|
|
repo.insert(
|
|
user_id="u-tot", prompt_tokens=999, generated_tokens=0, cost=9,
|
|
timestamp=_now() - timedelta(days=40),
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"bucket, expected",
|
|
[("all", (143, 0.9375)), ("direct", (117, 0.625)), ("agent", (26, 0.3125))],
|
|
)
|
|
def test_totals_per_bucket(self, pg_conn, bucket, expected):
|
|
repo = _repo(pg_conn)
|
|
self._seed(repo)
|
|
assert repo.usage_totals(user_id="u-tot", start=_now() - timedelta(days=1), bucket=bucket) == expected
|
|
|
|
def test_no_usage_is_zero(self, pg_conn):
|
|
assert _repo(pg_conn).usage_totals(user_id="nobody", start=_now() - timedelta(days=1)) == (0, 0.0)
|
|
|
|
def test_unknown_bucket_rejected(self, pg_conn):
|
|
with pytest.raises(ValueError):
|
|
_repo(pg_conn).usage_totals(user_id="u", start=_now(), bucket="nope")
|
|
|
|
|
|
class TestReassignApiKey:
|
|
def test_rewrites_and_preserves_rate_limit_window(self, pg_conn):
|
|
# Rotating an agent key must carry the running 24h usage window over
|
|
# to the new key so rotation cannot be used to reset the limit.
|
|
repo = _repo(pg_conn)
|
|
repo.insert(api_key="tu-old", prompt_tokens=20, generated_tokens=10)
|
|
repo.insert(api_key="tu-other", prompt_tokens=1, generated_tokens=1)
|
|
|
|
moved = repo.reassign_api_key(old_key="tu-old", new_key="tu-new")
|
|
assert moved == 1
|
|
|
|
window = dict(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1)
|
|
)
|
|
assert repo.sum_tokens_in_range(api_key="tu-old", **window) == 0
|
|
assert repo.sum_tokens_in_range(api_key="tu-new", **window) == 30
|
|
# Unrelated key untouched.
|
|
assert repo.sum_tokens_in_range(api_key="tu-other", **window) == 2
|
|
|
|
def test_noop_on_blank_keys(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
assert repo.reassign_api_key(old_key="", new_key="x") == 0
|
|
assert repo.reassign_api_key(old_key="x", new_key="") == 0
|
|
|
|
|
|
class TestModelId:
|
|
def _latest_model_id(self, conn, user_id):
|
|
return conn.execute(
|
|
text(
|
|
"SELECT model_id FROM token_usage "
|
|
"WHERE user_id = :u ORDER BY id DESC LIMIT 1"
|
|
),
|
|
{"u": user_id},
|
|
).scalar_one()
|
|
|
|
def test_persists_model_id(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(user_id="u-model", prompt_tokens=1, generated_tokens=1, model_id="gpt-4o")
|
|
assert self._latest_model_id(pg_conn, "u-model") == "gpt-4o"
|
|
|
|
def test_persists_byom_uuid_model_id(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
uuid_id = "11111111-1111-1111-1111-111111111111"
|
|
repo.insert(user_id="u-byom", prompt_tokens=1, generated_tokens=1, model_id=uuid_id)
|
|
assert self._latest_model_id(pg_conn, "u-byom") == uuid_id
|
|
|
|
def test_model_id_defaults_to_null(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(user_id="u-nomodel", prompt_tokens=1, generated_tokens=1)
|
|
assert self._latest_model_id(pg_conn, "u-nomodel") is None
|
|
|
|
|
|
class TestSumTokensInRange:
|
|
def test_sums_correctly(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(user_id="u1", prompt_tokens=10, generated_tokens=5)
|
|
repo.insert(user_id="u1", prompt_tokens=20, generated_tokens=10)
|
|
repo.insert(user_id="u2", prompt_tokens=100, generated_tokens=50)
|
|
total = repo.sum_tokens_in_range(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u1"
|
|
)
|
|
assert total == 45
|
|
|
|
def test_returns_zero_when_no_rows(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
total = repo.sum_tokens_in_range(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="nobody"
|
|
)
|
|
assert total == 0
|
|
|
|
def test_respects_time_range(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
old = _now() - timedelta(hours=48)
|
|
repo.insert(user_id="u1", prompt_tokens=100, generated_tokens=0, timestamp=old)
|
|
repo.insert(user_id="u1", prompt_tokens=10, generated_tokens=0)
|
|
total = repo.sum_tokens_in_range(
|
|
start=_now() - timedelta(hours=1), end=_now() + timedelta(minutes=1), user_id="u1"
|
|
)
|
|
assert total == 10
|
|
|
|
|
|
class TestCountInRange:
|
|
def test_counts_rows(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(user_id="u1", prompt_tokens=1, generated_tokens=1)
|
|
repo.insert(user_id="u1", prompt_tokens=1, generated_tokens=1)
|
|
repo.insert(user_id="u2", prompt_tokens=1, generated_tokens=1)
|
|
count = repo.count_in_range(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u1"
|
|
)
|
|
assert count == 2
|
|
|
|
def test_filters_by_api_key(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
repo.insert(api_key="k1", prompt_tokens=1, generated_tokens=1)
|
|
repo.insert(api_key="k2", prompt_tokens=1, generated_tokens=1)
|
|
count = repo.count_in_range(
|
|
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), api_key="k1"
|
|
)
|
|
assert count == 1
|
|
|
|
def test_distinct_request_id_collapses_multi_call_stream(self, pg_conn):
|
|
"""A multi-tool agent run produces N rows with the same
|
|
``request_id`` but counts as one user request."""
|
|
repo = _repo(pg_conn)
|
|
rid = "req-multi-1"
|
|
# 4 LLM calls within a single user request.
|
|
for _ in range(4):
|
|
repo.insert(
|
|
user_id="u-multi",
|
|
prompt_tokens=10,
|
|
generated_tokens=20,
|
|
request_id=rid,
|
|
)
|
|
# And a separate request from the same user.
|
|
repo.insert(
|
|
user_id="u-multi",
|
|
prompt_tokens=5,
|
|
generated_tokens=5,
|
|
request_id="req-multi-2",
|
|
)
|
|
count = repo.count_in_range(
|
|
start=_now() - timedelta(minutes=1),
|
|
end=_now() + timedelta(minutes=1),
|
|
user_id="u-multi",
|
|
)
|
|
assert count == 2
|
|
|
|
def test_excludes_side_channel_sources(self, pg_conn):
|
|
"""Title / compression / rag_condense / fallback rows don't tick
|
|
the request limit — only ``agent_stream`` rows count.
|
|
"""
|
|
repo = _repo(pg_conn)
|
|
for src in ("title", "compression", "rag_condense", "fallback"):
|
|
repo.insert(
|
|
user_id="u-side",
|
|
prompt_tokens=5,
|
|
generated_tokens=5,
|
|
source=src,
|
|
request_id="req-side-x",
|
|
)
|
|
# One real user request.
|
|
repo.insert(
|
|
user_id="u-side",
|
|
prompt_tokens=5,
|
|
generated_tokens=5,
|
|
source="agent_stream",
|
|
request_id="req-side-x",
|
|
)
|
|
count = repo.count_in_range(
|
|
start=_now() - timedelta(minutes=1),
|
|
end=_now() + timedelta(minutes=1),
|
|
user_id="u-side",
|
|
)
|
|
assert count == 1
|
|
|
|
def test_mixes_legacy_null_and_new_request_id_rows(self, pg_conn):
|
|
"""Pre-migration rows have ``request_id=NULL`` and are counted
|
|
one-per-row; new rows are DISTINCT'd. The two branches sum.
|
|
"""
|
|
repo = _repo(pg_conn)
|
|
# Two legacy rows (NULL request_id) — count as 2.
|
|
repo.insert(user_id="u-mix", prompt_tokens=1, generated_tokens=1)
|
|
repo.insert(user_id="u-mix", prompt_tokens=1, generated_tokens=1)
|
|
# One new request with 3 rows under the same id — counts as 1.
|
|
for _ in range(3):
|
|
repo.insert(
|
|
user_id="u-mix",
|
|
prompt_tokens=1,
|
|
generated_tokens=1,
|
|
request_id="req-mix",
|
|
)
|
|
count = repo.count_in_range(
|
|
start=_now() - timedelta(minutes=1),
|
|
end=_now() + timedelta(minutes=1),
|
|
user_id="u-mix",
|
|
)
|
|
assert count == 3
|
|
|
|
|
|
class TestBucketedTotals:
|
|
def test_day_bucket_sums_per_day(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
t1 = datetime(2026, 4, 10, 10, 0, tzinfo=timezone.utc)
|
|
t2 = datetime(2026, 4, 10, 23, 30, tzinfo=timezone.utc)
|
|
t3 = datetime(2026, 4, 11, 0, 15, tzinfo=timezone.utc)
|
|
repo.insert(
|
|
user_id="u-day", prompt_tokens=10, generated_tokens=5, cost=0.5, timestamp=t1
|
|
)
|
|
repo.insert(
|
|
user_id="u-day", prompt_tokens=20, generated_tokens=7, cost=0.25, timestamp=t2
|
|
)
|
|
repo.insert(user_id="u-day", prompt_tokens=1, generated_tokens=1, timestamp=t3)
|
|
rows = repo.bucketed_totals(
|
|
bucket_unit="day",
|
|
user_id="u-day",
|
|
timestamp_gte=datetime(2026, 4, 10, tzinfo=timezone.utc),
|
|
timestamp_lt=datetime(2026, 4, 12, tzinfo=timezone.utc),
|
|
)
|
|
# ``cached_tokens`` is None when no row in the bucket reported a cache
|
|
# breakdown: "the provider said nothing", not "no cache hits".
|
|
assert rows == [
|
|
{
|
|
"bucket": "2026-04-10",
|
|
"prompt_tokens": 30,
|
|
"generated_tokens": 12,
|
|
"cost": 0.75,
|
|
"cached_tokens": None,
|
|
"cache_eligible_prompt_tokens": 0,
|
|
},
|
|
{
|
|
"bucket": "2026-04-11",
|
|
"prompt_tokens": 1,
|
|
"generated_tokens": 1,
|
|
"cost": 0.0,
|
|
"cached_tokens": None,
|
|
"cache_eligible_prompt_tokens": 0,
|
|
},
|
|
]
|
|
|
|
def test_hour_bucket(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
t1 = datetime(2026, 4, 10, 10, 5, tzinfo=timezone.utc)
|
|
t2 = datetime(2026, 4, 10, 10, 50, tzinfo=timezone.utc)
|
|
t3 = datetime(2026, 4, 10, 11, 0, tzinfo=timezone.utc)
|
|
repo.insert(user_id="u-hour", prompt_tokens=1, generated_tokens=2, timestamp=t1)
|
|
repo.insert(user_id="u-hour", prompt_tokens=3, generated_tokens=4, timestamp=t2)
|
|
repo.insert(user_id="u-hour", prompt_tokens=5, generated_tokens=6, timestamp=t3)
|
|
rows = repo.bucketed_totals(bucket_unit="hour", user_id="u-hour")
|
|
buckets = {r["bucket"]: r for r in rows}
|
|
assert buckets["2026-04-10 10:00"]["prompt_tokens"] == 4
|
|
assert buckets["2026-04-10 10:00"]["generated_tokens"] == 6
|
|
assert buckets["2026-04-10 11:00"]["prompt_tokens"] == 5
|
|
|
|
def test_minute_bucket(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
t1 = datetime(2026, 4, 10, 10, 5, 15, tzinfo=timezone.utc)
|
|
t2 = datetime(2026, 4, 10, 10, 5, 45, tzinfo=timezone.utc)
|
|
repo.insert(user_id="u-min", prompt_tokens=1, generated_tokens=2, timestamp=t1)
|
|
repo.insert(user_id="u-min", prompt_tokens=3, generated_tokens=4, timestamp=t2)
|
|
rows = repo.bucketed_totals(bucket_unit="minute", user_id="u-min")
|
|
assert len(rows) == 1
|
|
assert rows[0]["bucket"] == "2026-04-10 10:05:00"
|
|
assert rows[0]["prompt_tokens"] == 4
|
|
assert rows[0]["generated_tokens"] == 6
|
|
|
|
def test_filters_by_api_key(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
t = datetime(2026, 4, 10, 10, 0, tzinfo=timezone.utc)
|
|
repo.insert(api_key="dash-key", prompt_tokens=1, generated_tokens=1, timestamp=t)
|
|
repo.insert(api_key="other-key", prompt_tokens=99, generated_tokens=99, timestamp=t)
|
|
rows = repo.bucketed_totals(bucket_unit="day", api_key="dash-key")
|
|
assert len(rows) == 1
|
|
assert rows[0]["prompt_tokens"] == 1
|
|
|
|
def test_rejects_invalid_bucket_unit(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
with pytest.raises(ValueError):
|
|
repo.bucketed_totals(bucket_unit="week")
|
|
|
|
def test_respects_timestamp_range(self, pg_conn):
|
|
repo = _repo(pg_conn)
|
|
inside = datetime(2026, 4, 10, 12, 0, tzinfo=timezone.utc)
|
|
outside = datetime(2026, 4, 9, 12, 0, tzinfo=timezone.utc)
|
|
repo.insert(user_id="u-range", prompt_tokens=10, generated_tokens=0, timestamp=inside)
|
|
repo.insert(user_id="u-range", prompt_tokens=99, generated_tokens=0, timestamp=outside)
|
|
rows = repo.bucketed_totals(
|
|
bucket_unit="day",
|
|
user_id="u-range",
|
|
timestamp_gte=datetime(2026, 4, 10, tzinfo=timezone.utc),
|
|
timestamp_lt=datetime(2026, 4, 11, tzinfo=timezone.utc),
|
|
)
|
|
assert len(rows) == 1
|
|
assert rows[0]["prompt_tokens"] == 10
|
|
|
|
|
|
class TestCacheBins:
|
|
def _row(self, pg_conn, user_id):
|
|
return pg_conn.execute(
|
|
text(
|
|
"SELECT cached_tokens, cache_write_tokens FROM token_usage "
|
|
"WHERE user_id = :u ORDER BY id DESC LIMIT 1"
|
|
),
|
|
{"u": user_id},
|
|
).fetchone()
|
|
|
|
def test_persists_cache_bins(self, pg_conn):
|
|
repo = TokenUsageRepository(pg_conn)
|
|
repo.insert(
|
|
user_id="cache-user",
|
|
prompt_tokens=1000,
|
|
generated_tokens=10,
|
|
cached_tokens=800,
|
|
cache_write_tokens=100,
|
|
)
|
|
assert tuple(self._row(pg_conn, "cache-user")) == (800, 100)
|
|
|
|
def test_cache_bins_default_to_null_not_zero(self, pg_conn):
|
|
"""NULL means "provider did not report"; 0 means "reported no cache
|
|
activity". Keeping them distinct is what makes a hit-rate query honest."""
|
|
repo = TokenUsageRepository(pg_conn)
|
|
repo.insert(user_id="cache-user-2", prompt_tokens=10, generated_tokens=1)
|
|
assert tuple(self._row(pg_conn, "cache-user-2")) == (None, None)
|