100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
"""Async batch logger for request logs."""
|
|
import asyncio
|
|
from typing import Optional, List
|
|
from dataclasses import dataclass
|
|
from ..database import db_connection
|
|
from ..config import LOG_BATCH_SIZE, LOG_FLUSH_INTERVAL
|
|
|
|
|
|
@dataclass
|
|
class LogEntry:
|
|
"""A single log entry."""
|
|
client_key_id: int
|
|
provider_id: int
|
|
model: str
|
|
prompt_tokens: Optional[int] = None
|
|
completion_tokens: Optional[int] = None
|
|
latency_ms: Optional[int] = None
|
|
success: bool = True
|
|
error_message: Optional[str] = None
|
|
|
|
|
|
class AsyncLogger:
|
|
"""Asynchronous batch logger for request logs."""
|
|
|
|
def __init__(self):
|
|
self._queue: asyncio.Queue[LogEntry] = asyncio.Queue()
|
|
self._task: Optional[asyncio.Task] = None
|
|
|
|
async def start(self):
|
|
"""Start the background flush task."""
|
|
if self._task is None:
|
|
self._task = asyncio.create_task(self._flush_loop())
|
|
|
|
async def stop(self):
|
|
"""Stop the background flush task."""
|
|
if self._task:
|
|
self._task.cancel()
|
|
try:
|
|
await self._task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._task = None
|
|
# Flush remaining logs
|
|
await self._flush_all()
|
|
|
|
async def log(self, entry: LogEntry):
|
|
"""Queue a log entry."""
|
|
await self._queue.put(entry)
|
|
|
|
async def _flush_loop(self):
|
|
"""Background loop to flush logs periodically."""
|
|
while True:
|
|
try:
|
|
await asyncio.sleep(LOG_FLUSH_INTERVAL)
|
|
await self._flush_batch()
|
|
except asyncio.CancelledError:
|
|
raise
|
|
|
|
async def _flush_batch(self):
|
|
"""Flush a batch of logs to database."""
|
|
entries: List[LogEntry] = []
|
|
|
|
# Collect up to batch_size entries
|
|
while len(entries) < LOG_BATCH_SIZE:
|
|
try:
|
|
entry = self._queue.get_nowait()
|
|
entries.append(entry)
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
if not entries:
|
|
return
|
|
|
|
# Insert batch
|
|
async with db_connection() as db:
|
|
await db.executemany(
|
|
"""INSERT INTO request_logs
|
|
(client_key_id, provider_id, model, prompt_tokens,
|
|
completion_tokens, latency_ms, success, error_message)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
[
|
|
(
|
|
e.client_key_id, e.provider_id, e.model,
|
|
e.prompt_tokens, e.completion_tokens,
|
|
e.latency_ms, e.success, e.error_message
|
|
)
|
|
for e in entries
|
|
]
|
|
)
|
|
await db.commit()
|
|
|
|
async def _flush_all(self):
|
|
"""Flush all remaining logs."""
|
|
while not self._queue.empty():
|
|
await self._flush_batch()
|
|
|
|
|
|
# Global instance
|
|
async_logger = AsyncLogger()
|