Add optional client-side rate limiting (max_rpm) and preserve partial results on parallel failure

[ojassharma7]

Description

Addresses the two still-open items of problem #2 in #358 (Gemini 429 RESOURCE_EXHAUSTED handling). The exponential backoff that issue also asked for already landed in 3aab86c, so this PR covers what remains:

• max_rpm — client-side throttle. A new optional GeminiLanguageModel parameter. When > 0, real-time requests are spaced by a thread-safe leaky bucket so at most max_rpm start per minute (shared across max_workers). This prevents 429 RESOURCE_EXHAUSTED on quota-limited tiers (e.g. the free tier's 15 RPM) rather than only retrying after the fact. Defaults to 0 (disabled) — no behavior change for existing callers. Does not apply to the Batch API.
• Preserve completed chunks on failure. The parallel infer() path raised on the first failed chunk, which abandoned the still-running futures and discarded every chunk that had already succeeded. It now lets the pool drain, then raises once with the completed work attached to InferenceRuntimeError (partial_results, aligned to the input prompts, plus failed_indices). The all-or-nothing raise is preserved, so annotation's zip(batch, outputs) alignment is unaffected; callers who want to recover already-done work can now read it off the exception.

This is a partial fix for #358 (chunking for very large documents is a separate design question, not touched here), so:

Related to #358

Type: Feature (the max_rpm throttle; the partial-result change also fixes a real loss-of-work bug).

How Has This Been Tested?

Unit tests in tests/gemini_ratelimit_partial_test.py cover the limiter's spacing (mocked clock), max_rpm wiring + validation, and partial-result preservation across the all-success, single-failure, and multi-failure cases.

$ python -m pytest tests/gemini_ratelimit_partial_test.py tests/gemini_retry_test.py

Differential-clean against the full suite: the only failures on my machine are pre-existing and environmental (the optional openai extra isn't installed, plus 3 plugin-packaging tests) — identical with and without this change — and the 10 new tests pass on top. pyink + isort + pylint --rcfile=.pylintrc / --rcfile=tests/.pylintrc (CI's exact configs) all clean.

Checklist:

• I have read and acknowledged Google's Open Source Code of conduct.
• I have read the Contributing page and am covered by the Google Individual CLA (signed for Forward Gemini generation settings #509).
• I have discussed my proposed solution with code owners in the linked issue(s) and we have agreed upon the general approach. — I flagged intent on Challenges extracting structured data from large Brazilian fund regulation PDFs (230K+ chars, 22 entity types) #358, and both items come straight from the issue's problem Do I need to set up an API URL to run langextract using a local model? #2; opening this as a concrete, tested proposal and happy to adjust the approach or split it if you'd prefer.
• I have made any needed documentation changes (the max_rpm constructor docstring; the retry-family params aren't documented outside the docstring either).
• I have added tests covering the changes.
• I have followed Google's Python Style Guide and ran pylint over the affected code.

[github-actions]

⚠️ Branch Update Required

Your branch is 1 commits behind main. Please update your branch to ensure CI checks run with the latest code:

git fetch origin main
git merge origin/main
git push

Note: Enable "Allow edits by maintainers" to allow automatic updates.

[ojassharma7]

Rebased onto main to clear the conflict with #534.

Worth flagging what the conflict actually was, since it is not visible in the diff. Textually it was just two additions landing in the same place (_no_text_diagnostic from #534, _RateLimiter here), and keeping both resolves cleanly. But that resolution quietly disabled this branch: #534 added an except InferenceRuntimeError: raise pass-through to the parallel result loop, and it runs before the draining logic here. Once _process_single_prompt wraps a failure, the first failing chunk short-circuits and the completed chunks are discarded again, which is the behaviour this branch exists to fix.

Failures are now collected uniformly and the pool is allowed to drain. What gets raised then depends on whether anything survived:

• Nothing completed: the provider's own InferenceRuntimeError is re-raised unchanged, so a blocked prompt or an exhausted token budget still reports its own reason instead of being buried a level down. The drained bookkeeping is attached to it.
• Some chunks completed: the aggregate error carries partial_results and failed_indices, keeps the first underlying error as original, and quotes it in the message.

That keeps both behaviours: the no-text diagnostic from #534 propagates intact when it is the whole story, and completed work survives when there is some.

Verification on the rebased tree:

python -m pytest tests/gemini_ratelimit_partial_test.py tests/gemini_retry_test.py -m "not live_api"
100 passed

The full suite is differential-clean against main: the only failures are the 3 in tests/create_provider_plugin_test.py, and those reproduce identically on unmodified main in the same environment. isort, pyink and pylint --rcfile=.pylintrc are clean on the changed files.

[github-actions]

⚠️ Branch Update Required

Your branch is 3 commits behind main. Please update your branch to ensure CI checks run with the latest code:

git fetch origin main
git merge origin/main
git push

Note: Enable "Allow edits by maintainers" to allow automatic updates.

[aksg87]

Thanks for putting this together. Gemini already has configurable retries with backoff for transient API errors, and custom providers can handle more specialized retry and rate-limiting policies.

I’m closing this PR for now and holding off on adding rate-limiting and partial-result recovery APIs. If a specific use case shows where the existing options fall short, feel free to open a focused issue. Happy to revisit and reopen this if needed.