Coverage for src/chat_limiter/limiter.py: 82%

1"""

2Core rate limiter implementation using PyrateLimiter.

3"""

5import asyncio

6import logging

7import time

8from collections.abc import AsyncIterator, Iterator

9from contextlib import asynccontextmanager, contextmanager

10from dataclasses import dataclass, field

11from typing import Any

13import httpx

14from pyrate_limiter import Duration, Limiter, Rate

15from tenacity import (

16 retry,

17 retry_if_exception_type,

18 stop_after_attempt,

19 wait_exponential,

20)

22from .adapters import get_adapter

23from .providers import (

24 Provider,

25 ProviderConfig,

26 RateLimitInfo,

27 detect_provider_from_url,

28 extract_rate_limit_info,

29 get_provider_config,

30)

31from .types import (

32 ChatCompletionRequest,

33 ChatCompletionResponse,

34 Message,

35 MessageRole,

36 detect_provider_from_model,

37)

39logger = logging.getLogger(__name__)

42@dataclass

43class LimiterState:

44 """Current state of the rate limiter."""

46 # Current limits (None if not yet discovered)

47 request_limit: int | None = None

48 token_limit: int | None = None

50 # Usage tracking

51 requests_used: int = 0

52 tokens_used: int = 0

54 # Timing

55 last_request_time: float = field(default_factory=time.time)

56 last_limit_update: float = field(default_factory=time.time)

58 # Rate limit info from last response

59 last_rate_limit_info: RateLimitInfo | None = None

61 # Adaptive behavior

62 consecutive_rate_limit_errors: int = 0

63 adaptive_backoff_factor: float = 1.0

66class ChatLimiter:

67 """

68 A Pythonic rate limiter for API calls supporting OpenAI, Anthropic, and OpenRouter.

70 Features:

71 - Automatic rate limit discovery and adaptation

72 - Sync and async support with context managers

73 - Intelligent retry logic with exponential backoff

74 - Token and request rate limiting

75 - Provider-specific optimizations

77 Example:

78 # High-level interface (recommended)

79 async with ChatLimiter.for_model("gpt-4o", api_key="sk-...") as limiter:

80 response = await limiter.chat_completion(

81 model="gpt-4o",

82 messages=[Message(role=MessageRole.USER, content="Hello!")]

83 )

85 # Low-level interface (for advanced users)

86 async with ChatLimiter(provider=Provider.OPENAI, api_key="sk-...") as limiter:

87 response = await limiter.request("POST", "/chat/completions", json=data)

88 """

90 def __init__(

91 self,

92 provider: Provider | None = None,

93 api_key: str | None = None,

94 base_url: str | None = None,

95 config: ProviderConfig | None = None,

96 http_client: httpx.AsyncClient | None = None,

97 sync_http_client: httpx.Client | None = None,

98 enable_adaptive_limits: bool = True,

99 enable_token_estimation: bool = True,

100 request_limit: int | None = None,

101 token_limit: int | None = None,

102 max_retries: int | None = None,

103 base_backoff: float | None = None,

104 timeout: float | None = None,

105 **kwargs: Any,

106 ):

107 """

108 Initialize the ChatLimiter.

109

110 Args:

111 provider: The API provider (OpenAI, Anthropic, OpenRouter)

112 api_key: API key for authentication

113 base_url: Base URL for API requests

114 config: Custom provider configuration

115 http_client: Custom async HTTP client

116 sync_http_client: Custom sync HTTP client

117 enable_adaptive_limits: Enable adaptive rate limit adjustment

118 enable_token_estimation: Enable token usage estimation

119 request_limit: Override request limit (if not provided, must be discovered from API)

120 token_limit: Override token limit (if not provided, must be discovered from API)

121 max_retries: Override max retries (defaults to 3 if not provided)

122 base_backoff: Override base backoff (defaults to 1.0 if not provided)

123 timeout: HTTP request timeout in seconds (defaults to 120.0 for better reliability)

124 **kwargs: Additional arguments passed to HTTP clients

125 """

126 # Determine provider and config

127 if config:

128 self.config = config

129 self.provider = config.provider

130 elif provider:

131 self.provider = provider

132 self.config = get_provider_config(provider)

133 elif base_url:

134 detected_provider = detect_provider_from_url(base_url)

135 if detected_provider:

136 self.provider = detected_provider

137 self.config = get_provider_config(detected_provider)

138 else:

139 raise ValueError(f"Could not detect provider from URL: {base_url}")

140 else:

141 raise ValueError("Must provide either provider, config, or base_url")

142

143 # Override base_url if provided

144 if base_url:

145 self.config.base_url = base_url

146

147 # Store configuration

148 self.api_key = api_key

149 self.enable_adaptive_limits = enable_adaptive_limits

150 self.enable_token_estimation = enable_token_estimation

151

152 # Store user-provided overrides

