feat: Support Vertex AI Embeddings and Gemini Cached Content Tokens - #2488
feat: Support Vertex AI Embeddings and Gemini Cached Content Tokens#2488Simba98 wants to merge 4851 commits into
Conversation
feat: EditTokenModal 中针对用户创建的 token 默认无限额度
feat: add environment variable switch for critical rate limit
…80p-image fix: trim suffix p for jimeng image model
…space fix: tag splitting by whitespace
…odisable fix(channel): 当没有可用密钥时返回错误而不是第一个密钥
… new sections for partners, acknowledgments, and deployment instructions
- 更新中文README.md中的语言链接 - 完全重写英文README.en.md，包含所有详细功能说明 - 完全重写法文README.fr.md，确保内容一致性 - 完全重写日文README.ja.md，提供完整的项目说明 所有语言版本现在具有： - 相同的结构和格式 - 一致的语言导航 - 完整的功能特性和部署指南 - 统一的环境变量配置说明
…ence 修复viduq2不支持参考生视频的问题
…annel feat: replicate channel flux model
fix GetChannelKey AdminAuth -> RootAuth
fix GetChannelKey AdminAuth -> RootAuth
…future adjustments
…-response-id fix: 模型设置增加针对Vertex渠道过滤content[].part[].functionResponse.id的选项，默认启用
- Replace legacy `docs.newapi.pro` paths with the new `/{lang}/docs/...` structure across all README translations
- Point key sections (installation, env vars, API, support, features) to their new locations
- Ensure language-specific links use the correct locale prefix (zh/en/ja) and keep FR aligned with English routes
Keep new-site links (/{lang}/docs/...) where matching pages exist in the current docs repo
Revert links that have no equivalent in the new docs to the legacy paths on doc.newapi.pro:
Google Gemini Chat
Midjourney-Proxy image docs
Suno music docs
Apply the same rule consistently across all README translations (zh/en/ja/fr)
…te-doc-links-new-routing 🔗 docs(readme): update documentation links to new site routing
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds cached-content token tracking from Gemini metadata into usage, routes Gemini embedding requests through a Vertex-compatible transform, introduces a Vertex embedding response type and handler that aggregates tokens, and adjusts quota accounting to subtract cached tokens for Gemini and VertexAi channels. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Relay as Relay (Vertex Adaptor)
participant VertexAPI as Vertex AI
participant Handler as vertexEmbeddingHandler
participant Quota as Quota Service
Client->>Relay: POST Gemini embedding request
Relay->>Relay: Read & parse Gemini embedding body
Relay->>Relay: Build Vertex `instances` (aggregate content, task_type/title, params)
Relay->>VertexAPI: POST Vertex-compatible embedding request
VertexAPI-->>Relay: JSON embedding response
Relay->>Handler: Route response to vertexEmbeddingHandler
Handler->>Handler: Unmarshal VertexEmbeddingResponse, aggregate token counts
Handler-->>Client: Stream original response body and attach usage metadata
Client->>Quota: Report usage (includes cached tokens)
Quota->>Quota: Subtract cached tokens for Gemini/VertexAi when consuming quota
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: Organization UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
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: 3
🧹 Nitpick comments (3)
relay/channel/vertex/relay-vertex.go (2)
63-67: Consider usingcommon.Unmarshalfor consistency.Other handlers in the codebase (e.g.,
GeminiEmbeddingHandlerinrelay-gemini.go) usecommon.Unmarshalinstead ofjson.Unmarshal. Using the common utility ensures consistent behavior across the codebase.🔎 Proposed fix
var vertexResponse VertexEmbeddingResponse - err = json.Unmarshal(responseBody, &vertexResponse) + err = common.Unmarshal(responseBody, &vertexResponse) if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) }
37-49: Anonymous nested structs reduce reusability.The
VertexEmbeddingResponseuses anonymous nested structs. While functional, extracting named types would improve code readability and enable reuse if these structures are needed elsewhere.This is a minor style consideration and can be addressed later if the response format needs to be referenced in other parts of the codebase.
relay/channel/vertex/adaptor.go (1)
375-421: Consider extracting duplicate instance-building logic.The code for building instances from embedding requests is duplicated between the batch and single request paths. Consider extracting this into a helper function to improve maintainability.
🔎 Example refactor
func buildEmbeddingInstance(content dto.GeminiChatContent, taskType, title string) map[string]interface{} { instance := make(map[string]interface{}) contentStr := "" for _, part := range content.Parts { if part.Text != "" { contentStr += part.Text } } instance["content"] = contentStr if taskType != "" { instance["task_type"] = taskType } if title != "" { instance["title"] = title } return instance }Then use it in both paths:
if err := json.Unmarshal(bodyBytes, &req); err != nil { return nil, fmt.Errorf("failed to unmarshal batch embedding request: %w", err) } for _, r := range req.Requests { - instance := make(map[string]interface{}) - content := "" - for _, part := range r.Content.Parts { - if part.Text != "" { - content += part.Text - } - } - instance["content"] = content - if r.TaskType != "" { - instance["task_type"] = r.TaskType - } - if r.Title != "" { - instance["title"] = r.Title - } + instance := buildEmbeddingInstance(r.Content, r.TaskType, r.Title) instances = append(instances, instance) }
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
dto/gemini.go(1 hunks)relay/channel/gemini/relay-gemini-native.go(1 hunks)relay/channel/gemini/relay-gemini.go(2 hunks)relay/channel/vertex/adaptor.go(5 hunks)relay/channel/vertex/relay-vertex.go(2 hunks)service/quota.go(1 hunks)
🧰 Additional context used
🧠 Learnings (5)
📓 Common learnings
Learnt from: RedwindA
Repo: QuantumNous/new-api PR: 1537
File: relay/gemini_handler.go:330-342
Timestamp: 2025-08-08T17:12:43.157Z
Learning: In the new-api repository, the `GeminiEmbeddingHandler` function in `relay/gemini_handler.go` is designed specifically for native Gemini embedding requests and therefore does not require the `ConvertGeminiRequest` step that is used in the chat handler. The embedding requests are already in the native Gemini format and don't need conversion.
Learnt from: 9Ninety
Repo: QuantumNous/new-api PR: 1273
File: relay/channel/gemini/relay-gemini.go:97-116
Timestamp: 2025-06-21T03:37:41.726Z
Learning: In relay/channel/gemini/relay-gemini.go, the thinking budget calculation logic (including the MaxOutputTokens multiplication) was introduced in PR #1247. PR #1273 focused specifically on decoupling the thoughts summary feature from thinking budget settings and did not modify the existing thinking budget behavior.
📚 Learning: 2025-08-08T17:12:43.157Z
Learnt from: RedwindA
Repo: QuantumNous/new-api PR: 1537
File: relay/gemini_handler.go:330-342
Timestamp: 2025-08-08T17:12:43.157Z
Learning: In the new-api repository, the `GeminiEmbeddingHandler` function in `relay/gemini_handler.go` is designed specifically for native Gemini embedding requests and therefore does not require the `ConvertGeminiRequest` step that is used in the chat handler. The embedding requests are already in the native Gemini format and don't need conversion.
Applied to files:
relay/channel/vertex/relay-vertex.gorelay/channel/vertex/adaptor.gorelay/channel/gemini/relay-gemini.go
📚 Learning: 2025-08-05T17:14:17.246Z
Learnt from: neotf
Repo: QuantumNous/new-api PR: 1511
File: setting/ratio_setting/model_ratio.go:118-123
Timestamp: 2025-08-05T17:14:17.246Z
Learning: Claude models handle "-thinking" variants differently from Gemini models. For Claude models, only the base model (without "-thinking") gets an entry in defaultModelRatio map. The "-thinking" variants rely on the Claude relay handler stripping the suffix using strings.TrimSuffix(textRequest.Model, "-thinking") before looking up the ratio, so they automatically use the base model's ratio.
Applied to files:
relay/channel/vertex/adaptor.go
📚 Learning: 2025-06-21T03:37:41.726Z
Learnt from: 9Ninety
Repo: QuantumNous/new-api PR: 1273
File: relay/channel/gemini/relay-gemini.go:97-116
Timestamp: 2025-06-21T03:37:41.726Z
Learning: In relay/channel/gemini/relay-gemini.go, the thinking budget calculation logic (including the MaxOutputTokens multiplication) was introduced in PR #1247. PR #1273 focused specifically on decoupling the thoughts summary feature from thinking budget settings and did not modify the existing thinking budget behavior.
Applied to files:
relay/channel/gemini/relay-gemini-native.gorelay/channel/gemini/relay-gemini.goservice/quota.go
📚 Learning: 2025-06-18T12:20:25.779Z
Learnt from: neotf
Repo: QuantumNous/new-api PR: 1120
File: service/quota.go:0-0
Timestamp: 2025-06-18T12:20:25.779Z
Learning: For OpenRouter integration: cacheTokens returned from upstream always belong to [0, promptTokens], meaning cacheTokens will never exceed the original promptTokens value. This constraint ensures that operations like `promptTokens -= cacheTokens` will not result in negative values.
Applied to files:
service/quota.go
🧬 Code graph analysis (2)
relay/channel/vertex/adaptor.go (3)
dto/gemini.go (2)
GeminiBatchEmbeddingRequest(423-425)GeminiEmbeddingRequest(391-397)common/json.go (2)
Unmarshal(9-11)Marshal(21-23)logger/logger.go (1)
LogDebug(69-76)
service/quota.go (1)
constant/channel.go (2)
ChannelTypeGemini(28-28)ChannelTypeVertexAi(41-41)
🔇 Additional comments (9)
relay/channel/gemini/relay-gemini-native.go (1)
46-46: Cached token propagation looks correct.The assignment follows the established pattern used for
ReasoningTokenson line 45, consistently propagating the newCachedContentTokenCountfield from the Gemini API response into the usage structure.relay/channel/gemini/relay-gemini.go (2)
1105-1105: Streaming handler cached token propagation is consistent.The assignment correctly propagates
CachedContentTokenCountwithin the streaming usage metadata update block, following the same pattern as theReasoningTokensassignment on line 1104.
1240-1240: Non-streaming handler cached token propagation is consistent.The assignment in the chat handler follows the same pattern as the streaming handler and native handler, ensuring cached tokens are properly tracked across all Gemini response paths.
dto/gemini.go (1)
348-355: DTO field addition for cached content token tracking looks good.The new
CachedContentTokenCountfield follows the existing naming conventions and type patterns inGeminiUsageMetadata. The JSON tagcachedContentTokenCountaligns with the Gemini API response format.relay/channel/vertex/relay-vertex.go (1)
74-82: Verify if Vertex embedding response should be converted to OpenAI format.The
GeminiEmbeddingHandlerinrelay-gemini.goconverts the Gemini embedding response to the OpenAI embedding format (dto.OpenAIEmbeddingResponse) before returning. This handler returns the raw Vertex response body viaIOCopyBytesGracefully.If clients expect a consistent OpenAI-compatible embedding response format, consider converting the Vertex response similarly:
openAIResponse := dto.OpenAIEmbeddingResponse{ Object: "list", Data: make([]dto.OpenAIEmbeddingResponseItem, 0, len(vertexResponse.Predictions)), Model: info.UpstreamModelName, } for i, prediction := range vertexResponse.Predictions { openAIResponse.Data = append(openAIResponse.Data, dto.OpenAIEmbeddingResponseItem{ Object: "embedding", Embedding: prediction.Embeddings.Values, Index: i, }) } openAIResponse.Usage = *usageIf the raw Vertex format is intentional for native Vertex API passthrough, this is fine as-is.
service/quota.go (1)
270-272: Remove this review. The Gemini/VertexAi API design guarantees thatcachedContentTokenCountis always a subset ofpromptTokenCountsince cached content forms a prefix within the prompt. No bounds check is needed.Likely an incorrect or invalid review comment.
relay/channel/vertex/adaptor.go (3)
4-4: LGTM!The new imports are properly utilized in the embedding transformation logic.
Also applies to: 14-14
231-233: LGTM!The condition correctly identifies embedding models and routes them to use the "predict" suffix, consistent with how imagen models are handled.
450-452: No issues found. Theinfo.RequestURLPathis properly populated fromc.Request.URL.String()before reaching the DoResponse function, and the routing logic correctly identifies and directs embedding requests tovertexEmbeddingHandler.
| if a.RequestMode == RequestModeGemini && strings.Contains(c.Request.URL.Path, "embed") { | ||
| bodyBytes, err := io.ReadAll(requestBody) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| var newBodyBytes []byte | ||
| vertexReq := make(map[string]interface{}) | ||
| instances := make([]interface{}, 0) | ||
|
|
||
| if info.IsGeminiBatchEmbedding { | ||
| var req dto.GeminiBatchEmbeddingRequest | ||
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | ||
| for _, r := range req.Requests { | ||
| instance := make(map[string]interface{}) | ||
| content := "" | ||
| for _, part := range r.Content.Parts { | ||
| if part.Text != "" { | ||
| content += part.Text | ||
| } | ||
| } | ||
| instance["content"] = content | ||
| if r.TaskType != "" { | ||
| instance["task_type"] = r.TaskType | ||
| } | ||
| if r.Title != "" { | ||
| instance["title"] = r.Title | ||
| } | ||
| instances = append(instances, instance) | ||
| } | ||
| } | ||
| } else { | ||
| var req dto.GeminiEmbeddingRequest | ||
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | ||
| instance := make(map[string]interface{}) | ||
| content := "" | ||
| for _, part := range req.Content.Parts { | ||
| if part.Text != "" { | ||
| content += part.Text | ||
| } | ||
| } | ||
| instance["content"] = content | ||
| if req.TaskType != "" { | ||
| instance["task_type"] = req.TaskType | ||
| } | ||
| if req.Title != "" { | ||
| instance["title"] = req.Title | ||
| } | ||
| instances = append(instances, instance) | ||
|
|
||
| if req.OutputDimensionality > 0 { | ||
| vertexReq["parameters"] = map[string]interface{}{ | ||
| "outputDimensionality": req.OutputDimensionality, | ||
| } | ||
| } | ||
| } | ||
| } | ||
| vertexReq["instances"] = instances | ||
| newBodyBytes, _ = json.Marshal(vertexReq) | ||
| requestBody = bytes.NewReader(newBodyBytes) | ||
| logger.LogDebug(c, "Vertex Embedding request body: "+string(newBodyBytes)) | ||
| } |
There was a problem hiding this comment.
Validate that instances array is non-empty.
After processing the request, there's no validation that the instances array contains at least one element. Sending an empty instances array would result in an invalid API request.
🔎 Proposed fix
}
vertexReq["instances"] = instances
+ if len(instances) == 0 {
+ return nil, errors.New("no valid embedding content found in request")
+ }
newBodyBytes, _ = json.Marshal(vertexReq)🤖 Prompt for AI Agents
In relay/channel/vertex/adaptor.go around lines 365 to 426, after building the
instances slice validate that it contains at least one element and fail fast if
not: if len(instances) == 0 return an error (or an HTTP 400 style error)
indicating "empty instances in embedding request" instead of marshaling and
forwarding an empty instances array; ensure the error is returned consistently
with the surrounding function's error handling and log a debug/warn message with
context before returning.
|
|
||
| if info.IsGeminiBatchEmbedding { | ||
| var req dto.GeminiBatchEmbeddingRequest | ||
| if err := json.Unmarshal(bodyBytes, &req); err == nil { |
There was a problem hiding this comment.
Handle unmarshaling errors explicitly.
The unmarshaling errors are silently ignored with err == nil checks. If unmarshaling fails, the code proceeds with an empty instances array, which could result in invalid API requests.
🔎 Proposed fix
if info.IsGeminiBatchEmbedding {
var req dto.GeminiBatchEmbeddingRequest
- if err := json.Unmarshal(bodyBytes, &req); err == nil {
+ if err := json.Unmarshal(bodyBytes, &req); err != nil {
+ return nil, fmt.Errorf("failed to unmarshal batch embedding request: %w", err)
+ }
for _, r := range req.Requests {
instance := make(map[string]interface{})
content := ""
for _, part := range r.Content.Parts {
if part.Text != "" {
content += part.Text
}
}
instance["content"] = content
if r.TaskType != "" {
instance["task_type"] = r.TaskType
}
if r.Title != "" {
instance["title"] = r.Title
}
instances = append(instances, instance)
}
- }
} else {
var req dto.GeminiEmbeddingRequest
- if err := json.Unmarshal(bodyBytes, &req); err == nil {
+ if err := json.Unmarshal(bodyBytes, &req); err != nil {
+ return nil, fmt.Errorf("failed to unmarshal embedding request: %w", err)
+ }
instance := make(map[string]interface{})
content := ""
for _, part := range req.Content.Parts {
if part.Text != "" {
content += part.Text
}
}
instance["content"] = content
if req.TaskType != "" {
instance["task_type"] = req.TaskType
}
if req.Title != "" {
instance["title"] = req.Title
}
instances = append(instances, instance)
if req.OutputDimensionality > 0 {
vertexReq["parameters"] = map[string]interface{}{
"outputDimensionality": req.OutputDimensionality,
}
}
- }
}Also applies to: 398-398
🤖 Prompt for AI Agents
In relay/channel/vertex/adaptor.go around lines 377 and 398, the json.Unmarshal
errors are currently ignored (using only err == nil) which lets execution
continue with empty/invalid data; change both blocks to explicitly check for err
!= nil, log the unmarshaling error with context, and return an appropriate HTTP
error response (e.g., 400 Bad Request or an internal error) instead of
proceeding, so invalid JSON does not result in invalid API requests.
| } | ||
| } | ||
| vertexReq["instances"] = instances | ||
| newBodyBytes, _ = json.Marshal(vertexReq) |
There was a problem hiding this comment.
Handle the Marshal error.
The error from json.Marshal is ignored. If marshaling fails, newBodyBytes will be nil or empty, leading to an invalid request.
🔎 Proposed fix
- newBodyBytes, _ = json.Marshal(vertexReq)
+ newBodyBytes, err = json.Marshal(vertexReq)
+ if err != nil {
+ return nil, fmt.Errorf("failed to marshal vertex embedding request: %w", err)
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| newBodyBytes, _ = json.Marshal(vertexReq) | |
| newBodyBytes, err = json.Marshal(vertexReq) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to marshal vertex embedding request: %w", err) | |
| } |
🤖 Prompt for AI Agents
In relay/channel/vertex/adaptor.go around line 423, json.Marshal's error is
currently ignored when marshaling vertexReq into newBodyBytes; update the code
to capture the returned error, handle it (e.g., log/processLogger.Error or
return the error up the call stack), and abort building/sending the request when
marshaling fails. Concretely: assign both newBodyBytes and err from
json.Marshal, check if err != nil, then return an appropriate error or respond
with a failure (and log the marshaling error) instead of proceeding with a
nil/empty body.
There was a problem hiding this comment.
Pull request overview
This PR adds support for Vertex AI embeddings and Gemini cached content tokens to improve API compatibility and token accounting accuracy.
- Added support for transforming Gemini embedding requests to Vertex AI's prediction API format
- Implemented cached content token tracking in Gemini responses for accurate quota calculation
- Added vertex embedding response handler with token counting from API statistics
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| service/quota.go | Added cached token deduction logic for Gemini and Vertex AI channels in quota calculation |
| relay/channel/vertex/relay-vertex.go | Added vertex embedding response handler with token count extraction from predictions |
| relay/channel/vertex/adaptor.go | Implemented embedding request transformation from Gemini format to Vertex AI format with proper routing logic |
| relay/channel/gemini/relay-gemini.go | Added cached content token tracking to both streaming and non-streaming chat handlers |
| relay/channel/gemini/relay-gemini-native.go | Added cached content token tracking to native Gemini text generation handler |
| dto/gemini.go | Added CachedContentTokenCount field to GeminiUsageMetadata structure for tracking cached tokens |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | ||
| for _, r := range req.Requests { | ||
| instance := make(map[string]interface{}) | ||
| content := "" | ||
| for _, part := range r.Content.Parts { | ||
| if part.Text != "" { | ||
| content += part.Text | ||
| } | ||
| } | ||
| instance["content"] = content | ||
| if r.TaskType != "" { | ||
| instance["task_type"] = r.TaskType | ||
| } | ||
| if r.Title != "" { | ||
| instance["title"] = r.Title | ||
| } | ||
| instances = append(instances, instance) | ||
| } | ||
| } | ||
| } else { | ||
| var req dto.GeminiEmbeddingRequest | ||
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | ||
| instance := make(map[string]interface{}) | ||
| content := "" | ||
| for _, part := range req.Content.Parts { | ||
| if part.Text != "" { | ||
| content += part.Text | ||
| } | ||
| } | ||
| instance["content"] = content | ||
| if req.TaskType != "" { | ||
| instance["task_type"] = req.TaskType | ||
| } | ||
| if req.Title != "" { | ||
| instance["title"] = req.Title | ||
| } | ||
| instances = append(instances, instance) | ||
|
|
||
| if req.OutputDimensionality > 0 { | ||
| vertexReq["parameters"] = map[string]interface{}{ | ||
| "outputDimensionality": req.OutputDimensionality, | ||
| } |
There was a problem hiding this comment.
The error from json.Unmarshal is silently ignored with 'if err == nil'. If unmarshaling fails, the code continues with an empty instances array, which will result in an empty vertexReq being sent to the API. This could lead to API errors or unexpected behavior. Consider returning an error if unmarshaling fails instead of silently proceeding.
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | |
| for _, r := range req.Requests { | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range r.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if r.TaskType != "" { | |
| instance["task_type"] = r.TaskType | |
| } | |
| if r.Title != "" { | |
| instance["title"] = r.Title | |
| } | |
| instances = append(instances, instance) | |
| } | |
| } | |
| } else { | |
| var req dto.GeminiEmbeddingRequest | |
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range req.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if req.TaskType != "" { | |
| instance["task_type"] = req.TaskType | |
| } | |
| if req.Title != "" { | |
| instance["title"] = req.Title | |
| } | |
| instances = append(instances, instance) | |
| if req.OutputDimensionality > 0 { | |
| vertexReq["parameters"] = map[string]interface{}{ | |
| "outputDimensionality": req.OutputDimensionality, | |
| } | |
| if err := json.Unmarshal(bodyBytes, &req); err != nil { | |
| return nil, err | |
| } | |
| for _, r := range req.Requests { | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range r.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if r.TaskType != "" { | |
| instance["task_type"] = r.TaskType | |
| } | |
| if r.Title != "" { | |
| instance["title"] = r.Title | |
| } | |
| instances = append(instances, instance) | |
| } | |
| } else { | |
| var req dto.GeminiEmbeddingRequest | |
| if err := json.Unmarshal(bodyBytes, &req); err != nil { | |
| return nil, err | |
| } | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range req.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if req.TaskType != "" { | |
| instance["task_type"] = req.TaskType | |
| } | |
| if req.Title != "" { | |
| instance["title"] = req.Title | |
| } | |
| instances = append(instances, instance) | |
| if req.OutputDimensionality > 0 { | |
| vertexReq["parameters"] = map[string]interface{}{ | |
| "outputDimensionality": req.OutputDimensionality, |
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | ||
| for _, r := range req.Requests { | ||
| instance := make(map[string]interface{}) | ||
| content := "" | ||
| for _, part := range r.Content.Parts { | ||
| if part.Text != "" { | ||
| content += part.Text | ||
| } | ||
| } | ||
| instance["content"] = content | ||
| if r.TaskType != "" { | ||
| instance["task_type"] = r.TaskType | ||
| } | ||
| if r.Title != "" { | ||
| instance["title"] = r.Title | ||
| } | ||
| instances = append(instances, instance) | ||
| } | ||
| } | ||
| } else { | ||
| var req dto.GeminiEmbeddingRequest | ||
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | ||
| instance := make(map[string]interface{}) | ||
| content := "" | ||
| for _, part := range req.Content.Parts { | ||
| if part.Text != "" { | ||
| content += part.Text | ||
| } | ||
| } | ||
| instance["content"] = content | ||
| if req.TaskType != "" { | ||
| instance["task_type"] = req.TaskType | ||
| } | ||
| if req.Title != "" { | ||
| instance["title"] = req.Title | ||
| } | ||
| instances = append(instances, instance) | ||
|
|
||
| if req.OutputDimensionality > 0 { | ||
| vertexReq["parameters"] = map[string]interface{}{ | ||
| "outputDimensionality": req.OutputDimensionality, | ||
| } |
There was a problem hiding this comment.
The error from json.Unmarshal is silently ignored with 'if err == nil'. If unmarshaling fails, the code continues with an empty instances array, which will result in an empty vertexReq being sent to the API. This could lead to API errors or unexpected behavior. Consider returning an error if unmarshaling fails instead of silently proceeding.
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | |
| for _, r := range req.Requests { | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range r.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if r.TaskType != "" { | |
| instance["task_type"] = r.TaskType | |
| } | |
| if r.Title != "" { | |
| instance["title"] = r.Title | |
| } | |
| instances = append(instances, instance) | |
| } | |
| } | |
| } else { | |
| var req dto.GeminiEmbeddingRequest | |
| if err := json.Unmarshal(bodyBytes, &req); err == nil { | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range req.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if req.TaskType != "" { | |
| instance["task_type"] = req.TaskType | |
| } | |
| if req.Title != "" { | |
| instance["title"] = req.Title | |
| } | |
| instances = append(instances, instance) | |
| if req.OutputDimensionality > 0 { | |
| vertexReq["parameters"] = map[string]interface{}{ | |
| "outputDimensionality": req.OutputDimensionality, | |
| } | |
| if err := json.Unmarshal(bodyBytes, &req); err != nil { | |
| return nil, fmt.Errorf("failed to unmarshal GeminiBatchEmbeddingRequest: %w", err) | |
| } | |
| for _, r := range req.Requests { | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range r.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if r.TaskType != "" { | |
| instance["task_type"] = r.TaskType | |
| } | |
| if r.Title != "" { | |
| instance["title"] = r.Title | |
| } | |
| instances = append(instances, instance) | |
| } | |
| } else { | |
| var req dto.GeminiEmbeddingRequest | |
| if err := json.Unmarshal(bodyBytes, &req); err != nil { | |
| return nil, fmt.Errorf("failed to unmarshal GeminiEmbeddingRequest: %w", err) | |
| } | |
| instance := make(map[string]interface{}) | |
| content := "" | |
| for _, part := range req.Content.Parts { | |
| if part.Text != "" { | |
| content += part.Text | |
| } | |
| } | |
| instance["content"] = content | |
| if req.TaskType != "" { | |
| instance["task_type"] = req.TaskType | |
| } | |
| if req.Title != "" { | |
| instance["title"] = req.Title | |
| } | |
| instances = append(instances, instance) | |
| if req.OutputDimensionality > 0 { | |
| vertexReq["parameters"] = map[string]interface{}{ | |
| "outputDimensionality": req.OutputDimensionality, |
| } | ||
| } | ||
| vertexReq["instances"] = instances | ||
| newBodyBytes, _ = json.Marshal(vertexReq) |
There was a problem hiding this comment.
The error returned from json.Marshal is ignored using the blank identifier. If marshaling fails, newBodyBytes will be nil, which could cause a panic or unexpected behavior when creating the bytes.NewReader. Consider handling this error appropriately by returning it to the caller.
| newBodyBytes, _ = json.Marshal(vertexReq) | |
| newBodyBytes, err = json.Marshal(vertexReq) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to marshal vertex embedding request: %w", err) | |
| } |
seefs001
added
the
ready to merge
|
真的很需要这个特性 |
|
这个功能现在得到支持了吗？pr为何被关闭了？ |
|
pr因为main分支force push被自动关闭了，目前支持了缓存token读取，但是Vertex AI Embedding还不支持 |
This PR introduces two key improvements to the Vertex AI and Gemini channels:
Summary by CodeRabbit
New Features
Bug Fixes
✏️ Tip: You can customize this high-level summary in your review settings.