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

Support CDI MemoryIdProvider #523

Closed
sberyozkin opened this issue Apr 29, 2024 · 22 comments
Closed

Support CDI MemoryIdProvider #523

sberyozkin opened this issue Apr 29, 2024 · 22 comments

Comments

@sberyozkin
Copy link
Contributor

sberyozkin commented Apr 29, 2024

Right now, the memory id can be supported with @MemoryId or it is set to the connection id with the WebSockets Next integration.

I'd like to be able to use the current SecurityIdentity instance. Or the current JWT token's unique subject claim value.

I'm not 100% sure what is the best way to do it.
My proposal is to update Quarkus LangChain4j code which checks what memory id is if no @MemoryId is set, but before falling back to the WS connection id.
Here it will check with Arc if the correct scope MemoryIdProvider is registered, if yes, get it and retrieve an instance of Quarkus LangChain4j MemoryIdentifier interface from it and use as a chat memory id, or to keep it simple, just return String.

For example, one demo can have a request scoped MemoryIdProvider registered which will have JsonWebToken injected and return its token subject claim value, ensuring a verified user identity subject value is used as a memory id.

How does it sound ? I'm happy to start looking at this enhancement request once it is agreed upon

@geoand
Copy link
Collaborator

geoand commented Apr 29, 2024

I think this is a very interesting idea.

I do have a question however: do you expect this to behave any differently than how it does for Websockets or REST?

@sberyozkin
Copy link
Contributor Author

As far as I understand, with REST, if @MemoryId is used, it means that somehow the client using Quarkus LangChain4j has to manage that memory id and pass it to Quarkus, and with WebSockets Next, this id is implicitly set up from the WS connection id.

If this provider is introduced then if the client is say HTML based client which does not use Web Sockets, and authenticates to Quarkus with Google, then every time the client accesses Quarkus, the memory id will be calculated from SecurityIdentity. It will also work exactly the same for Web Sockets Next once quarkusio/quarkus#40312 is sorted out - if this provider is available, the security identity will get associated with the WS channel and instead of the connection id, it will be a security identity id, or may be a pair of the connection id and security identity id.

Did I understand your question correctly ?

@geoand
Copy link
Collaborator

geoand commented Apr 29, 2024

As far as I understand, with REST, if @memoryid is used, it means that somehow the client using Quarkus LangChain4j has to manage that memory id and pass it to Quarkus, and with WebSockets Next, this id is implicitly set up from the WS connectio

Actually with REST you have the same idea as with Websockets - i.e. we use the active equest itself

@geoand
Copy link
Collaborator

geoand commented Apr 29, 2024

if this provider is available, the security identity will get associated with the WS channel and instead of the connection id, it will be a security identity id, or may be a pair of the connection id and security identity id.

Sure, but what is the advantage of doing that?

@sberyozkin
Copy link
Contributor Author

sberyozkin commented Apr 29, 2024

I see, if it also works the similar way for both REST and WS, then, if the provider is available, then it is checked, if not fallback to the way Quarkus LangChain4j does it out of the box for REST and WebSockets. I can experiment with a chat bot demo which does not use WebSockets, as a POC.

Sure, but what is the advantage of doing that?

AFAIK, with Web Sockets, a connection id does not guarantee that it came from the user who started the WS upgrade, this is why Quarkus users want for example @RolesAllowed("admin") to work in the on-message WS calls, not only in the on-start calls.

But may be it will be much more relevant with just REST, you mentioned the memory id was also calculated similarly for REST, can you clarify how is it done to correlate things between multiple calls ? FYI, with the provider, the way it will work across calls is that for example an OIDC or Form authentication session cookie will help to create the security identity, etc.

I guess this provider, if introduced, may offer options for users to make more sophisticated correlations, even without the security in place

@geoand
Copy link
Collaborator

geoand commented Apr 29, 2024

I think there is a misunderstanding here (which is totally reasonable). The fact that a specific ID is used is not strictly tied to how long memory is available.
In your case, that means if that even if you did have something like the user id, it does not mean that you automatically have the memory of all the previous interactions.
What is means is that we guarantee that the memory is available for the duration of the specific "request" (and remove the memory if the proper scope is used when that "request" terminates).

@sberyozkin
Copy link
Contributor Author

In your case, that means if that even if you did have something like the user id, it does not mean that you automatically have the memory of all the previous interactions.

