truss 0.11.2rc1__tar.gz → 0.11.2rc2__tar.gz

{truss-0.11.2rc1 → truss-0.11.2rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.2rc1
+Version: 0.11.2rc2
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co

{truss-0.11.2rc1 → truss-0.11.2rc2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.11.2rc1"
+version = "0.11.2rc2"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "no-reply@baseten.co" },

{truss-0.11.2rc1 → truss-0.11.2rc2}/truss-chains/truss_chains/deployment/code_gen.py

@@ -546,7 +546,7 @@ async def websocket(self, websocket: fastapi.WebSocket) -> None:
         src=src,
         imports={
             "import fastapi",
-            "import starlette.websockets",
+            "from starlette.websockets import WebSocketState",
             "from truss_chains.remote_chainlet import utils",
         },
     )

{truss-0.11.2rc1 → truss-0.11.2rc2}/truss-chains/truss_chains/remote_chainlet/utils.py

@@ -39,7 +39,7 @@ if TYPE_CHECKING:
 try:
     import prometheus_client
 except ImportError:
-    logging.warning("Optional `prometheus_client` is not installed. ")
+    logging.warning("Optional `prometheus_client` is not installed.")
 
     class _NoOpMetric:
         def labels(self, *args: object, **kwargs: object) -> "_NoOpMetric":
@@ -65,6 +65,26 @@ except ImportError:
             return _NoOpMetric()
 
 
+try:
+    from fastapi import WebSocketDisconnect
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketDisconnect, in case local environment doesn't have
+    # fastapi.
+    class WebSocketDisconnect(Exception):  # type: ignore[no-redef]
+        def __init__(self, code: int, reason: Optional[str] = None):
+            super().__init__()
+            self.code = code
+            self.reason = reason
+
+
+try:
+    from starlette.websockets import WebSocketState
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketState, in case local environment doesn't have starlette.
+    class WebSocketState:  # type: ignore[no-redef]
+        CONNECTED = "connected"
+
+
 T = TypeVar("T")
 
 _LockT = TypeVar("_LockT", bound=Union[threading.Lock, asyncio.Lock])
@@ -588,17 +608,12 @@ class WebsocketWrapperFastAPI:
         await self._websocket.close(code=code, reason=reason)
 
     async def receive(self) -> Union[str, bytes]:
-        try:
-            import fastapi
-        except ImportError:
-            raise utils.make_optional_import_error("fastapi")
-
         message = await self._websocket.receive()
 
         if message.get("type") == "websocket.disconnect":
             # NB(nikhil): Mimics FastAPI `_raise_on_disconnect`, since otherwise the user has no
             # way of detecting that the client disconnected.
-            raise fastapi.WebSocketDisconnect(message["code"], message.get("reason"))
+            raise WebSocketDisconnect(message["code"], message.get("reason"))
         elif message.get("text"):
             return typing.cast(str, message["text"])
         else:
@@ -632,11 +647,6 @@ class WebsocketWrapperFastAPI:
         return self._websocket.iter_json()
 
     async def is_connected(self) -> bool:
-        try:
-            from starlette.websockets import WebSocketState
-        except ImportError:
-            raise utils.make_optional_import_error("starlette.websockets")
-
         # NB(nikhil): This isn't a foolproof mechanism for detecting whether a websocket
         # connection is actually alive, ping/pong messages are best suited for that. However,
         # as a heuristic to determine if a message is safe to send, this can do a pretty good job.