47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
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
|
|
|
|
T = TypeVar("T", bound=BaseModel)
|
|
|
|
|
|
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]:
|
|
|
|
if base_uri.endswith("/"):
|
|
base_uri = base_uri[:-1]
|
|
|
|
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:
|
|
trace = response.json()
|
|
except Exception:
|
|
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")
|