| from __future__ import annotations |
| import asyncio, json, time |
| from app.services.storage import async_session_maker |
| from app.models.request_log import RequestLog |
|
|
| _buffer: list[dict] = [] |
| _lock: asyncio.Lock = asyncio.Lock() |
| _flush_task: asyncio.Task | None = None |
| FLUSH_INTERVAL = 2.0 |
| FLUSH_BATCH_SIZE = 100 |
|
|
|
|
| async def log(entry: dict) -> None: |
| async with _lock: |
| _buffer.append(entry) |
| if len(_buffer) >= FLUSH_BATCH_SIZE: |
| asyncio.create_task(flush()) |
|
|
|
|
| async def flush() -> None: |
| async with _lock: |
| entries = _buffer[:] |
| _buffer.clear() |
| if not entries: |
| return |
| try: |
| async with async_session_maker() as session: |
| for e in entries: |
| session.add(RequestLog( |
| request_id=e.get("request_id", ""), |
| provider=e.get("provider", ""), |
| model=e.get("model", ""), |
| endpoint=e.get("endpoint", ""), |
| key_id=e.get("key_id"), |
| hub_key_id=e.get("hub_key_id"), |
| tokens_in=e.get("tokens_in", 0), |
| tokens_out=e.get("tokens_out", 0), |
| cost=e.get("cost", 0), |
| latency_ms=e.get("latency_ms", 0), |
| status_code=e.get("status_code", 200), |
| error=e.get("error"), |
| cached=e.get("cached", 0), |
| cache_type=e.get("cache_type"), |
| streaming=e.get("streaming", 0), |
| fallback_count=e.get("fallback_count", 0), |
| ip=e.get("ip"), |
| session_id=e.get("session_id"), |
| injection_flagged=e.get("injection_flagged", 0), |
| has_images=e.get("has_images", 0), |
| ab_variant=e.get("ab_variant"), |
| cache_read_tokens=e.get("cache_read_tokens", 0), |
| cache_write_tokens=e.get("cache_write_tokens", 0), |
| )) |
| await session.commit() |
| except Exception as exc: |
| print(f"[batch_logger] Flush error: {exc}") |
|
|
|
|
| async def _periodic_flush(): |
| while True: |
| await asyncio.sleep(FLUSH_INTERVAL) |
| async with _lock: |
| if _buffer: |
| await flush() |
|
|
|
|
| def start() -> None: |
| global _flush_task |
| if _flush_task is None or _flush_task.done(): |
| _flush_task = asyncio.create_task(_periodic_flush()) |
|
|
