diff --git a/src/great_ai/great_ai/deploy/great_ai.py b/src/great_ai/great_ai/deploy/great_ai.py index d438926..0f306e6 100644 --- a/src/great_ai/great_ai/deploy/great_ai.py +++ b/src/great_ai/great_ai/deploy/great_ai.py @@ -95,7 +95,7 @@ class GreatAI(Generic[T]): func: Optional[Callable[..., T]] = None, *, version: str = "0.0.1", - return_raw_result: bool =False, + return_raw_result: bool = False, disable_rest_api: bool = False, disable_docs: bool = False, disable_dashboard: bool = False, @@ -112,7 +112,9 @@ class GreatAI(Generic[T]): ), ) - instance = GreatAI[T](func, version=version, return_raw_result=return_raw_result) + instance = GreatAI[T]( + func, version=version, return_raw_result=return_raw_result + ) if not disable_rest_api: instance._bootstrap_rest_api( diff --git a/src/great_ai/great_ai/remote/call_remote_great_ai.py b/src/great_ai/great_ai/remote/call_remote_great_ai.py index cd03251..6b3e436 100644 --- a/src/great_ai/great_ai/remote/call_remote_great_ai.py +++ b/src/great_ai/great_ai/remote/call_remote_great_ai.py @@ -1,5 +1,5 @@ import asyncio -from typing import Any, Mapping, Optional, TypeVar +from typing import Any, Mapping, Optional, Type, TypeVar from pydantic import BaseModel @@ -9,9 +9,14 @@ from .call_remote_great_ai_async import call_remote_great_ai_async logger = get_logger("call_remote_great_ai") -T = TypeVar("T", BaseModel) +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 + base_uri: str, + data: Mapping[str, Any], + retry_count: int = 4, + model_class: Optional[Type[T]] = None, ) -> Trace[T]: try: asyncio.get_running_loop() @@ -22,7 +27,7 @@ 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, model_class=model_class ) return asyncio.run(future) diff --git a/src/great_ai/great_ai/remote/call_remote_great_ai_async.py b/src/great_ai/great_ai/remote/call_remote_great_ai_async.py index 95dbbb2..7604be7 100644 --- a/src/great_ai/great_ai/remote/call_remote_great_ai_async.py +++ b/src/great_ai/great_ai/remote/call_remote_great_ai_async.py @@ -8,10 +8,14 @@ from .remote_call_error import RemoteCallError http: Optional[HttpClient] = None -T = TypeVar("T", BaseModel) +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 + 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: @@ -27,7 +31,7 @@ async def call_remote_great_ai_async( try: if model_class is not None: - response['output'] = model_class.parse_obj(response['output']) + response["output"] = model_class.parse_obj(response["output"]) return Trace.parse_obj(response) except Exception: raise RemoteCallError("Could not parse response")