LangChain-compatible chat model wrapper with transparent key switching. Each :meth:`invoke` / :meth:`ainvoke` call acquires a key from the pool, builds a :class:`~langchain_openai.ChatOpenAI` instance on the fly, and releases the key when done. On rate-limit errors the wrapper releases
| 395 | |
| 396 | |
| 397 | class PooledChatModel: |
| 398 | """LangChain-compatible chat model wrapper with transparent key switching. |
| 399 | |
| 400 | Each :meth:`invoke` / :meth:`ainvoke` call acquires a key from the pool, |
| 401 | builds a :class:`~langchain_openai.ChatOpenAI` instance on the fly, and |
| 402 | releases the key when done. On rate-limit errors the wrapper releases |
| 403 | the key with ``success=False``, picks a different key, and retries. |
| 404 | |
| 405 | Parameters |
| 406 | ---------- |
| 407 | pool : |
| 408 | An :class:`ApiKeyPool` with at least one configured key. |
| 409 | max_tokens : |
| 410 | ``max_completion_tokens`` passed to each ``ChatOpenAI`` instance. |
| 411 | timeout : |
| 412 | Request timeout in seconds passed to each ``ChatOpenAI`` instance. |
| 413 | max_retries : |
| 414 | Maximum number of key-switch retries on rate-limit errors before |
| 415 | giving up. |
| 416 | """ |
| 417 | |
| 418 | def __init__( |
| 419 | self, |
| 420 | pool: ApiKeyPool, |
| 421 | *, |
| 422 | max_tokens: int = 4096, |
| 423 | timeout: float = 30.0, |
| 424 | max_retries: int = _MAX_RATE_LIMIT_RETRIES, |
| 425 | ) -> None: |
| 426 | self._pool = pool |
| 427 | self._max_tokens = max_tokens |
| 428 | self._timeout = timeout |
| 429 | self._max_retries = max_retries |
| 430 | |
| 431 | # -- Public API ----------------------------------------------------------- |
| 432 | |
| 433 | def invoke(self, prompt: str) -> object: |
| 434 | """Synchronous invoke with automatic key switching on rate-limit.""" |
| 435 | return self._invoke_with_retry(prompt) |
| 436 | |
| 437 | async def ainvoke(self, prompt: str) -> object: |
| 438 | """Async invoke with automatic key switching on rate-limit.""" |
| 439 | return await self._ainvoke_with_retry(prompt) |
| 440 | |
| 441 | # -- Internal ------------------------------------------------------------- |
| 442 | |
| 443 | def _invoke_with_retry(self, prompt: str) -> object: |
| 444 | """Sync retry loop — acquire slot, call LLM, release, retry on 429.""" |
| 445 | last_exception: Exception | None = None |
| 446 | |
| 447 | for attempt in range(self._max_retries + 1): |
| 448 | key = self._pool.acquire() |
| 449 | llm = self._build_llm(key) |
| 450 | try: |
| 451 | result = llm.invoke(prompt) |
| 452 | self._pool.release(key, success=True) |
| 453 | if attempt > 0: |
| 454 | self._pool.record_retry_success() |
no outgoing calls