fix: gemini batch embedding token not counted

[creamlike1024]

Summary by CodeRabbit

• New Features
  • Added support for batch embedding requests in Gemini integration, enabling efficient processing of multiple embedding requests simultaneously.

[coderabbitai]

Walkthrough

A new public handler GetAndValidateGeminiBatchEmbeddingRequest is introduced to validate batch embedding requests. The Gemini routing logic is refactored to distinguish between single embedding (":embedContent") and batch embedding (":batchEmbedContents") requests, with each routed to its respective validation handler.

Changes

Cohort / File(s)	Summary
Gemini batch embedding request validation relay/helper/valid_request.go	Adds new public handler GetAndValidateGeminiBatchEmbeddingRequest() to unmarshal and validate GeminiBatchEmbeddingRequest. Refactors routing logic to route ":batchEmbedContents" to the new handler and ":embedContent" to the existing GetAndValidateGeminiEmbeddingRequest().

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Poem

🐰 A batch of embeddings hops into the warren,
No longer mingled with their single-bound sparren!
The handler now splits them with care and precision,
Each path gets its route—a rabbit's wise vision! 🥕

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Title Check	⚠️ Warning	The PR title "fix: gemini batch embedding token not counted" describes a problem statement or objective rather than clearly summarizing the actual implementation changes. The raw summary indicates the primary changes are adding a new validation handler (GetAndValidateGeminiBatchEmbeddingRequest) and adjusting request routing logic to distinguish batch embedding requests from single embedding requests. While batch embeddings are mentioned in both the title and changes, the title suggests a fix for token counting functionality without clarifying that the implementation is about request routing and validation. A teammate scanning the repository history would not immediately understand that this PR introduces new validation and routing logic based solely on the title.	Consider revising the title to more accurately reflect the implementation details, such as "feat: add batch embedding request validation and routing" or "feat: distinguish batch and single embedding requests with separate handlers". This would give reviewers and future readers a clearer understanding of what code changes the PR introduces, rather than only stating the problem being addressed.
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch fix-gemini-batch-embedding-token-count

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

relay/helper/valid_request.go (1)

325-332: Consider adding validation for the Requests field.

The handler follows the pattern of GetAndValidateGeminiEmbeddingRequest, but consider adding defensive validation to ensure the Requests array is not empty, similar to how GetAndValidateRerankRequest validates its fields (lines 93-98).

Apply this diff to add validation:

 func GetAndValidateGeminiBatchEmbeddingRequest(c *gin.Context) (*dto.GeminiBatchEmbeddingRequest, error) {
 	request := &dto.GeminiBatchEmbeddingRequest{}
 	err := common.UnmarshalBodyReusable(c, request)
 	if err != nil {
 		return nil, err
 	}
+	if len(request.Requests) == 0 {
+		return nil, errors.New("requests is required and cannot be empty")
+	}
 	return request, nil
 }

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between dd77516 and e2d3b46.

📒 Files selected for processing (1)

• relay/helper/valid_request.go (2 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

relay/helper/valid_request.go (3)

dto/request_common.go (1)

• Request (8-12)

dto/gemini.go (1)

• GeminiBatchEmbeddingRequest (381-383)

common/gin.go (1)

• UnmarshalBodyReusable (32-53)

🔇 Additional comments (2)

relay/helper/valid_request.go (2)

25-31: ---

Routing and batch handler implementation is correct; token counting aggregation verified.

The routing logic properly separates :embedContent and :batchEmbedContents requests into distinct handlers. The new GetAndValidateGeminiBatchEmbeddingRequest function follows the established pattern, and the GeminiBatchEmbeddingRequest DTO correctly aggregates token counts across all batch requests in its GetTokenCountMeta() implementation (iterating over the Requests slice). This integration with relay/gemini_handler.go ensures token counting works correctly for batch operations, addressing the PR's stated fix.

---

325-332: No issues found — GeminiBatchEmbeddingRequest properly implements the Request interface.

The verification confirms that dto.GeminiBatchEmbeddingRequest correctly implements the Request interface with all required methods:

• IsStream() returns false for non-streamed batch requests
• GetTokenCountMeta() (lines 390–401) aggregates token count metadata from nested requests by iterating through r.Requests and combining their text
• SetModelName() (lines 404–410) propagates the model name to nested requests

The token counting functionality for batch embeddings is properly implemented.