Move files

great_ai/remote/__init__.py

@@ -0,0 +1,4 @@
+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.py

@@ -0,0 +1,34 @@
+import asyncio
+from typing import Any, Mapping, Optional, Type, TypeVar
+
+from pydantic import BaseModel
+
+from great_ai.utilities import get_logger
+
+from ..views import Trace
+from .call_remote_great_ai_async import call_remote_great_ai_async
+
+logger = get_logger("call_remote_great_ai")
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def call_remote_great_ai(
+    base_uri: str,
+    data: Mapping[str, Any],
+    retry_count: int = 4,
+    model_class: Optional[Type[T]] = None,
+) -> Trace[T]:
+    try:
+        asyncio.get_running_loop()
+        raise Exception(
+            f"Already running in an event loop, you have to call `{call_remote_great_ai_async.__name__}`"
+        )
+    except RuntimeError:
+        pass
+
+    future = call_remote_great_ai_async(
+        base_uri=base_uri, data=data, retry_count=retry_count, model_class=model_class
+    )
+
+    return asyncio.run(future)

great_ai/remote/call_remote_great_ai_async.py

@@ -0,0 +1,37 @@
+from typing import Any, Mapping, Optional, Type, TypeVar
+
+from pydantic import BaseModel
+
+from ..views import Trace
+from .http_client import HttpClient
+from .remote_call_error import RemoteCallError
+
+http: Optional[HttpClient] = None
+
+T = TypeVar("T", bound=BaseModel)
+
+
+async def call_remote_great_ai_async(
+    base_uri: str,
+    data: Mapping[str, Any],
+    retry_count: int = 4,
+    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
+    )
+
+    try:
+        if model_class is not None:
+            response["output"] = model_class.parse_obj(response["output"])
+        return Trace.parse_obj(response)
+    except Exception:
+        raise RemoteCallError("Could not parse response")

great_ai/remote/http_client.py

@@ -0,0 +1,60 @@
+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()

great_ai/remote/remote_call_error.py

@@ -0,0 +1,2 @@
+class RemoteCallError(Exception):
+    pass