fix: handle Gemini batch embedding count mismatch (avoid zip crash)#926
fix: handle Gemini batch embedding count mismatch (avoid zip crash)#926Stromweld wants to merge 2 commits into
Conversation
Batch Gemini embedding calls previously assumed the SDK always returns exactly one embedding per input item in order, then used zip(batch, response.embeddings, strict=True). When Gemini returns fewer embeddings than input items (e.g. empty/whitespace items dropped, API quirks), this crashed with 'zip() argument 2 is shorter than argument 1'. Now the batch call verifies embeddings/inputs cardinality explicitly. On a match, the efficient batch path is used as before. On a mismatch, we fall back to embedding each item individually (preserving order), requiring exactly one embedding per individual call. If an individual call itself returns zero or multiple embeddings, we raise ValueError with a clear diagnostic instead of silently misaligning or dropping items. Signed-off-by: Corey Hemminger <hemminger@hotmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
WalkthroughGemini batch embedding now validates that response cardinality matches the input batch, preserves item order, and falls back to per-item calls when it does not. Invalid counts or missing embedding values raise ChangesGemini batch alignment
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant EmbeddingClient
participant GeminiAPI
participant ItemEmbeddingCalls
EmbeddingClient->>GeminiAPI: request batch embeddings
GeminiAPI-->>EmbeddingClient: return batch embeddings
EmbeddingClient->>ItemEmbeddingCalls: retry items when cardinality mismatches
ItemEmbeddingCalls->>GeminiAPI: request single-item embeddings
GeminiAPI-->>ItemEmbeddingCalls: return one embedding per item
ItemEmbeddingCalls-->>EmbeddingClient: return aligned mappings
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/embedding_client.py (1)
484-489: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSequential per-item fallback calls add up to N round-trips.
The fallback loops over every item in
batch(up tomax_batch_size=100) with sequentialawaitcalls. Since this comment block itself notes the mismatch "has been observed" in practice (not purely theoretical), this sequential path can add meaningful latency when triggered. Considerasyncio.gatherto issue the per-item calls concurrently, preserving output order (gather returns results in input order regardless of completion order) — but note the test double'ssingle_callslist order is completion-order dependent, sotests/llm/test_embedding_client.py's assertions on call order may need to account for that if this is changed.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/embedding_client.py` around lines 484 - 489, Update the per-item fallback in the batch embedding flow to issue Gemini requests concurrently with asyncio.gather while preserving batch output order. Keep each item mapped to its corresponding response, and adjust affected tests such as single_calls order assertions to account for completion-order variation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/embedding_client.py`:
- Around line 462-466: Update both the happy-path loop using zip and the
per-item fallback to raise ValueError when embedding.values is None or empty,
matching embed()'s “No embedding returned from Gemini API” behavior. Do not
silently skip the item; only validate and store embeddings when values are
present.
---
Nitpick comments:
In `@src/embedding_client.py`:
- Around line 484-489: Update the per-item fallback in the batch embedding flow
to issue Gemini requests concurrently with asyncio.gather while preserving batch
output order. Keep each item mapped to its corresponding response, and adjust
affected tests such as single_calls order assertions to account for
completion-order variation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d957eaef-21ef-4f0e-898e-ddfe4e2f05bd
📒 Files selected for processing (2)
src/embedding_client.pytests/llm/test_embedding_client.py
…ently dropping CodeRabbit flagged that both the cardinality-matched batch (zip) path and the per-item fallback path silently omitted an item from the result dict when embedding.values was falsy (None or empty list). This reintroduced the exact silent-misalignment risk this PR was created to eliminate. Now both paths raise a ValueError with diagnostic context (text_id, chunk_index, and which path) when a returned embedding has no values, matching the diagnostic style already used for the wrong-count case. Also adds regression test coverage: - fallback call returning multiple embeddings for a single item raises - batch path with one falsy-values embedding in an otherwise cardinality-matched response raises - fallback path with a falsy-values single embedding raises Signed-off-by: Corey Hemminger <hemminger@hotmail.com>
|
Finn-loop review of 153f55a CI: not configured (mergeStateStatus=BLOCKED, mergeable=MERGEABLE, but ReviewPR #926 fixes a 1. Must fix before merge
2. Should fix soonNone. The code changes are well-scoped:
3. Safe to mergeNo — CI gate is missing. Automated review of the code itself finds no defects, security issues, or AC gaps. A human must verify lint/typecheck/tests pass and resolve the stale branch-protection rule before merging. Notes
|
2 participants
Problem
Batch Gemini embedding calls in
src/embedding_client.pydid:This assumed Gemini always returns exactly one embedding per input item, in order. While the SDK docstring claims this, it does not always hold in practice (possible causes: empty/whitespace items dropped, batch size >100, other API quirks). When the counts didn't match,
zip(..., strict=True)raisedValueError: zip() argument 2 is shorter than argument 1, crashing the whole batch.Fix
len(response.embeddings) != len(batch)before zipping.ValueErrorwith a clear diagnostic (text_id/chunk_index) instead of silently dropping or misaligning data.Tests
Added to
tests/llm/test_embedding_client.py:ValueErrorThis is a separate, independent fix from BUG1 (PR #925) — no overlap with that branch/commit.
Created with AI assistance — Hermes Agent v0.17.0, model Claude Sonnet 4.6.
Summary by CodeRabbit