I see what I mean. But I think in the secure Quarkus appplication, the lifespan of the "request" (I'm assuming here that it involves everything for example in a given WS session) will match the security identity lifespan.

Let me consider this example. I authenticate to Quarkus via Keycloak and keep accessing Quarkus, for as long as the OIDC session is active, Quarkus will see the same user id and this very same user id will keep the user specific interactions keyed correctly for LangChain4j - this should be the case with or without WS being used. If the OIDC user logs out, the security event is sent and it can help to clean the corresponding memory (though I haven't though yet about the clean up specifics in the logout case).

It would be useful for tracing as well, which user asked what. Right now, even if I'm biased :-), I think it will be beneficial to associate the security related data with the memory id.

@sberyozkin
Copy link
Contributor Author

Can you please clarify your earlier comment,

Actually with REST you have the same idea as with Websockets - i.e. we use the active equest itself

How does it work across multiple requests ?

@sberyozkin
Copy link
Contributor Author

sberyozkin commented Apr 29, 2024

By the way, another reason I came up with this proposal, is that, after looking at the langchain4j code, I can see the memory id is passed with Metadata and it can probably impact the way the custom augmentors provide a user specific content. As long as we can assert/confirm with tests the the security identity name remains the same during a given ChatBot session, it should help with the user specific document selection.

@geoand
Copy link
Collaborator

geoand commented Apr 29, 2024

Quarkus will see the same user id and this very same user id will keep the user specific interactions keyed correctly for LangChain4j - this should be the case with or without WS being used

This would only be the case if we knew the lifecycle of the user id.
My point is that REST requests and WS requests work automatically because the AI service itself also has the proper scope (meaning we can clear out memory when it goes away).

@geoand
Copy link
Collaborator

geoand commented Apr 29, 2024

How does it work across multiple requests ?

It does not, that's part of what I am failing to explain 😄.
Our automatic memory handling does not ensure that subsequent requests from the same user see the previous memory items.
If users want that, they need to handle memory manually (which is not too hard)

@sberyozkin
Copy link
Contributor Author

Hi @geoand

I think I understand that the way @SessionScope works supports the correct memory management at that langchain4j level. My understanding the CDI session scope is bound to the lifecycle of the HTTP session.

I see how each user's chat bot session avoids clashing with another user's session in this case.

And it is exactly the same case for example, for a typical secure OIDC session, we don't depend there on the CDI session scope, manage it in a different way, but the idea is the same.

So getting back to your question, what is the advantage of trying to use a user name or some other security related identifier, which can work in in the latter case, with the memory id provider. I think it will let introduce the chat bot experience beyond the first entry, where the user has authenticated and authorized, where it is a user-specific, personalized experience. Take a DB SQL query example mentioned at the Insights call yesterday - if a principal name is [email protected] then a user specific SQL query can be run on the local DB to do the user specific RAG.

With @SessionScope alone, it works technically, but does not give an info about the current user. And I'm not sure how it aligns with request-scoped Quarkus security, what are security characteristics of such a case.

IMHO, it is worth giving it a try. Someone will do it sooner or later in some form, the only question IMHO is where it will be done.

Like I said, if agreed in principle, I can start planning a draft PR to show how it may work, and learn a few things along the way.
I'm thinking of the following demo:

  • Quarkus endpoint authenticates the user to Google or one of the other providers, the OIDC session is live, no WebSockets are used (or may they are used but the OIDC token is attached to it)
  • User principal is passed as a memory identifier (or, to be considered, in some cases, it can be a pair of the user name and the automatically allocated memory id, to ensure it is unique to a specific authenticate user session, lets say in the future the custom memory id provider will be able to tap into the automatic memory management code to have a memory id issued, etc)
  • The injestion is done with from the DB which contains doc entries for alice, bob, etc
  • The augmentation fetches an entry for the currently authenticated user only, say, alice

Something along these lines. The user logs out, the provider can catch a security event and clear the Chat memory somehow. Or may be CDI session scope can be used along the way since I expect the lifespan of this session scope and the OIDC session be more or less equal. This will have to be figured out.

I think all I'll need to get started is to know where I can put this provider check in a pure REST case, to start with, just before the automatic memory management decides on allocating the memory id itself.

I guess I may be still missing a few details :-), but some kind of association between the the ChatBot session and the current security identity will be useful IMHO.

Please think about it with the rest of the Quarkus langchain4j team :-), I can try to help if this enhancement request can be of interest...

@geoand
Copy link
Collaborator

geoand commented Apr 30, 2024

Like I said, if agreed in principle, I can start planning a draft PR to show how it may work, and learn a few things along the way.

Sounds good to me!

@sberyozkin
Copy link
Contributor Author

Thanks @geoand, I'll give it a try, make take me a bit of time to get something real produced with a few other issues in the mix, but I do commit to it. Cheers

@geoand
Copy link
Collaborator

geoand commented Apr 30, 2024

👌

@sberyozkin
Copy link
Contributor Author

At the moment I'm looking at the fraud detection demo which does not require a custom memory id, but once we resolve the WS-Next securty integration issue, it would work nicely with the secure variation of a demo like csv-chat-bot

@sberyozkin
Copy link
Contributor Author

Hi @geoand I've found DefaultMemoryIdProvider and I can register a custom implementation in #539, with a low priority, and it picks up the current identity.
So may be then this issue is invalid ? My only minor concern is that the security based one must go at the -1 priority to take over not only RequestScopeStateDefaultMemoryIdProvider but also WebSocketConnectionDefaultMemoryIdProvider which is at the default priority. But may be in the future, DefaultMemoryIdProvider can get discovered as a CDI bean as well.

But in any case, I'm not sure now MemoryIdProvider is needed if DefaultMemorIdProvider is already available, what do you think ?

@sberyozkin
Copy link
Contributor Author

Let me work on updating #539 to show how one might want to use it

@geoand
Copy link
Collaborator

geoand commented May 20, 2024

But in any case, I'm not sure now MemoryIdProvider is needed if DefaultMemorIdProvider is already available, what do you think ?

I am not sure exactly which classes you mean since we only have DefaultMemoryIdProvider :)

@sberyozkin
Copy link
Contributor Author

I am not sure exactly which classes you mean since we only have DefaultMemoryIdProvider :)

I meant that I was thinking of adding a new interface, MemoryIdProvider, in this issue, but since DefaultMemoryIdProvider is available, duplicating it with another interface seems unnecessary.

@sberyozkin
Copy link
Contributor Author

Yeah, I think I should close it for now, sorry I did not do my homework earlier :-)

@sberyozkin sberyozkin closed this as not planned Won't fix, can't repro, duplicate, stale May 20, 2024
@geoand
Copy link
Collaborator

geoand commented May 20, 2024

NP!

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

No branches or pull requests

2 participants