153 self._user_request_limit = request_limit

154 self._user_token_limit = token_limit

155 self._user_max_retries = max_retries or 3 # Default to 3 if not provided

156 self._user_base_backoff = base_backoff or 1.0 # Default to 1.0 if not provided

157 self._user_timeout = timeout or 120.0 # Default to 120 seconds for better reliability

158

159 # Determine initial limits (user override, config default, or None for discovery)

160 initial_request_limit = (

161 request_limit or

162 self.config.default_request_limit or

163 None

164 )

165 initial_token_limit = (

166 token_limit or

167 self.config.default_token_limit or

168 None

169 )

170

171 # Initialize state - will be None if no defaults and no discovery yet

172 self.state = LimiterState(

173 request_limit=initial_request_limit,

174 token_limit=initial_token_limit,

175 )

176

177 # Flag to track if we need to discover limits

178 self._limits_discovered = initial_request_limit is not None and initial_token_limit is not None

179

180 # Initialize HTTP clients

181 self._init_http_clients(http_client, sync_http_client, **kwargs)

182

183 # Initialize rate limiters

184 self._init_rate_limiters()

185

186 # Context manager state

187 self._async_context_active = False

188 self._sync_context_active = False

189

190 # Verbose mode (can be set by batch processor)

191 self._verbose_mode = False

192

193 @classmethod

194 def for_model(

195 cls,

196 model: str,

197 api_key: str | None = None,

198 provider: str | Provider | None = None,

199 use_dynamic_discovery: bool = True,

200 request_limit: int | None = None,

201 token_limit: int | None = None,

202 max_retries: int | None = None,

203 base_backoff: float | None = None,

204 timeout: float | None = None,

205 **kwargs: Any,

206 ) -> "ChatLimiter":

207 """

208 Create a ChatLimiter instance automatically detecting the provider from the model name.

209

210 Args:

211 model: The model name (e.g., "gpt-4o", "claude-3-sonnet-20240229")

212 api_key: API key for the provider. If None, will be read from environment variables

213 (OPENAI_API_KEY, ANTHROPIC_API_KEY, OPENROUTER_API_KEY)

214 provider: Override provider detection. Can be "openai", "anthropic", "openrouter",

215 or Provider enum. If None, will be auto-detected from model name

216 use_dynamic_discovery: Whether to query live APIs for model availability (default: True).

217 Requires appropriate API keys to be available. Falls back to

218 hardcoded model lists when disabled or when API calls fail.

219 **kwargs: Additional arguments passed to ChatLimiter

220

221 Returns:

222 Configured ChatLimiter instance

223

224 Raises:

225 ValueError: If provider cannot be determined from model name or API key not found

226

227 Example:

228 # Auto-detect provider with dynamic discovery (default behavior)

229 async with ChatLimiter.for_model("gpt-4o") as limiter:

230 response = await limiter.simple_chat("gpt-4o", "Hello!")

231

232 # Override provider detection

233 async with ChatLimiter.for_model("custom-model", provider="openai") as limiter:

234 response = await limiter.simple_chat("custom-model", "Hello!")

235

236 # Disable dynamic discovery to use only hardcoded model lists

237 async with ChatLimiter.for_model("gpt-4o", use_dynamic_discovery=False) as limiter:

238 response = await limiter.simple_chat("gpt-4o", "Hello!")

239 """

240 import os

241

242 # Determine provider

243 if provider is not None:

244 # Use provided provider

245 if isinstance(provider, str):

246 provider_enum = Provider(provider)

247 else:

248 provider_enum = provider

249 provider_name = provider_enum.value

250 else:

251 # Auto-detect from model name

252 # If dynamic discovery is requested, we need to collect API keys first

253 api_keys_for_discovery = {}

254 if use_dynamic_discovery:

255 # Collect available API keys from environment

256 env_var_map = {

257 "openai": "OPENAI_API_KEY",

258 "anthropic": "ANTHROPIC_API_KEY",

259 "openrouter": "OPENROUTER_API_KEY"

260 }

261

262 for provider_key, env_var in env_var_map.items():

263 key_value = os.getenv(env_var)

264 if key_value:

265 api_keys_for_discovery[provider_key] = key_value

266

267 # Try dynamic discovery first to get more detailed information

268 discovery_result = None

269 if use_dynamic_discovery and api_keys_for_discovery:

270 from .models import detect_provider_from_model_sync

271 discovery_result = detect_provider_from_model_sync(model, api_keys_for_discovery)

272 detected_provider = discovery_result.found_provider

273 else:

274 detected_provider = detect_provider_from_model(model, use_dynamic_discovery, api_keys_for_discovery)

275

276 if not detected_provider:

277 discovery_msg = " with dynamic API discovery" if use_dynamic_discovery else ""

278 error_msg = f"Could not determine provider from model '{model}'{discovery_msg}. "

279

280 # Add detailed information about available models if we have discovery results

281 if discovery_result and discovery_result.get_total_models_found() > 0:

282 error_msg += f"\n\nFound {discovery_result.get_total_models_found()} models across providers:\n"

283 for provider_name, models in discovery_result.get_all_models().items():

284 error_msg += f" {provider_name}: {len(models)} models\n"

285 for example in sorted(list(models)):

286 error_msg += f" - {example}\n"

287 error_msg += "\nPlease check the model name or specify the provider explicitly using the 'provider' parameter."

288 else:

289 error_msg += "Please specify the provider explicitly using the 'provider' parameter."

290

291 # Add information about discovery errors if any

292 if discovery_result and discovery_result.errors:

293 error_msg += f"\n\nDiscovery errors encountered:\n"

294 for provider_name, error in discovery_result.errors.items():

295 error_msg += f" {provider_name}: {error}\n"

296

297 raise ValueError(error_msg)

298 assert detected_provider is not None # Help MyPy understand type narrowing

299 provider_name = detected_provider

300 provider_enum = Provider(provider_name)

301

302 # Determine API key

303 if api_key is None:

304 # Try to get from environment variables

305 env_var_map = {

306 "openai": "OPENAI_API_KEY",

307 "anthropic": "ANTHROPIC_API_KEY",

308 "openrouter": "OPENROUTER_API_KEY"

309 }

310

311 env_var_name: str | None = env_var_map.get(provider_name)

312 if env_var_name:

313 api_key = os.getenv(env_var_name)

314 if not api_key:

315 raise ValueError(

316 f"API key not provided and {env_var_name} environment variable not set. "

317 f"Please provide api_key parameter or set {env_var_name} environment variable."

318 )

319 else:

320 raise ValueError(

321 f"Unknown provider '{provider_name}'. Cannot determine environment variable for API key."

322 )

323

324 return cls(

325 provider=provider_enum,

326 api_key=api_key,

327 request_limit=request_limit,

328 token_limit=token_limit,

329 max_retries=max_retries,

330 base_backoff=base_backoff,

331 timeout=timeout,

332 **kwargs

333 )

334

335 def _init_http_clients(

336 self,

337 http_client: httpx.AsyncClient | None,

338 sync_http_client: httpx.Client | None,

339 **kwargs: Any,

340 ) -> None:

341 """Initialize HTTP clients with proper headers."""

342 # Prepare headers

343 headers = {

344 "User-Agent": f"chat-limiter/0.1.0 ({self.provider.value})",

345 }

346

347 # Add provider-specific headers

348 if self.api_key:

349 if self.provider == Provider.OPENAI:

350 headers["Authorization"] = f"Bearer {self.api_key}"

351 elif self.provider == Provider.ANTHROPIC:

352 headers["x-api-key"] = self.api_key

353 headers["anthropic-version"] = "2023-06-01"

354 elif self.provider == Provider.OPENROUTER:

355 headers["Authorization"] = f"Bearer {self.api_key}"

356 headers["HTTP-Referer"] = "https://github.com/your-repo/chat-limiter"

357

358 # Merge with user-provided headers

359 if "headers" in kwargs:

360 headers.update(kwargs["headers"])

361 kwargs["headers"] = headers

362

363 # Initialize clients

364 if http_client:

365 self.async_client = http_client

366 else:

367 self.async_client = httpx.AsyncClient(

368 base_url=self.config.base_url,

369 timeout=httpx.Timeout(self._user_timeout), # Configurable timeout

370 **kwargs,

371 )

372

373 if sync_http_client:

374 self.sync_client = sync_http_client

375 else:

376 self.sync_client = httpx.Client(

377 base_url=self.config.base_url,

378 timeout=httpx.Timeout(self._user_timeout), # Configurable timeout

379 **kwargs,

380 )

381

382 def _init_rate_limiters(self) -> None:

383 """Initialize PyrateLimiter instances."""

384 # Only initialize if we have limits

385 if self.state.request_limit is None or self.state.token_limit is None:

386 # Cannot initialize rate limiters without limits

387 # This will be called again after limits are discovered

388 self.request_limiter = None

389 self.token_limiter = None

390 self._effective_request_limit = None

391 self._effective_token_limit = None

392 return

393

394 # Calculate effective limits with buffer

395 effective_request_limit = int(self.state.request_limit * self.config.request_buffer_ratio)

396 effective_token_limit = int(self.state.token_limit * self.config.token_buffer_ratio)

397

398 # Request rate limiter

399 self.request_limiter = Limiter(

400 Rate(

401 effective_request_limit,

402 Duration.MINUTE,

403 )

404 )

405

406 # Token rate limiter

407 self.token_limiter = Limiter(

408 Rate(

409 effective_token_limit,

410 Duration.MINUTE,

411 )

412 )

413

414 # Store effective limits for logging

