feat: Add batch API support for entity extraction (Gemini + OpenAI) - #2901
Closed
acorretti wants to merge 14 commits into
Closed
feat: Add batch API support for entity extraction (Gemini + OpenAI)#2901acorretti wants to merge 14 commits into
acorretti wants to merge 14 commits into
Conversation
Introduce a provider-agnostic BatchProvider ABC for submitting LLM completion requests as batch jobs, with data types for request/response handling and job lifecycle management. Add GeminiBatchProvider that implements the abstraction using Google's Gemini Batch Prediction API (client.aio.batches), with inlined request submission, polling, and result retrieval.
Add batch_provider, llm_batch_timeout, and llm_batch_poll_interval fields for opt-in batch API ingestion support. These flow through global_config to the extraction pipeline.
When batch_provider is set in global_config, extract_entities() takes an alternate path that pre-generates all prompts, checks the LLM cache, submits cache misses as a single batch via BatchProvider, and optionally runs a second batch for gleaning. Failed rows fall back to the live API. Includes _merge_gleaning_results() and _await_batch_with_cancellation() helpers shared between batch and live paths.
Add gemini_batch as a recognized LLM/embedding binding that reuses Gemini's live API functions for queries while additionally creating a GeminiBatchProvider for batch entity extraction during ingestion.
Users can now select gemini_batch as an LLM or embedding provider in the interactive setup flow. It shares Gemini's configuration prompts (API key, endpoint, model) and default values.
Document the LLM_BINDING=gemini_batch option with LLM_BATCH_TIMEOUT and LLM_BATCH_POLL_INTERVAL settings.
Store {job_id, submitted_keys} in llm_response_cache after submitting
a batch. On re-entry, if a persisted job exists and the chunk keys
match the current cache misses, resume polling instead of resubmitting.
Clear persisted state once the job reaches a terminal state.
OpenAIBatchProvider implements BatchProvider using OpenAI's Batch API (JSONL file upload -> batch creation -> polling -> results download). Wire openai_batch as a selectable binding through config, server, setup wizard, and env.example.
When a batch fails with a token limit error (detected via error_code from the provider), automatically split the unresolved requests in half and resubmit as smaller sub-batches. This handles both submission- time rejections and post-submission validation failures (like OpenAI's enqueued token limit). Persistence updated to store multiple job IDs.
19 offline tests covering: - _merge_gleaning_results: new/existing entities, description length comparison, edges, None handling - Batch extraction: all-cached, all-miss, mixed cache, per-row failure fallback, entire batch failure fallback, gleaning as second batch - Auto-split on token limit errors (recursive halving) - Cache key compatibility between batch and live paths - Job persistence: save after submission, resume on matching keys, discard on mismatched keys - Pipeline cancellation cancels provider-side job
The split condition checked status.total > 1, but OpenAI doesn't populate request_counts when rejecting a batch for token limits (total=0). Removed the total check and instead guard against infinite splitting by checking len(unresolved) > 1. Updated test to match real OpenAI behavior (total=0 on rejection).
When a sub-batch failed with token_limit_exceeded, the split logic scanned all cache_misses instead of just the failed job's requests, causing duplicate sub-batches. Now each job ID maps to its specific requests via job_requests dict, so splitting only affects the failed job's requests.
…ting The "Enqueued token limit reached" error means the org-wide batch queue is full, not that the individual batch is too large. Splitting only created more queued batches, worsening the problem. Now waits for the queue to drain (2x poll interval) then resubmits the same batch unchanged.
Author
|
Hi @danielaskdd — thanks for triaging this last week. No rush, I can see you're busy. No pressure on timing. Thanks for maintaining LightRAG. |
Collaborator
|
Thank you for your interest and contributions to LightRAG. Given the significant changes introduced in version 1.5, we will review this PR following the official release of v1.5. |
|
@acorretti, dude you are a hero! I really needed this feature, and randomly stumbling on this PR was a great surprise. Great work! |
|
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds batch API support for LightRAG's entity extraction pipeline. Depending on the provider, this may lead to reduced ingestion costs by submitting all extraction prompts as a single batch job instead of individual concurrent calls. The tradeoff is higher latency — batch jobs may queue for minutes to hours before processing starts.
Two new LLM bindings --
gemini_batchandopenai_batch-- use their respective provider batch APIs for entity extraction during ingestion while keeping queries on the live API for real-time responses.Architecture
batch_provider.py): Provider-agnostic ABC withsubmit_completion_batch(),get_job_status(),get_results(),cancel_job(), and a default polling implementationgemini_batch.py): Usesclient.aio.batches.create()with inlined requestsopenai_batch.py): Uses JSONL file upload -> Files API -> Batch API -> results downloadoperate.py): Bypasses the per-chunk semaphore, pre-generates all prompts, checks the LLM cache, submits cache misses as a batch, and handles gleaning as a second batch roundKey Features
llm_response_cache. If the server restarts mid-batch, re-ingesting the same document resumes polling instead of resubmittingRelated Issues
Closes #2210
Changes Made
New Files
lightrag/llm/batch_provider.py-- BatchProvider ABC + BatchRequest/BatchResponse/BatchJobStatus data typeslightrag/llm/gemini_batch.py-- GeminiBatchProvider implementationlightrag/llm/openai_batch.py-- OpenAIBatchProvider implementationtests/test_batch_extraction.py-- 19 offline unit tests (merge logic, cache compatibility, persistence, auto-split, cancellation)docs/BatchAPIImplementationPlan.md-- Design document with architecture, retry strategies, and edge casesModified Files
lightrag/operate.py-- Batch extraction path, auto-split logic, job persistence, polling helperslightrag/lightrag.py--batch_provider,llm_batch_timeout,llm_batch_poll_intervalconfig fields;_build_global_config()helper for safeasdict()handlinglightrag/api/lightrag_server.py--gemini_batchandopenai_batchbinding wiring, validation, provider instantiationlightrag/api/config.py-- Batch binding arg registration and default hostsscripts/setup/setup.sh-- Batch bindings in setup wizardenv.example-- Batch configuration examplesConfiguration
Checklist
Additional Notes
EMBEDDING_BINDINGis independent ofLLM_BINDING-- set it to match your existing vector storeBatchProviderabstraction is designed for extensibility -- adding new providers requires implementing 4 methods