Skip to content

Return all embeddings for batched OpenAI inputs - #71989

Open
YAshhh29 wants to merge 1 commit into
apache:mainfrom
YAshhh29:fix-openai-batched-embeddings
Open

Return all embeddings for batched OpenAI inputs#71989
YAshhh29 wants to merge 1 commit into
apache:mainfrom
YAshhh29:fix-openai-batched-embeddings

Conversation

@YAshhh29

@YAshhh29 YAshhh29 commented Aug 23, 2026

Copy link
Copy Markdown
Contributor

Why

The OpenAI Embeddings API accepts batches, but OpenAIHook.create_embeddings() returned only response.data[0]. Every successful batch result after the first item was silently discarded.

What changed

  • Preserve the existing single-input return shape.
  • Return every batch embedding ordered by the response index.
  • Update the operator return type and user documentation.
  • Add regressions for string and token-array batches.

Tests

  • Focused batch regression test against this branch: 2 passed.
  • The same regression test against the unmodified implementation: 2 failed.
  • ruff format and ruff check --fix: passed.
  • git diff --check: passed.

The full provider suite was not run locally because Airflow does not support native Windows; CI provides the Linux suite result.


Was generative AI tooling used to co-author this PR?
  • Yes - GitHub Copilot

Generated-by: GitHub Copilot following the guidelines


Drafted-by: GitHub Copilot (no human review before posting)

The SDK returns one indexed embedding per batch item, but the hook exposed only the first result. Batch callers therefore received a plausible-looking response with no indication that every later input had been discarded.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant