feat: add context-based header and path overrides with key selection skipping - #728
Conversation
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR introduces context-aware path resolution and header handling across the Bifrost provider framework. It adds new context keys for request customization, modifies utility functions to accept context parameters, and replaces hardcoded endpoint paths with dynamic context-derived path construction throughout provider implementations. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/bifrost.go (1)
2330-2340: Guard skip-key path against nil contexts.
(*ctx).Value(...)is invoked without first ensuring bothctxand*ctxare non-nil. WhenselectKeyFromProviderForModelis called with a nil context (e.g., upstream passescontext.Context(nil)), this dereference panics before we even reach provider code. Fold the skip-key branch into the existing context guard so both the direct-key override and the skip flag only run when a real context is present.- if ctx != nil { - key, ok := (*ctx).Value(schemas.BifrostContextKeyDirectKey).(schemas.Key) - if ok { - return key, nil - } - } - - if skipKeySelection, ok := (*ctx).Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skipKeySelection && isKeySkippingAllowed(providerKey) { - return schemas.Key{}, nil - } + if ctx != nil && *ctx != nil { + if key, ok := (*ctx).Value(schemas.BifrostContextKeyDirectKey).(schemas.Key); ok { + return key, nil + } + if skipKeySelection, ok := (*ctx).Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skipKeySelection && isKeySkippingAllowed(providerKey) { + return schemas.Key{}, nil + } + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
core/bifrost.go(1 hunks)core/providers/anthropic.go(9 hunks)core/providers/azure.go(3 hunks)core/providers/bedrock.go(3 hunks)core/providers/cohere.go(7 hunks)core/providers/gemini.go(8 hunks)core/providers/groq.go(3 hunks)core/providers/mistral.go(4 hunks)core/providers/ollama.go(5 hunks)core/providers/openai.go(22 hunks)core/providers/openrouter.go(7 hunks)core/providers/parasail.go(3 hunks)core/providers/sgl.go(6 hunks)core/providers/utils.go(3 hunks)core/providers/vertex.go(3 hunks)core/schemas/bifrost.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (5)
core/bifrost.go (2)
core/schemas/bifrost.go (1)
BifrostContextKeySkipKeySelection(107-107)core/schemas/account.go (1)
Key(8-17)
core/providers/utils.go (1)
core/schemas/bifrost.go (2)
BifrostContextKeyExtraHeaders(108-108)BifrostContextKeyURLPath(109-109)
core/providers/anthropic.go (1)
core/schemas/providers/anthropic/models.go (1)
ToAnthropicListModelsURL(11-43)
core/providers/gemini.go (1)
core/schemas/providers/gemini/models.go (1)
ToGeminiListModelsURL(11-27)
core/providers/cohere.go (1)
core/schemas/providers/cohere/models.go (1)
ToCohereListModelsURL(10-34)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (8)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (7)
core/providers/vertex.go (1)
182-182: LGTM! Context-aware header handling consistently applied.The addition of
ctxparameter tosetExtraHeadersHTTPcalls is consistently applied across all HTTP request methods (ListModels, ChatCompletion, and handleVertexEmbedding). This enables dynamic header injection while maintaining the existing request flow.Also applies to: 397-397, 704-704
core/providers/groq.go (1)
75-75: LGTM! Context-aware path resolution properly implemented.URL construction now uses
getPathFromContext(ctx, defaultPath)for all major endpoints (models, chat completions, and streaming). This enables runtime endpoint customization while maintaining backward compatibility through default paths.Also applies to: 164-164, 183-183
core/providers/parasail.go (1)
75-75: LGTM! Context-aware paths consistently applied.Path resolution properly uses
getPathFromContext(ctx, defaultPath)for ListModels, ChatCompletion, and ChatCompletionStream endpoints.Also applies to: 101-101, 120-120
core/providers/openrouter.go (1)
74-74: LGTM! Comprehensive context-aware refactoring.All endpoints now support context-based customization:
- Header injection via
setExtraHeaders(ctx, ...)- Dynamic path resolution via
getPathFromContext(ctx, path)The changes cover all request types including the alpha responses endpoint, maintaining consistency across the provider.
Also applies to: 76-76, 123-123, 140-140, 156-156, 175-175, 191-191, 206-206
core/providers/openai.go (1)
87-87: LGTM! Context-aware request construction thoroughly implemented.All OpenAI endpoints now support context-based customization:
- Headers:
setExtraHeaders(ctx, ...)andsetExtraHeadersHTTP(ctx, ...)- Paths:
getPathFromContext(ctx, defaultPath)This comprehensive refactoring enables dynamic request modification across all operations (models, completions, chat, responses, embeddings, speech, transcriptions) while maintaining backward compatibility.
Also applies to: 109-109, 165-165, 202-202, 257-257, 328-328, 510-510, 549-549, 612-612, 683-683, 867-867, 906-906, 967-967, 1040-1040, 1193-1193, 1234-1234, 1309-1309, 1311-1311, 1399-1399, 1560-1560, 1562-1562, 1643-1643, 1662-1662
core/providers/ollama.go (1)
76-86: LGTM! Context-aware paths consistently applied across all Ollama endpoints.All endpoints (ListModels, TextCompletion, TextCompletionStream, ChatCompletion, ChatCompletionStream, Embedding) now use
getPathFromContext(ctx, defaultPath)for dynamic path resolution. The reformatting at lines 76-86 also improves code readability.Also applies to: 94-94, 111-111, 127-127, 146-146, 187-187
core/providers/cohere.go (1)
113-113: LGTM! Context-aware refactoring comprehensively applied to Cohere provider.All Cohere v2 endpoints now support context-based customization:
- Header injection:
setExtraHeaders(ctx, ...)andsetExtraHeadersHTTP(ctx, ...)- Dynamic paths:
getPathFromContext(ctx, path)for/v1/models,/v2/chat, and/v2/embedThe changes cover both synchronous and streaming operations, maintaining consistency across the entire provider.
Also applies to: 116-116, 233-233, 235-235, 311-311, 336-336, 593-593, 618-618, 785-785, 787-787
| responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/complete"), key.Value) | ||
| if err != nil { |
There was a problem hiding this comment.
Undefined variable in text completion request.
jsonData no longer exists in this scope, so this call doesn’t compile. Pass the marshaling payload we actually have (reqBody), letting completeRequest handle the JSON encoding internally.
- responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/complete"), key.Value)
+ responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/complete"), key.Value)📝 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.
| responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/complete"), key.Value) | |
| if err != nil { | |
| responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/complete"), key.Value) | |
| if err != nil { |
🤖 Prompt for AI Agents
In core/providers/anthropic.go around lines 273 to 274, the call to
completeRequest uses the undefined variable jsonData; replace jsonData with the
existing marshaling payload reqBody so completeRequest receives the correct
payload (it will handle JSON encoding internally) and ensure reqBody is in scope
for that call.
| responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value) | ||
| if err != nil { |
There was a problem hiding this comment.
Same undefined variable hits chat completions.
jsonData is still referenced after being removed, which breaks the build. Hand the structured reqBody to completeRequest instead.
- responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value)
+ responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value)📝 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.
| responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value) | |
| if err != nil { | |
| responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value) | |
| if err != nil { |
🤖 Prompt for AI Agents
In core/providers/anthropic.go around lines 325 to 326, the call to
completeRequest still references the removed variable jsonData which breaks the
build; replace that argument with the structured reqBody (the request body
variable used for chat completions) so completeRequest receives the proper
payload, and ensure the function signature/serialization expectations match
(marshal inside completeRequest or pass the correct typed object) so types
compile.
| responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value) | ||
| if err != nil { |
There was a problem hiding this comment.
Responses path still references removed jsonData.
Here too we hand an undefined identifier into completeRequest, so the code won’t compile. Forward reqBody.
- responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value)
+ responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value)📝 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.
| responseBody, latency, err := provider.completeRequest(ctx, jsonData, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value) | |
| if err != nil { | |
| responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), key.Value) | |
| if err != nil { |
🤖 Prompt for AI Agents
In core/providers/anthropic.go around lines 604 to 605, the call to
provider.completeRequest passes the removed/undefined identifier jsonData
causing a compile error; replace jsonData with the existing reqBody variable
when calling completeRequest (i.e., forward reqBody as the request body
argument) so the function receives the correct payload and the file compiles.
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), bytes.NewReader(jsonData)) | ||
| if err != nil { |
There was a problem hiding this comment.
Streaming request still uses jsonData.
The byte reader should wrap jsonBody, the buffer we just marshaled. As written the compiler can’t resolve jsonData.
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), bytes.NewReader(jsonData))
+ req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), bytes.NewReader(jsonBody))📝 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.
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), bytes.NewReader(jsonData)) | |
| if err != nil { | |
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/messages"), bytes.NewReader(jsonBody)) | |
| if err != nil { |
🤖 Prompt for AI Agents
In core/providers/anthropic.go around lines 663 to 664, the streaming request
incorrectly passes bytes.NewReader(jsonData) where jsonData doesn't exist;
change the request to read from the marshaled buffer instead (e.g. use
bytes.NewReader(jsonBody.Bytes()) or otherwise pass jsonBody as the request
body) so the http.NewRequestWithContext uses the actual JSON buffer you just
created.
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonData)) | ||
| if err != nil { |
There was a problem hiding this comment.
Replace jsonData with the existing jsonBody.
jsonData is not defined in this scope, so this line breaks compilation. Use the marshalled jsonBody bytes you already built above.
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonData))
+ req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonBody))🤖 Prompt for AI Agents
In core/providers/gemini.go around lines 377 to 378, the
http.NewRequestWithContext call uses an undefined variable jsonData which breaks
compilation; replace jsonData with the already-marshalled []byte jsonBody (i.e.,
use bytes.NewReader(jsonBody)) so the request body uses the correct byte slice
and the code compiles.
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonData)) | ||
| if err != nil { |
There was a problem hiding this comment.
Reuse jsonBody instead of undefined jsonData.
Same issue here—the code references jsonData, which does not exist, so this path won’t compile.
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonData))
+ req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonBody))📝 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.
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonData)) | |
| if err != nil { | |
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/models/"+request.Model+":streamGenerateContent?alt=sse"), bytes.NewReader(jsonBody)) | |
| if err != nil { |
🤖 Prompt for AI Agents
In core/providers/gemini.go around lines 655-656, the request is using an
undefined jsonData variable which causes a compile error; replace the undefined
jsonData with the existing marshaled byte slice variable (jsonBody) or ensure
you marshal the payload into a []byte named jsonBody before this call, then pass
bytes.NewReader(jsonBody) to http.NewRequestWithContext so the code compiles and
sends the intended JSON body.
|
|
||
| // Set any extra headers from network config | ||
| setExtraHeaders(req, provider.networkConfig.ExtraHeaders, nil) | ||
| setExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
🧩 Analysis chain
Verify Line 81 for consistency.
Context-aware handling is properly implemented for headers (line 79) and most endpoints (lines 139, 158, 201). However, Line 81 still uses a hardcoded path:
req.SetRequestURI(provider.networkConfig.BaseURL + "/v1/models")Verify whether this should also use getPathFromContext(ctx, "/v1/models") for consistency with the other endpoints.
Also applies to: 139-139, 158-158, 201-201
🏁 Script executed:
#!/bin/bash
# Check if other providers consistently use getPathFromContext for all endpoints
rg -n "SetRequestURI.*BaseURL.*\+\s*\"/v1/" core/providers/Length of output: 94
[/run_scripts]
🏁 Script executed:
#!/bin/bash
# Inspect the specific lines mentioned in the review comment
echo "=== Lines 79-81 (around the flagged line) ==="
sed -n '79,81p' core/providers/mistral.go
echo ""
echo "=== Lines 139, 158, 201 (allegedly using context-aware handling) ==="
sed -n '139p;158p;201p' core/providers/mistral.go
echo ""
echo "=== Check for getPathFromContext usage ==="
rg -n "getPathFromContext" core/providers/mistral.goLength of output: 787
Refactor line 81 to use context-aware path handling for consistency.
Line 81 uses a hardcoded path while all other endpoints (lines 139, 158, 201) use getPathFromContext() for context-aware handling. Update:
req.SetRequestURI(provider.networkConfig.BaseURL + getPathFromContext(ctx, "/v1/models"))🤖 Prompt for AI Agents
In core/providers/mistral.go around line 79, replace the hardcoded "/v1/models"
path when setting the request URI with a context-aware path using
getPathFromContext(ctx, "/v1/models") so the request uses
provider.networkConfig.BaseURL + getPathFromContext(ctx, "/v1/models"); keep the
rest of the call (SetRequestURI and BaseURL) unchanged to match how other
endpoints handle context-aware paths.
|
|
||
| // Create HTTP request for streaming | ||
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+"/v1/audio/speech", bytes.NewReader(jsonBody)) | ||
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/audio/speech"), bytes.NewReader(jsonData)) |
There was a problem hiding this comment.
Fix undefined variable reference.
Line 1380 references jsonData which is not defined in this scope. Should be jsonBody (defined at line 1366).
Apply this diff:
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/audio/speech"), bytes.NewReader(jsonData))
+ req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/audio/speech"), bytes.NewReader(jsonBody))📝 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.
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/audio/speech"), bytes.NewReader(jsonData)) | |
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.networkConfig.BaseURL+getPathFromContext(ctx, "/v1/audio/speech"), bytes.NewReader(jsonBody)) |
🤖 Prompt for AI Agents
In core/providers/openai.go around line 1380, the request creation uses an
undefined variable jsonData; replace that reference with the correctly defined
jsonBody (defined at line 1366) so the http.NewRequestWithContext call reads the
body from jsonBody instead of jsonData; ensure no other references to jsonData
exist in this scope.
1 participant
Summary
Add support for context-based request customization, allowing dynamic modification of API requests through context values.
Changes
BifrostContextKeySkipKeySelectionBifrostContextKeyExtraHeadersBifrostContextKeyURLPathto dynamically change API endpointsType of change
Affected areas
How to test
Test the new context-based customization features:
Breaking changes
Related issues
N/A
Security considerations
Checklist
docs/contributing/README.mdand followed the guidelines