47 lines
2.4 KiB
Python
47 lines
2.4 KiB
Python
from pathlib import Path
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
sys.path.insert(0, r'C:\kefu\wechat_rpa')
|
|
from admin_backend import Database
|
|
|
|
with tempfile.TemporaryDirectory(prefix='model-usage-query-test-') as temporary:
|
|
db = Database(Path(temporary) / 'usage.db')
|
|
db.migrate()
|
|
_, since, until = db._model_usage_window(7)
|
|
with db.connect() as connection:
|
|
connection.execute('BEGIN')
|
|
connection.executemany('''INSERT INTO model_calls
|
|
(tenant_id,task_id,created_at) VALUES (?,?,?)''',
|
|
[('tenant-a', f'task-{i}', since) for i in range(20000)])
|
|
connection.executemany('''INSERT INTO model_usage_events
|
|
(event_id,task_id,tenant_id,purpose,role,attempt,status,latency_ms,created_at,total_tokens)
|
|
VALUES (?,?,?,'chat','answer',1,'success',1,?,100)''',
|
|
[(f'event-{i}', f'task-{i}', 'tenant-a', since) for i in range(10000)])
|
|
queries = {
|
|
'tenant_period': (
|
|
'SELECT SUM(total_tokens) FROM model_usage_events WHERE tenant_id=? AND created_at>=? AND created_at<?',
|
|
['tenant-a', since, until]),
|
|
'legacy_coverage': ('''SELECT COUNT(*) FROM model_calls AS calls
|
|
WHERE datetime(calls.created_at)>=datetime(?) AND datetime(calls.created_at)<datetime(?)
|
|
AND calls.tenant_id=? AND NOT EXISTS (
|
|
SELECT 1 FROM model_usage_events AS usage
|
|
WHERE calls.task_id != '' AND usage.task_id != '' AND usage.task_id=calls.task_id
|
|
AND (usage.tenant_id=calls.tenant_id OR
|
|
(usage.tenant_id IN ('','default') AND calls.tenant_id IN ('','default'))))''',
|
|
[since, until, 'tenant-a']),
|
|
}
|
|
report = {}
|
|
with db.connect() as connection:
|
|
for name, (query, parameters) in queries.items():
|
|
plan = [dict(row) for row in connection.execute('EXPLAIN QUERY PLAN ' + query, parameters)]
|
|
started = time.perf_counter()
|
|
result = connection.execute(query, parameters).fetchone()[0]
|
|
report[name] = {'plan': plan, 'result': result, 'elapsed_ms': round((time.perf_counter()-started)*1000, 2)}
|
|
assert report['legacy_coverage']['result'] == 10000
|
|
assert report['tenant_period']['result'] == 1000000
|
|
Path(__file__).with_name('query-plans.json').write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8')
|
|
print(json.dumps(report, ensure_ascii=False))
|