Skip to content

Commit

Permalink
feat: implement a gcp auth chain for vertex-ai clients (#1328)
Browse files Browse the repository at this point in the history
Implement support for GCP application default credentials and everything
else that the `gcp_auth` crate supports.

Also make `BAML_LOG=debug` useful for customers debugging the GCP auth
decision graph.

In theory I could probably fork `gcp_auth` and make it work in wasm (the
compile errors suggest to me that the underlying rand/ring
implementations just need to get swapped out for wasm-compatible ones to
make it work) but it was more time-bounded to just reuse the existing
implementation. (We should strongly look into using `gcp_auth`, though,
or at least upgrading to using the `oauth2` crate to handle the
exchange.)

Tested on cloud run: 

![image](https://github.com/user-attachments/assets/b02016d4-8dda-48d4-993c-b62642e22d82)

Tested with `gcloud auth application-default login`:
<img width="924" alt="image"
src="https://github.com/user-attachments/assets/a08abb81-fba5-43e9-b5b0-d939fbac8b9a"
/>

Tested with `gcloud auth print-access-token` fallback:
<img width="928" alt="image"
src="https://github.com/user-attachments/assets/e9c85ae7-ffc6-4d42-be48-b2d15095e112"
/>

Fixes #1312 

<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Implement GCP authentication chain for Vertex AI clients using
`gcp_auth` crate, update documentation, and modify build/test scripts.
> 
>   - **Authentication**:
> - Implement GCP authentication chain for Vertex AI clients using
`gcp_auth` crate in `vertex.rs`.
> - Introduce `UnresolvedGcpAuthStrategy` and `ResolvedGcpAuthStrategy`
enums for handling different auth strategies.
> - Add `VertexAuth` struct in `std_auth.rs` and `wasm_auth.rs` for
standard and WASM environments.
>   - **Documentation**:
> - Update `vertex.mdx` to include new authentication strategies and
examples.
>   - **Build and Test**:
> - Update `tools/build` to include new test paths and commands for
Vertex AI.
> - Modify `.github/workflows/build-python-release.reusable.yaml` to
track new branches.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 7674f18. It will automatically
update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->
  • Loading branch information
sxlijin authored Jan 18, 2025
1 parent 1cbb268 commit 6dfa23b
Show file tree
Hide file tree
Showing 17 changed files with 775 additions and 353 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/build-python-release.reusable.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ name: BAML Release - Build Python
on:
workflow_call: {}
push:
branches: [manylinux-x86, aaron-fix]
branches: [sam/gcp-auth-nope]

concurrency:
# suffix is important to prevent a concurrency deadlock with the calling workflow
Expand Down
56 changes: 55 additions & 1 deletion engine/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 5 additions & 0 deletions engine/baml-lib/llm-client/src/clients/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@ pub mod round_robin;
pub mod vertex;

#[derive(Debug)]
/// `Meta` is a generic carrying span information, so that if it comes from a .baml file,
/// we can trace it back to the original location in said .baml file. In dynamic clients,
/// though, we can't do that, so we just pass in `()`.
pub enum UnresolvedClientProperty<Meta> {
OpenAI(openai::UnresolvedOpenAI<Meta>),
Anthropic(anthropic::UnresolvedAnthropic<Meta>),
Expand Down Expand Up @@ -121,6 +124,8 @@ impl<Meta: Clone> UnresolvedClientProperty<Meta> {
}

impl crate::ClientProvider {
/// This codepath is common to both static clients (defined in a .baml file) and dynamic clients
/// (defined in user's python/ts code).
pub fn parse_client_property<Meta: Clone>(
&self,
properties: PropertyHandler<Meta>,
Expand Down
Loading

0 comments on commit 6dfa23b

Please sign in to comment.