This commit is contained in:
Andras Schmelczer 2022-07-01 21:10:38 +02:00
parent 6a50f87abe
commit 076f676f4f
3 changed files with 20 additions and 9 deletions

View file

@ -95,7 +95,7 @@ class GreatAI(Generic[T]):
func: Optional[Callable[..., T]] = None, func: Optional[Callable[..., T]] = None,
*, *,
version: str = "0.0.1", version: str = "0.0.1",
return_raw_result: bool =False, return_raw_result: bool = False,
disable_rest_api: bool = False, disable_rest_api: bool = False,
disable_docs: bool = False, disable_docs: bool = False,
disable_dashboard: 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: if not disable_rest_api:
instance._bootstrap_rest_api( instance._bootstrap_rest_api(

View file

@ -1,5 +1,5 @@
import asyncio import asyncio
from typing import Any, Mapping, Optional, TypeVar from typing import Any, Mapping, Optional, Type, TypeVar
from pydantic import BaseModel 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") logger = get_logger("call_remote_great_ai")
T = TypeVar("T", BaseModel) T = TypeVar("T", bound=BaseModel)
def call_remote_great_ai( 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]: ) -> Trace[T]:
try: try:
asyncio.get_running_loop() asyncio.get_running_loop()
@ -22,7 +27,7 @@ def call_remote_great_ai(
pass pass
future = call_remote_great_ai_async( 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) return asyncio.run(future)

View file

@ -8,10 +8,14 @@ from .remote_call_error import RemoteCallError
http: Optional[HttpClient] = None http: Optional[HttpClient] = None
T = TypeVar("T", BaseModel) T = TypeVar("T", bound=BaseModel)
async def call_remote_great_ai_async( 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]: ) -> Trace[T]:
global http global http
if http is None: if http is None:
@ -27,7 +31,7 @@ async def call_remote_great_ai_async(
try: try:
if model_class is not None: 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) return Trace.parse_obj(response)
except Exception: except Exception:
raise RemoteCallError("Could not parse response") raise RemoteCallError("Could not parse response")