支持gimini格式,优化debug日志
This commit is contained in:
parent
e726f11bad
commit
4de6db13f9
16 changed files with 1783 additions and 55 deletions
72
utils/usage_tracker.py
Normal file
72
utils/usage_tracker.py
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
"""用量统计 — 内存聚合
|
||||
|
||||
按模型名聚合请求数、token 用量等统计数据。
|
||||
重启后重置,适合轻量监控场景。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
|
||||
class _ModelStats:
|
||||
__slots__ = ('request_count', 'input_tokens', 'output_tokens', 'first_seen', 'last_seen')
|
||||
|
||||
def __init__(self):
|
||||
self.request_count = 0
|
||||
self.input_tokens = 0
|
||||
self.output_tokens = 0
|
||||
self.first_seen = time.time()
|
||||
self.last_seen = time.time()
|
||||
|
||||
|
||||
class UsageTracker:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._stats: dict[str, _ModelStats] = {}
|
||||
self._start_time = time.time()
|
||||
|
||||
def record(
|
||||
self,
|
||||
model: str,
|
||||
usage: dict[str, Any] | None = None,
|
||||
*,
|
||||
input_key: str = 'prompt_tokens',
|
||||
output_key: str = 'completion_tokens',
|
||||
) -> None:
|
||||
"""记录一次请求的用量。"""
|
||||
with self._lock:
|
||||
if model not in self._stats:
|
||||
self._stats[model] = _ModelStats()
|
||||
s = self._stats[model]
|
||||
s.request_count += 1
|
||||
s.last_seen = time.time()
|
||||
if usage:
|
||||
s.input_tokens += usage.get(input_key, 0) or 0
|
||||
s.output_tokens += usage.get(output_key, 0) or 0
|
||||
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""返回所有模型的聚合统计。"""
|
||||
with self._lock:
|
||||
result = {}
|
||||
for model, s in self._stats.items():
|
||||
result[model] = {
|
||||
'request_count': s.request_count,
|
||||
'input_tokens': s.input_tokens,
|
||||
'output_tokens': s.output_tokens,
|
||||
'total_tokens': s.input_tokens + s.output_tokens,
|
||||
}
|
||||
return {
|
||||
'uptime_seconds': int(time.time() - self._start_time),
|
||||
'models': result,
|
||||
}
|
||||
|
||||
def reset(self) -> None:
|
||||
with self._lock:
|
||||
self._stats.clear()
|
||||
self._start_time = time.time()
|
||||
|
||||
|
||||
usage_tracker = UsageTracker()
|
||||
Loading…
Add table
Add a link
Reference in a new issue