def get_websocket_callbacks(chain: Chain, websocket: WebSocket) -> list[Callable]:
    """Get websocket callbacks for a LangChain instance.

    Note: This function might not support all LangChain
    chain and agent types. Please open an issue on GitHub to
    request support for a specific type.

    Args:
        chain: A LangChain instance.
        websocket: A WebSocket instance.
    """
    callbacks = []

    if "source_documents" in chain.output_keys:
        callbacks.append(SourceDocumentsWebSocketCallbackHandler(websocket=websocket))

    if len(set(chain.output_keys) - {"source_documents"}) > 1:
        logger.warning(
            f"""multiple output keys found: {set(chain.output_keys) - {'source_documents'}}.

        Only the first output key will be used for sending tokens. For more complex websocket logic, define the endpoint function manually.
        """
        )

    if isinstance(chain, AgentExecutor):
        callbacks.append(FinalTokenWebSocketCallbackHandler(websocket=websocket))
    else:
        callbacks.extend(
            [
                TokenWebSocketCallbackHandler(
                    output_key=chain.output_keys[0], websocket=websocket
                ),
            ]
        )

    return callbacks
Get websocket callbacks for a LangChain instance.

    Note: This function might not support all LangChain
    chain and agent types. Please open an issue on GitHub to
    request support for a specific type.

    Args:
        chain: A LangChain instance.
        websocket: A WebSocket instance.
    
lanarky/adapters/langchain/utils.py