apache/beam · error · IllegalStateException

Number of responses must match number of inputs

Error message

Number of responses must match number of inputs

What it means

GeminiModelHandler.request() asserts that the request function returned exactly one response per input element. If the GeminiRequestFunction produces a response list whose size differs from the input list, the batch invariant is broken and the handler fails with this IllegalStateException rather than emitting misaligned PredictionResults.

Solutions

  1. Fix the custom GeminiRequestFunction so it returns exactly one OutputT per InputT in the same order.
  2. If inputs may be skipped, keep list alignment by returning a placeholder/error OutputT for failed inputs instead of omitting them.
  3. Add a size assertion inside the request function to catch misalignment close to the source.

Example fix

// before
return responses; // may have fewer entries than inputs
// after
if (responses.size() != inputs.size()) { /* pad or fix alignment */ }
return responses;
Defensive patterns

Strategy: validation

Validate before calling

List<OutputT> responses = requestFn.apply(modelName, inputs); if (responses.size() != inputs.size()) { /* realign or pad before calling handler.request */ }

Try / catch

try { handler.request(batch); } catch (IllegalStateException e) { /* log batch sizes and realign request fn output */ }

Prevention

When it happens

Trigger: Calling modelHandler.request(inputs) (directly or via the results() DoFn) with a custom GeminiRequestFunction whose apply() returns fewer or more OutputT elements than the List<InputT> passed in.

Common situations: A custom request function that drops inputs the model refused to answer, batches internally and returns per-batch results, or returns a single aggregated response for the whole batch. Also seen when the model response parsing flattens multiple choices into one.

Understand the failure class

Background: "This is a bug, please report it": internal invariant violations, unreachable panics, and SNH errors explained — this error's family across 47 libraries.

Related errors


AI-assisted analysis of apache/beam@12126d8942 (2026-09-13). Data as JSON: /api/errors/77f01120a003d258. Report an issue: GitHub.

Appendix: source

Thrown at sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java:74

      Client.Builder builder = Client.builder();
      if (parameters.getProject() != null && parameters.getLocation() != null) {
        builder.vertexAI(true).project(parameters.getProject()).location(parameters.getLocation());
      } else if (parameters.getProject() != null || parameters.getLocation() != null) {
        throw new IllegalArgumentException(
            "Project and location must both be provided if one is provided");
      }
      this.client = builder.build();
    }
  }

  @Override
  public Iterable<PredictionResult<InputT, OutputT>> request(List<InputT> input) {
    try {
      GeminiRequestFunction<InputT, OutputT> requestFn = modelParameters.getRequestFn();
      List<OutputT> responses = requestFn.apply(modelParameters.getModelName(), input, client);

      if (responses.size() != input.size()) {
        throw new IllegalStateException("Number of responses must match number of inputs");
      }

      List<PredictionResult<InputT, OutputT>> results = new ArrayList<>();
      for (int i = 0; i < input.size(); i++) {
        results.add(PredictionResult.create(input.get(i), responses.get(i)));
      }
      return results;
    } catch (Exception e) {
      throw new RuntimeException("Error during Gemini inference request", e);
    }
  }
}

View on GitHub (pinned to 12126d8942)