Skip to content

feat(provider): add Google VertexAI support#24

Merged
nirga merged 15 commits intotraceloop:mainfrom
onyedikachi-david:feat/add-vertexai-provider
May 16, 2025
Merged

feat(provider): add Google VertexAI support#24
nirga merged 15 commits intotraceloop:mainfrom
onyedikachi-david:feat/add-vertexai-provider

Conversation

@onyedikachi-david
Copy link
Copy Markdown
Contributor

@onyedikachi-david onyedikachi-david commented Dec 30, 2024 •
edited
Loading

Add Google VertexAI Provider Support

This PR adds support for Google VertexAI models (Gemini) as a new provider in the Hub, allowing users to route their LLM requests to Google's models through our unified API interface.

Definition of Done

  • Implementation of VertexAI provider
    • Chat completions endpoint with Gemini models
    • Text completions endpoint
    • Embeddings endpoint with textembedding-gecko
    • Tool/function calling support
    • Streaming support
    • Multi-modal support
  • Unit tests
    • Provider configuration tests
    • Request/response conversion tests
    • Error handling tests
    • Model format conversion tests
  • Integration tests
    • Chat completions test with recorded responses
    • Completions test with recorded responses
    • Embeddings test with recorded responses
    • Tool calling test with recorded responses
    • Test cassettes for offline testing
    • Quota retry mechanism
  • Documentation
    • Configuration examples in README
    • Authentication methods (API key and service account)
    • Model support documentation
    • Usage examples with OpenAI SDK
  • Configuration
    • Example configuration in config-example.yaml
    • Support for both API key and service account auth
    • Required parameters (project_id)
    • Optional parameters (location, credentials_path)
  • Error Handling
    • Proper status code mapping from Google API
    • Informative error messages
    • Quota limit handling with configurable retries
    • Authentication error handling
  • Code Review and Approval

Changes Made

  1. Added new VertexAI provider implementation

    • Full support for Gemini models (chat, completion)
    • Text embeddings with textembedding-gecko
    • Tool/function calling with proper mapping
    • Streaming support for real-time responses
    • Multi-modal capabilities for image+text inputs

  2. Added comprehensive test suite

    • Unit tests for core functionality
    • Integration tests with recorded responses
    • Test cassettes for offline testing
    • Quota retry mechanism for rate limits
    • Both auth methods covered

  3. Updated documentation

    • Detailed VertexAI setup instructions
    • Authentication configuration guide
    • Model compatibility list
    • Usage examples with OpenAI SDK
    • Configuration parameter reference

  4. Added robust configuration support

    • API key authentication
    • Service account authentication (JSON key file)
    • Project ID configuration
    • Regional endpoint configuration
    • Default values for optional parameters

Testing

  • Unit tests: cargo test
  • Integration tests with credentials: 
    # Service Account Auth
export VERTEXAI_CREDENTIALS_PATH="../credentials/vertexai-key.json"

# Record new responses
RECORD_MODE=1 cargo test

# Replay mode (default)
cargo test
  • Test retry configuration: 
    export RETRY_DELAY=60  # Seconds between retries

Security Considerations

  • Credentials stored outside repository
  • Test cassettes cleaned of sensitive data
  • Support for both API key and service account auth
  • Environment variable support for credentials
  • No hardcoded sensitive values

Notes

  • Default location is us-central1 (configurable)
  • Automatic retry on quota limits (configurable)
  • Test cassettes provided for offline development
  • Compatible with existing OpenAI SDK clients

Fixes #19
/claim #19

@algora-pbc algora-pbc bot mentioned this pull request Dec 30, 2024
Closed
8 tasks
@onyedikachi-david
cargo fmt fix
aa628ef 
Signed-off-by: David Anyatonwu <davidanyatonwu@gmail.com>
@CLAassistant
Copy link
Copy Markdown

CLAassistant commented Dec 31, 2024 •
edited
Loading

CLA assistant check
All committers have signed the CLA.

@onyedikachi-david
feat(vertexai): support both api key and service and robust tests
5af4770 
Signed-off-by: David Anyatonwu <davidanyatonwu@gmail.com>
@onyedikachi-david
fix(vertexai): enforce required project_id and improve endpoint forma…
b26628c 
…tting

- Changed project_id retrieval to use expect for mandatory parameter.
- Updated location retrieval to use unwrap_or for default value.
- Modified endpoint formatting to dynamically include location in the URL for both chat and embeddings requests.
- Refactored test provider setup to use constants for project_id and location.
@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented Dec 31, 2024 •
edited
Loading

