ajndkr / lanarky

The web framework for building LLM microservices
https://lanarky.ajndkr.com/
MIT License
975 stars 74 forks source link

feat: add Option to Accept or Not Accept WebSocket Connection #86

Closed Tooflex closed 1 year ago

Tooflex commented 1 year ago

Description

As a developer, I want an option to control whether to accept a WebSocket connection or not so I can prevent unnecessary "websocket.accept" messages after the connection has already been accepted.

Currently, my application is encountering an issue where it expects a "websocket.send" or "websocket.close" message after the WebSocket connection has been accepted, but it is receiving another "websocket.accept" message. This is likely due to the connection.connect() method trying to accept the WebSocket connection again, which is not allowed.

To address this issue, this pull request introduces an option to control whether to accept a WebSocket connection or not. This will provide more flexibility in managing WebSocket connections and help prevent the aforementioned issue.

For instance, in the case where the first message from the client is used to validate a JWT token, we can use this new option to prevent the connection.connect() method from accepting the WebSocket connection again after it has already been accepted:

@router.websocket("/chat")
async def chat_websocket_endpoint(
    websocket: WebSocket,
    chain: ConversationChain = Depends(conversation_chain),
):
    # Accept the connection
    await websocket.accept()

    # Wait for the client to send the initial authentication message (jwt token)
    token = await websocket.receive_text()

    # Verify the token and get the user
    user = await get_current_user_from_token(token)

    if not user:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    # If the user is authenticated, proceed with the connection
    connection = WebsocketConnection.from_chain(chain=chain, websocket=websocket)
    await connection.connect(accept_connection=False)

Acceptance Criteria