Use the correct import

Author: igalshilman

python/restate/__init__.py

@@ -81,7 +81,7 @@ def test_harness(
     @asynccontextmanager
     async def create_client(
         ingress: str, headers: typing.Optional[dict] = None
-    ) -> typing.AsyncGenerator[RestateClient]:
+    ) -> typing.AsyncGenerator[RestateClient, None]:
         """a dummy client constructor to raise ImportError. Install restate-sdk[client] to use this feature"""
         raise ImportError("Install restate-sdk[client] to use this feature")
         yield  # type: ignore

python/restate/client.py

@@ -308,7 +308,9 @@ async def generic_send(
 
 
 @asynccontextmanager
-async def create_client(ingress: str, headers: typing.Optional[dict] = None) -> typing.AsyncGenerator[RestateClient]:
+async def create_client(
+    ingress: str, headers: typing.Optional[dict] = None
+) -> typing.AsyncGenerator[RestateClient, None]:
     """
     Create a new Restate client.
     """

python/restate/types.py

@@ -17,7 +17,7 @@
 
 from dataclasses import dataclass
 
-from restate.client import RestateClient
+from restate.client_types import RestateClient
 
 
 @dataclass