@galkleinman This PR is ready for review.

running 9 tests
test providers::vertexai::tests::test_chat_completions_with_api_key ... ignored, Requires valid API key which is not available yet
test providers::vertexai::provider::tests::test_gemini_request_conversion ... ok
test providers::vertexai::provider::tests::test_gemini_response_conversion ... ok
test providers::vertexai::provider::tests::test_provider_new_missing_project_id - should panic ... ok
test providers::vertexai::provider::tests::test_provider_new ... ok
test providers::vertexai::tests::test_chat_completions ... ok
test providers::vertexai::tests::test_chat_completions_with_tools ... ok
test providers::vertexai::tests::test_embeddings ... ok
test providers::vertexai::tests::test_completions ... ok

test result: ok. 8 passed; 0 failed; 1 ignored; 0 measured; 0 filtered out; finished in 3.32s

     Running unittests src/main.rs (target/debug/deps/hub-e2c6417f88e881cd)

running 0 tests

test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

   Doc-tests hub

running 0 tests

test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s
image

@onyedikachi-david
feat(vertexai): add location validation and unit comprehensive tests
a381664 
- Introduced a new method `validate_location` to sanitize and validate location input, defaulting to "us-central1" if invalid.
- Updated the provider initialization to utilize the new location validation method.
- Added extensive unit tests for the provider, covering various scenarios including location validation, request conversion, and handling of empty messages.
- Ensured that invalid characters in location parameters are filtered out correctly.
- Enhanced tests to verify the precedence of API key over credentials path in configuration.

This commit improves the robustness of the VertexAIProvider by ensuring valid location formats and enhancing test coverage.
nirga
nirga reviewed Jan 2, 2025
View reviewed changes
Copy link
Copy Markdown
Member

@nirga nirga left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @onyedikachi-david - left a lot of comments. You're code isn't following idiomatic rust constructs and it's not following other providers already implemented in this repo. You're also missing integration (black box) tests which tests the whole system - not just the provider.

src/providers/vertexai/provider.rs Outdated
.collect();

if sanitized.is_empty() {
"us-central1".to_string() // Default if invalid
Copy link
Copy Markdown
Member

@nirga nirga Jan 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why?

Copy link
Copy Markdown
Contributor Author

@onyedikachi-david onyedikachi-david Jan 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for say the user didn't configure one

src/providers/vertexai/provider.rs
}
}

#[cfg(test)]
Copy link
Copy Markdown
Member

@nirga nirga Jan 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should not be here, it belongs to the test file

src/providers/vertexai/provider.rs Outdated

}

#[cfg(test)]
Copy link
Copy Markdown
Member

@nirga nirga Jan 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is it here?

@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented Jan 22, 2025

addressed your reviews @nirga

@nirga
Copy link
Copy Markdown
Member

nirga commented May 2, 2025

Hey @onyedikachi-david - can you sign the CLA on this?

@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented May 2, 2025

Hi, @nirga, I just did, I think it's requesting for yours also, since you made merge.

✅ onyedikachi-david
❌ nirga

@nirga
Copy link
Copy Markdown
Member

nirga commented May 2, 2025

I'll fix that @onyedikachi-david
Can you fix the formatting issues?

@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented May 2, 2025

Done @nirga

@nirga nirga force-pushed the feat/add-vertexai-provider branch from ba4a72d to 79a9fbe Compare May 8, 2025 13:19
nirga
nirga reviewed May 8, 2025
View reviewed changes
Copy link
Copy Markdown
Member

@nirga nirga left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @onyedikachi-david - I rebased but it seems build / test is now failing. Can you check?
Also, you haven't addressed my comment regarding test code in verrtexai/provider.rs

@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented May 8, 2025

Thanks @onyedikachi-david - I rebased but it seems build / test is now failing. Can you check? Also, you haven't addressed my comment regarding test code in verrtexai/provider.rs

Oh, totally forgot, let me check.

@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented May 12, 2025

Hello, @nirga, I have addressed all and did some (alot 😌) refactoring also.

hub/Users/onyedikachi on  feat/add-vertexai-provider is 📦 0.3.0 via 🦀 1.88.0 using ☁️  default/heavenya took 2.4s 
⇣6% ➜ cargo test                                                                                            
   Compiling hub v0.3.0 (/Users/onyedikachi/Documents/codes/algora-bounties/hub)
    Finished `test` profile [unoptimized + debuginfo] target(s) in 7.86s
     Running unittests src/lib.rs (target/debug/deps/hub-9b56fda26de774a9)