415 self._effective_request_limit = effective_request_limit

416 self._effective_token_limit = effective_token_limit

417

418 async def __aenter__(self) -> "ChatLimiter":

419 """Async context manager entry."""

420 if self._async_context_active:

421 raise RuntimeError(

422 "ChatLimiter is already active as an async context manager"

423 )

424

425 self._async_context_active = True

426

427 # Discover rate limits if supported

428 if self.config.supports_dynamic_limits:

429 await self._discover_rate_limits()

430

431 # Print rate limit information if verbose mode is enabled

432 if self._verbose_mode:

433 self._print_rate_limit_info()

434

435 return self

436

437 async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object) -> None:

438 """Async context manager exit."""

439 self._async_context_active = False

440 await self.async_client.aclose()

441

442 def __enter__(self) -> "ChatLimiter":

443 """Sync context manager entry."""

444 if self._sync_context_active:

445 raise RuntimeError(

446 "ChatLimiter is already active as a sync context manager"

447 )

448

449 self._sync_context_active = True

450

451 # Discover rate limits if supported

452 if self.config.supports_dynamic_limits:

453 self._discover_rate_limits_sync()

454

455 # Print rate limit information if verbose mode is enabled

456 if self._verbose_mode:

457 self._print_rate_limit_info()

458

459 return self

460

461 def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object) -> None:

462 """Sync context manager exit."""

463 self._sync_context_active = False

464 self.sync_client.close()

465

466 async def _discover_rate_limits(self) -> None:

467 """Discover current rate limits from the API."""

468 try:

469 if self.provider == Provider.OPENROUTER and self.config.auth_endpoint:

470 # OpenRouter uses a special auth endpoint

471 response = await self.async_client.get(self.config.auth_endpoint)

472 response.raise_for_status()

473

474 data = response.json()

475 # Update limits based on response

476 # This is a simplified version - actual implementation would parse the response

477 logger.info(f"Discovered OpenRouter limits: {data}")

478

479 else:

480 # For other providers, we'll discover limits on first request

481 if self._verbose_mode:

482 print(f"Rate limit discovery will happen on first request for {self.provider.value}")

483 logger.info(

484 f"Rate limit discovery will happen on first request for {self.provider.value}"

485 )

486

487 except Exception as e:

488 logger.warning(f"Failed to discover rate limits: {e}")

489

490 def _discover_rate_limits_sync(self) -> None:

491 """Sync version of rate limit discovery."""

492 try:

493 if self.provider == Provider.OPENROUTER and self.config.auth_endpoint:

494 response = self.sync_client.get(self.config.auth_endpoint)

495 response.raise_for_status()

496

497 data = response.json()

498 logger.info(f"Discovered OpenRouter limits: {data}")

499 else:

500 logger.info(

501 f"Rate limit discovery will happen on first request for {self.provider.value}"

502 )

503

504 except Exception as e:

505 logger.warning(f"Failed to discover rate limits: {e}")

506

507 def _update_rate_limits(self, rate_limit_info: RateLimitInfo) -> None:

508 """Update rate limits based on response headers."""

509 updated = False

510 was_uninitialized = self.state.request_limit is None or self.state.token_limit is None

511

512 # Update request limits

513 if (

514 rate_limit_info.requests_limit

515 and rate_limit_info.requests_limit != self.state.request_limit

516 ):

517 old_limit = self.state.request_limit

518 self.state.request_limit = rate_limit_info.requests_limit

519 updated = True

520 if was_uninitialized:

521 message = f"Discovered request limit: {self.state.request_limit} req/min"

522 if self._verbose_mode:

523 print(message)

524 logger.info(message)

525 else:

526 message = f"Updated request limit: {old_limit} -> {self.state.request_limit} req/min"

527 if self._verbose_mode:

528 print(message)

529 logger.info(message)

530

531 # Update token limits

532 if (

533 rate_limit_info.tokens_limit

534 and rate_limit_info.tokens_limit != self.state.token_limit

535 ):

536 old_limit = self.state.token_limit

537 self.state.token_limit = rate_limit_info.tokens_limit

538 updated = True

539 if was_uninitialized:

540 message = f"Discovered token limit: {self.state.token_limit} tokens/min"

541 if self._verbose_mode:

542 print(message)

543 logger.info(message)

544 else:

545 message = f"Updated token limit: {old_limit} -> {self.state.token_limit} tokens/min"

546 if self._verbose_mode:

547 print(message)

548 logger.info(message)

549

550 if updated:

551 # Reinitialize rate limiters with new limits

552 self._init_rate_limiters()

553

554 # Update limits_discovered flag if both limits are now available

555 if self.state.request_limit is not None and self.state.token_limit is not None:

556 self._limits_discovered = True

557

558 if was_uninitialized:

559 message = "Rate limiters initialized after discovery"

560 if self._verbose_mode:

561 print(message)

