feat: vertex batches api - #4152
Merged
Merged
Conversation
|
|
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds Vertex AI Batch Prediction support: exported Vertex batch types, converters between Vertex and Bifrost shapes, BatchCreate/List/Retrieve/Cancel/Delete/Results implementations with GCS IO, GenAI routing, and native Vertex Python integration tests. ChangesVertex AI Batch Prediction Integration
Sequence DiagramsequenceDiagram
participant Client
participant VertexProvider
participant VertexAPI
participant GCS
Client->>VertexProvider: BatchCreate(request)
alt inline requests
VertexProvider->>VertexProvider: Convert items to JSONL
VertexProvider->>GCS: Upload staged JSONL
GCS-->>VertexProvider: Input URI
end
VertexProvider->>VertexAPI: POST /batchPredictionJobs
VertexAPI-->>VertexProvider: BatchPredictionJob
VertexProvider-->>Client: BifrostBatchCreateResponse
Client->>VertexProvider: BatchResults(batch_id)
VertexProvider->>VertexAPI: GET /batchPredictionJobs/{id}
VertexAPI-->>VertexProvider: outputInfo.gcsOutputDirectory
VertexProvider->>GCS: List objects by prefix
GCS-->>VertexProvider: Object pages
loop each prediction JSONL file
VertexProvider->>GCS: Download object (alt=media)
GCS-->>VertexProvider: JSONL bytes
VertexProvider->>VertexProvider: Parse lines and map results
end
VertexProvider-->>Client: BifrostBatchResultsResponse
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
TejasGhatte
marked this pull request as ready for review
June 8, 2026 11:02
Contributor
Confidence Score: 4/5Safe to merge for non-production use; the cancel route path ambiguity is minor and only matters if clients omit the required The core batch operations are well-structured and the multi-key fan-out pattern is consistent with the rest of the provider. The two previously-flagged correctness issues (BatchResults always using the first key, and unbounded GCS pagination) have been addressed or are known open items. The one new finding is the cancel HTTP route registered at POST .../batchPredictionJobs/{batch_id} rather than with a :cancel literal suffix — a bare POST without transports/bifrost-http/integrations/genai.go — cancel route path registration; core/providers/vertex/vertex.go — GCS listing loop remains unbounded (noted in a prior review). Important Files Changed
Reviews (6): Last reviewed commit: "feat: vertex batches api" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/vertex/vertex.go`:
- Around line 3161-3168: The loop currently swallows malformed JSONL lines;
instead, catch the unmarshal error inside the loop that reads into
vertexBatchOutputLine and append a parse error entry to the response
ExtraFields.ParseErrors (matching schemas.BifrostResponseExtraFields) recording
the raw line and/or the unmarshal error, then continue; update the code around
the loop that produces vertexBatchOutputLine to reference the response's
ExtraFields (or create one if missing) and push a descriptive entry for each
skipped line so callers can detect parse failures rather than silently losing
rows.
- Around line 2736-2745: The new Vertex batch/GCS request code creates fasthttp
requests (e.g., the block that acquires req/resp and sets URI, method,
content-type, Authorization, and body) but never applies
NetworkConfig.ExtraHeaders; update these request creation sites to call
providerUtils.SetExtraHeaders(req.Header, cfg.NetworkConfig.ExtraHeaders) (or
the equivalent helper used elsewhere) immediately after setting the auth header
and before sending the request so custom routing/quota/org headers are
forwarded; apply the same change to the other similar blocks referenced (the
blocks around the batch/job/GCS request code at the other mentioned locations)
to keep behavior consistent with existing Vertex request paths.
- Around line 2757-2759: The error-return branches that check resp.StatusCode()
!= fasthttp.StatusOK must evict the cached token source when the status is 401
or 403 before returning; call removeVertexClient(ctx, projectID) (or the
equivalent cache eviction helper used for vertexTokenSourcePool) immediately
when resp.StatusCode() is fasthttp.StatusUnauthorized or
fasthttp.StatusForbidden, then proceed to return the enriched error (e.g., the
paths that call parseVertexJobAPIError and providerUtils.EnrichError should
first remove the cached client). Apply the same change to the other similar
branches referenced (the ones using parseVertexJobAPIError or GCS error parsing)
so that poisoned/rotated ADC or impersonated credentials are cleared from
vertexTokenSourcePool on auth failures.
- Around line 3116-3143: BatchResults currently uses only keys[0] which breaks
lookups when the batch belongs to a different key; update BatchResults to probe
each candidate key in the provided keys slice: for each key call
vertexGetBatchJob(ctx, key, request.BatchID), skip non-fatal errors and
continue, then on a successful job that has OutputInfo.GcsOutputDirectory set
parse it with parseGCSURI, obtain authHeader via gcsGetAuthHeader(key), list
objects with gcsListAllObjects and return the results; only return an operation
error if all keys fail (aggregate/track errors or return the last meaningful
error), and stop iterating as soon as a successful GCS listing is obtained.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: 50e9aed2-424b-4f42-a02f-316a7e569042
📒 Files selected for processing (3)
core/providers/vertex/batch.gocore/providers/vertex/types.gocore/providers/vertex/vertex.go
TejasGhatte
force-pushed
the
06-08-feat_vertex_batches_api
branch
from
June 8, 2026 20:24
75f5eda to
05d2e26
Compare
TejasGhatte
force-pushed
the
06-03-feat_vertex_files_api
branch
from
June 8, 2026 20:24
67c8118 to
f6ba673
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (2)
core/providers/vertex/vertex.go (2)
3262-3274:⚠️ Potential issue | 🟠 Major | ⚡ Quick winEvict the cached Vertex token source on GCS 401/403 responses.
These helper branches still return immediately on auth failures without calling
removeVertexClient(...). After ADC or impersonated credentials are rotated/revoked,BatchResultscan keep reusing the poisoned token source until some other Vertex path happens to clear it.🛡️ Suggested fix
if statusCode != fasthttp.StatusOK { + if statusCode == fasthttp.StatusUnauthorized || statusCode == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } apiErr = parseGCSAPIError(resp.Body(), statusCode, "list") } @@ if resp.StatusCode() != fasthttp.StatusOK { + if resp.StatusCode() == fasthttp.StatusUnauthorized || resp.StatusCode() == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } return nil, parseGCSAPIError(resp.Body(), resp.StatusCode(), "content download") }Also applies to: 3302-3304
🤖 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/vertex/vertex.go` around lines 3262 - 3274, When a GCS call returns 401 or 403 we must evict the cached Vertex token source so BatchResults won't reuse a poisoned credential; update the error handling in the GCS list/other response branches (the block checking statusCode != fasthttp.StatusOK and the similar branch around lines 3302-3304) to call removeVertexClient(...) before returning (i.e., when parseGCSAPIError reports an auth error), then continue to release request/response and return the API error as before; ensure you call removeVertexClient with the same identifiers used when the client was cached so the token source is cleared for future calls.
3199-3201:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winDon't silently drop malformed batch-result rows.
continuehere returns a partial result set with no signal that rows were lost. Record these decode failures inExtraFields.ParseErrors(or another surfaced field on the response) so callers can tell “no result” from “corrupt result file”.🤖 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/vertex/vertex.go` around lines 3199 - 3201, The loop that unmarshals rawLine into VertexBatchOutputLine silently skips malformed rows; instead capture these decode failures by appending a descriptive entry (including err.Error() and identifying info such as the rawLine or row index) to the response's ExtraFields.ParseErrors (or the existing ExtraFields map) so callers can detect corrupt rows. Locate the unmarshalling block around VertexBatchOutputLine and rawLine, replace the bare continue with code that creates/initializes ExtraFields.ParseErrors if needed and appends the parse error string, then continue to the next line; ensure you use the same response object/variable that collects batch results so the error is surfaced with the result set.
🤖 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/vertex/vertex.go`:
- Around line 2719-2728: The code sets
ctx.SetValue(schemas.BifrostContextKeyPassthroughExtraParams, true) before
calling providerUtils.CheckContextAndGetRequestBody which mutates the shared
BifrostContext in place; save the previous value of
schemas.BifrostContextKeyPassthroughExtraParams from ctx, set the flag as
currently done, then defer restoring the original value immediately after
setting it so that after calling providerUtils.CheckContextAndGetRequestBody
(and ToVertexBatchCreateRequest) the context is returned to its prior state and
passthrough extra_params won’t leak into retries/fallbacks.
In `@tests/integrations/python/tests/test_google.py`:
- Around line 2971-2974: The Vertex native tests (e.g., test_vertex_batch_get
and the other tests at the noted ranges) are incorrectly gated by
skip_if_no_api_key("vertex"); remove that decorator so the tests rely only on
the native-Vertex skip helper (skip_if_no_vertex_native_batch or similar)
because get_vertex_job_service_client() uses AnonymousCredentials() and Bifrost
supplies credentials from the gateway; update the decorators on
test_vertex_batch_get and the sibling tests (lines referenced: ~2971, ~3001,
~3015, ~3052 ranges) to drop skip_if_no_api_key("vertex") and keep only the
native Vertex skip helper.
- Around line 2940-2943: The cleanup currently calls
client.delete_batch_prediction_job(name=job_name) and returns immediately, which
can leave the job if the long-running delete fails; change each call
(client.delete_batch_prediction_job) to capture the returned LRO, wait for
completion (e.g., call result() or wait() on the operation) inside the try
block, and handle/log any exceptions raised during the wait so cleanup only
succeeds after the delete operation finishes; update all occurrences referenced
(the calls around job_name at the three noted locations).
- Around line 2951-2959: Tests are using create_google_batch_json_content
(developer-batches {"key":..., "request":..."}) but the Vertex provider expects
native Vertex JSONL lines of {"request": {...}} with custom_id carried via
request labels in core/providers/vertex/batch.go; replace the usage of
create_google_batch_json_content in the test (calls to stage_vertex_batch_input
and build_vertex_batch_prediction_job) with a new helper (e.g.,
create_vertex_batch_json_content) that emits one JSON-per-line where each line
is {"request": {...}} containing contents/config as Vertex expects, and update
the other occurrences in the same test file that call
create_google_batch_json_content accordingly so the staged input matches
Vertex's contract.
- Around line 3007-3013: The test currently only iterates
client.list_batch_prediction_jobs and can pass if the iterator is empty; modify
the test to create a real batch prediction job first (use whatever helper or API
used elsewhere in the suite to create a Vertex batch job), capture its returned
resource name, then call
client.list_batch_prediction_jobs(parent=self._vertex_parent()) and assert that
at least one returned job has name == the created job's resource name (instead
of merely counting up to 10); keep the existing loop structure (variables
count/job) but add the creation step before it and a boolean or direct assert
that the created job is found in the listing.
---
Duplicate comments:
In `@core/providers/vertex/vertex.go`:
- Around line 3262-3274: When a GCS call returns 401 or 403 we must evict the
cached Vertex token source so BatchResults won't reuse a poisoned credential;
update the error handling in the GCS list/other response branches (the block
checking statusCode != fasthttp.StatusOK and the similar branch around lines
3302-3304) to call removeVertexClient(...) before returning (i.e., when
parseGCSAPIError reports an auth error), then continue to release
request/response and return the API error as before; ensure you call
removeVertexClient with the same identifiers used when the client was cached so
the token source is cleared for future calls.
- Around line 3199-3201: The loop that unmarshals rawLine into
VertexBatchOutputLine silently skips malformed rows; instead capture these
decode failures by appending a descriptive entry (including err.Error() and
identifying info such as the rawLine or row index) to the response's
ExtraFields.ParseErrors (or the existing ExtraFields map) so callers can detect
corrupt rows. Locate the unmarshalling block around VertexBatchOutputLine and
rawLine, replace the bare continue with code that creates/initializes
ExtraFields.ParseErrors if needed and appends the parse error string, then
continue to the next line; ensure you use the same response object/variable that
collects batch results so the error is surfaced with the result set.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: 362be416-cca2-447b-89a1-5320700ba9b3
📒 Files selected for processing (7)
core/providers/vertex/batch.gocore/providers/vertex/types.gocore/providers/vertex/vertex.gotests/integrations/python/config.ymltests/integrations/python/tests/test_google.pytests/integrations/python/tests/utils/common.pytests/integrations/python/tests/utils/config_loader.py
TejasGhatte
force-pushed
the
06-08-feat_vertex_batches_api
branch
from
June 9, 2026 05:26
05d2e26 to
3a3ee1d
Compare
TejasGhatte
force-pushed
the
06-03-feat_vertex_files_api
branch
from
June 9, 2026 05:26
f6ba673 to
12bdf28
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (6)
tests/integrations/python/tests/test_google.py (4)
2940-2943:⚠️ Potential issue | 🟡 Minor | 💤 Low valueDelete operation returns an LRO that should be awaited for reliable cleanup.
The
delete_batch_prediction_job()method returns a long-running operation. If the LRO fails after the call returns, the test passes but leaves an orphaned job behind. Consider waiting for the delete operation to complete.Suggested fix
try: - client.delete_batch_prediction_job(name=job_name) + op = client.delete_batch_prediction_job(name=job_name) + op.result(timeout=300) except Exception as e: print(f"Cleanup info: Could not delete job: {e}")🤖 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 `@tests/integrations/python/tests/test_google.py` around lines 2940 - 2943, The cleanup currently calls client.delete_batch_prediction_job(name=job_name) but does not wait for the long-running operation to finish, risking orphaned jobs; change the code to capture the returned operation (op = client.delete_batch_prediction_job(name=job_name)) and wait for completion (e.g., op.result(timeout=...) or op.wait()) and handle exceptions from the result call so failures during deletion are logged/raised; update the try/except around client.delete_batch_prediction_job and the operation result to reference the operation variable and job_name for clear diagnostics.
2951-2959:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftInput JSONL uses Gemini Developer batches format instead of Vertex batch format.
create_google_batch_json_content()produces{"key": ..., "request": ...}(Gemini Developer batches format), but the Vertex batch prediction provider incore/providers/vertex/batch.goexpects{"request": {...}}withcustom_idcarried via request labels.Consider creating a dedicated
create_vertex_batch_json_content()helper that produces the Vertex-native format:def create_vertex_batch_json_content(num_requests: int = 2) -> str: requests_list = [] for i in range(num_requests): prompt = BATCH_INLINE_PROMPTS[i % len(BATCH_INLINE_PROMPTS)] line = { "request": { "contents": [{"role": "user", "parts": [{"text": prompt}]}], } } requests_list.append(json.dumps(line)) return "\n".join(requests_list)🤖 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 `@tests/integrations/python/tests/test_google.py` around lines 2951 - 2959, The tests call create_google_batch_json_content() which emits Gemini Developer batch lines {"key":..., "request":...} but the Vertex batch handler in core/providers/vertex/batch.go expects each line to be {"request": {...}} with any custom_id carried via request labels; change the test to produce Vertex-native JSONL by adding a new helper create_vertex_batch_json_content(num_requests:int=2) and use it where build_vertex_batch_prediction_job/stage_vertex_batch_input are used in the failing test; implement create_vertex_batch_json_content to iterate requests, wrap the prompt in {"request": {"contents":[{"role":"user","parts":[{"text":...}]}]}} and return newline-joined JSON strings so the Vertex provider receives the expected shape.
3094-3095:⚠️ Potential issue | 🟡 Minor | 💤 Low valueWait for the delete LRO to complete.
Same as the cleanup helper issue —
delete_batch_prediction_job()returns an LRO. Consider calling.result()on the returned operation to ensure the delete completes before the test ends.- client.delete_batch_prediction_job(name=job.name) + op = client.delete_batch_prediction_job(name=job.name) + op.result(timeout=300) print(f"Success: Deleted Vertex batch job {job.name}")🤖 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 `@tests/integrations/python/tests/test_google.py` around lines 3094 - 3095, The test calls client.delete_batch_prediction_job(name=job.name) but ignores the long-running operation it returns; change the cleanup to capture the returned operation from client.delete_batch_prediction_job(...) and call .result() on it to wait for completion before proceeding (ensure you still log/print the successful deletion referencing job.name after the .result() completes).
2971-2974:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRemove
@skip_if_no_api_key("vertex")decorator from native Vertex tests.These tests use
get_vertex_job_service_client()which creates a client withAnonymousCredentials(). The docstring correctly notes that Bifrost supplies the real Vertex credentials from its key config. Therefore, gating onVERTEX_API_KEYenv var will skip valid test environments where the gateway is configured correctly but the local env var is intentionally absent.The tests should rely only on
skip_if_no_vertex_native_batch()which checks for GCS configuration.- `@skip_if_no_api_key`("vertex") def test_vertex_batch_get(self, test_config):Also applies to: 3001-3004, 3032-3033, 3069-3070
🤖 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 `@tests/integrations/python/tests/test_google.py` around lines 2971 - 2974, Remove the `@skip_if_no_api_key`("vertex") decorator from native Vertex test functions that call get_vertex_job_service_client() (e.g., test_vertex_batch_get) because those clients use AnonymousCredentials and rely on gateway-supplied credentials; leave the skip_if_no_vertex_native_batch() check in place. Update each affected test to only use skip_if_no_vertex_native_batch() and remove the `@skip_if_no_api_key`("vertex") line from other native Vertex tests listed (the additional occurrences noted in the review) so they no longer gate on the VERTEX_API_KEY env var.core/providers/vertex/vertex.go (2)
3202-3205:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winExpose malformed JSONL rows instead of silently dropping them.
Skipping bad lines with
continuereturns a partial result set with no signal that data was lost. Record each parse failure inExtraFields.ParseErrorsand keep going.🤖 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/vertex/vertex.go` around lines 3202 - 3205, When sonic.Unmarshal(rawLine, &line) fails inside the VertexBatchOutputLine parsing loop, do not silently continue; instead append a parse-failure entry to the container's ExtraFields.ParseErrors (include the rawLine and the error message) and then continue processing the next line. Update the code around the sonic.Unmarshal call so parse errors are recorded on the relevant result object (use the existing ExtraFields.ParseErrors field) while still skipping the malformed row from contributing a valid VertexBatchOutputLine.
3231-3284:⚠️ Potential issue | 🟠 Major | ⚡ Quick winEvict the cached token source on GCS 401/403 responses.
gcsGetAuthHeaderreuses the cached token source, but these helper error paths return immediately on GCS auth failures. After credential rotation or revocation,BatchResultswill keep reusing the poisoned source until some other endpoint happens to clear it. The root cause is that these helpers no longer have access to the key/auth-credentials needed to callremoveVertexClient(...).Also applies to: 3288-3311
🤖 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/vertex/vertex.go` around lines 3231 - 3284, When GCS returns 401/403 the code should evict the poisoned cached token source so subsequent calls won't reuse it; update gcsListAllObjects to detect statusCode == 401 || statusCode == 403 (after parseGCSAPIError returns apiErr) and call the provider cache-eviction routine (e.g. removeVertexClient or a new provider.evictCachedTokenSource helper) with the same client/key identity used when caching the token source, then proceed to return the apiErr; apply the same eviction logic to the other GCS helper referenced around lines 3288-3311 (the functions that call gcsGetAuthHeader and parseGCSAPIError) so token sources are cleared on auth failures.
🤖 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/vertex/vertex.go`:
- Around line 2638-2645: When a file input is provided the code currently only
checks presence (hasFileInput) but allows any opaque string; add upfront
validation to ensure request.InputFileID is a valid GCS URI (starts with
"gs://", non-empty path, and preferably ends with .jsonl or at least has an
object name) before calling the Vertex batch API. In the block guarded by
hasFileInput (where NewBifrostOperationError is returned for mutual/exclusive
cases), parse/validate request.InputFileID and return
providerUtils.NewBifrostOperationError with a clear message if it is not a
well-formed gs:// URI so invalid client input is rejected early. Ensure you
reference request.InputFileID and keep the existing mutual-exclusion checks
(hasInlineRequests) intact.
In `@tests/integrations/python/config.yml`:
- Around line 168-173: The new batch model mappings (batch_create, batch_inline,
batch_file_upload, batch_list, batch_retrieve, batch_cancel) in
tests/integrations/python/config.yml are unused because the Vertex batch
scenarios (provider_scenarios.vertex.batch_*) are disabled; update the
provider_scenarios.vertex.batch_create/batch_inline/batch_file_upload/batch_list/batch_retrieve/batch_cancel
flags to true to enable config-driven batch test coverage, or alternatively
remove the unused batch_* model mappings if you intentionally want to keep
Vertex batch scenarios disabled; ensure you change the
provider_scenarios.vertex.batch_* booleans (the exact names shown) to reflect
the intended test behavior.
In `@transports/bifrost-http/integrations/genai.go`:
- Around line 758-784: The Vertex list route is losing pagination because
extractVertexBatchPathParams (used as PreCallback in the RouteConfig for
collectionPath) only stamps provider/batch IDs and does not copy HTTP query
pagination into the BifrostBatchListRequest; update either
extractVertexBatchPathParams or the RouteConfig's BatchRequestConverter to read
pageSize and pageToken from the incoming request context and set them on the
created *schemas.BifrostBatchListRequest (the object returned by
GetRequestTypeInstance and consumed in BatchRequestConverter) so that
BifrostBatchListRequest.pageSize/pageToken are preserved for Vertex paginated
calls.
- Around line 859-862: The BatchDeleteResponseConverter currently returns a
synthetic {"done": true} and discards the provider response; change it to return
the provider raw response so native Vertex delete operation metadata is
preserved (i.e., return the supplied resp value from
BatchDeleteResponseConverter instead of the synthetic map). Update the converter
implementation in BatchDeleteResponseConverter to pass through resp (or
resp.Operation if the other batch converters expose the nested operation) so the
transport remains compatible with provider-native request/response shapes and
retains operation metadata.
---
Duplicate comments:
In `@core/providers/vertex/vertex.go`:
- Around line 3202-3205: When sonic.Unmarshal(rawLine, &line) fails inside the
VertexBatchOutputLine parsing loop, do not silently continue; instead append a
parse-failure entry to the container's ExtraFields.ParseErrors (include the
rawLine and the error message) and then continue processing the next line.
Update the code around the sonic.Unmarshal call so parse errors are recorded on
the relevant result object (use the existing ExtraFields.ParseErrors field)
while still skipping the malformed row from contributing a valid
VertexBatchOutputLine.
- Around line 3231-3284: When GCS returns 401/403 the code should evict the
poisoned cached token source so subsequent calls won't reuse it; update
gcsListAllObjects to detect statusCode == 401 || statusCode == 403 (after
parseGCSAPIError returns apiErr) and call the provider cache-eviction routine
(e.g. removeVertexClient or a new provider.evictCachedTokenSource helper) with
the same client/key identity used when caching the token source, then proceed to
return the apiErr; apply the same eviction logic to the other GCS helper
referenced around lines 3288-3311 (the functions that call gcsGetAuthHeader and
parseGCSAPIError) so token sources are cleared on auth failures.
In `@tests/integrations/python/tests/test_google.py`:
- Around line 2940-2943: The cleanup currently calls
client.delete_batch_prediction_job(name=job_name) but does not wait for the
long-running operation to finish, risking orphaned jobs; change the code to
capture the returned operation (op =
client.delete_batch_prediction_job(name=job_name)) and wait for completion
(e.g., op.result(timeout=...) or op.wait()) and handle exceptions from the
result call so failures during deletion are logged/raised; update the try/except
around client.delete_batch_prediction_job and the operation result to reference
the operation variable and job_name for clear diagnostics.
- Around line 2951-2959: The tests call create_google_batch_json_content() which
emits Gemini Developer batch lines {"key":..., "request":...} but the Vertex
batch handler in core/providers/vertex/batch.go expects each line to be
{"request": {...}} with any custom_id carried via request labels; change the
test to produce Vertex-native JSONL by adding a new helper
create_vertex_batch_json_content(num_requests:int=2) and use it where
build_vertex_batch_prediction_job/stage_vertex_batch_input are used in the
failing test; implement create_vertex_batch_json_content to iterate requests,
wrap the prompt in {"request":
{"contents":[{"role":"user","parts":[{"text":...}]}]}} and return newline-joined
JSON strings so the Vertex provider receives the expected shape.
- Around line 3094-3095: The test calls
client.delete_batch_prediction_job(name=job.name) but ignores the long-running
operation it returns; change the cleanup to capture the returned operation from
client.delete_batch_prediction_job(...) and call .result() on it to wait for
completion before proceeding (ensure you still log/print the successful deletion
referencing job.name after the .result() completes).
- Around line 2971-2974: Remove the `@skip_if_no_api_key`("vertex") decorator from
native Vertex test functions that call get_vertex_job_service_client() (e.g.,
test_vertex_batch_get) because those clients use AnonymousCredentials and rely
on gateway-supplied credentials; leave the skip_if_no_vertex_native_batch()
check in place. Update each affected test to only use
skip_if_no_vertex_native_batch() and remove the `@skip_if_no_api_key`("vertex")
line from other native Vertex tests listed (the additional occurrences noted in
the review) so they no longer gate on the VERTEX_API_KEY env var.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: 46f139ec-8c19-4e66-a65a-efa1b2d9d1ff
📒 Files selected for processing (9)
core/providers/vertex/batch.gocore/providers/vertex/types.gocore/providers/vertex/vertex.gocore/schemas/batch.gotests/integrations/python/config.ymltests/integrations/python/tests/test_google.pytests/integrations/python/tests/utils/common.pytests/integrations/python/tests/utils/config_loader.pytransports/bifrost-http/integrations/genai.go
TejasGhatte
force-pushed
the
06-08-feat_vertex_batches_api
branch
from
June 9, 2026 09:09
3a3ee1d to
60a8d9a
Compare
TejasGhatte
force-pushed
the
06-03-feat_vertex_files_api
branch
from
June 9, 2026 09:09
12bdf28 to
4c8b9d5
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 11
♻️ Duplicate comments (6)
tests/integrations/python/tests/test_google.py (2)
2940-2943:⚠️ Potential issue | 🟠 Major | ⚡ Quick winWait for Vertex delete LRO completion before treating cleanup as successful.
delete_batch_prediction_job(...)is fire-and-forget here; if the LRO fails after return, tests still pass and jobs can leak.Suggested fix
- client.delete_batch_prediction_job(name=job_name) + delete_op = client.delete_batch_prediction_job(name=job_name) + delete_op.result(timeout=300)- client.delete_batch_prediction_job(name=job.name) + delete_op = client.delete_batch_prediction_job(name=job.name) + delete_op.result(timeout=300)- client.delete_batch_prediction_job(name=job.name) + delete_op = client.delete_batch_prediction_job(name=job.name) + delete_op.result(timeout=300)In google-cloud-aiplatform Python client, does JobServiceClient.delete_batch_prediction_job return a long-running operation that should be awaited with operation.result()?Also applies to: 3064-3067, 3094-3095
🤖 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 `@tests/integrations/python/tests/test_google.py` around lines 2940 - 2943, The cleanup currently calls client.delete_batch_prediction_job(name=job_name) fire-and-forget; change it to capture the returned long-running operation, call operation.result() to wait for completion, and handle any exceptions from result() so failures don't silently allow leaked jobs — update the code paths using client.delete_batch_prediction_job (and the similar calls at the other locations) to await the LRO completion and log or raise if operation.result() errors.
2971-2971:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRemove
@skip_if_no_api_key("vertex")from native Vertex batch tests.These tests already call
skip_if_no_vertex_native_batch()and use the native Vertex client path; this extra gate can skip valid environments and hide regressions.Suggested fix
- `@skip_if_no_api_key`("vertex") def test_vertex_batch_get(self, test_config): @@ - `@skip_if_no_api_key`("vertex") def test_vertex_batch_list(self, test_config): @@ - `@skip_if_no_api_key`("vertex") def test_vertex_batch_cancel(self, test_config): @@ - `@skip_if_no_api_key`("vertex") def test_vertex_batch_delete(self, test_config):Also applies to: 3001-3001, 3032-3032, 3069-3069
🤖 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 `@tests/integrations/python/tests/test_google.py` at line 2971, Remove the redundant `@skip_if_no_api_key`("vertex") decorator from the native Vertex batch tests that already call skip_if_no_vertex_native_batch() and exercise the native Vertex client path; specifically, find occurrences of the decorator (skip_if_no_api_key) used on tests that call skip_if_no_vertex_native_batch() and delete that decorator so those tests rely only on skip_if_no_vertex_native_batch() to control skipping.core/providers/vertex/vertex.go (4)
3219-3299:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPass the key into the GCS helpers so 401/403 can evict the cached token source.
vertexTokenSourcePoolis supposed to be cleared on auth failures, butgcsListAllObjectsandgcsDownloadObjectonly receive the bearer string. When GCS returns 401/403 fromBatchResultsorFile*, these helpers cannot callremoveVertexClient(...), so retries keep reusing the poisoned source.♻️ Suggested fix pattern
-func (provider *VertexProvider) gcsListAllObjects(ctx *schemas.BifrostContext, authHeader, bucket, prefix string) ([]gcsObjectMetadata, *schemas.BifrostError) { +func (provider *VertexProvider) gcsListAllObjects(ctx *schemas.BifrostContext, key schemas.Key, authHeader, bucket, prefix string) ([]gcsObjectMetadata, *schemas.BifrostError) { @@ - if statusCode != fasthttp.StatusOK { + if statusCode != fasthttp.StatusOK { + if statusCode == fasthttp.StatusUnauthorized || statusCode == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } apiErr = parseGCSAPIError(resp.Body(), statusCode, "list") }Apply the same change to
gcsDownloadObjectand update the call sites.🤖 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/vertex/vertex.go` around lines 3219 - 3299, gcsListAllObjects and gcsDownloadObject only accept the bearer string, so on 401/403 they cannot evict the cached token source; modify both functions to accept the vertex client key/identifier (the same key used by vertexTokenSourcePool/removeVertexClient), propagate that extra parameter from all callers (e.g. BatchResults/File* call sites), and call removeVertexClient(key) when parseGCSAPIError indicates a 401/403 before returning the error so the poisoned token source gets cleared and retries obtain a fresh token.
2759-2778:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPropagate provider response headers on the new batch/file success paths.
These returns bypass the existing Vertex pattern of publishing filtered upstream headers. Most never set
ExtraFields.ProviderResponseHeaders; the upload paths set it but still skipctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders, ...). That drops request-id/quota metadata for the new batch and GCS routes.♻️ Suggested fix pattern
+ headers := providerUtils.ExtractProviderResponseHeaders(resp) + ctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders, headers) + result := &schemas.BifrostBatchCreateResponse{ ... ExtraFields: schemas.BifrostResponseExtraFields{ - Latency: time.Since(startTime).Milliseconds(), + Latency: time.Since(startTime).Milliseconds(), + ProviderResponseHeaders: headers, }, }Apply the same pattern to each new success return in this batch/GCS surface.
As per coding guidelines, provider changes should preserve response/error metadata. Based on learnings, filtered provider headers should be surfaced through both
ExtraFields.ProviderResponseHeadersandschemas.BifrostContextKeyProviderResponseHeaders.Also applies to: 2905-2913, 2934-2941, 3048-3056, 3114-3121, 3209-3215, 3517-3531, 3581-3596, 3688-3698, 3771-3785, 3846-3853, 3920-3927
🤖 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/vertex/vertex.go` around lines 2759 - 2778, The new-batch/GCS success paths (e.g., where result is constructed and fields like ExtraFields.RawRequest/RawResponse are set) are not propagating filtered provider response headers; update each success return to set result.ExtraFields.ProviderResponseHeaders from the filtered headers and also call ctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders, filteredHeaders) so upstream metadata (request-id/quota) is preserved; locate the construction of result (symbol: result := &schemas.BifrostBatchCreateResponse) and after any RawRequest/RawResponse handling add the ProviderResponseHeaders assignment and ctx.SetValue call following the same pattern used in the upload paths.Sources: Coding guidelines, Learnings
3186-3193:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winDon’t silently drop malformed prediction rows.
The
continueon unmarshal failure makesBatchResultslook complete even when one of thepredictions*.jsonlrows was lost. Record each skipped row inExtraFields.ParseErrors(or fail the request) so callers can detect partial corruption.🤖 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/vertex/vertex.go` around lines 3186 - 3193, The loop over bytes.Split(content, []byte("\n")) currently skips malformed lines on sonic.Unmarshal failure; instead, capture each parse failure into the BatchResults ExtraFields.ParseErrors so callers can detect partial corruption: when sonic.Unmarshal(rawLine, &line) returns an error, append an entry including the rawLine (or its trimmed string) and the error message to the container used for BatchResults.ExtraFields.ParseErrors (create the slice/map if missing) rather than silently continuing, and only skip adding that line to the successful results; use the existing VertexBatchOutputLine, rawLine, err and the BatchResults/ExtraFields structures to locate and populate the parse-errors metadata.
2637-2680:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winValidate both GCS URIs before creating the batch job.
InputFileIDis only treated as “present or absent”, so an opaque string orgs://bucketwith no object name is forwarded upstream.output_folder.urlis only parsed in the inline-upload branch, so the file-input path can still send a non-gs://destination to Vertex.♻️ Suggested fix
if !hasFileInput && !hasInlineRequests { return nil, providerUtils.NewBifrostOperationError("either input_file_id (gs:// JSONL URI) or requests is required for Vertex batch API", nil) } + if hasFileInput { + _, objectKey, err := parseGCSURI(request.InputFileID) + if err != nil || strings.Trim(objectKey, "/") == "" { + return nil, providerUtils.NewBifrostOperationError("input_file_id must be a gs:// URI pointing to a JSONL object", nil) + } + } @@ if outputURI == "" { return nil, providerUtils.NewBifrostOperationError("output_folder.url (gs:// prefix) is required for Vertex batch API", nil) } + if _, _, err := parseGCSURI(outputURI); err != nil { + return nil, providerUtils.NewBifrostOperationError("output_folder.url must be a gs:// prefix", nil) + }As per coding guidelines, validate all untrusted input before provider calls.
🤖 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/vertex/vertex.go` around lines 2637 - 2680, Validate both request.InputFileID and request.OutputFolder.URL as proper GCS URIs before creating the Vertex batch job: parse and check request.OutputFolder.URL (trimmed) right after you derive outputURI and return a NewBifrostOperationError if parseGCSURI fails, and if request.InputFileID != "" call parseGCSURI(inputFileID) (or equivalent validation) to ensure it includes a bucket and object path and return a NewBifrostOperationError on failure; move or add these validations before any calls that build the job (e.g., before vertexBatchJobsBaseURL usage and before the inline upload branch that calls vertexConvertRequestsToJSONL) so untrusted URIs are always validated.Source: Coding guidelines
🤖 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/vertex/batch.go`:
- Around line 157-173: To preserve native Vertex input/output configs, change
ToVertexBatchCreateRequest so it does not unconditionally collapse to jsonl+GCS;
instead detect and copy the incoming Vertex-style InputConfig/OutputConfig from
the schemas.BifrostBatchCreateRequest when present (e.g., if the request already
contains BigquerySource/BigqueryDestination or a non-"jsonl"
Instances/PredictionsFormat) and populate VertexBatchCreateRequest.InputConfig
and .OutputConfig directly rather than overwriting with hard-coded
VertexBatchInputConfig/VertexBatchOutputConfig; update similar logic in the
other converter range (the other conversion block around lines 291-359) to
mirror this behavior, keeping use of the existing symbols
VertexBatchCreateRequest, VertexBatchInputConfig, VertexBatchOutputConfig,
VertexGcsSource, VertexGcsDestination, InputFileID and OutputFolder to locate
where to copy/preserve the native fields.
- Around line 157-160: In ToVertexBatchCreateRequest, avoid dereferencing the
optional pointer request.Model directly; instead initialize model to empty
string, check if request.Model != nil before assigning model = *request.Model,
and only apply the prefix logic (adding "publishers/google/models/") if model is
non-empty and doesn't contain "/"; this prevents a nil-pointer panic while
preserving existing behavior when Model is provided.
In `@core/providers/vertex/vertex_test.go`:
- Around line 32-47: The test enables file/batch scenarios unconditionally even
though fileStorageConfig and batchOutputFolder are only set when
VERTEX_GCS_BUCKET is present; update the test to be deterministic by gating
those scenarios: either set the VERTEX_GCS_BUCKET/ VERTEX_GCS_PREFIX env vars in
the test setup so fileStorageConfig and batchOutputFolder are populated, or wrap
the file/batch scenario assertions so they only run when fileStorageConfig !=
nil && batchOutputFolder != nil (or when os.Getenv("VERTEX_GCS_BUCKET") != "").
Reference the variables fileStorageConfig, batchOutputFolder and the env var
VERTEX_GCS_BUCKET to locate and fix the code paths.
In `@core/providers/vertex/vertex.go`:
- Line 3837: Remove the stray stdout print in FileDelete: delete the
fmt.Println(resp.StatusCode()) call and instead log the response using the
component's structured logger (e.g., use the existing logger or processLogger
used in this file) or handle the status code via the function's normal
error/response flow so deletes don't print directly to stdout; locate the call
inside the FileDelete function and replace it with a structured log or
conditional error handling of resp.StatusCode().
In `@core/schemas/batch.go`:
- Line 82: Add the new display_name field into the HTTP parser and request
struct wiring so client-supplied names are not left in ExtraParams: update
batchCreateParamsKnownFields (in transports/bifrost-http/handlers/inference.go)
to include "display_name", parse it into
schemas.BifrostBatchCreateRequest.DisplayName (populate the pointer when
present) and ensure the Vertex conversion path uses the correct target key (map
schemas.BifrostBatchCreateRequest.DisplayName to Vertex's "displayName") rather
than forwarding the raw "display_name" key.
In `@tests/integrations/python/tests/test_google.py`:
- Around line 155-157: The current header merge lets callers override the
routing header by doing headers = {"x-model-provider": provider} then
headers.update(extra_headers); change the merge so extra_headers cannot replace
"x-model-provider" — e.g., when combining extra_headers into headers (or
building the final headers), explicitly preserve headers["x-model-provider"]
based on the local provider variable and ignore or remove that key from
extra_headers before update; locate the header construction around the headers
and extra_headers variables in the test (headers = {"x-model-provider":
provider} and headers.update(extra_headers)) and ensure the x-model-provider
value always comes from provider.
In `@tests/integrations/python/tests/utils/common.py`:
- Around line 2684-2687: The default filename generation using int(time.time())
is collision-prone; update the filename assignment (the block that sets filename
when filename is None, which feeds into blob_name and uses cfg["prefix"]) to use
a collision-resistant identifier such as uuid.uuid4().hex (or a timestamp plus
uuid for readability), and add the corresponding import (uuid) so each test gets
a unique object name and avoids overwrites.
- Around line 2643-2655: The helper that reads VERTEX_CREDENTIALS currently
returns None on invalid values which lets stage_vertex_batch_input() silently
fall back to ADC; change it to fail fast: when raw is set but
os.path.isfile(raw) is False, raise a clear exception indicating the file is
missing; when json.loads(raw) raises ValueError/TypeError, raise a clear
exception indicating malformed JSON; ensure the exceptions surface instead of
returning None so callers (e.g. stage_vertex_batch_input) cannot silently use
ambient credentials; update the logic around raw, os.path.isfile(raw),
json.loads(raw), service_account.Credentials.from_service_account_file and
service_account.Credentials.from_service_account_info to raise these errors.
In `@tests/integrations/python/tests/utils/config_loader.py`:
- Line 26: The "vertex" key is being advertised as a routable integration (the
mapping entry "vertex": "vertex") which causes callers like list_integrations()
and get_integration_url("vertex") to expect a configured endpoint; either remove
"vertex" from the public integrations mapping in config_loader (so it is not
returned by list_integrations/get_integration_url) or add a corresponding
bifrost.endpoints.vertex entry to the runtime config.yml; alternatively, if you
need Vertex only for model selection, move the "vertex" entry into the
provider-only lookup used by model selection rather than the integrations map so
list_integrations()/get_integration_url() won't try to resolve an endpoint.
In `@transports/bifrost-http/integrations/genai.go`:
- Around line 832-835: The BatchCancelResponseConverter currently returns an
empty map for cancel responses; update it to mirror the pattern used in the
create/list/retrieve converters by checking resp.ExtraFields.RawResponse and
returning that when present (e.g., if resp != nil && resp.ExtraFields != nil &&
resp.ExtraFields.RawResponse != nil return resp.ExtraFields.RawResponse, nil),
otherwise fall back to returning an empty map as before; modify the
BatchCancelResponseConverter function in the integration (the converter for
schemas.BifrostBatchCancelResponse in the genai transport) to implement this
passthrough.
- Around line 876-898: The code extracts batchID without validation and then
assigns it to retrieve/cancel/delete requests; replicate the guard used in
extractGeminiBatchIDFromPath by validating batchID is non-empty before setting
r.BatchID for *schemas.BifrostBatchRetrieveRequest,
*schemas.BifrostBatchCancelRequest, and *schemas.BifrostBatchDeleteRequest—if
batchID is empty return an appropriate error (400/invalid path param) rather
than assigning an empty string to r.BatchID so provider calls cannot proceed
with an invalid ID.
---
Duplicate comments:
In `@core/providers/vertex/vertex.go`:
- Around line 3219-3299: gcsListAllObjects and gcsDownloadObject only accept the
bearer string, so on 401/403 they cannot evict the cached token source; modify
both functions to accept the vertex client key/identifier (the same key used by
vertexTokenSourcePool/removeVertexClient), propagate that extra parameter from
all callers (e.g. BatchResults/File* call sites), and call
removeVertexClient(key) when parseGCSAPIError indicates a 401/403 before
returning the error so the poisoned token source gets cleared and retries obtain
a fresh token.
- Around line 2759-2778: The new-batch/GCS success paths (e.g., where result is
constructed and fields like ExtraFields.RawRequest/RawResponse are set) are not
propagating filtered provider response headers; update each success return to
set result.ExtraFields.ProviderResponseHeaders from the filtered headers and
also call ctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders,
filteredHeaders) so upstream metadata (request-id/quota) is preserved; locate
the construction of result (symbol: result :=
&schemas.BifrostBatchCreateResponse) and after any RawRequest/RawResponse
handling add the ProviderResponseHeaders assignment and ctx.SetValue call
following the same pattern used in the upload paths.
- Around line 3186-3193: The loop over bytes.Split(content, []byte("\n"))
currently skips malformed lines on sonic.Unmarshal failure; instead, capture
each parse failure into the BatchResults ExtraFields.ParseErrors so callers can
detect partial corruption: when sonic.Unmarshal(rawLine, &line) returns an
error, append an entry including the rawLine (or its trimmed string) and the
error message to the container used for BatchResults.ExtraFields.ParseErrors
(create the slice/map if missing) rather than silently continuing, and only skip
adding that line to the successful results; use the existing
VertexBatchOutputLine, rawLine, err and the BatchResults/ExtraFields structures
to locate and populate the parse-errors metadata.
- Around line 2637-2680: Validate both request.InputFileID and
request.OutputFolder.URL as proper GCS URIs before creating the Vertex batch
job: parse and check request.OutputFolder.URL (trimmed) right after you derive
outputURI and return a NewBifrostOperationError if parseGCSURI fails, and if
request.InputFileID != "" call parseGCSURI(inputFileID) (or equivalent
validation) to ensure it includes a bucket and object path and return a
NewBifrostOperationError on failure; move or add these validations before any
calls that build the job (e.g., before vertexBatchJobsBaseURL usage and before
the inline upload branch that calls vertexConvertRequestsToJSONL) so untrusted
URIs are always validated.
In `@tests/integrations/python/tests/test_google.py`:
- Around line 2940-2943: The cleanup currently calls
client.delete_batch_prediction_job(name=job_name) fire-and-forget; change it to
capture the returned long-running operation, call operation.result() to wait for
completion, and handle any exceptions from result() so failures don't silently
allow leaked jobs — update the code paths using
client.delete_batch_prediction_job (and the similar calls at the other
locations) to await the LRO completion and log or raise if operation.result()
errors.
- Line 2971: Remove the redundant `@skip_if_no_api_key`("vertex") decorator from
the native Vertex batch tests that already call skip_if_no_vertex_native_batch()
and exercise the native Vertex client path; specifically, find occurrences of
the decorator (skip_if_no_api_key) used on tests that call
skip_if_no_vertex_native_batch() and delete that decorator so those tests rely
only on skip_if_no_vertex_native_batch() to control skipping.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: 6b028bc5-0e6f-40f2-be02-6ae8c177cc5e
📒 Files selected for processing (13)
core/internal/llmtests/account.gocore/internal/llmtests/batch.gocore/providers/vertex/batch.gocore/providers/vertex/types.gocore/providers/vertex/vertex.gocore/providers/vertex/vertex_test.gocore/schemas/batch.gotests/integrations/python/config.ymltests/integrations/python/tests/test_google.pytests/integrations/python/tests/utils/common.pytests/integrations/python/tests/utils/config_loader.pytransports/bifrost-http/integrations/genai.gotransports/bifrost-http/integrations/openai.go
TejasGhatte
force-pushed
the
06-03-feat_vertex_files_api
branch
from
June 9, 2026 11:28
4c8b9d5 to
1315a19
Compare
TejasGhatte
force-pushed
the
06-08-feat_vertex_batches_api
branch
from
June 9, 2026 11:28
60a8d9a to
e0d50c5
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/vertex/vertex.go`:
- Around line 3723-3726: parseGCSURI currently accepts URIs with no object path
(e.g., "gs://bucket") which yields an empty objectKey and causes invalid GCS /o/
requests; create a new parser (e.g., parseGCSObjectURI) that returns an error
when the object path is empty and replace usages in the file-object APIs (the
retrieve, delete and content handlers that call parseGCSURI) to call this new
parser; on parse failure return providerUtils.NewBifrostOperationError with the
parse error message so invalid file_id inputs are rejected before any provider
call.
- Around line 2651-2659: The file-input (input_file_id) code path must validate
output_folder.url the same as the inline path: after computing outputURI from
request.OutputFolder.URL (the existing outputURI string and empty check), call
parseGCSURI(outputURI) and handle any parsing error by returning
providerUtils.NewBifrostOperationError with the parse error before forwarding to
Vertex; ensure this validation is added in the branch that handles
request.InputFileID (or the function that processes file inputs) so opaque
strings never reach Vertex.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: c70dd9d7-db0a-48c0-80e4-f0dc858d603f
📒 Files selected for processing (14)
core/internal/llmtests/account.gocore/internal/llmtests/batch.gocore/providers/vertex/batch.gocore/providers/vertex/types.gocore/providers/vertex/vertex.gocore/providers/vertex/vertex_test.gocore/schemas/batch.gotests/integrations/python/config.ymltests/integrations/python/tests/test_google.pytests/integrations/python/tests/utils/common.pytests/integrations/python/tests/utils/config_loader.pytransports/bifrost-http/handlers/inference.gotransports/bifrost-http/integrations/genai.gotransports/bifrost-http/integrations/openai.go
💤 Files with no reviewable changes (7)
- tests/integrations/python/tests/utils/config_loader.py
- tests/integrations/python/config.yml
- transports/bifrost-http/handlers/inference.go
- tests/integrations/python/tests/test_google.py
- transports/bifrost-http/integrations/genai.go
- tests/integrations/python/tests/utils/common.py
- transports/bifrost-http/integrations/openai.go
Contributor
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 2
🤖 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/vertex/vertex.go`:
- Around line 3723-3726: parseGCSURI currently accepts URIs with no object path
(e.g., "gs://bucket") which yields an empty objectKey and causes invalid GCS /o/
requests; create a new parser (e.g., parseGCSObjectURI) that returns an error
when the object path is empty and replace usages in the file-object APIs (the
retrieve, delete and content handlers that call parseGCSURI) to call this new
parser; on parse failure return providerUtils.NewBifrostOperationError with the
parse error message so invalid file_id inputs are rejected before any provider
call.
- Around line 2651-2659: The file-input (input_file_id) code path must validate
output_folder.url the same as the inline path: after computing outputURI from
request.OutputFolder.URL (the existing outputURI string and empty check), call
parseGCSURI(outputURI) and handle any parsing error by returning
providerUtils.NewBifrostOperationError with the parse error before forwarding to
Vertex; ensure this validation is added in the branch that handles
request.InputFileID (or the function that processes file inputs) so opaque
strings never reach Vertex.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: c70dd9d7-db0a-48c0-80e4-f0dc858d603f
📒 Files selected for processing (14)
core/internal/llmtests/account.gocore/internal/llmtests/batch.gocore/providers/vertex/batch.gocore/providers/vertex/types.gocore/providers/vertex/vertex.gocore/providers/vertex/vertex_test.gocore/schemas/batch.gotests/integrations/python/config.ymltests/integrations/python/tests/test_google.pytests/integrations/python/tests/utils/common.pytests/integrations/python/tests/utils/config_loader.pytransports/bifrost-http/handlers/inference.gotransports/bifrost-http/integrations/genai.gotransports/bifrost-http/integrations/openai.go
💤 Files with no reviewable changes (7)
- tests/integrations/python/tests/utils/config_loader.py
- tests/integrations/python/config.yml
- transports/bifrost-http/handlers/inference.go
- tests/integrations/python/tests/test_google.py
- transports/bifrost-http/integrations/genai.go
- tests/integrations/python/tests/utils/common.py
- transports/bifrost-http/integrations/openai.go
🛑 Comments failed to post (2)
core/providers/vertex/vertex.go (2)
2651-2659:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winValidate
output_folder.urlin the file-input path too.Right now this is only checked for non-empty. Inline requests eventually hit
parseGCSURI(outputURI), but theinput_file_idpath forwards any opaque string to Vertex and turns a deterministic client error into an upstream 4xx.♻️ Suggested fix
if outputURI == "" { return nil, providerUtils.NewBifrostOperationError("output_folder.url (gs:// prefix) is required for Vertex batch API", nil) } + if _, _, parseErr := parseGCSURI(outputURI); parseErr != nil { + return nil, providerUtils.NewBifrostOperationError(parseErr.Error(), nil) + }As per coding guidelines, validate all untrusted input before provider calls.
🤖 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/vertex/vertex.go` around lines 2651 - 2659, The file-input (input_file_id) code path must validate output_folder.url the same as the inline path: after computing outputURI from request.OutputFolder.URL (the existing outputURI string and empty check), call parseGCSURI(outputURI) and handle any parsing error by returning providerUtils.NewBifrostOperationError with the parse error before forwarding to Vertex; ensure this validation is added in the branch that handles request.InputFileID (or the function that processes file inputs) so opaque strings never reach Vertex.Source: Coding guidelines
3723-3726:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winRequire an object path for the file-object APIs.
parseGCSURI("gs://bucket")succeeds with an emptyobjectKey, so these paths build/o/requests against GCS instead of rejecting an invalidfile_id. Please add a dedicated object-URI parser and reuse it across retrieve/delete/content.♻️ Suggested fix
+func parseGCSObjectURI(uri string) (bucket, objectKey string, err error) { + bucket, objectKey, err = parseGCSURI(uri) + if err != nil { + return "", "", err + } + if objectKey == "" { + return "", "", fmt.Errorf("invalid GCS URI %q: object name is required", uri) + } + return bucket, objectKey, nil +}- bucket, objectKey, parseErr := parseGCSURI(request.FileID) + bucket, objectKey, parseErr := parseGCSObjectURI(request.FileID)As per coding guidelines, validate all untrusted input before provider calls.
Also applies to: 3810-3813, 3876-3879
🤖 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/vertex/vertex.go` around lines 3723 - 3726, parseGCSURI currently accepts URIs with no object path (e.g., "gs://bucket") which yields an empty objectKey and causes invalid GCS /o/ requests; create a new parser (e.g., parseGCSObjectURI) that returns an error when the object path is empty and replace usages in the file-object APIs (the retrieve, delete and content handlers that call parseGCSURI) to call this new parser; on parse failure return providerUtils.NewBifrostOperationError with the parse error message so invalid file_id inputs are rejected before any provider call.Source: Coding guidelines
Collaborator
Merge activity
|
Pratham-Mishra04
changed the base branch from
06-03-feat_vertex_files_api
to
graphite-base/4152
June 9, 2026 13:58
Pratham-Mishra04
force-pushed
the
06-08-feat_vertex_batches_api
branch
from
June 9, 2026 14:00
e0d50c5 to
7b7ebfc
Compare
18 tasks
akshaydeo
pushed a commit
that referenced
this pull request
Jun 12, 2026
## Summary Implements full Vertex AI batch prediction support, replacing the previous stub implementations that returned `UnsupportedOperationError` for all batch operations. The Vertex AI Batch Prediction API uses GCS for input/output and `BatchPredictionJob` resources rather than OpenAI-style file IDs, requiring a dedicated mapping layer. ## Changes - **`BatchCreate`**: Accepts either a `gs://` input file URI or inline request items. Inline items are serialized to Vertex-format JSONL (with `custom_id` embedded in request labels via `bifrost_custom_id`) and uploaded to GCS before the job is submitted. The GCS output prefix is resolved from `extra_params["output_uri"]` or derived from `extra_params["gcs_bucket"]`/`["gcs_prefix"]`. Vertex-native fields (`modelParameters`, `labels`, `encryptionSpec`, etc.) are passed through via `extra_params`. - **`BatchList`**: Paginates across all configured keys using `SerialListHelper`, since batch jobs are scoped to a project/region. Each key's native Vertex `pageToken` is re-encoded into the Bifrost cursor. - **`BatchRetrieve`**: Tries each key in turn until the job is found, since a job ID is only resolvable within the project/region that created it. - **`BatchCancel`** / **`BatchDelete`**: Same multi-key fan-out pattern as retrieve. - **`BatchResults`**: Fetches the job to locate its GCS output directory, lists all `predictions-*.jsonl` files, downloads and parses each line, and recovers `custom_id` from the echoed request labels. - **`vertexJobStateToBatchStatus`**: Maps Vertex `JOB_STATE_*` values to Bifrost `BatchStatus` constants. - **`ToVertexBatchCreateRequest`**: Maps a Bifrost batch create request to a `VertexBatchCreateRequest`, stripping Bifrost control keys (`output_uri`, `gcs_bucket`, `gcs_prefix`, `job_name`) before forwarding `extra_params` to Vertex. - **`vertexConvertRequestsToJSONL`**: Converts inline `BatchRequestItem` entries to Vertex JSONL, injecting `custom_id` into each request's `labels` map without mutating the caller's data. - **GCS helpers** (`gcsListAllObjects`, `gcsDownloadObject`): Added to support paginated object listing and raw object download needed by `BatchResults`. - **Batch prediction types**: Added `vertexBatchPredictionJob`, `VertexBatchCreateRequest`, `vertexBatchJobListResponse`, `vertexBatchOutputLine`, and supporting config/stats structs. A regional Vertex key (e.g. `us-central1`) is required; `global` is explicitly rejected since the Batch Prediction API does not support it. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [ ] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/providers/vertex/... ``` To exercise end-to-end: 1. Configure a Vertex key with a regional `VertexKeyConfig` (e.g. `us-central1`) and a GCS bucket accessible by the service account. 2. Call `BatchCreate` with either `input_file_id` (a `gs://` JSONL URI) or inline `requests`, and set `extra_params["gcs_bucket"]` or `extra_params["output_uri"]`. 3. Poll `BatchRetrieve` until the job reaches `completed`. 4. Call `BatchResults` to retrieve per-request responses with `custom_id` round-tripped correctly. 5. Verify `BatchList`, `BatchCancel`, and `BatchDelete` against jobs in the target project/region. ## Breaking changes - [ ] Yes - [x] No ## Related issues ## Security considerations - GCS and Vertex API calls are authenticated via the existing `gcsGetAuthHeader` helper (service account credentials from the key config). No credentials are logged or returned in responses. - The `bifrost_custom_id` label is user-supplied and echoed back from Vertex; callers should treat it as untrusted input if used downstream. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Full Vertex AI Batch Prediction: create, list, retrieve, cancel, delete, and fetch results with status/timestamp mapping, structured errors, display_name support, and GCS upload/download helpers. * Inline JSONL support: converts inline requests, injects/preserves custom IDs, stages to GCS, and parses prediction JSONL outputs. * Vertex-native HTTP routes for batch operations and response passthrough. * **Tests** * End-to-end Vertex batch integration tests, GCS staging helpers, and test config mappings. * **Chores** * Schema/config updates to enable Vertex batch workflows and request handling. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
18 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements full Vertex AI batch prediction support, replacing the previous stub implementations that returned
UnsupportedOperationErrorfor all batch operations. The Vertex AI Batch Prediction API uses GCS for input/output andBatchPredictionJobresources rather than OpenAI-style file IDs, requiring a dedicated mapping layer.Changes
BatchCreate: Accepts either ags://input file URI or inline request items. Inline items are serialized to Vertex-format JSONL (withcustom_idembedded in request labels viabifrost_custom_id) and uploaded to GCS before the job is submitted. The GCS output prefix is resolved fromextra_params["output_uri"]or derived fromextra_params["gcs_bucket"]/["gcs_prefix"]. Vertex-native fields (modelParameters,labels,encryptionSpec, etc.) are passed through viaextra_params.BatchList: Paginates across all configured keys usingSerialListHelper, since batch jobs are scoped to a project/region. Each key's native VertexpageTokenis re-encoded into the Bifrost cursor.BatchRetrieve: Tries each key in turn until the job is found, since a job ID is only resolvable within the project/region that created it.BatchCancel/BatchDelete: Same multi-key fan-out pattern as retrieve.BatchResults: Fetches the job to locate its GCS output directory, lists allpredictions-*.jsonlfiles, downloads and parses each line, and recoverscustom_idfrom the echoed request labels.vertexJobStateToBatchStatus: Maps VertexJOB_STATE_*values to BifrostBatchStatusconstants.ToVertexBatchCreateRequest: Maps a Bifrost batch create request to aVertexBatchCreateRequest, stripping Bifrost control keys (output_uri,gcs_bucket,gcs_prefix,job_name) before forwardingextra_paramsto Vertex.vertexConvertRequestsToJSONL: Converts inlineBatchRequestItementries to Vertex JSONL, injectingcustom_idinto each request'slabelsmap without mutating the caller's data.gcsListAllObjects,gcsDownloadObject): Added to support paginated object listing and raw object download needed byBatchResults.vertexBatchPredictionJob,VertexBatchCreateRequest,vertexBatchJobListResponse,vertexBatchOutputLine, and supporting config/stats structs.A regional Vertex key (e.g.
us-central1) is required;globalis explicitly rejected since the Batch Prediction API does not support it.Type of change
Affected areas
How to test
go test ./core/providers/vertex/...To exercise end-to-end:
VertexKeyConfig(e.g.us-central1) and a GCS bucket accessible by the service account.BatchCreatewith eitherinput_file_id(ags://JSONL URI) or inlinerequests, and setextra_params["gcs_bucket"]orextra_params["output_uri"].BatchRetrieveuntil the job reachescompleted.BatchResultsto retrieve per-request responses withcustom_idround-tripped correctly.BatchList,BatchCancel, andBatchDeleteagainst jobs in the target project/region.Breaking changes
Related issues
Security considerations
gcsGetAuthHeaderhelper (service account credentials from the key config). No credentials are logged or returned in responses.bifrost_custom_idlabel is user-supplied and echoed back from Vertex; callers should treat it as untrusted input if used downstream.Checklist
docs/contributing/README.mdand followed the guidelinesSummary by CodeRabbit
New Features
Tests
Chores