Improve remote call

great_ai/remote/call_remote_great_ai.py

@@ -17,6 +17,7 @@ def call_remote_great_ai(
     base_uri: str,
     data: Mapping[str, Any],
     retry_count: int = 4,
+    timeout_in_seconds: Optional[int] = 300,
     model_class: Optional[Type[T]] = None,
 ) -> Trace[T]:
     try:
@@ -28,7 +29,11 @@ def call_remote_great_ai(
         pass
 
     future = call_remote_great_ai_async(
-        base_uri=base_uri, data=data, retry_count=retry_count, model_class=model_class
+        base_uri=base_uri,
+        data=data,
+        retry_count=retry_count,
+        timeout_in_seconds=timeout_in_seconds,
+        model_class=model_class,
     )
 
     return asyncio.run(future)

great_ai/remote/call_remote_great_ai_async.py

@@ -13,7 +13,7 @@ async def call_remote_great_ai_async(
     base_uri: str,
     data: Mapping[str, Any],
     retry_count: int = 4,
-    timeout_in_seconds: int = 300,
+    timeout_in_seconds: Optional[int] = 300,
     model_class: Optional[Type[T]] = None,
 ) -> Trace[T]:
 
@@ -24,6 +24,8 @@ async def call_remote_great_ai_async(
         base_uri = f"{base_uri}/predict"
 
     transport = httpx.AsyncHTTPTransport(retries=retry_count)
+
+    try:
         async with httpx.AsyncClient(
             transport=transport, timeout=timeout_in_seconds
         ) as client:
@@ -32,6 +34,8 @@ async def call_remote_great_ai_async(
                 response.raise_for_status()
             except Exception:
                 raise RemoteCallError("Unexpected status code")
+    except Exception as e:
+        raise RemoteCallError from e
 
     try:
         trace = response.json()