feat(ai-proxy): add support for vertex ai provider #1590
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the WalkthroughThe pull request introduces updates to the Changes
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![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
🧹 Outside diff range and nitpick comments (12)
go/ai-proxy/proxy/provider/vertex.go (1)
44-56: Ensure consistent error handling inModifyResponseIn the
ModifyResponsemethod, whenmodifyResponseBodyfails, the error is logged and returned. However, the error logging could include more context for better debugging.Enhance the error log to include additional context, such as the response status code:
err := in.modifyResponseBody(r) if err != nil { - klog.ErrorS(err, "failed to map response body") + klog.ErrorS(err, "failed to map response body", "statusCode", r.StatusCode) return err }This provides more information when reviewing logs.
go/ai-proxy/proxy/provider/ollama.go (1)
Line range hint
10-19: Consider adding documentation for NewOllamaProxySince this is part of a broader authentication standardization effort, it would be helpful to add documentation explaining why certain parameters are nil and what configuration options are available for Ollama.
Add godoc comments explaining the parameters and their purpose:
+// NewOllamaProxy creates a new Ollama translation proxy. +// host: The Ollama server host URL +// Returns: A TranslationProxy implementation for Ollama func NewOllamaProxy(host string) (api.TranslationProxy, error) {go/ai-proxy/main.go (1)
Line range hint
19-23: Consider enhancing error messaging for credentials issuesThe current error message "Could not create proxy" could be more specific about credential-related failures to help with troubleshooting.
p, err := proxy.NewOllamaTranslationProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials()) if err != nil { - klog.ErrorS(err, "Could not create proxy") + klog.ErrorS(err, "Failed to create proxy", "provider", args.Provider()) os.Exit(1) }go/ai-proxy/api/vertex/vertex.go (2)
8-14: Consider adding URL construction helpersThe endpoint URL contains placeholders (
${PROJECT_ID}and${LOCATION}) that need runtime substitution. Consider adding:
- A helper method to construct the URL with proper validation
- Documentation about required environment variables
Example helper method:
func (e Endpoint) BuildURL(projectID, location string) (string, error) { if projectID == "" || location == "" { return "", fmt.Errorf("projectID and location are required") } url := string(e) url = strings.Replace(url, "${PROJECT_ID}", projectID, 1) url = strings.Replace(url, "${LOCATION}", location, 1) return url, nil }
16-22: Add documentation for error structureThe
ErrorResponsestruct matches Vertex AI's error format, but it would be helpful to add documentation explaining:
- The expected error format from Vertex AI
- Example error scenarios
Add documentation like:
// ErrorResponse represents the error format returned by Vertex AI API. // Example error: // { // "error": { // "code": 400, // "message": "Invalid request", // "status": "INVALID_ARGUMENT" // } // }go/ai-proxy/api/provider.go (1)
63-68: Consider using error returns instead of panicWhile the implementation is consistent with existing patterns, using panic for path validation might be too aggressive for runtime errors. Consider refactoring to return errors instead, allowing callers to handle missing paths gracefully.
-func ToProviderAPIPath(target Provider, path string) string { +func ToProviderAPIPath(target Provider, path string) (string, error) { switch target { case ProviderOllama: - return path + return path, nil case ProviderVertex: targetPath, exists := ollamaToVertex[path] if !exists { - panic(fmt.Sprintf("path %s not registered for provider %s", path, target)) + return "", fmt.Errorf("path %s not registered for provider %s", path, target) } - return targetPath + return targetPath, nil } - panic(fmt.Sprintf("invalid provider: %s", target)) + return "", fmt.Errorf("invalid provider: %s", target) }go/ai-proxy/proxy/provider/base.go (1)
98-98: LGTM: Consider adding documentationThe constructor changes are well-implemented. Consider adding a comment explaining the expected format and behavior of the mapping function for future maintainers.
Add documentation above the constructor:
+// newBaseTranslationProxy creates a new proxy with optional path mapping. +// The mapping function, if provided, is used to transform API paths using os.Expand +// format (e.g., $name -> mapped value). It should return sanitized paths suitable +// for use in URLs. func newBaseTranslationProxy(Also applies to: 113-113
go/ai-proxy/args/args.go (2)
27-30: Consider enhancing provider configuration robustness.A few suggestions to improve the provider configuration:
- Consider reordering providers in the help text to match their declaration order in the code
- Add validation for the service account file path existence when the Vertex provider is selected
Here's a suggested improvement for the provider flag description:
- argProvider = pflag.String("provider", defaultProvider.String(), "Provider name. Must be one of: ollama, openai, vertex. Defaults to 'ollama' type API.") + argProvider = pflag.String("provider", defaultProvider.String(), "Provider name. Must be one of: vertex, openai, ollama. Defaults to 'ollama' type API.")
73-86: Consider refactoring for better readability and error clarity.While the logic is correct, the function could be more readable with early returns and more specific error messages.
Here's a suggested refactor:
func ProviderCredentials() string { + provider := Provider() + + if provider == defaultProvider { + return "" + } + + if provider == api.ProviderOpenAI && len(*argProviderToken) > 0 { + return *argProviderToken + } + + if provider == api.ProviderVertex && len(*argProviderServiceAccount) > 0 { + return *argProviderServiceAccount + } + + var missing string + switch provider { + case api.ProviderOpenAI: + missing = "token" + case api.ProviderVertex: + missing = "service account" + } + panic(fmt.Errorf("provider %s must be provided when %s provider is used", missing, provider)) - if len(*argProviderToken) > 0 && Provider() == api.ProviderOpenAI { - return *argProviderToken - } - - if len(*argProviderServiceAccount) > 0 && Provider() == api.ProviderVertex { - return *argProviderServiceAccount - } - - if Provider() == defaultProvider { - return "" - } - - panic(fmt.Errorf("provider credentials must be provided when %s provider is used", Provider())) }go/ai-proxy/go.mod (2)
Line range hint
3-3: Fix incorrect Go versionThe specified Go version
1.22.8does not exist yet (as of November 2024). The latest stable version is 1.21.x.-go 1.22.8 +go 1.21.4
54-59: Consider structured logging integration with OpenTelemetryThe addition of OpenTelemetry dependencies suggests an investment in observability. To maximize the value of these dependencies:
- Ensure consistent trace context propagation across service boundaries
- Consider implementing structured logging that correlates with trace IDs
- Define meaningful metrics for the Vertex AI integration
Would you like assistance in implementing a logging middleware that integrates with OpenTelemetry?
charts/ai-proxy/values.yaml (1)
23-24: Security improvement: Removed hardcoded default tokenGood security practice to:
- Remove the hardcoded default token
- Set both sensitive fields (
tokenandserviceAccount) to null by defaultHowever, the new
serviceAccountfield needs documentation.Consider adding comments to document:
- The purpose of the serviceAccount field
- The expected format (JSON key file content?)
- Whether this is specifically for Vertex AI authentication
secrets: token: ~ + # Service account credentials in JSON format, required when using the 'vertex' provider serviceAccount: ~
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
⛔ Files ignored due to path filters (1)
go/ai-proxy/go.sumis excluded by!**/*.sum
📒 Files selected for processing (14)
charts/ai-proxy/Chart.yaml(1 hunks)charts/ai-proxy/templates/secrets.yaml(1 hunks)charts/ai-proxy/values.yaml(1 hunks)go/ai-proxy/api/provider.go(5 hunks)go/ai-proxy/api/vertex/vertex.go(1 hunks)go/ai-proxy/args/args.go(2 hunks)go/ai-proxy/go.mod(2 hunks)go/ai-proxy/main.go(1 hunks)go/ai-proxy/proxy/provider/base.go(6 hunks)go/ai-proxy/proxy/provider/ollama.go(1 hunks)go/ai-proxy/proxy/provider/openai.go(2 hunks)go/ai-proxy/proxy/provider/vertex.go(1 hunks)go/ai-proxy/proxy/proxy.go(1 hunks)go/ai-proxy/test/helpers/common.go(1 hunks)
🔇 Additional comments (24)
go/ai-proxy/proxy/provider/vertex.go (1)
61-63: Ensure correct endpoint expansion in modifyRequestBody
Using os.Expand with in.mapping on vertex.EndpointChat relies on accurate variable substitution. If vertex.EndpointChat contains unexpected patterns, this could lead to incorrect endpoint URLs.
Verify that vertex.EndpointChat contains the correct patterns and that in.mapping handles all required variables.
Run the following script to check for variable patterns:
This ensures that all placeholders are intentional and correctly mapped.
✅ Verification successful
Variable substitution in vertex.EndpointChat is correctly implemented
The implementation is safe and correct:
vertex.EndpointChatuses standard${PROJECT_ID}and${LOCATION}patterns- These variables are properly defined as constants
EnvProjectIDandEnvLocation - The
mappingfunction inNewVertexProxycorrectly handles both variables:- Tries to use credentials.ProjectID first for PROJECT_ID
- Falls back to environment variables with proper error logging
- Handles LOCATION from environment variables with error logging
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for variable patterns in vertex.EndpointChat
# Expected: Should contain only known environment variables
grep -E '\$\{?[A-Za-z_][A-Za-z0-9_]*\}?' go/ai-proxy/api/vertex/endpoints.go
Length of output: 147
Script:
#!/bin/bash
# Let's first find the correct path for vertex endpoints
fd "endpoints.go" --type f
# Then check the content of the file containing vertex endpoints
fd "endpoints.go" --type f --exec cat {}
Length of output: 69
Script:
#!/bin/bash
# Let's check the vertex.go file itself to understand the endpoint definition
cat go/ai-proxy/proxy/provider/vertex.go
# Also search for any other files that might contain vertex endpoint definitions
rg -l "EndpointChat"
Length of output: 3077
Script:
#!/bin/bash
# Let's check the vertex.go file in the api package to see the endpoint definition
cat go/ai-proxy/api/vertex/vertex.go
# Also check the mapping variables used
rg "EnvProjectID|EnvLocation" go/ai-proxy/api/vertex/vertex.go
Length of output: 1020
charts/ai-proxy/templates/secrets.yaml (1)
10-10: Consider making the service account key conditional
Since this is being added for Vertex AI provider support, consider making this key conditional based on the provider type to maintain clean configuration.
+ {{- if eq .Values.config.provider "vertex" }}
PLRL_PROVIDER_SERVICE_ACCOUNT: {{ .Values.secrets.serviceAccount | b64enc | quote }}
+ {{- end }}Let's verify if the provider configuration exists in the values file:
✅ Verification successful
Let me gather more information about the provider configuration and service account usage.
The conditional service account configuration is appropriate
The values.yaml file shows that:
- Provider configuration exists with
vertexas a valid option:# One of: ollama, openai, vertex - Service account is already defined as an optional value:
serviceAccount: ~
Making the service account secret conditional on the vertex provider is a good architectural suggestion since:
- It's only needed for Vertex AI
- It keeps the configuration clean for other providers
- The values file already supports both the provider selection and service account configuration
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for provider configuration in values.yaml
# Expected: Find provider configuration that would control this conditional
rg -A 5 "provider:" charts/ai-proxy/values.yaml
Length of output: 147
Script:
#!/bin/bash
# Check for any vertex-related configuration or service account references
rg -i "vertex|serviceaccount" charts/ai-proxy/values.yaml
# Check for provider-related configuration in other yaml files
fd -e yaml -e yml -x rg -l "provider:|vertex"
Length of output: 1613
go/ai-proxy/proxy/provider/ollama.go (1)
12-12: Verify if nil arguments are intentional
The code passes nil for all parameters to newBaseTranslationProxy. While this might be intentional if Ollama doesn't require credentials or additional configuration, it would be good to verify this is the expected behavior.
Let's check how other providers handle these parameters:
✅ Verification successful
Let me gather more information about the base proxy implementation to understand if nil arguments are indeed valid for Ollama.
Passing nil arguments is correct for Ollama implementation
The code review reveals that passing nil arguments is intentional and correct for the Ollama provider:
- The
newBaseTranslationProxyimplementation shows that whennilis passed formodifyRequestandmodifyResponse, it falls back to using the base implementation's methods:
if modifyRequest == nil {
baseProxy.proxy.Rewrite = baseProxy.ModifyRequest
}
if modifyResponse == nil {
baseProxy.proxy.ModifyResponse = baseProxy.ModifyResponse
}- Unlike OpenAI (which needs token auth) and Vertex (which needs request/response modifications), Ollama is designed to work with the base implementation's behavior without any custom modifications.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check how other providers implement newBaseTranslationProxy
# Expected: Find patterns of parameter usage across different providers
# Search for all calls to newBaseTranslationProxy
rg "newBaseTranslationProxy\(" --type go
# Check if Ollama specifically requires any credentials or config
rg -B 5 -A 5 "type OllamaProxy struct" --type go
Length of output: 1242
Script:
#!/bin/bash
# Check the implementation of newBaseTranslationProxy to understand parameter usage
rg -B 5 -A 15 "func newBaseTranslationProxy" --type go
# Check if there are any Ollama-specific request/response modifications
rg -A 5 "ModifyRequest|ModifyResponse" --type go
Length of output: 5832
go/ai-proxy/proxy/proxy.go (2)
10-10: LGTM: Function signature change is appropriate.
The change from token to credentials better represents the generic nature of authentication across different providers.
Let's verify the consistency of this change across the codebase:
#!/bin/bash
# Description: Verify that all callers have been updated to use credentials
# Check for any remaining references to ProviderToken
rg -l "ProviderToken"
# Check for the new credential parameter usage
rg -l "ProviderCredentials"
# Look for any direct token references that might need updating
ast-grep --pattern 'NewOllamaTranslationProxy($_, $_, $_)'15-17: Verify provider implementations and error handling consistency.
The provider handling looks good, but let's ensure consistency across implementations.
Let's verify the provider implementations:
Consider these architectural suggestions:
- Ensure all providers implement consistent error handling patterns
- Consider adding provider capability validation at initialization
- Document credential format requirements for each provider in the interface
✅ Verification successful
Let me gather more specific information about the provider implementations and error handling patterns.
Let me check one more aspect about the interface implementation and error handling in the base proxy.
Based on the gathered information, I can now provide the final response:
Provider implementations show consistent error handling and initialization patterns ✓
All provider implementations (OpenAI, Vertex, and Ollama) demonstrate consistent patterns:
- Uniform error propagation from base implementation
- Consistent initialization error handling
- Proper implementation of request/response modification methods
- Structured inheritance from baseTranslationProxy
The architectural patterns are well-maintained across all providers with proper error handling, credential validation, and consistent proxy initialization. No issues found.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify provider implementations and interface compliance
# Check for provider implementations
fd -e go . -x rg -l "NewVertexProxy|NewOpenAIProxy|NewOllamaProxy"
# Verify consistent error handling patterns
rg -A 2 "fmt.Errorf" "go/ai-proxy/proxy/provider/"
# Look for interface implementation
ast-grep --pattern 'type $_Proxy struct {
$$$
}'
Length of output: 316
Script:
#!/bin/bash
# Check provider struct definitions and their methods
rg -A 10 "type \w+Proxy struct" go/ai-proxy/proxy/provider/
# Look for error handling patterns in each provider
echo "=== OpenAI Provider ==="
rg "return.*err|fmt.Errorf|errors\." go/ai-proxy/proxy/provider/openai.go
echo -e "\n=== Vertex Provider ==="
rg "return.*err|fmt.Errorf|errors\." go/ai-proxy/proxy/provider/vertex.go
echo -e "\n=== Ollama Provider ==="
rg "return.*err|fmt.Errorf|errors\." go/ai-proxy/proxy/provider/ollama.go
# Check interface implementation methods
rg -A 2 "func \(.*Proxy\) \w+" go/ai-proxy/proxy/provider/
Length of output: 6354
Script:
#!/bin/bash
# Check the interface definition and base implementation
echo "=== Interface Definition ==="
rg "type TranslationProxy interface" -A 10 go/ai-proxy/proxy/api/
echo "=== Base Error Handling ==="
rg "func.*error" -A 3 go/ai-proxy/proxy/provider/base.go
# Check initialization error handling
echo "=== Initialization Error Handling ==="
rg "func New\w+Proxy" -A 5 go/ai-proxy/proxy/provider/
Length of output: 2182
go/ai-proxy/main.go (1)
19-19: Verify provider compatibility with new credentials system
The change from ProviderToken() to ProviderCredentials() looks good and aligns with the Vertex AI provider addition. However, let's verify the compatibility across all providers.
✅ Verification successful
Let me gather more information about the provider implementations to verify the credentials handling.
Let me check one more thing to see the implementation of NewOpenAIProxy and NewVertexProxy functions.
Based on the gathered information, I can now provide the final response:
Change from ProviderToken to ProviderCredentials is properly implemented
The verification shows that the credentials system is correctly implemented across all providers:
- OpenAI provider continues to use token-based auth (
ProviderTokenis mapped to credentials for backward compatibility) - Vertex AI provider uses service account credentials
- Ollama provider doesn't require credentials (passes nil)
Each provider implementation properly handles its specific authentication needs through the new credentials parameter:
NewOpenAIProxy(target, token string)NewVertexProxy(target, serviceAccount string)NewOllamaProxy(host string)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if all providers support the new credentials system
# Expected: Find all provider implementations and their credential handling
# Look for provider struct definitions and their creation
ast-grep --pattern 'type $provider struct {
$$$
}'
# Check credential usage across providers
rg -A 5 "ProviderCredentials|NewOllamaTranslationProxy"
Length of output: 11443
Script:
#!/bin/bash
# Check provider implementations and their credential handling
rg -A 10 "type \w+Proxy struct"
# Check NewOpenAIProxy and NewVertexProxy implementations
ast-grep --pattern 'func New(OpenAI|Vertex)Proxy($$$) {
$$$
}'
# Look for provider interface definition
ast-grep --pattern 'type Provider interface {
$$$
}'
Length of output: 2836
Script:
#!/bin/bash
# Check NewOpenAIProxy and NewVertexProxy function implementations
rg -A 10 "func New(OpenAI|Vertex)Proxy"
# Check Provider type definition and related constants
rg -A 5 "type Provider"
Length of output: 7802
go/ai-proxy/api/vertex/vertex.go (2)
1-7: LGTM: Clean package structure and imports
The package structure follows Go conventions, and imports are minimal and purposeful.
10-14: Verify completeness of Vertex AI endpoints
Currently, only the chat completions endpoint is defined. Verify if other Vertex AI endpoints are needed for full functionality.
✅ Verification successful
Current endpoint implementation is sufficient
After reviewing the implementation, the single chat completions endpoint is appropriate because:
- The VertexProxy implementation shows that it's specifically designed to translate OpenAI-compatible chat completion requests to Vertex AI format
- The codebase only handles chat completion functionality, with request/response mapping focused on this specific use case
- No other Vertex AI endpoints are being used or required by the current implementation
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for Vertex AI endpoint usage across the codebase
rg -l "vertex.*endpoint"
# Search for potential missing endpoint definitions
ast-grep --pattern 'const $_ = $$$endpoint$$$'
Length of output: 75
Script:
#!/bin/bash
# Search for Vertex AI related files and their content
rg -l "vertex" --type go
# Show content of vertex.go to understand the implementation
cat go/ai-proxy/api/vertex/vertex.go
# Search for any API calls or endpoint usage
rg "projects.*locations.*endpoints" --type go -A 3
Length of output: 1476
Script:
#!/bin/bash
# Check the vertex provider implementation to understand endpoint usage
cat go/ai-proxy/proxy/provider/vertex.go
# Look for any model or completion related structs/interfaces
rg "type.*Completion" --type go -A 5
rg "interface.*Provider" --type go -A 10
Length of output: 19793
charts/ai-proxy/Chart.yaml (2)
24-24: Document version changes in release notes
The app version bump to v1.1.0 aligns well with the addition of new functionality (Vertex AI provider). However, ensure that:
- The changes are documented in the release notes
- The version matches the actual application container image tag
Let's verify the version consistency:
#!/bin/bash
# Description: Check version references across the chart
# Check for version references in values.yaml
rg "v1\.1\.0" charts/ai-proxy/
# Check for image tag references
fd values.yaml charts/ai-proxy -x cat {} | grep -A 5 "image:"18-18: Consider if a major version bump is necessary
A major version bump (0.1.0 -> 1.0.0) typically indicates breaking changes that affect backward compatibility. Adding a new provider functionality seems more aligned with a minor version increment (0.2.0) unless there are breaking changes not apparent from the context.
Let's check for potential breaking changes:
If this is indeed a breaking change, please ensure:
- The changes are documented in CHANGELOG.md
- Migration steps are provided for existing deployments
- Release notes clearly explain the impact
go/ai-proxy/api/provider.go (5)
8-8: LGTM: Import follows project conventions
The vertex package import is correctly placed and follows the project's package structure.
25-26: LGTM: Provider conversion case follows existing pattern
The Vertex provider case is correctly implemented and maintains consistency with other provider conversions.
36-36: LGTM: Provider constant follows conventions
The ProviderVertex constant is correctly defined and follows the established naming pattern.
25-26: Verify test coverage for Vertex provider implementation
Please ensure that test cases are added/updated to cover:
- ToProvider function with Vertex provider
- Path mappings for Vertex endpoints
- Error cases in ToProviderAPIPath
#!/bin/bash
# Check for existing test coverage
rg -l "TestToProvider|TestToProviderAPIPath" --type go
# Look for Vertex-related test cases
rg "ProviderVertex" --type go -lAlso applies to: 36-36, 47-49, 63-68
47-49: Verify completeness of endpoint mappings
The mapping currently only includes the chat endpoint. Please verify if other Ollama endpoints need to be supported for the Vertex provider.
✅ Verification successful
The previous script didn't show Ollama endpoints due to incorrect path. Let me search again with the correct path.
Let me try one more time with the correct path structure.
Endpoint mapping is complete
The mapping is correct as implemented. From the verification results:
- Ollama only exposes one endpoint:
EndpointChat = "/api/chat" - Vertex AI also has one corresponding endpoint:
EndpointChat = "/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/endpoints/openapi/chat/completions" - The current mapping in
ollamaToVertexcorrectly maps the only available endpoint.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check for all available Ollama endpoints that might need mapping
rg -A 1 "Endpoint[A-Za-z]+ =" go/ai-proxy/api/ollama/
# Check for all available Vertex endpoints
rg -A 1 "Endpoint[A-Za-z]+ =" go/ai-proxy/api/vertex/
Length of output: 620
Script:
#!/bin/bash
# Check for all available Ollama endpoints
rg -A 1 "Endpoint[A-Za-z]+ =" ollama/
Length of output: 87
Script:
#!/bin/bash
# Check for all available Ollama endpoints
rg -A 1 "Endpoint[A-Za-z]+ =" go/ai-proxy/api/ollama/
# Also check the imports to ensure we're looking at the right files
rg -l "package ollama" go/
Length of output: 252
go/ai-proxy/proxy/provider/openai.go (1)
39-39: LGTM! Error message accuracy improved
The error message now correctly indicates that it failed to map the response body instead of the request body, which is more accurate given this is in the ModifyResponse method.
go/ai-proxy/test/helpers/common.go (1)
19-19: Verify test coverage for the new credentials handling
The change from ProviderToken to ProviderCredentials aligns with the PR objective of adding Vertex AI support. However, ensure that:
- Test coverage exists for different credential scenarios (token-based, key-based, etc.)
- Backward compatibility is maintained for existing providers
Let's verify the test coverage:
✅ Verification successful
Credential handling is properly implemented with backward compatibility
The verification shows that the credential handling is well-implemented:
-
The
ProviderCredentials()function inargs/args.gomaintains backward compatibility:- For OpenAI: Uses
argProviderToken(via--provider-tokenflag orPROVIDER_TOKENenv var) - For Vertex AI: Uses
argProviderServiceAccount(via--provider-service-accountflag orPROVIDER_SERVICE_ACCOUNTenv var)
- For OpenAI: Uses
-
The change in
test/helpers/common.gois correct as it uses the unifiedProviderCredentials()method which internally handles the appropriate credential type based on the provider.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check test coverage for credential handling scenarios
# Look for test cases covering different credential scenarios
rg -A 5 "TestSetupServer|TestNewOllamaTranslationProxy" go/ai-proxy/test/
# Check for existing provider token usage to ensure backward compatibility
ast-grep --pattern 'ProviderToken()'
Length of output: 114
Script:
#!/bin/bash
# Let's try a broader search to understand the test coverage and credential handling
# Search for test files in the ai-proxy directory
fd ".*_test.go" go/ai-proxy/
# Look for credential-related test cases with more context
rg -B 2 -A 5 "Provider.*Credentials|Provider.*Token" go/ai-proxy/
# Check for credential-related interfaces or structs
ast-grep --pattern 'interface {
$$$
ProviderCredentials() $_
$$$
}'
# Check for actual implementations
ast-grep --pattern 'func ($_ $_) ProviderCredentials() $_'
Length of output: 3487
go/ai-proxy/proxy/provider/base.go (2)
5-5: LGTM: Clean additions for new functionality
The new imports and struct field are well-organized and properly typed for the intended functionality.
Also applies to: 11-11, 29-29
82-89: Verify path mapping security
Using os.Expand for path mapping could be risky if the path contains special characters or if the mapping function doesn't properly sanitize its output. Consider adding path validation.
Let's check for any potentially dangerous path manipulations in the codebase:
✅ Verification successful
Let me gather more context about the path handling and validation mechanisms.
Let me check one more aspect about the path handling and mapping implementation.
Path mapping implementation is secure
The current implementation is safe because:
- The path variable is first processed through
ToProviderAPIPathwhich uses a strict mapping of predefined endpoints - The
os.Expandis only used with controlled mapping functions that handle specific environment variables - The path mapping is applied to well-defined API endpoints (like
vertex.EndpointChatandollama.EndpointChat)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for path-related function calls that might need validation
rg -t go "os\.Expand|url\.Parse|path\." --no-heading
# Look for existing path validation functions
ast-grep --pattern 'func $NAME($ARGS) {
$$$
validate$$$Path
$$$
}'
Length of output: 945
Script:
#!/bin/bash
# Check how ToProviderAPIPath is implemented and what validation it does
ast-grep --pattern 'func ToProviderAPIPath($ARGS) $RET {
$$$
}'
# Look for path sanitization or validation in the codebase
rg -A 5 "func.*Path" --no-heading
# Check the mapping usage context
rg -A 3 "mapping.*:=" --no-heading
Length of output: 12068
Script:
#!/bin/bash
# Check the implementation of ToProviderAPIPath and related path mappings
rg -A 10 "ollamaToOpenAI" --no-heading
# Look for any path sanitization or URL encoding
rg -A 5 "url\.Path|url\.Query|url\.Encode" --no-heading
# Check the actual usage of os.Expand in the codebase
rg -B 3 -A 3 "os\.Expand" --no-heading
Length of output: 2345
go/ai-proxy/args/args.go (2)
18-19: LGTM: Environment variable constants follow consistent naming pattern.
The new envProviderServiceAccount constant follows the established naming pattern and purpose.
73-73: Verify the impact of renaming ProviderToken to ProviderCredentials.
The function has been renamed from ProviderToken to ProviderCredentials. Let's verify all callers have been updated.
✅ Verification successful
Function rename from ProviderToken to ProviderCredentials has been properly updated across all references
The verification shows that there are no remaining references to the old function name ProviderToken, and all three occurrences of ProviderCredentials are using the new name consistently:
go/ai-proxy/args/args.go: Function definitiongo/ai-proxy/main.go: Function usage in maingo/ai-proxy/test/helpers/common.go: Function usage in test helpers
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for any remaining references to ProviderToken
rg "ProviderToken\(" --type go
# Search for current usage of ProviderCredentials
rg "ProviderCredentials\(" --type go
Length of output: 405
go/ai-proxy/go.mod (2)
63-72: Review outdated dependency versions
Some Google API dependencies appear to be using older versions:
golang.org/x/oauth2 v0.23.0(newer versions available)google.golang.org/grpc v1.67.1(appears to be a future version, current stable is v1.59.x)
#!/bin/bash
# Description: Check for latest versions and compatibility of Google API dependencies
# Check latest versions
for pkg in "golang.org/x/oauth2" "google.golang.org/grpc" "google.golang.org/api"; do
echo "Checking $pkg..."
go list -m -versions "$pkg" 2>/dev/null || echo "Package not found locally"
done
# Verify gRPC compatibility matrix
curl -s https://raw.githubusercontent.com/grpc/grpc-go/master/README.md | grep -A 5 "Supported Go Versions"16-23: Verify Google Cloud dependencies versions
The addition of Google Cloud dependencies is appropriate for Vertex AI integration. However, let's verify these versions for security and compatibility.
✅ Verification successful
Google Cloud dependencies are up-to-date and properly configured
The verification shows that:
- All Google Cloud dependencies are using recent versions (e.g.,
cloud.google.com/go v0.116.0,cloud.google.com/go/vertexai v0.13.2) - No replace directives are present that could affect these versions
- No obvious security vulnerabilities were found in the dependencies
- All dependencies are marked as indirect, indicating they are properly managed through transitive dependencies
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for latest versions and security advisories of Google Cloud dependencies
# Check for latest versions on pkg.go.dev
for pkg in "cloud.google.com/go" "cloud.google.com/go/aiplatform" "cloud.google.com/go/vertexai"; do
echo "Checking $pkg..."
curl -s "https://pkg.go.dev/$pkg" | grep -A 1 "Latest version"
done
# Check for security advisories
gh api graphql -f query='
{
securityVulnerabilities(first: 5, ecosystem: GO, package: "cloud.google.com/go") {
nodes {
advisory {
summary
severity
publishedAt
}
vulnerableVersionRange
firstPatchedVersion {
identifier
}
}
}
}'
Length of output: 1076
Script:
#!/bin/bash
# Let's try a different approach to verify the dependencies
# Check go.mod for all Google Cloud dependencies
echo "Current Google Cloud dependencies:"
grep "cloud.google.com/" go/ai-proxy/go.mod
# Check if there are any replace directives that might affect these versions
echo -e "\nReplace directives:"
grep "replace" go/ai-proxy/go.mod -A 1
# Check if there are any known vulnerabilities in the dependencies
echo -e "\nChecking for vulnerability patterns in go.sum:"
grep -i "CVE\|vulnerability\|security" go/ai-proxy/go.sum || echo "No obvious vulnerability patterns found"
# List all direct dependencies to see if versions are consistent
echo -e "\nAll direct dependencies:"
grep -v "indirect" go/ai-proxy/go.mod | grep "^[[:space:]]*[a-z]"
Length of output: 1302
charts/ai-proxy/values.yaml (1)
17-17: Verify provider configuration validation
The addition of "vertex" as a provider option aligns with the PR objective. However, we should ensure that the application properly validates this configuration.
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (3)
charts/ai-proxy/templates/secrets.yaml (1)
12-14: Consider adding validation and documentation for service account credentials.While the implementation is correct, since this is for Vertex AI credentials which are typically JSON service account keys, consider:
- Adding validation to ensure the value is valid JSON
- Adding a comment documenting the expected format
- Considering size limitations of the secret
Add validation and documentation:
{{- with .Values.secrets.serviceAccount }} + {{- if not (fromJson .) }} + {{- fail "PLRL_PROVIDER_SERVICE_ACCOUNT must be a valid JSON service account key" }} + {{- end }} + # Vertex AI service account key in JSON format PLRL_PROVIDER_SERVICE_ACCOUNT: {{ . | b64enc | quote }} {{- end }}go/ai-proxy/proxy/provider/base.go (1)
103-103: Add parameter documentation for the mapping functionConsider adding a comment to document the
mappingparameter in the constructor function signature to explain its purpose and expected behavior.func newBaseTranslationProxy( target string, provider api.Provider, modifyRequest func(*httputil.ProxyRequest), modifyResponse func(*http.Response) error, - mapping func(string) string, + // mapping is an optional function to transform API paths and expand environment variables + mapping func(string) string, ) (*baseTranslationProxy, error) {Also applies to: 118-118
go/ai-proxy/proxy/provider/vertex.go (1)
58-80: Consider adding input validation and logging for unhandled endpointsWhile the implementation is solid, consider these enhancements:
- Add debug logging for unhandled endpoints
- Validate request/response body before transformation
Example enhancement:
func (in *VertexProxy) modifyRequestBody(r *httputil.ProxyRequest) error { endpoint := r.Out.URL.Path + klog.V(4).InfoS("modifying request body", "endpoint", endpoint) switch endpoint { case os.Expand(vertex.EndpointChat, in.mapping): + if r.Out.Body == nil { + return fmt.Errorf("empty request body for chat completion") + } return replaceRequestBody(r, openai.ToChatCompletionRequest) } + klog.V(4).InfoS("no modification needed for endpoint", "endpoint", endpoint) return nil }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (4)
charts/ai-proxy/Chart.yaml(1 hunks)charts/ai-proxy/templates/secrets.yaml(1 hunks)go/ai-proxy/proxy/provider/base.go(7 hunks)go/ai-proxy/proxy/provider/vertex.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- charts/ai-proxy/Chart.yaml
🧰 Additional context used
🪛 yamllint
charts/ai-proxy/templates/secrets.yaml
[warning] 10-10: wrong indentation: expected 0 but found 2
(indentation)
🔇 Additional comments (7)
charts/ai-proxy/templates/secrets.yaml (2)
9-11: LGTM! Proper conditional handling of the token secret.
The implementation correctly uses the with block to handle nil values and properly encodes the secret.
🧰 Tools
🪛 yamllint
[warning] 10-10: wrong indentation: expected 0 but found 2
(indentation)
Line range hint 1-14: Verify completeness of Vertex AI configuration.
Let's ensure all required Vertex AI authentication configurations are included.
✅ Verification successful
Let me gather more specific information about the Vertex AI configuration requirements.
Let me try a different approach to find the configuration requirements.
Let me check one more crucial piece to verify the provider implementation.
Vertex AI configuration is complete and correctly implemented
The secrets configuration in the Helm chart is complete and aligns with the Vertex AI provider implementation. The code shows that:
PLRL_PROVIDER_SERVICE_ACCOUNTis used for Google Cloud authentication- The provider automatically extracts
PROJECT_IDfrom the service account JSON LOCATIONis obtained from the environment- No additional secret configurations are required
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for Vertex AI related configuration patterns in the codebase
echo "Searching for Vertex AI configuration patterns..."
rg -l "vertex" --type go
# Look for authentication related code
echo "Checking authentication related code..."
ast-grep --pattern 'type $_ struct {
$$$
serviceAccount string
$$$
}'
Length of output: 502
Script:
#!/bin/bash
# Check Vertex AI provider implementation for required configuration
echo "Checking Vertex AI provider configuration..."
rg -A 10 "type Config struct" go/ai-proxy/api/vertex/vertex.go
# Check provider initialization
echo "Checking provider initialization..."
rg -A 10 "NewProvider" go/ai-proxy/api/vertex/vertex.go
# Check environment variable usage
echo "Checking environment variable usage..."
rg "os\.Getenv.*PLRL_PROVIDER" go/ai-proxy/api/vertex/vertex.go
Length of output: 453
Script:
#!/bin/bash
# Look for the Vertex AI provider implementation
echo "Checking Vertex AI provider file content..."
cat go/ai-proxy/api/vertex/vertex.go
# Check for environment variable references in provider setup
echo "Checking provider setup in args..."
cat go/ai-proxy/args/args.go | grep -A 10 -B 10 "vertex"
Length of output: 2445
Script:
#!/bin/bash
# Check the provider implementation for required configuration
echo "Checking provider implementation..."
cat go/ai-proxy/proxy/provider/vertex.go
# Check deployment settings
echo "Checking deployment settings..."
cat go/controller/api/v1alpha1/deploymentsettings_types.go | grep -A 10 -B 10 "vertex"
Length of output: 4776
🧰 Tools
🪛 yamllint
[warning] 10-10: wrong indentation: expected 0 but found 2
(indentation)
go/ai-proxy/proxy/provider/base.go (3)
5-5: LGTM: Well-documented new field with appropriate imports
The new mapping field is properly documented and its type aligns with the os.Expand function's requirements.
Also applies to: 11-11, 29-31
57-57: LGTM: Proper gzip error handling implemented
The gzip handling implementation properly checks for errors and follows consistent patterns with the brotli case. The previous review concerns about error handling have been addressed.
Also applies to: 68-73
87-90: Verify the path mapping functionality
The implementation looks correct, but let's verify that the path mapping works as expected with environment variables.
✅ Verification successful
Based on the search results, I can now generate a final response as we have sufficient evidence about the path mapping implementation:
Path mapping implementation is correctly implemented and used
The code shows proper implementation of path mapping functionality:
- The mapping function is used correctly with
os.Expandin both base.go and vertex.go - There's a concrete example in vertex.go showing how the mapping function is implemented to replace variables like
vertex.EnvProjectID - The mapping is properly defined as an optional function parameter (
mapping func(string) string)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for path mapping usage patterns in the codebase
# Look for environment variable usage in paths
rg -A 2 'os\.Expand.*\$\{.*\}'
# Look for path mapping function implementations
ast-grep --pattern 'func $_(string) string {
$$$
}'
Length of output: 88
Script:
#!/bin/bash
# Let's try a broader search to understand the path mapping usage
# Search for os.Expand usage without specific pattern
rg -A 3 'os\.Expand'
# Search for mapping-related code
rg -A 3 'mapping.*string'
# Look for test cases that might demonstrate path mapping
rg -A 5 'Test.*Path.*Map'
# Search for string map declarations that might be used for path mapping
rg 'map\[string\]string'
Length of output: 6159
go/ai-proxy/proxy/provider/vertex.go (2)
19-23: LGTM! Clean struct definition with proper separation of concerns
The VertexProxy struct follows good practices by:
- Embedding the base proxy implementation
- Only storing the essential
TokenSourcefor authentication - Not storing sensitive information like service account credentials
44-56: LGTM! Good error handling in ModifyResponse
The implementation properly:
- Propagates errors from the base implementation
- Handles response body modification errors
- Returns appropriate error responses
floreks
deleted the
sebastian/prod-2866-add-vertex-ai-support-for-ai-proxy
branch
Added support for vertex ai provider to the ai proxy.
Test Plan
Locally + dev env + basic unit test
Checklist
Summary by CodeRabbit
New Features
Improvements
Configuration Changes
values.yaml, including new provider selection and service account management.Bug Fixes
Documentation