562 # Print updated rate limit info after discovery

563 self._print_rate_limit_info()

564 logger.info(message)

565

566 # Store the rate limit info

567 self.state.last_rate_limit_info = rate_limit_info

568 self.state.last_limit_update = time.time()

569

570 def _estimate_tokens(self, request_data: dict[str, Any]) -> int:

571 """Estimate token usage from request data."""

572 if not self.enable_token_estimation:

573 return 0

574

575 # Simple token estimation

576 # This is a placeholder - real implementation would use tiktoken or similar

577 if "messages" in request_data:

578 text = ""

579 for message in request_data["messages"]:

580 if isinstance(message, dict) and "content" in message:

581 text += str(message["content"])

582

583 # Rough estimation: 1 token ≈ 4 characters

584 return len(text) // 4

585

586 return 0

587

588 @asynccontextmanager

589 async def _acquire_rate_limits(

590 self, estimated_tokens: int = 0

591 ) -> AsyncIterator[None]:

592 """Acquire rate limits before making a request."""

593 # Check if rate limiters are initialized

594 if self.request_limiter is None or self.token_limiter is None:

595 # Limits not yet discovered - this request will help discover them

596 logger.info("Rate limits not yet discovered, proceeding without rate limiting for discovery")

597 else:

598 # Wait for request rate limit

599 await asyncio.to_thread(self.request_limiter.try_acquire, "request")

600

601 # Wait for token rate limit if we have token estimation and limiters are initialized

602 if estimated_tokens > 0 and self.token_limiter is not None and self._effective_token_limit is not None:

603 # Check if request is too large for bucket capacity

604 if estimated_tokens > self._effective_token_limit:

605 # Log warning for large requests

606 logger.warning(

607 f"Request estimated at {estimated_tokens} tokens exceeds bucket capacity "

608 f"of {self._effective_token_limit} tokens. This may cause delays."

609 )

610 # For very large requests, we'll split the acquisition

611 # Acquire tokens in chunks to avoid bucket overflow

612 remaining_tokens = estimated_tokens

613 while remaining_tokens > 0:

614 chunk_size = min(remaining_tokens, self._effective_token_limit // 2)

615 await asyncio.to_thread(self.token_limiter.try_acquire, "token", chunk_size)

616 remaining_tokens -= chunk_size

617 if remaining_tokens > 0:

618 # Brief pause to let bucket refill

619 await asyncio.sleep(0.1)

620 else:

621 # Normal acquisition for smaller requests

622 await asyncio.to_thread(self.token_limiter.try_acquire, "token", estimated_tokens)

623

624 try:

625 yield

626 finally:

627 # Update usage tracking

628 self.state.requests_used += 1

629 self.state.tokens_used += estimated_tokens

630 self.state.last_request_time = time.time()

631

632 @contextmanager

633 def _acquire_rate_limits_sync(self, estimated_tokens: int = 0) -> Iterator[None]:

634 """Sync version of rate limit acquisition."""

635 # Check if rate limiters are initialized

636 if self.request_limiter is None or self.token_limiter is None:

637 # Limits not yet discovered - this request will help discover them

638 logger.info("Rate limits not yet discovered, proceeding without rate limiting for discovery")

639 else:

640 # Wait for request rate limit

641 self.request_limiter.try_acquire("request")

642

643 # Wait for token rate limit if we have token estimation and limiters are initialized

644 if estimated_tokens > 0 and self.token_limiter is not None and self._effective_token_limit is not None:

645 # Check if request is too large for bucket capacity

646 if estimated_tokens > self._effective_token_limit:

647 # Log warning for large requests

648 logger.warning(

649 f"Request estimated at {estimated_tokens} tokens exceeds bucket capacity "

650 f"of {self._effective_token_limit} tokens. This may cause delays."

651 )

652 # For very large requests, we'll split the acquisition

653 # Acquire tokens in chunks to avoid bucket overflow

654 remaining_tokens = estimated_tokens

655 while remaining_tokens > 0:

656 chunk_size = min(remaining_tokens, self._effective_token_limit // 2)

657 self.token_limiter.try_acquire("token", chunk_size)

658 remaining_tokens -= chunk_size

659 if remaining_tokens > 0:

660 # Brief pause to let bucket refill

661 time.sleep(0.1)

662 else:

663 # Normal acquisition for smaller requests

664 self.token_limiter.try_acquire("token", estimated_tokens)

665

666 try:

667 yield

668 finally:

669 # Update usage tracking

670 self.state.requests_used += 1

671 self.state.tokens_used += estimated_tokens

672 self.state.last_request_time = time.time()

673

674 def _get_retry_decorator(self):

675 """Get retry decorator with user-configured parameters."""

676 return retry(

677 stop=stop_after_attempt(self._user_max_retries),

678 wait=wait_exponential(multiplier=self._user_base_backoff, min=1, max=60),

679 retry=retry_if_exception_type((httpx.HTTPStatusError, httpx.RequestError, httpx.ReadTimeout, httpx.ConnectTimeout)),

680 )

681

682 async def request(

683 self,

684 method: str,

685 url: str,

686 *,

687 json: dict[str, Any] | None = None,

688 **kwargs: Any,

689 ) -> httpx.Response:

690 """Wrapper that applies retry decorator dynamically."""

691 try:

692 return await self._get_retry_decorator()(self._request_impl)(method, url, json=json, **kwargs)

693 except Exception as e:

694 # Check if this is a retry error wrapping a timeout

695 if hasattr(e, 'last_attempt') and e.last_attempt and e.last_attempt.exception():

696 original_exception = e.last_attempt.exception()

697 if isinstance(original_exception, (httpx.ReadTimeout, httpx.ConnectTimeout)):

698 # Enhance timeout error with helpful information

699 timeout_info = (

700 f"\n💡 Timeout Error Help:\n"

701 f" Current timeout: {self._user_timeout}s\n"

702 f" To increase timeout, use: ChatLimiter.for_model('{self.provider.value}', timeout={int(self._user_timeout + 60)})\n"

703 f" Or reduce batch concurrency if processing multiple requests\n"

704 f" Retries attempted: {self._user_max_retries}\n"

705 )

706 raise type(original_exception)(str(original_exception) + timeout_info) from e

707

708 # For direct timeout errors (shouldn't happen due to retry decorator but just in case)

709 if isinstance(e, (httpx.ReadTimeout, httpx.ConnectTimeout)):

710 timeout_info = (

711 f"\n💡 Timeout Error Help:\n"

712 f" Current timeout: {self._user_timeout}s\n"

713 f" To increase timeout, use: ChatLimiter.for_model('{self.provider.value}', timeout={int(self._user_timeout + 60)})\n"

714 f" Or reduce batch concurrency if processing multiple requests\n"

715 )

716 raise type(e)(str(e) + timeout_info) from e

717

718 # Re-raise any other exceptions unchanged

719 raise

720

721 async def _request_impl(

722 self,

723 method: str,

724 url: str,

725 *,

726 json: dict[str, Any] | None = None,

727 **kwargs: Any,

728 ) -> httpx.Response:

729 """

730 Make an async HTTP request with rate limiting.

731

732 Args:

733 method: HTTP method (GET, POST, etc.)

734 url: URL or path for the request

735 json: JSON data to send

736 **kwargs: Additional arguments passed to httpx

737

738 Returns:

739 HTTP response

740

741 Raises:

742 httpx.HTTPStatusError: For HTTP error responses

743 httpx.RequestError: For request errors

744 """

745 if not self._async_context_active:

746 raise RuntimeError("ChatLimiter must be used as an async context manager")

747

748 # Estimate tokens if we have JSON data

749 estimated_tokens = self._estimate_tokens(json or {})

750

751 # Acquire rate limits

752 async with self._acquire_rate_limits(estimated_tokens):

753 # Make the request

754 response = await self.async_client.request(method, url, json=json, **kwargs)

755

756 # Extract rate limit info

757 rate_limit_info = extract_rate_limit_info(

758 dict(response.headers), self.config

759 )

760

761 # Update our rate limits

762 if self.enable_adaptive_limits:

763 self._update_rate_limits(rate_limit_info)

764

765 # Handle rate limit errors

766 if response.status_code == 429:

767 self.state.consecutive_rate_limit_errors += 1

768 if rate_limit_info.retry_after:

769 await asyncio.sleep(rate_limit_info.retry_after)

770 else:

771 # Exponential backoff

772 backoff = self.config.base_backoff * (

773 2**self.state.consecutive_rate_limit_errors

774 )

775 await asyncio.sleep(min(backoff, self.config.max_backoff))

776

777 response.raise_for_status()

778 else:

779 # Reset consecutive errors on success

780 self.state.consecutive_rate_limit_errors = 0

781

782 return response

783

784 def request_sync(

785 self,

786 method: str,

787 url: str,

788 *,

789 json: dict[str, Any] | None = None,

790 **kwargs: Any,

791 ) -> httpx.Response:

792 """Wrapper that applies retry decorator dynamically."""

793 # For sync, we need to use the sync version of retry

794 retry_decorator = retry(

795 stop=stop_after_attempt(self._user_max_retries),

796 wait=wait_exponential(multiplier=self._user_base_backoff, min=1, max=60),

797 retry=retry_if_exception_type((httpx.HTTPStatusError, httpx.RequestError, httpx.ReadTimeout, httpx.ConnectTimeout)),

798 )

799 try:

800 return retry_decorator(self._request_sync_impl)(method, url, json=json, **kwargs)

801 except (httpx.ReadTimeout, httpx.ConnectTimeout) as e:

802 # Enhance timeout error with helpful information

803 timeout_info = (

804 f"\n💡 Timeout Error Help:\n"

805 f" Current timeout: {self._user_timeout}s\n"

806 f" To increase timeout, use: ChatLimiter.for_model('{self.provider.value}', timeout={int(self._user_timeout + 60)})\n"

807 f" Or reduce batch concurrency if processing multiple requests\n"

808 )

809 raise type(e)(str(e) + timeout_info) from e

810

811 def _request_sync_impl(

812 self,

813 method: str,

814 url: str,

815 *,

816 json: dict[str, Any] | None = None,

817 **kwargs: Any,

818 ) -> httpx.Response:

819 """

820 Make a sync HTTP request with rate limiting.

821

822 Args:

823 method: HTTP method (GET, POST, etc.)

824 url: URL or path for the request

825 json: JSON data to send

826 **kwargs: Additional arguments passed to httpx

827

828 Returns:

829 HTTP response

830

831 Raises:

832 httpx.HTTPStatusError: For HTTP error responses

833 httpx.RequestError: For request errors

834 """

835 if not self._sync_context_active:

836 raise RuntimeError("ChatLimiter must be used as a sync context manager")

837

838 # Estimate tokens if we have JSON data

839 estimated_tokens = self._estimate_tokens(json or {})

840

841 # Acquire rate limits

842 with self._acquire_rate_limits_sync(estimated_tokens):

843 # Make the request

844 response = self.sync_client.request(method, url, json=json, **kwargs)

845

846 # Extract rate limit info

847 rate_limit_info = extract_rate_limit_info(

848 dict(response.headers), self.config

849 )

850

851 # Update our rate limits

852 if self.enable_adaptive_limits:

853 self._update_rate_limits(rate_limit_info)

854

855 # Handle rate limit errors

856 if response.status_code == 429:

857 self.state.consecutive_rate_limit_errors += 1

858 if rate_limit_info.retry_after:

859 time.sleep(rate_limit_info.retry_after)

860 else:

861 # Exponential backoff

862 backoff = self.config.base_backoff * (

863 2**self.state.consecutive_rate_limit_errors

864 )

865 time.sleep(min(backoff, self.config.max_backoff))

866

867 response.raise_for_status()

868 else:

869 # Reset consecutive errors on success

870 self.state.consecutive_rate_limit_errors = 0

871

872 return response

873

874 def get_current_limits(self) -> dict[str, Any]:

875 """Get current rate limit information."""

876 return {

877 "provider": self.provider.value,

878 "request_limit": self.state.request_limit,

879 "token_limit": self.state.token_limit,

880 "requests_used": self.state.requests_used,

881 "tokens_used": self.state.tokens_used,

882 "last_request_time": self.state.last_request_time,

883 "last_limit_update": self.state.last_limit_update,

884 "consecutive_rate_limit_errors": self.state.consecutive_rate_limit_errors,

885 }

886

887 def reset_usage_tracking(self) -> None:

888 """Reset usage tracking counters."""

889 self.state.requests_used = 0

890 self.state.tokens_used = 0

891 self.state.consecutive_rate_limit_errors = 0

892

893 # High-level chat completion methods

894

895 async def chat_completion(

896 self,

897 model: str,

898 messages: list[Message],

899 max_tokens: int | None = None,

900 temperature: float | None = None,

901 top_p: float | None = None,

902 stop: str | list[str] | None = None,

903 stream: bool = False,

904 **kwargs: Any,

905 ) -> ChatCompletionResponse:

906 """

907 Make a high-level chat completion request.

908

909 Args:

910 model: The model to use for completion

911 messages: List of messages in the conversation

912 max_tokens: Maximum tokens to generate

913 temperature: Sampling temperature

914 top_p: Top-p sampling parameter

915 stop: Stop sequences

916 stream: Whether to stream the response

917 **kwargs: Additional provider-specific parameters

918

919 Returns:

920 ChatCompletionResponse with the completion result

921

922 Raises:

923 ValueError: If provider cannot be determined from model

924 httpx.HTTPStatusError: For HTTP error responses

925 httpx.RequestError: For request errors

926 """

927 if not self._async_context_active:

928 raise RuntimeError("ChatLimiter must be used as an async context manager")

929

930 # Create request object

931 request = ChatCompletionRequest(

932 model=model,

933 messages=messages,

934 max_tokens=max_tokens,

935 temperature=temperature,

936 top_p=top_p,

937 stop=stop,

938 stream=stream,

939 **kwargs

940 )

941

942 # Get the appropriate adapter

943 adapter = get_adapter(self.provider)

944

945 # Format the request for the provider

946 formatted_request = adapter.format_request(request)

947

948 # Make the HTTP request

949 response = await self.request(

950 "POST",

951 adapter.get_endpoint(),

952 json=formatted_request

953 )

954

955 # Parse the response

956 response_data = response.json()

957 return adapter.parse_response(response_data, request)

958

959 def chat_completion_sync(

960 self,

961 model: str,

962 messages: list[Message],

963 max_tokens: int | None = None,

964 temperature: float | None = None,

965 top_p: float | None = None,

966 stop: str | list[str] | None = None,

967 stream: bool = False,

968 **kwargs: Any,

969 ) -> ChatCompletionResponse:

970 """

971 Make a synchronous high-level chat completion request.

972

973 Args:

974 model: The model to use for completion

975 messages: List of messages in the conversation

976 max_tokens: Maximum tokens to generate

977 temperature: Sampling temperature

978 top_p: Top-p sampling parameter

979 stop: Stop sequences

980 stream: Whether to stream the response

981 **kwargs: Additional provider-specific parameters

982

983 Returns:

984 ChatCompletionResponse with the completion result

985

986 Raises:

987 ValueError: If provider cannot be determined from model

988 httpx.HTTPStatusError: For HTTP error responses

989 httpx.RequestError: For request errors

990 """

991 if not self._sync_context_active:

992 raise RuntimeError("ChatLimiter must be used as a sync context manager")

993

994 # Create request object

995 request = ChatCompletionRequest(

996 model=model,

997 messages=messages,

998 max_tokens=max_tokens,

999 temperature=temperature,

1000 top_p=top_p,

1001 stop=stop,

1002 stream=stream,

1003 **kwargs

1004 )

1005

1006 # Get the appropriate adapter

1007 adapter = get_adapter(self.provider)

1008

1009 # Format the request for the provider

1010 formatted_request = adapter.format_request(request)

1011

1012 # Make the HTTP request

1013 response = self.request_sync(

1014 "POST",

1015 adapter.get_endpoint(),

1016 json=formatted_request

1017 )

1018

1019 # Parse the response

1020 response_data = response.json()

1021 return adapter.parse_response(response_data, request)

1022

1023 # Convenience methods for different message types

1024

1025 async def simple_chat(

1026 self,

1027 model: str,

1028 prompt: str,

1029 max_tokens: int | None = None,

1030 temperature: float | None = None,

1031 **kwargs: Any,

1032 ) -> str:

1033 """

1034 Simple chat completion that returns just the text response.

1035

1036 Args:

1037 model: The model to use

1038 prompt: The user prompt

1039 max_tokens: Maximum tokens to generate

1040 temperature: Sampling temperature

1041 **kwargs: Additional parameters

1042

1043 Returns:

1044 The text response from the model

1045 """

1046 messages = [Message(role=MessageRole.USER, content=prompt)]

1047 response = await self.chat_completion(

1048 model=model,

1049 messages=messages,

1050 max_tokens=max_tokens,

1051 temperature=temperature,

1052 **kwargs

1053 )

1054

1055 if response.choices:

1056 return response.choices[0].message.content

1057 return ""

1058

1059 def simple_chat_sync(

1060 self,

1061 model: str,

1062 prompt: str,

1063 max_tokens: int | None = None,

1064 temperature: float | None = None,

1065 **kwargs: Any,

1066 ) -> str:

1067 """

1068 Simple synchronous chat completion that returns just the text response.

1069

1070 Args:

1071 model: The model to use

1072 prompt: The user prompt

1073 max_tokens: Maximum tokens to generate

1074 temperature: Sampling temperature

1075 **kwargs: Additional parameters

1076

1077 Returns:

1078 The text response from the model

1079 """

1080 messages = [Message(role=MessageRole.USER, content=prompt)]

1081 response = self.chat_completion_sync(

1082 model=model,

1083 messages=messages,

1084 max_tokens=max_tokens,

1085 temperature=temperature,

1086 **kwargs

1087 )

1088

1089 if response.choices:

1090 return response.choices[0].message.content

1091 return ""

1092

1093 def set_verbose_mode(self, verbose: bool) -> None:

1094 """Set verbose mode for detailed logging."""

1095 self._verbose_mode = verbose

1096

1097 def _print_rate_limit_info(self) -> None:

1098 """Print current rate limit configuration."""

1099 print(f"\n=== Rate Limit Configuration for {self.provider.value.title()} ===")

1100 print(f"Provider: {self.provider.value}")

1101 print(f"Base URL: {self.config.base_url}")

1102

1103 # Handle None values for limits

1104 if self.state.request_limit is not None:

1105 effective_req = self._effective_request_limit or "not calculated"

1106 print(f"Request Limit: {self.state.request_limit}/minute (effective: {effective_req}/minute)")

1107 else:

1108 print("Request Limit: Not yet discovered (will be fetched from API)")

1109

1110 if self.state.token_limit is not None:

1111 effective_tok = self._effective_token_limit or "not calculated"

1112 print(f"Token Limit: {self.state.token_limit}/minute (effective: {effective_tok}/minute)")