fix: image search for genai search tool - #5647
Conversation
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
📝 WalkthroughSummary by CodeRabbit
WalkthroughGemini search conversions now preserve localization coordinates, text/image search types, image grounding metadata, merged multi-source citations, and streaming parity. Responses schemas carry the added fields, with round-trip tests covering these paths. A separate test validates cached-content token accounting. ChangesGemini search and grounding
Cached token accounting
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Gemini
participant BifrostResponses
participant StreamState
Gemini->>BifrostResponses: convert search tools and grounding metadata
BifrostResponses->>StreamState: buffer grounding annotations
StreamState->>BifrostResponses: emit merged citations and image sources
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
|
|
![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)
core/providers/gemini/responses.go (1)
3772-3805: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
urlToIndexMapcollision risk between web and image chunks sharing the same page URL.
urlToIndexMapis keyed bysource.URLalone (line 3804, pre-existing). The new image branch (3786-3801) writesGroundingChunkImage.SourceURI = source.URL— the same URL namespace as web chunks. If a web citation and an image citation attribute to the same page URL, the later chunk overwrites the earlier map entry, and any annotation citing that URL gets silently misattributed to the wrong chunk (and possibly the wrong chunk type) later in the function.♻️ Suggested fix: key by URL+kind instead of URL alone
- groundingChunks = append(groundingChunks, chunk) - urlToIndexMap[source.URL] = int32(len(groundingChunks) - 1) + groundingChunks = append(groundingChunks, chunk) + key := source.URL + if chunk.Image != nil { + key += "|image" + } + urlToIndexMap[key] = int32(len(groundingChunks) - 1)(and apply the matching key construction where
urlToIndexMapis looked up for annotations, based on whether the annotation source was an image source)🤖 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 `@core/providers/gemini/responses.go` around lines 3772 - 3805, Update the grounding-chunk mapping around urlToIndexMap to use a composite key containing the source URL and chunk kind, distinguishing web sources from image sources. Apply the identical key construction when annotation lookups occur, using whether the annotation references an image source, so web and image chunks sharing a page URL resolve to their respective chunk indices.
🤖 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 `@core/providers/gemini/responses.go`:
- Around line 2490-2515: Update groundingChunkSource to return the existing
zero-value source and false immediately when chunk is nil, before accessing
chunk.Web or chunk.Image. Keep the current URL, title, image URL, domain, and
unsupported-chunk handling unchanged for non-nil inputs; the callers should
continue receiving the same boolean-based skip behavior.
---
Nitpick comments:
In `@core/providers/gemini/responses.go`:
- Around line 3772-3805: Update the grounding-chunk mapping around urlToIndexMap
to use a composite key containing the source URL and chunk kind, distinguishing
web sources from image sources. Apply the identical key construction when
annotation lookups occur, using whether the annotation references an image
source, so web and image chunks sharing a page URL resolve to their respective
chunk indices.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 68661c67-2bd0-4da0-a96d-34f8f724c8c7
📒 Files selected for processing (5)
core/providers/gemini/payload_ordering_test.gocore/providers/gemini/responses.gocore/providers/gemini/types.gocore/providers/gemini/websearchstreamstate_test.gocore/schemas/responses.go
af8a7ee to
c012275
Compare
ce522ef to
4b9cd7b
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/providers/gemini/responses.go (1)
3775-3808: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winURL-keyed chunk map can misattribute citations between web and image chunks sharing the same page URL.
urlToIndexMapismap[string]int32, so when a web-type source and an image-type source share the sameURL(the image's containing page can legitimately equal a separate web citation's URL), the second insert silently overwrites the first. Later, annotation lookup only ever resolves to the last chunk stored for that URL:var chunkIndices []int32 if annotation.URL != nil { if chunkIdx, exists := urlToIndexMap[*annotation.URL]; exists { chunkIndices = []int32{chunkIdx} } }This can attribute a plain web citation to the image chunk (or vice versa), showing incorrect source metadata (e.g.
ImageURL/Domain) for that citation. Before image grounding was added, all chunks wereWeb-type so a URL collision was a harmless duplicate; now it can cross type boundaries.🛡️ Proposed fix
- urlToIndexMap := make(map[string]int32) // Map URL to chunk index for annotation processing + urlToIndexMap := make(map[string][]int32) // Map URL to all matching chunk indices for annotation processing @@ groundingChunks = append(groundingChunks, chunk) - urlToIndexMap[source.URL] = int32(len(groundingChunks) - 1) + urlToIndexMap[source.URL] = append(urlToIndexMap[source.URL], int32(len(groundingChunks)-1)) @@ var chunkIndices []int32 if annotation.URL != nil { - if chunkIdx, exists := urlToIndexMap[*annotation.URL]; exists { - chunkIndices = []int32{chunkIdx} - } + if idxs, exists := urlToIndexMap[*annotation.URL]; exists { + chunkIndices = idxs + } }Also applies to: 3838-3844
🤖 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 `@core/providers/gemini/responses.go` around lines 3775 - 3808, Update the grounding chunk index bookkeeping in the source loop and the annotation URL lookup so duplicate page URLs do not overwrite one another across web and image chunks. Preserve every matching chunk index, then resolve annotations using the appropriate source type or URL rather than always selecting the last entry; keep citation attribution tied to the correct web or image metadata.
🤖 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.
Outside diff comments:
In `@core/providers/gemini/responses.go`:
- Around line 3775-3808: Update the grounding chunk index bookkeeping in the
source loop and the annotation URL lookup so duplicate page URLs do not
overwrite one another across web and image chunks. Preserve every matching chunk
index, then resolve annotations using the appropriate source type or URL rather
than always selecting the last entry; keep citation attribution tied to the
correct web or image metadata.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 1678cee6-fef2-4980-8552-25c6f3e68ce2
📒 Files selected for processing (5)
core/providers/gemini/payload_ordering_test.gocore/providers/gemini/responses.gocore/providers/gemini/types.gocore/providers/gemini/websearchstreamstate_test.gocore/schemas/responses.go
🚧 Files skipped from review as they are similar to previous changes (4)
- core/providers/gemini/payload_ordering_test.go
- core/schemas/responses.go
- core/providers/gemini/types.go
- core/providers/gemini/websearchstreamstate_test.go
Merge activity
|
The base branch was changed.
4b9cd7b to
c13f703
Compare
3 participants
Summary
Extends the Gemini provider's Google Search grounding support to cover image search:
searchTypes(web + image), image grounding chunks, image search queries, and search localization vialatLng. Also fixes a bug where multi-source grounding supports were fanned out into one support per cited chunk instead of being regrouped by segment.Changes
SearchTypes,WebSearch, andImageSearchtypes totypes.go, with camelCase/snake_caseUnmarshalJSONsupport.GoogleSearchnow carries aSearchTypesfield.GroundingChunkImageto represent image-search grounding chunks (withsourceUri,imageUri,title,domain).GroundingChunknow includes anImagefield alongsideWeb.ImageSearchQueriestoGroundingMetadatato keep image queries separate from web queries.groundingChunkSourcehelper that normalizes both web and image chunks into aResponsesWebSearchToolCallActionSearchSource, replacing scattered inline chunk-to-source conversions throughoutresponses.go.search_content_typesonResponsesToolWebSearchnow maps"text"→WebSearchand"image"→ImageSearchwhen converting to/from Gemini'ssearchTypes.latLngfromtoolConfig.retrievalConfig) is now round-tripped throughResponsesToolWebSearchUserLocation.Latitude/Longitudeso it survives the Gemini → Bifrost Responses → Gemini hop.ImageQueriesadded toResponsesWebSearchToolCallActionandImageURL/Domainadded toResponsesWebSearchToolCallActionSearchSourceto carry image-specific grounding data through the neutral schema.buildGroundingMetadataFromWebSearch: annotations were previously emitting oneGroundingSupportper(segment, chunk)pair. They are now regrouped by segment key so multi-source supports correctly list all cited chunk indices rather than duplicating the segment.emitWebSearchFromGroundingMetadataandemitAnnotationsFromGroundingSupportsupdated to handle image chunks viagroundingChunkSourceand to propagateImageQueries.Type of change
Affected areas
How to test
go test ./core/providers/gemini/...New tests cover:
TestGeminiGoogleSearchToolRoundTrip— verifiessearchTypes,excludeDomains,timeRangeFilter, andlatLngsurvive a full Gemini → Bifrost Responses → Gemini round trip.TestGeminiGoogleSearchToolRoundTripSnakeCase— same for snake_case input; asserts unselected search types are not synthesized.TestGeminiImageGroundingRoundTrip— verifies image chunks keep their asset URL and domain, and thatimageSearchQueriesstay separate fromwebSearchQueries.TestGeminiGroundedRoundTripMergesMultiSourceSupports— non-streaming path: four supports over four chunks, three citing two sources, must come back as four supports (not seven).TestGeminiGroundedStreamRoundTripMergesMultiSourceSupports— streaming counterpart of the above.TestToBifrostCountTokensResponseCachedContent— locks cached-token accounting so cache modalities are never folded into the breakdown a second time.Breaking changes
Related issues
Security considerations
No auth, secrets, or PII implications.
latLngcoordinates flow through the existing tool config path and are not persisted.Checklist
docs/contributing/README.mdand followed the guidelines