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

Update vertex-ai-gemini to check OIDC TokenCredential #626

Conversation

sberyozkin
Copy link
Contributor

Related to #610

@sberyozkin
Copy link
Contributor Author

OK, it works now, got:

Beans, Servlets, JBoss galore,
EJBs dance 'cross Java's shore,
XML deploys, wars galore 
A web of code to explore! 

(Author: Gemini) 

Let me update README with the security considerations section.

@sberyozkin sberyozkin force-pushed the vertex-ai-gemini-token-credential branch from c9986ef to c2bfa61 Compare May 27, 2024 17:45
@sberyozkin sberyozkin marked this pull request as ready for review May 27, 2024 17:45
@sberyozkin sberyozkin requested a review from a team as a code owner May 27, 2024 17:45
@sberyozkin
Copy link
Contributor Author

@geoand @jmartisk It all works now nicely, please review.

The user authenticates to Google by accessing http://localhost:8080 and the token acquired and authorized during the authentication process is used to access the Vertex AI Gemini endpoint.
README has a security consideration section detailing this process.

I've also added a configuration property, keeping the new option enabled by default as it gives a user specific access to the Gemini API out of the box, since setting up default credentials can not be user-unique if used by more than one user.

I plan to follow up with a similar demo for Azure OpenAI

Thanks

@sberyozkin sberyozkin force-pushed the vertex-ai-gemini-token-credential branch from c2bfa61 to 52bdad5 Compare May 27, 2024 18:11
@sberyozkin
Copy link
Contributor Author

I update the demo's user message to request only a short poem about Java. Requesting both short and funny poem for several times gets on the wrong side of Vertex AI, it starts giving some risk scores...

@sberyozkin
Copy link
Contributor Author

By the way, this injected TokenCredential will also pick up an incoming bearer access token, as long as Quarkus OIDC can verify it (example, SPA authenticates with Google and passes the token to Quarkus PoemResource endpoint).

@sberyozkin
Copy link
Contributor Author

What can make a very cool demo is to have both vertex ai gemini and azure openai providers used in the same demo but make them accessible to either Azure or Google authenticated users, thus also demonstrating OIDC multitenancy

@sberyozkin sberyozkin force-pushed the vertex-ai-gemini-token-credential branch from 52bdad5 to a41c061 Compare May 28, 2024 14:42
@sberyozkin sberyozkin force-pushed the vertex-ai-gemini-token-credential branch from a41c061 to 35a91a8 Compare May 28, 2024 17:47
@sberyozkin
Copy link
Contributor Author

I decided to disable picking up a security identity token by default, it is easy to enable with a property and users should make a decision if it is what they'd like to do

@geoand
Copy link
Collaborator

geoand commented Jun 3, 2024

I like it!

I have the same comment as with #634

@sberyozkin
Copy link
Contributor Author

Thanks @geoand, this one can not be because AuthProvider interface is OpenAiRestApi bound, this is what I implied in #634.
So we can have an oidc auth provider module which registers AuthProvider but for it to work across multiple providers like Azure OpenAI, Vertex Gemini, AuthProvider should be in the core module...
Sorry if I've missed something, I'm on PTO, heading to the ⛱️

@geoand
Copy link
Collaborator

geoand commented Jun 3, 2024

Have fun!

We'll take this one up again when you are back

@sberyozkin
Copy link
Contributor Author

sberyozkin commented Jun 11, 2024

Thanks @geoand, I'm back now.

So, let me refer you to the way OIDC token propagation works in Quarkus in general. For example, with rest-client-oidc-token-propagation. This is essentially a single REST client filter which collects the current token from the injected TokenCredential and propagates it as a bearer access token to the downstream target.

For example, let's say we have an OIDC secured FrontendResource which uses some REST client to access an OIDC secured DownstreamResource. I authenticate to the FrontendResource with Keycloak/Auth0/Google/Azure/whatever other SSO provider, REST client, being annotated with @AccessToken, has my OIDC session's access token picked up and propagated to the DownstreamEndpoint where this token is verified as a bearer access token.

An important point is that adding rest-client-oidc-token-propagation is enough to propagate this token, no matter which SSO provider was used to authenticate me.

