Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat (provider/azure): add support for azure managed identity. #3911

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

ShervK
Copy link

@ShervK ShervK commented Nov 27, 2024

Fixes: #3124

Similar to my last PR, #3463 but updated to v4 and cleaned up the code. Had to close the previous PR as it was too far behind and easier to start over than rebase.

  • added option for a token provider to be sent
  • added new tests
  • added to the docs

Made it a fetch wrapper to grab a new token on new requests, similar to how OpenAI package does it.

Please let me know what else I can add or change.

@ShervK ShervK force-pushed the azure-managed-identity branch from 0294700 to 4804e03 Compare November 27, 2024 06:43
@ShervK ShervK marked this pull request as ready for review November 27, 2024 06:45
@lgrammel
Copy link
Collaborator

Thanks for the PR. I would prefer a different design for passing in auth that's more extensible.

cc @shaper ideally this would align with the design we end up with for google vertex so it's similar across providers

@ShervK
Copy link
Author

ShervK commented Nov 27, 2024

Thanks for the feedback, I didn't realize there were plans to implement token based auth for google vertex as well.

One idea that comes to mind is setting up a type-based approach for providers that have multiple auth methods, something like

type AuthenticationMethod = 
  | { type: 'apiKey'; value: string }
  | { type: 'bearer'; tokenProvider: () => Promise<string> }
  | { type: 'custom'; getAuth: () => Promise<Record<string, string>> };

You could still have apiKey as an option and default to that if its provided, as to not cause breaking changes. Would this be more along the lines of what you're looking for?

@shaper
Copy link
Contributor

shaper commented Nov 27, 2024

Thanks for the feedback, I didn't realize there were plans to implement token based auth for google vertex as well.

One idea that comes to mind is setting up a type-based approach for providers that have multiple auth methods, something like

type AuthenticationMethod = 
  | { type: 'apiKey'; value: string }
  | { type: 'bearer'; tokenProvider: () => Promise<string> }
  | { type: 'custom'; getAuth: () => Promise<Record<string, string>> };

You could still have apiKey as an option and default to that if its provided, as to not cause breaking changes. Would this be more along the lines of what you're looking for?

Hi, thanks for the contribution, I will take a closer look at it!

A question: all of the above seem effectively implementable if we expose a way to optionally add to request headers. We have one today via headers in provider settings today, however, it isn't async, which can limit what one might need to do (e.g. fetching/refreshing an auth token with its own async ops needed). We could add a new one that allows async.

That has a nice feel from a simplicity of API surface perspective as it's one thing to think about, though it's also admittedly a little raw, as not all devs or docs want to think about or deal with talking about headers vs. the auth primitives/methods.

A concern could be any performance impact/overhead from calling the dev-provided headers population (do we do it once and cache it, or on each request, or catch auth errors and try refreshing by calling it again, ...). I expect we already have this auth-lifetime issue in some cases, and I wonder whether folks are having to catch auth timeout/failure and re-create the provider with updated auth in some cases.

Stepping back -- today I theorize you actually can already do what you want by calling your identity token provider before you create the Azure provider, getting your token, and then using the existing headers settings. The issues I can see you'd face, per the above, are (1) a little less semantically-relevant API for Azure (others would have to know to do similarly, we'd want to enhance docs/examples), and (2) any token timeout/refresh issues. Is this accurate, am I missing anything, what else comes to mind?

It would be great to hear your thoughts on the above.

@ShervK
Copy link
Author

ShervK commented Nov 28, 2024

@shaper Those are great points and you nailed the main challenges here.

I actually thought about using getHeaders for this PR at first, but the async nature of token providers is what stopped me. Since getHeaders is used consistently across providers, I didn’t want to break a standard just for this feature. That said, an async getHeaders would be a great way to clean this up and keep things consistent across providers if thats an option.

Initially, I had token auth working by setting it up in the headers option at the request level. I ran into issues with token lifetimes when using headers at the provider level since the provider had to be configured ahead of time (I'm using customProvider). As the project grew and we started using more providers, it got trickier to manage for the reasons you brought up (caching, validating, refetching). That’s when I started looking into how client libraries (like OpenAI, Anthropic, etc.) handle token-based auth. Most of them seem to fetch a fresh token on each request, which avoids those lifetime issues.
Right around that time, you rolled out the custom fetch feature for providers so I used that to fetch a new token on each request, mimicking the behaviour of those libraries.

You can achieve this with what's available in the SDK right now, but it's not as intuitive as other libraries, and like you said, "a little raw". I was hoping to help bridge that gap, but I agree that the docs could address this just as well. I’ve actually seen a few comments from people saying they can’t use the SDK because of its missing native token-based auth support for their cloud provider. I’ve also heard similar feedback in enterprise settings. So, even if this PR doesn’t move forward, adding an example in the docs for token-based auth would go a long way toward helping folks out.

I’m happy to tweak this further or help with the docs if that’s the way to go

@ShervK
Copy link
Author

ShervK commented Dec 5, 2024

Hi @shaper

I saw the PRs for the Resolvable utility type and Vertex AI auth changed get merged. I'll get to work on changing this PR to use those standards.

I saw that the Vertex auth implementation used the google-auth-library package and added authentication options to the provider settings, is this something you want to see in the Azure provider? @azure/identity supports a lot of different credential methods so native support would cause the provider to balloon in size.
An easier approach would be updating the docs with a section about using getBearerTokenProvider function included in @azure/identity to retrieve the token using the user's preferred credential method and adding the token to the Resolvable Header.

I can support both implementations, let me know which sounds better.

@shaper
Copy link
Contributor

shaper commented Dec 5, 2024

Hi @shaper

I saw the PRs for the Resolvable utility type and Vertex AI auth changed get merged. I'll get to work on changing this PR to use those standards.

I saw that the Vertex auth implementation used the google-auth-library package and added authentication options to the provider settings, is this something you want to see in the Azure provider? @azure/identity supports a lot of different credential methods so native support would cause the provider to balloon in size. An easier approach would be updating the docs with a section about using getBearerTokenProvider function included in @azure/identity to retrieve the token using the user's preferred credential method and adding the token to the Resolvable Header.

I can support both implementations, let me know which sounds better.

The easier approach and keeping the provider lighter weight sounds better to me. If you put together a draft PR we'd be happy to look at it and discuss, that would be great, thanks!

There were two reasons we kept the google-auth-library for vertex:

  • backwards-compatibility (folks were already using it) which is something we care about a lot
  • DX, it's really nice having the one-line import of the default vertex provider instance "just work"

The last point is a down side of the simpler approach. I think you'll end up with something that's like at least 4 lines or so wherever someone creates a provider instance.

Let's see what it looks like when you try it? My only other idea, if we want to pursue the goal of lighter-weight providers, is to have a separate contrib-type package where small hooks like this can live and folks could leverage that way, and we can push heavyweight deps to those other more specific packages that depend on the provider. I am not saying we should do this yet, it's just the only thing that comes to mind. Maybe you'll find something else syntactic-sugar-wise etc.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Support azure authentication with managed identity
3 participants