chainlit with custom endpoint

[cesi08]

I have to use chainlit without using the ui, I'm using the version 1.3.2
my goal is to define three endpoints (one for login, one for instantiating an agent (related to user session), and one for invoking the agent also related to user session)

When I set the agent and then retrieve it on the endpoint that invokes it, it arrives always None as if the session was not maintained between the custom endpoints that I define.
Do you have any advice?

from chainlit.auth import authenticate_user
from chainlit.context import init_http_context
from chainlit.data.sql_alchemy import SQLAlchemyDataLayer
from chainlit.data.storage_clients import AzureStorageClient
from chainlit.element import ElementBased
from chainlit.server import app
from fastapi import Depends, Request
from fastapi.responses import HTMLResponse, JSONResponse

class LLMRequest(BaseModel):
    message: str

@cl.password_auth_callback
def auth_callback(username: str, password: str):
      return cl.User(
          identifier="test",
          metadata={"role": "admin", "provider": "credentials"},
      )


@app.post("/set_user_agent")
async def set_user_agent(
    current_user: Annotated[Union[cl.User], Depends(authenticate_user)]
):
    init_http_context(user=current_user)
    agent = create_agent()
    cl.user_session.set("agent", agent)
    print(f"Agent created for user {current_user.identifier}: {agent}")
    stored_agent = cl.user_session.get("agent")
    print(f"Stored agent for user {current_user.identifier}: {stored_agent}")
    return JSONResponse({"message": "Agent created"})


@app.post("/endpoint_message")
async def endpoint_message(
    request: LLMRequest,
    current_user: Annotated[Union[cl.User], Depends(authenticate_user)],
):
    init_http_context(user=current_user)
    print(cl.user_session.get("agent"))
    print(current_user)
    agent_user = cl.user_session.get("agent")
    result = agent_user.invoke(input={"input": request.message})
    return JSONResponse(result["output"])

@dosu

[dosubot]

I found a similar issue that might be relevant to your situation. In a closed issue titled Problem with user session variables, it was noted that you need to be in a Chainlit context to use user_session. This means calling the code within specific functions like on_chat_start, langchain_factory, or action_callback. Each tab/window of the Chainlit app will have its own user session, so ensure you don't have multiple tabs/windows open if you want to avoid multiple sessions being triggered.

For your specific case, to maintain user session between custom endpoints in Chainlit version 1.3.2, you can use the UserSession class to store and retrieve user-specific data between calls. Here’s a brief guide:

1. Storing Data in the Session:
   Use the set method to store data in the user session:

   cl.user_session.set('agent', your_agent_object)

2. Retrieving Data from the Session:
   Use the get method to retrieve data from the user session:

   agent = cl.user_session.get('agent')
   if agent is None:
       # Handle the case where the agent is not found
       pass

Additionally, ensure that the init_http_context function is correctly initializing the session context for each request. This function creates a new HTTPSession object with a unique session ID and thread ID, which is crucial for maintaining session continuity ^[1][2].

If you are using Chainlit with FastAPI, you can mount your Chainlit app on an existing FastAPI app and use Chainlit APIs within your FastAPI endpoints. This setup ensures that the session state is maintained across different endpoints, and the agent is correctly stored and retrieved for each user ^[3].

To continue talking to Dosu, mention @dosu.