|
|
|
|
|
|
| """Request log persistence and OpenAI stream logging helpers."""
|
|
|
| from __future__ import annotations
|
|
|
| import json
|
| import time
|
| from typing import Any, AsyncGenerator, Dict, Optional, Tuple
|
|
|
| from app.services.request_log_dao import get_request_log_dao
|
| from app.utils.format import format_compact_number
|
| from app.utils.logger import logger, log_exception
|
| from app.utils.request_source import RequestSourceInfo
|
| from app.utils.token_pool import get_token_pool
|
| from app.utils.utlis import mask_token
|
|
|
| CACHE_CREATION_FLOOR = 1024
|
| CACHE_STRIDE = 128
|
|
|
|
|
| def _empty_usage() -> Dict[str, int]:
|
| return {
|
| "input_tokens": 0,
|
| "output_tokens": 0,
|
| "cache_creation_tokens": 0,
|
| "cache_read_tokens": 0,
|
| "total_tokens": 0,
|
| }
|
|
|
|
|
| def _coerce_int_or_none(value: Any) -> Optional[int]:
|
| try:
|
| if value is None:
|
| return None
|
| return int(value)
|
| except (TypeError, ValueError):
|
| return None
|
|
|
|
|
| def _coerce_int(value: Any, default: int = 0) -> int:
|
| coerced = _coerce_int_or_none(value)
|
| return default if coerced is None else coerced
|
|
|
|
|
| def _get_first_present(
|
| mapping: Dict[str, Any],
|
| *keys: str,
|
| ) -> Tuple[Optional[Any], Optional[str]]:
|
| for key in keys:
|
| if key in mapping and mapping[key] is not None:
|
| return mapping[key], key
|
| return None, None
|
|
|
|
|
| def _estimate_cache_creation_tokens(
|
| input_tokens: int,
|
| cache_read_tokens: int,
|
| ) -> int:
|
| if input_tokens < CACHE_CREATION_FLOOR:
|
| return 0
|
| cacheable = (input_tokens // CACHE_STRIDE) * CACHE_STRIDE
|
| return max(0, cacheable - cache_read_tokens)
|
|
|
|
|
| def _merge_openai_usage(
|
| current: Dict[str, int],
|
| update: Dict[str, int],
|
| ) -> Dict[str, int]:
|
| merged = dict(current)
|
|
|
| for key in (
|
| "input_tokens",
|
| "output_tokens",
|
| "cache_creation_tokens",
|
| "cache_read_tokens",
|
| ):
|
| value = _coerce_int(update.get(key))
|
| if value > 0:
|
| merged[key] = value
|
|
|
| total_tokens = _coerce_int(update.get("total_tokens"))
|
| if total_tokens > 0:
|
| merged["total_tokens"] = total_tokens
|
| else:
|
| merged["total_tokens"] = merged["input_tokens"] + merged["output_tokens"]
|
|
|
| return merged
|
|
|
|
|
| def extract_openai_usage(response: Dict[str, Any]) -> Dict[str, int]:
|
| """Extract OpenAI-compatible usage, with cache fallback estimation."""
|
| usage = response.get("usage") or {}
|
| prompt_details = usage.get("prompt_tokens_details") or {}
|
| input_details = usage.get("input_token_details") or {}
|
|
|
| raw, _ = _get_first_present(usage, "prompt_tokens", "input_tokens")
|
| input_tokens = _coerce_int(raw)
|
|
|
| raw, _ = _get_first_present(usage, "completion_tokens", "output_tokens")
|
| output_tokens = _coerce_int(raw)
|
|
|
| raw, _ = _get_first_present(usage, "total_tokens")
|
| total_tokens = _coerce_int(raw)
|
| if total_tokens <= 0:
|
| total_tokens = input_tokens + output_tokens
|
|
|
| raw, source = _get_first_present(
|
| usage,
|
| "cache_read_input_tokens",
|
| "cache_read_tokens",
|
| "cached_tokens",
|
| )
|
| if not source:
|
| raw, source = _get_first_present(
|
| prompt_details,
|
| "cached_tokens",
|
| "cache_read_tokens",
|
| "cache_read_input_tokens",
|
| )
|
| if not source:
|
| raw, source = _get_first_present(
|
| input_details,
|
| "cached_tokens",
|
| "cache_read_input_tokens",
|
| "cache_read_tokens",
|
| )
|
| cache_read_tokens = max(0, _coerce_int(raw))
|
|
|
| raw, source = _get_first_present(
|
| usage,
|
| "cache_creation_input_tokens",
|
| "cache_creation_tokens",
|
| )
|
| if not source:
|
| raw, source = _get_first_present(
|
| prompt_details,
|
| "cache_creation_tokens",
|
| "cache_creation_input_tokens",
|
| )
|
| if not source:
|
| raw, source = _get_first_present(
|
| input_details,
|
| "cache_creation_input_tokens",
|
| "cache_creation_tokens",
|
| )
|
| reported_cache_creation = _coerce_int_or_none(raw)
|
|
|
| if reported_cache_creation is not None:
|
| cache_creation_tokens = max(0, reported_cache_creation)
|
| else:
|
| cache_creation_tokens = _estimate_cache_creation_tokens(
|
| input_tokens,
|
| cache_read_tokens,
|
| )
|
|
|
| return {
|
| "input_tokens": input_tokens,
|
| "output_tokens": output_tokens,
|
| "cache_creation_tokens": cache_creation_tokens,
|
| "cache_read_tokens": cache_read_tokens,
|
| "total_tokens": total_tokens,
|
| }
|
|
|
|
|
| async def write_request_log(
|
| *,
|
| provider: str,
|
| model: str,
|
| source_info: RequestSourceInfo,
|
| auth_token: Optional[str] = None,
|
| upstream_auth_token: Optional[str] = None,
|
| success: bool,
|
| started_at: float,
|
| status_code: int = 200,
|
| first_token_time: float = 0.0,
|
| input_tokens: int = 0,
|
| output_tokens: int = 0,
|
| cache_creation_tokens: int = 0,
|
| cache_read_tokens: int = 0,
|
| total_tokens: Optional[int] = None,
|
| error_message: Optional[str] = None,
|
| ) -> None:
|
| """Persist a request log entry without breaking request handling."""
|
| duration = max(0.0, time.perf_counter() - started_at)
|
|
|
| status_tag = "OK" if success else "ERR"
|
| masked_auth_token = mask_token(auth_token)
|
| masked_upstream_auth_token = mask_token(upstream_auth_token)
|
| auth_segment = f" | Auth: {masked_auth_token}" if masked_auth_token else ""
|
| upstream_auth_segment = (
|
| f" | Upstream Auth: {masked_upstream_auth_token}"
|
| if masked_upstream_auth_token
|
| else ""
|
| )
|
| logger.info(
|
| "[%s] [%s] %s%s%s | In: %s | Out: %s | %.2fs",
|
| status_tag,
|
| provider,
|
| model,
|
| auth_segment,
|
| upstream_auth_segment,
|
| format_compact_number(input_tokens),
|
| format_compact_number(output_tokens),
|
| duration,
|
| )
|
|
|
| try:
|
| dao = get_request_log_dao()
|
| dao.add_log_nowait(
|
| provider=provider,
|
| endpoint=source_info.endpoint,
|
| source=source_info.source,
|
| protocol=source_info.protocol,
|
| client_name=source_info.client_name,
|
| auth_token=auth_token,
|
| upstream_auth_token=upstream_auth_token,
|
| model=model,
|
| status_code=status_code,
|
| success=success,
|
| duration=duration,
|
| first_token_time=first_token_time,
|
| input_tokens=input_tokens,
|
| output_tokens=output_tokens,
|
| cache_creation_tokens=cache_creation_tokens,
|
| cache_read_tokens=cache_read_tokens,
|
| total_tokens=total_tokens,
|
| error_message=error_message,
|
| )
|
| except Exception as exc:
|
| log_exception(logger, "Failed to write request log")
|
|
|
|
|
| def _openai_response_has_output(response: Dict[str, Any]) -> bool:
|
| """Check whether a complete (non-stream) OpenAI response has visible output."""
|
| choices = response.get("choices") if isinstance(response, dict) else None
|
| if not choices:
|
| return False
|
| message = choices[0].get("message") or {}
|
| if message.get("content") or message.get("reasoning_content"):
|
| return True
|
| if message.get("tool_calls"):
|
| return True
|
| return False
|
|
|
|
|
| def _sse_chunk_has_output(chunk_text: str) -> bool:
|
| """Check whether a formatted SSE chunk string has visible output content."""
|
| for line in chunk_text.splitlines():
|
| if not line.startswith("data: "):
|
| continue
|
| payload_text = line[6:].strip()
|
| if not payload_text or payload_text == "[DONE]":
|
| continue
|
| try:
|
| payload = json.loads(payload_text)
|
| except json.JSONDecodeError:
|
| continue
|
| if isinstance(payload, dict):
|
| return _openai_payload_has_output(payload)
|
| return False
|
|
|
|
|
| def _openai_payload_has_output(payload: Dict[str, Any]) -> bool:
|
| """Treat only visible content/tool deltas as first-token output."""
|
| choice = ((payload.get("choices") or [{}])[0]) if isinstance(payload, dict) else {}
|
| delta = choice.get("delta") or {}
|
|
|
| if delta.get("content") or delta.get("reasoning_content"):
|
| return True
|
| if delta.get("tool_calls"):
|
| return True
|
| return False
|
|
|
|
|
| def _iter_sse_payloads(chunk: str):
|
| for line in chunk.splitlines():
|
| if not line.startswith("data: "):
|
| continue
|
| payload_text = line[6:].strip()
|
| if payload_text:
|
| yield payload_text
|
|
|
|
|
| async def wrap_openai_stream_with_logging(
|
| stream: AsyncGenerator[str, None],
|
| *,
|
| provider: str,
|
| model: str,
|
| source_info: RequestSourceInfo,
|
| auth_token: Optional[str] = None,
|
| upstream_auth_token: Optional[str] = None,
|
| started_at: float,
|
| ) -> AsyncGenerator[str, None]:
|
| """Wrap an OpenAI-compatible SSE stream and persist completion metadata."""
|
| success = True
|
| status_code = 200
|
| error_message: Optional[str] = None
|
| first_token_time = 0.0
|
| has_output = False
|
| usage = _empty_usage()
|
|
|
| try:
|
| async for chunk in stream:
|
| for payload_text in _iter_sse_payloads(chunk):
|
| if payload_text == "[DONE]":
|
| continue
|
|
|
| try:
|
| payload = json.loads(payload_text)
|
| except json.JSONDecodeError:
|
| payload = None
|
|
|
| if not isinstance(payload, dict):
|
| continue
|
|
|
| if "error" in payload:
|
| success = False
|
| error = payload.get("error") or {}
|
| error_message = error.get("message") or "Unknown stream error"
|
| status_code = _coerce_int(error.get("code"), 500)
|
| continue
|
|
|
| if _openai_payload_has_output(payload):
|
| has_output = True
|
| if not first_token_time:
|
| first_token_time = max(0.0, time.perf_counter() - started_at)
|
|
|
| if payload.get("usage"):
|
| usage = _merge_openai_usage(
|
| usage,
|
| extract_openai_usage(payload),
|
| )
|
|
|
| yield chunk
|
| except Exception as exc:
|
| success = False
|
| status_code = 500
|
| error_message = str(exc)
|
| raise
|
| finally:
|
| if success and not has_output:
|
| success = False
|
| error_message = error_message or "Empty response: no content in stream"
|
| token_pool = get_token_pool()
|
| if token_pool and upstream_auth_token:
|
| await token_pool.record_token_failure(
|
| upstream_auth_token,
|
| Exception(error_message),
|
| )
|
| await write_request_log(
|
| provider=provider,
|
| model=model,
|
| source_info=source_info,
|
| auth_token=auth_token,
|
| upstream_auth_token=upstream_auth_token,
|
| success=success,
|
| started_at=started_at,
|
| status_code=status_code,
|
| first_token_time=first_token_time,
|
| input_tokens=usage["input_tokens"],
|
| output_tokens=usage["output_tokens"],
|
| cache_creation_tokens=usage["cache_creation_tokens"],
|
| cache_read_tokens=usage["cache_read_tokens"],
|
| total_tokens=usage["total_tokens"],
|
| error_message=error_message,
|
| )
|
|
|