running 29 tests
test providers::bedrock::test::titan_tests::test_titan_provider_new ... ok
test providers::bedrock::test::antropic_tests::test_bedrock_provider_new ... ok
test providers::bedrock::test::ai21_tests::test_ai21_provider_new ... ok
test providers::bedrock::test::ai21_tests::test_ai21_provider_chat_completions ... ok
test providers::bedrock::test::antropic_tests::test_bedrock_provider_chat_completions ... ok
test providers::bedrock::test::titan_tests::test_chat_completions ... ok
test providers::vertexai::tests::test_chat_completions_with_api_key ... ignored, Requires valid API key which is not available yet
test providers::bedrock::test::ai21_tests::test_ai21_provider_completions ... ok
test providers::vertexai::tests::test_empty_message_handling ... ok
test providers::bedrock::test::titan_tests::test_embeddings ... ok
test providers::vertexai::tests::test_gemini_request_conversion ... ok
test providers::vertexai::tests::test_gemini_request_with_array_content ... ok
test providers::vertexai::tests::test_gemini_request_with_system_message ... ok
test providers::vertexai::tests::test_gemini_response_conversion ... ok
test providers::vertexai::tests::test_gemini_request_with_tools ... ok
test providers::vertexai::tests::test_gemini_response_with_tool_calls ... ok
test providers::vertexai::tests::test_generation_config_limits ... ok
test providers::vertexai::tests::test_location_validation ... ok
test providers::vertexai::tests::test_invalid_location_format - should panic ... ok
test providers::vertexai::tests::test_provider_new_missing_project_id - should panic ... ok
test providers::vertexai::tests::test_response_error_mapping ... ok
test providers::vertexai::tests::test_tool_choice_none ... ok
test providers::vertexai::tests::test_auth_config_credentials_only ... ok
test providers::vertexai::tests::test_provider_new ... ok
test providers::vertexai::tests::test_auth_config_precedence ... ok
test providers::vertexai::tests::test_completions - should panic ... ok
test providers::vertexai::tests::test_chat_completions_with_tools ... ok
test providers::vertexai::tests::test_embeddings ... ok
test providers::vertexai::tests::test_chat_completions ... ok

test result: ok. 28 passed; 0 failed; 1 ignored; 0 measured; 0 filtered out; finished in 2.90s

     Running unittests src/main.rs (target/debug/deps/hub-c43e446858e3fa57)

running 0 tests

test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

   Doc-tests hub

running 0 tests

test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s


hub/Users/onyedikachi on  feat/add-vertexai-provider is 📦 0.3.0 via 🦀 1.88.0 using ☁️  default/heavenya took 14.1s 
⇣6% ➜

@nirga
Copy link
Copy Markdown
Member

nirga commented May 13, 2025

Thanks @onyedikachi-david - can you also lint / fix formatting issues?

@onyedikachi-david
style: fmt
2975a95 
Signed-off-by: David Anyatonwu <davidanyatonwu@gmail.com>
@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented May 14, 2025

Done @nirga

nirga
nirga reviewed May 14, 2025
View reviewed changes
Copy link
Copy Markdown
Member

@nirga nirga left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @onyedikachi-david! Looks like tests are failing though

@nirga
Copy link
Copy Markdown
Member

nirga commented May 15, 2025

@onyedikachi-david and now formatting errors again 😅

@onyedikachi-david
chore: fmt
cc9775f 
Signed-off-by: David Anyatonwu <davidanyatonwu@gmail.com>
@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented May 15, 2025

@nirga Sorry about that, always forgetting to fmt. just did

@nirga
Copy link
Copy Markdown
Member

nirga commented May 15, 2025

@onyedikachi-david still failing

@onyedikachi-david
Copy link
Copy Markdown
Contributor Author

onyedikachi-david commented May 15, 2025

@onyedikachi-david still failing

Fixed, it was clippy warnings this time

@nirga nirga merged commit 2501921 into traceloop:main May 16, 2025
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@nirga nirga nirga approved these changes

Assignees

No one assigned

Labels

🙋 Bounty claim

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Feature Request: Add Google VertexAI Provider Support

3 participants

@onyedikachi-david @CLAassistant @nirga