The situation with this PR (and with #634) is identical as far as the token propagation is concerned - I login with Google to the frontend resources which must propagate a token to the downstream Vertex Gemini (or I login with Azure and then have a token propagated to Azure OpenAI - though I haven't made it confirmed it working yet).

But using OpenAiRestApi#ApiProvider does not work for Vertex Gemini, so the idea of having a single module, which if added, makes it work for both Azure OpenAI and Vertex Gemini is unfortunately not realizable, and I don't think model specific modules enabling such a token propagations should be shipped either.

I see that what I did in these 2 PRs might appear as inflexible, direct dependency on quarkus-security, and some tweaking with TokenCredential, but at least it is transparent to OIDC users, it just works for them and it can be disabled on a per model basis, where even with the OIDC enabled, no automatic token propagation is done.

If you prefer both of these PRs be re-implemented in terms of ApiProvider, essentially, I'll just end up opening a single PR which will add say an oidc-api-provider extension, then please move ApiProvider to the core, such that both Azure OpenAI, Vertex Gemini, and all other model providers which require token propagation, can depend on it. Otherwise I propose to keep these 2 PRs in the current form.

Perhaps a 3rd option is to add a REST client filter to the demo which propagates the token, but as implied in the docs, this is less flexible as one can't restrict to specific models.

I don't have a strong opinion, let me know please what you prefer, thanks

@geoand
Copy link
Collaborator

geoand commented Jun 12, 2024

Thanks for the input @sberyozkin.

Just to make sure we are on the same page, when you say

OpenAiRestApi#ApiProvider

I assume you mean OpenAiRestApi#AuthProvider ?

@sberyozkin
Copy link
Contributor Author

@geoand Thanks, yes, apologies for the typo, meant OpenAiRestApi#AuthProvider

@geoand
Copy link
Collaborator

geoand commented Jun 13, 2024

Cool thanks.

So here is how I understand this, please correct me if I'm wrong:

We move AuthProvider to the core (and probably call it RestAuthProvider or something). The OIDC implementation would simply grab the OIDC token and put it in the headers of the request.
The tricky parts becomes how to enable this propagation automatically...
I think that a good solution would to do that when both the following conditions are met:

  • The user has also added quarkus-rest-client-oidc-token-propagation
  • The user has opted in by using a specific property

WDYT?

@sberyozkin
Copy link
Contributor Author

Hi @geoand

That looks like it can work but I was thinking that we won't use quarkus-rest-client-oidc-token-propagation, instead this project would have an extension like rest-auth-provider-token-propagation, which would ship an implementation of RestAuthProvider which has TokenCredential injected and then interested users would add it to their project and it will enable the propagation automatically (for quarkus-oidc or quarkus-smallrye-jwt users) by simply including this extension, unconditionally. I may've missed your point here about the conditional dependencies though :-)

quarkus-rest-client-oidc-token-propagation ships a ClientRequestFilter which is added to all or specific REST clients, but I'm not sure if we could control at the level at whichrest-auth-provider-token-propagation can be controlled.

With quarkus-rest-client-oidc-token-propagation, one can use @AccessToken to bind the token propagation filter to a specific REST client only.

So I'm imagining, that if we had rest-auth-provider-token-propagation, then it would eventually support binding the shipped RestAuthProvider to specific models only in a multi-model setup.

What do you think ?

@geoand
Copy link
Collaborator

geoand commented Jun 13, 2024

That sounds fine with me.

What I really want is this to be an opt in for users and to be easy to use.
Furthermore I don't want other (non related) REST Clients to be affected.

@sberyozkin
Copy link
Contributor Author

@geoand,

What I really want is this to be an opt in for users and to be easy to use.

Sure, for example, in the demo introduced with this PR, users just add quarkus-langchain4j-model-auth-provider (I'm already thinking about a different name for the provider :-), ModelAuthProvider), and this Quarkus OIDC aware provider is detected by whatever model provider or providers which are currently used and which are aware of ModelAuthProvider.

I guess finer details like what if I use both Vertex Gemini and Azure Openai but only the latter is OK to pick up the logged in user's token, provided by one of the TokenCredential aware Quarkus extensions, while the former requires an API key can be figured as we move along, but I'm pretty sure that building on top of the current, possibly renamed AuthProvider, is what can help with supporting such variations in the future (may be by adding quarkus-langchain4j-model-auth-provider as conditional dependencies to such model providers or by allowing to refer from the quarkus-langchain4j-model-auth-provider configuration to specific models, etc)

For the final demo I have in mind though, once it also works for Azure OpenAI, I plan to have both providers propagating Azure and Google specific tokens respectively with a single quarkus-langchain4j-model-auth-provider dep.

Furthermore I don't want other (non related) REST Clients to be affected.

Definitely, won't be a problem.

So how do we proceed ? Can I suggest for @csotiriou to look at such a move (OpenAiRestApi.AuthProvider -> ModelAuthProvider or similarly named somewhere in the core and update the common openai code), since it will impact his work, if agreed in principle, and I can then follow up with a similar update for Vertex gemini and also add an OIDC specific auth provider extension ?

@geoand
Copy link
Collaborator

geoand commented Jun 13, 2024

So how do we proceed ? Can I suggest for @csotiriou to look at such a move (OpenAiRestApi.AuthProvider -> ModelAuthProvider or similarly named somewhere in the core and update the common openai code), since it will impact his work, if agreed in principle, and I can then follow up with a similar update for Vertex gemini and also add an OIDC specific auth provider extension ?

That makes sense to me

@csotiriou
Copy link
Contributor

@geoand @sberyozkin hello, just saw this conversation. So, in short I understand that I would need to

  • move the AuthProvider logic from my previous PR in the Core
  • Update the OpenAI modules to take advantage of it.
    and then, after this PR is fine and merged, then a subsequent PR possibly from @sberyozkin could use this to provide OIDC for Vertex and more?

@geoand
Copy link
Collaborator

geoand commented Jun 17, 2024

Yeah, that sounds correct :)

@sberyozkin
Copy link
Contributor Author

@csotiriou Indeed, I'll be happy to add a module shipping a Quarkus Security API specific auth provider, and rework this PR to use that new module instead, after you get a chance to move AuthProvider up to the core and update the Azure OpenAI code.
Later, I can rework #634 to only update the demo shipped with this PR to support a multiple OIDC provider authentication alongside multiple remote models
thanks

@sberyozkin
Copy link
Contributor Author

Closing in favor of #708

@sberyozkin sberyozkin closed this Jun 27, 2024
@sberyozkin sberyozkin deleted the vertex-ai-gemini-token-credential branch June 27, 2024 14:35
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.

None yet

3 participants