quarkiverse / quarkus-langchain4j

Quarkus Langchain4j extension
https://docs.quarkiverse.io/quarkus-langchain4j/dev/index.html
Apache License 2.0
116 stars 65 forks source link

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

Closed sberyozkin closed 5 days ago

sberyozkin commented 1 month ago

Related to #610

sberyozkin commented 1 month ago

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 commented 1 month ago

@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 commented 1 month ago

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 commented 1 month ago

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 commented 1 month ago

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 commented 1 month ago

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 commented 1 month ago

I like it!

I have the same comment as with https://github.com/quarkiverse/quarkus-langchain4j/pull/634

sberyozkin commented 4 weeks ago

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 commented 4 weeks ago

Have fun!

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

sberyozkin commented 3 weeks ago

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 commented 3 weeks ago

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 commented 2 weeks ago

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

geoand commented 2 weeks ago

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:

WDYT?

sberyozkin commented 2 weeks ago

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 commented 2 weeks ago

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 commented 2 weeks ago

@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 commented 2 weeks ago

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 commented 2 weeks ago

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

geoand commented 2 weeks ago

Yeah, that sounds correct :)

sberyozkin commented 2 weeks ago

@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 commented 5 days ago

Closing in favor of #708