feat: Add .env.example configuration file and integrate Redis service into Docker Compose for improved task management. Update desktop application to support document indexing and extraction workflows, enhancing user experience with a new extraction wizard and stepper components. Revise README and documentation for clarity on setup and usage.
This commit is contained in:
erik
@@ -0,0 +1,156 @@
|
||||
"""Rate limiter distribuido para llamadas a OpenAI.
|
||||
|
||||
Usa Redis como backend para compartir el estado entre múltiples workers.
|
||||
Implementa el algoritmo de token bucket con ventana deslizante para respetar
|
||||
los límites de RPM (requests per minute) y TPM (tokens per minute) de OpenAI.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Límites por defecto para OpenAI gpt-4o (tier 1).
|
||||
# Sobreescribibles con variables de entorno vía config.
|
||||
DEFAULT_RPM = 500 # requests per minute
|
||||
DEFAULT_TPM = 800_000 # tokens per minute
|
||||
|
||||
|
||||
class DistributedRateLimiter:
|
||||
"""Token bucket distribuido respaldado por Redis.
|
||||
|
||||
Un solo objeto se comparte entre todos los workers de un proceso.
|
||||
Múltiples réplicas del worker comparten el estado en Redis.
|
||||
|
||||
Uso:
|
||||
limiter = DistributedRateLimiter(redis_url="redis://redis:6379/1")
|
||||
async with limiter.acquire(estimated_tokens=4000):
|
||||
response = await openai_client.chat.completions.create(...)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis_url: str,
|
||||
rpm: int = DEFAULT_RPM,
|
||||
tpm: int = DEFAULT_TPM,
|
||||
key_prefix: str = "ips:rl",
|
||||
) -> None:
|
||||
self._redis_url = redis_url
|
||||
self.rpm = rpm
|
||||
self.tpm = tpm
|
||||
self._prefix = key_prefix
|
||||
self._client: Optional[aioredis.Redis] = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _get_client(self) -> aioredis.Redis:
|
||||
if self._client is None:
|
||||
async with self._lock:
|
||||
if self._client is None:
|
||||
self._client = aioredis.from_url(
|
||||
self._redis_url,
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=5,
|
||||
)
|
||||
return self._client
|
||||
|
||||
async def _sliding_window_count(self, key: str, window_s: float = 60.0) -> int:
|
||||
"""Cuenta eventos en la ventana deslizante usando un Sorted Set de Redis."""
|
||||
client = await self._get_client()
|
||||
now = time.monotonic()
|
||||
# Eliminar eventos expirados
|
||||
await client.zremrangebyscore(key, "-inf", now - window_s)
|
||||
return await client.zcard(key)
|
||||
|
||||
async def _record_event(self, key: str, tokens: int = 0, window_s: float = 60.0) -> None:
|
||||
"""Registra un evento en el sorted set. El score es el timestamp."""
|
||||
client = await self._get_client()
|
||||
now = time.monotonic()
|
||||
member = f"{now}:{tokens}"
|
||||
pipe = client.pipeline()
|
||||
pipe.zadd(key, {member: now})
|
||||
pipe.expire(key, int(window_s) + 5)
|
||||
await pipe.execute()
|
||||
|
||||
async def _token_sum_in_window(self, key: str, window_s: float = 60.0) -> int:
|
||||
"""Suma los tokens de todos los eventos en la ventana."""
|
||||
client = await self._get_client()
|
||||
now = time.monotonic()
|
||||
await client.zremrangebyscore(key, "-inf", now - window_s)
|
||||
members = await client.zrange(key, 0, -1)
|
||||
total = 0
|
||||
for m in members:
|
||||
parts = m.split(":", 1)
|
||||
if len(parts) == 2:
|
||||
try:
|
||||
total += int(parts[1])
|
||||
except ValueError:
|
||||
pass
|
||||
return total
|
||||
|
||||
async def acquire(self, estimated_tokens: int = 2000, max_wait_s: float = 120.0) -> None:
|
||||
"""Espera hasta que haya capacidad disponible en RPM y TPM.
|
||||
|
||||
Args:
|
||||
estimated_tokens: Tokens que se espera consumir (input + output).
|
||||
max_wait_s: Tiempo máximo de espera antes de lanzar TimeoutError.
|
||||
"""
|
||||
rpm_key = f"{self._prefix}:rpm"
|
||||
tpm_key = f"{self._prefix}:tpm"
|
||||
start = time.monotonic()
|
||||
sleep_s = 0.5
|
||||
|
||||
while True:
|
||||
rpm_count = await self._sliding_window_count(rpm_key)
|
||||
tpm_used = await self._token_sum_in_window(tpm_key)
|
||||
|
||||
rpm_ok = rpm_count < self.rpm
|
||||
tpm_ok = (tpm_used + estimated_tokens) <= self.tpm
|
||||
|
||||
if rpm_ok and tpm_ok:
|
||||
await self._record_event(rpm_key, tokens=0)
|
||||
await self._record_event(tpm_key, tokens=estimated_tokens)
|
||||
return
|
||||
|
||||
elapsed = time.monotonic() - start
|
||||
if elapsed >= max_wait_s:
|
||||
raise TimeoutError(
|
||||
f"Rate limiter: espera maxima alcanzada ({max_wait_s}s). "
|
||||
f"rpm={rpm_count}/{self.rpm} tpm={tpm_used}/{self.tpm}"
|
||||
)
|
||||
|
||||
reason = "RPM" if not rpm_ok else "TPM"
|
||||
logger.debug(
|
||||
"Rate limiter: esperando por %s (rpm=%d/%d tpm=%d/%d) — sleep %.1fs",
|
||||
reason, rpm_count, self.rpm, tpm_used, self.tpm, sleep_s,
|
||||
)
|
||||
await asyncio.sleep(sleep_s)
|
||||
sleep_s = min(sleep_s * 1.5, 10.0)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
|
||||
# Instancia global compartida por todos los workers del proceso.
|
||||
# Se inicializa de forma lazy en el primer uso.
|
||||
_limiter: Optional[DistributedRateLimiter] = None
|
||||
|
||||
|
||||
def get_rate_limiter() -> DistributedRateLimiter:
|
||||
"""Devuelve el rate limiter global, inicializándolo si es necesario."""
|
||||
global _limiter
|
||||
if _limiter is None:
|
||||
from domain.config import settings
|
||||
_limiter = DistributedRateLimiter(
|
||||
redis_url=settings.redis_url,
|
||||
rpm=settings.openai_rpm_limit,
|
||||
tpm=settings.openai_tpm_limit,
|
||||
)
|
||||
return _limiter
|
||||
Reference in New Issue
Block a user