bedrock.websockets.endpoint_with_websocket

 1import json  # pragma: unit
 2from typing import Any  # pragma: unit
 3
 4from bedrock.endpoints.dto.bedrock_response import BedrockResponse, BedrockErrorResponse  # pragma: unit
 5from bedrock.endpoints.endpoint import Endpoint  # pragma: unit
 6from bedrock.log import log_config  # pragma: unit
 7
 8log = log_config("WebsocketEndpoint")  # pragma: unit
 9
10
11class WebsocketEndpoint(Endpoint):  # pragma: unit
12    """
13    Your websocket handler class should inherit from this and be placed inside `./app/websocket_endpoints/`.
14    This introduces the `on_message` method which gets called when a message is received on the WebSocket connection.
15    """
16
17    def _handle_with_context(self, event, context=None) -> tuple[int, list[Any]]:  # pragma: no cover
18        return self._handle_websocket_event_with_context(event, context)
19
20    def _handle_websocket_event_with_context(self, event, context=None) -> tuple[int, list[Any]]:  # pragma: unit
21        try:
22            return self.on_message(event, context)
23        except Exception as e:
24            log.error(f"Error in _handle_websocket_event_with_context: {type(e).__name__}: {e}.")
25            return 500, BedrockErrorResponse({"error": "Error in _handle_websocket_event_with_context"}, event).as_json()
26
27    def on_message(self, event: dict, context: dict) -> tuple[int, BedrockResponse]:  # pragma: no cover
28        """
29        Override this method to handle websocket creation events in your own custom way.
30        By default, this gets called by `_handle_websocket_event_with_context`
31        """
32        pass