Spaces:
Sleeping
Sleeping
File size: 6,326 Bytes
732b14f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | from __future__ import annotations
import asyncio
import random
import threading
import time
from typing import Awaitable, Callable, TypeVar
import structlog
from tenacity import (
AsyncRetrying,
RetryCallState,
retry_if_exception_type,
)
from app.config import settings
try:
# openai v1: RateLimitError lives in the top-level module
from openai import RateLimitError
except Exception: # pragma: no cover
RateLimitError = Exception # type: ignore[assignment,misc]
T = TypeVar("T")
_semaphore: asyncio.Semaphore | None = None
_sync_semaphore: threading.Semaphore | None = None
# Configure structlog once on import. This is intentionally minimal: JSON to
# stdout with enough fields for the latency optimisation work.
structlog.configure(
processors=[
structlog.processors.TimeStamper(fmt="iso_8601", utc=True),
structlog.processors.add_log_level,
structlog.processors.StackInfoRenderer(),
structlog.processors.format_exc_info,
structlog.processors.JSONRenderer(),
],
logger_factory=structlog.stdlib.LoggerFactory(),
cache_logger_on_first_use=True,
)
_log = structlog.get_logger(__name__)
def _get_semaphore() -> asyncio.Semaphore:
global _semaphore
if _semaphore is None:
_semaphore = asyncio.Semaphore(int(settings.max_concurrent_llm_calls))
return _semaphore
def _get_sync_semaphore() -> threading.Semaphore:
global _sync_semaphore
if _sync_semaphore is None:
_sync_semaphore = threading.Semaphore(int(settings.max_concurrent_llm_calls))
return _sync_semaphore
def _exp_backoff_with_jitter(retry_state: RetryCallState) -> float:
"""Exponential backoff with +/-20% jitter."""
# attempt_number starts at 1 for the first retry.
attempt = retry_state.attempt_number
base = 1.0 * (2 ** (attempt - 1))
base = min(base, 60.0)
jitter = base * random.uniform(-0.2, 0.2)
return max(0.0, base + jitter)
def make_cache_hit_slot() -> list[bool | None]:
"""Mutable slot callers fill after ``log_openai_cache_usage`` inside ``call``."""
return [None]
def _resolved_cache_hit(
cache_hit: bool | None,
cache_hit_out: list[bool | None] | None,
) -> bool | None:
if cache_hit_out is not None and len(cache_hit_out) > 0:
return cache_hit_out[0]
return cache_hit
async def throttled_llm_call(
*,
phase: str,
section_id: str | None,
call: Callable[[], Awaitable[T]],
cache_hit: bool | None = None,
cache_hit_out: list[bool | None] | None = None,
) -> T:
"""Wrap an LLM call with:
- global concurrency semaphore
- exponential backoff with jitter on OpenAI RateLimitError
- structured JSON logs via structlog
"""
sem = _get_semaphore()
start = time.perf_counter()
last_exc: BaseException | None = None
async with sem:
retrying = AsyncRetrying(
retry=retry_if_exception_type(RateLimitError),
wait=_exp_backoff_with_jitter,
reraise=False,
)
async for attempt in retrying:
try:
result = await call()
duration_ms = int((time.perf_counter() - start) * 1000)
hit = _resolved_cache_hit(cache_hit, cache_hit_out)
_log.info(
event="llm_call",
duration_ms=duration_ms,
phase=phase,
section_id=section_id,
cache_hit=hit,
attempt=attempt.retry_state.attempt_number,
)
return result
except BaseException as exc: # noqa: BLE001
last_exc = exc
duration_ms = int((time.perf_counter() - start) * 1000)
hit = _resolved_cache_hit(cache_hit, cache_hit_out)
_log.warning(
event="llm_call_rate_limited",
duration_ms=duration_ms,
phase=phase,
section_id=section_id,
cache_hit=hit,
attempt=attempt.retry_state.attempt_number,
exc_type=type(exc).__name__,
)
# If we got here, reraise=False and we exhausted retries.
duration_ms = int((time.perf_counter() - start) * 1000)
hit = _resolved_cache_hit(cache_hit, cache_hit_out)
_log.error(
event="llm_call_exhausted_retries",
duration_ms=duration_ms,
phase=phase,
section_id=section_id,
cache_hit=hit,
)
if last_exc is None:
raise RuntimeError("LLM call exhausted retries without exception")
raise last_exc
def throttled_sync_llm_call(
*,
phase: str,
section_id: str | None,
call: Callable[[], T],
) -> T:
"""Sync OpenAI calls with global concurrency cap and rate-limit retries."""
from tenacity import Retrying, retry_if_exception_type
sem = _get_sync_semaphore()
start = time.perf_counter()
with sem:
retrying = Retrying(
retry=retry_if_exception_type(RateLimitError),
wait=_exp_backoff_with_jitter,
reraise=True,
)
for attempt in retrying:
with attempt:
try:
result = call()
except RateLimitError:
duration_ms = int((time.perf_counter() - start) * 1000)
_log.warning(
event="llm_call_rate_limited",
duration_ms=duration_ms,
phase=phase,
section_id=section_id,
cache_hit=None,
attempt=attempt.retry_state.attempt_number,
exc_type="RateLimitError",
)
raise
duration_ms = int((time.perf_counter() - start) * 1000)
_log.info(
event="llm_call",
duration_ms=duration_ms,
phase=phase,
section_id=section_id,
cache_hit=None,
attempt=attempt.retry_state.attempt_number,
)
return result
raise RuntimeError("unreachable")
|