tannnmayy opened a new pull request, #40914: URL: https://github.com/apache/beam/pull/40914
**Title:** [RunInference] Send one Gemini request per input so every element gets a result **Commit message:** [RunInference] Send one Gemini request per input so every element gets a result **Body:** Fixes #40913 `GeminiModelHandler` sent a whole batch of prompts in one `generate_content` call. google-genai turns a list of strings into a single user message, so Gemini returned one response for the batch, and `request()` wrapped it in a one-item list, so only the first element of each batch got a `PredictionResult`. The other elements were dropped without an error, and the one result answered the merged message. This PR: - sends one `generate_content` request per input in `generate_from_string` and `generate_image_from_strings_and_images`, dispatched on a bounded thread pool (at most 32 in flight per batch), following the `OpenAIModelHandler` approach from #40239; - makes `request()` check that the request function returned one response per input and raise a `ValueError` otherwise, so a mismatch can't silently drop elements again. A single response is still accepted for a single-element batch; - documents that a custom `request_fn` must return one response per input, in order, and adds a CHANGES.md entry. Testing: - New tests in `gemini_inference_test.py`: one request per prompt, concurrency with ordering, the concurrency cap, error propagation, the image request function, the response-count checks, and a pipeline test with batches of 5. 9 of the 10 new tests fail without the fix. - `gemini_inference_test.py` passes on Linux (Python 3.14) and Windows (Python 3.11): 21 tests. ruff, yapf and pyrefly report no issues. - Real Gemini API (`gemini-2.5-flash`), batches forced to 3: before the fix, 3 prompts gave 1 result (`4 Paris Blue`); with the fix, 6 prompts gave 6 correct results. The `gemini_text_classification.py` example still works. Open questions from the issue: raise vs. warn on a count mismatch, sharing the fan-out helper with `openai_inference.py`, and a default `max_batch_size`. R: @jrmccluskey -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
