Replace aiohttp with httpx

great_ai/__init__.py

@@ -4,9 +4,10 @@ __version__ = "0.0.12"
 
 from .context import configure
 from .deploy import GreatAI
-from .exceptions import (
+from .errors import (
     ArgumentValidationError,
     MissingArgumentError,
+    RemoteCallError,
     WrongDecoratorOrderError,
 )
 from .models import save_model, use_model
@@ -15,13 +16,9 @@ from .output_views import (
     MultiLabelClassificationOutput,
     RegressionOutput,
 )
-from .parameters import log_metric, parameter
+from .parameters.log_metric import log_metric
+from .parameters.parameter import parameter
 from .persistence import MongodbDriver, ParallelTinyDbDriver, TracingDatabaseDriver
-from .remote import (
-    HttpClient,
-    RemoteCallError,
-    call_remote_great_ai,
-    call_remote_great_ai_async,
-)
+from .remote import call_remote_great_ai, call_remote_great_ai_async
 from .tracing import add_ground_truth, delete_ground_truth, query_ground_truth
 from .views import RouteConfig, Trace

great_ai/__main__.py

@@ -17,7 +17,7 @@ from watchdog.observers import Observer
 from great_ai.constants import SERVER_NAME
 from great_ai.context import _is_in_production_mode
 from great_ai.deploy import GreatAI
-from great_ai.exceptions import ArgumentValidationError, MissingArgumentError
+from great_ai.errors import ArgumentValidationError, MissingArgumentError
 from great_ai.utilities import get_logger
 
 from .parse_arguments import parse_arguments

great_ai/errors/__init__.py

@@ -1,3 +1,4 @@
 from .argument_validation_error import ArgumentValidationError
 from .missing_argument_error import MissingArgumentError
+from .remote_call_error import RemoteCallError
 from .wrong_decorator_order_error import WrongDecoratorOrderError

great_ai/helper/assert_function_is_not_finalised.py

@@ -1,6 +1,6 @@
 from typing import Callable
 
-from ..exceptions import WrongDecoratorOrderError
+from ..errors import WrongDecoratorOrderError
 from .get_function_metadata_store import get_function_metadata_store
 
 

great_ai/parameters/__init__.py

@@ -1,3 +0,0 @@
-from .automatically_decorate_parameters import automatically_decorate_parameters
-from .log_metric import log_metric
-from .parameter import parameter

great_ai/parameters/parameter.py

@@ -3,7 +3,7 @@ from typing import Any, Callable, Dict, TypeVar, cast
 
 from typeguard import check_type
 
-from ..exceptions import ArgumentValidationError
+from ..errors import ArgumentValidationError
 from ..helper import get_arguments, get_function_metadata_store
 from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
 from ..tracing.tracing_context import TracingContext

great_ai/remote/__init__.py

@@ -1,4 +1,2 @@
 from .call_remote_great_ai import call_remote_great_ai
 from .call_remote_great_ai_async import call_remote_great_ai_async
-from .http_client import HttpClient
-from .remote_call_error import RemoteCallError

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")

great_ai/remote/http_client.py

@@ -1,60 +0,0 @@
-import logging
-from asyncio import sleep
-from typing import Any, Mapping, Optional
-
-import aiohttp
-
-from .remote_call_error import RemoteCallError
-
-logger = logging.getLogger("http")
-
-
-class HttpClient:
-    timeout_seconds: int = 600
-    wait_between_retries_seconds: float = 5
-
-    def __init__(
-        self,
-    ) -> None:
-        timeout = aiohttp.ClientTimeout(total=self.timeout_seconds)
-
-        self._session = aiohttp.ClientSession(
-            raise_for_status=False,
-            timeout=timeout,
-        )
-
-    async def post(
-        self,
-        url: str,
-        data: Mapping[str, Any],
-        retry_count: int = 0,
-        expected_status: Optional[int] = None,
-        **kwargs: Any,
-    ) -> Any:
-        for i in range(retry_count + 1):
-            try:
-                async with self._session.post(url, json=data, **kwargs) as r:
-                    if (
-                        expected_status is not None and r.status != expected_status
-                    ) or r.status >= 500:
-                        response_text = await r.text()
-                        raise ValueError(
-                            f"Found not-expected status code: {r.status}, response is: {response_text}"
-                        )
-                    try:
-                        return await r.json()
-                    except Exception:
-                        raise RemoteCallError(
-                            "JSON parsing failed",
-                        )
-            except Exception as e:
-                if retry_count - i > 1:
-                    logger.warning(
-                        f"Request failed ({e}), {retry_count - i - 1} retries left",
-                    )
-                    await sleep(self.wait_between_retries_seconds)
-
-        raise RemoteCallError(f"Request has failed too many ({retry_count + 1}) times")
-
-    async def close(self) -> None:
-        await self._session.close()