Replace aiohttp with httpx

great_ai/remote/call_remote_great_ai_async.py

@@ -1,12 +1,10 @@
 from typing import Any, Mapping, Optional, Type, TypeVar
 
+import httpx
 from pydantic import BaseModel
 
+from ..errors.remote_call_error import RemoteCallError
 from ..views import Trace
-from .http_client import HttpClient
-from .remote_call_error import RemoteCallError
-
-http: Optional[HttpClient] = None
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -15,23 +13,35 @@ async def call_remote_great_ai_async(
     base_uri: str,
     data: Mapping[str, Any],
     retry_count: int = 4,
+    timeout_in_seconds: int = 60,
     model_class: Optional[Type[T]] = None,
 ) -> Trace[T]:
-    global http
-    if http is None:
-        http = HttpClient()
 
     if base_uri.endswith("/"):
         base_uri = base_uri[:-1]
 
-    url = f"{base_uri}/predict/"
-    response = await http.post(
-        url=url, data=data, retry_count=retry_count, expected_status=200
-    )
+    if not base_uri.endswith("/predict"):
+        base_uri = f"{base_uri}/predict"
+
+    transport = httpx.AsyncHTTPTransport(retries=retry_count)
+    async with httpx.AsyncClient(
+        transport=transport, timeout=timeout_in_seconds
+    ) as client:
+        response = await client.post(base_uri, json=data)
+        try:
+            response.raise_for_status()
+        except Exception:
+            raise RemoteCallError("Unexpected status code")
 
     try:
-        if model_class is not None:
-            response["output"] = model_class.parse_obj(response["output"])
-        return Trace.parse_obj(response)
+        trace = response.json()
     except Exception:
-        raise RemoteCallError("Could not parse response")
+        raise RemoteCallError(
+            f"JSON parsing failed {response.text}",
+        )
+    try:
+        if model_class is not None:
+            trace["output"] = model_class.parse_obj(trace["output"])
+        return Trace.parse_obj(trace)
+    except Exception:
+        raise RemoteCallError("Could not parse Trace")