Add support for raw results

This commit is contained in:
Andras Schmelczer 2022-07-01 20:31:35 +02:00
parent 06cd7de379
commit 1a13c3f32e
No known key found for this signature in database
GPG key ID: 0EA1BC97D0AB076E

View file

@ -39,7 +39,7 @@ T = TypeVar("T")
class GreatAI(Generic[T]): class GreatAI(Generic[T]):
def __init__(self, func: Callable[..., Any], version: str): def __init__(self, func: Callable[..., Any], version: str, return_raw_result: bool):
func = automatically_decorate_parameters(func) func = automatically_decorate_parameters(func)
get_function_metadata_store(func).is_finalised = True get_function_metadata_store(func).is_finalised = True
@ -53,7 +53,7 @@ class GreatAI(Generic[T]):
) as t: ) as t:
result = func(*args, **kwargs) result = func(*args, **kwargs)
output = t.finalise(output=result) output = t.finalise(output=result)
return output return result if return_raw_result else output
self._cached_func = lru_cache(get_context().prediction_cache_size)( self._cached_func = lru_cache(get_context().prediction_cache_size)(
func_in_tracing_context func_in_tracing_context
@ -83,6 +83,7 @@ class GreatAI(Generic[T]):
@staticmethod @staticmethod
def create( def create(
version: str, version: str,
return_raw_result: bool,
disable_rest_api: bool, disable_rest_api: bool,
disable_docs: bool, disable_docs: bool,
disable_dashboard: bool, disable_dashboard: bool,
@ -94,6 +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,
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,
@ -103,13 +105,14 @@ class GreatAI(Generic[T]):
Callable[[Callable[..., T]], GreatAI[T]], Callable[[Callable[..., T]], GreatAI[T]],
partial( partial(
GreatAI.create, GreatAI.create,
disable_http=disable_rest_api, return_raw_result=return_raw_result,
disable_rest_api=disable_rest_api,
disable_docs=disable_docs, disable_docs=disable_docs,
disable_dashboard=disable_dashboard, disable_dashboard=disable_dashboard,
), ),
) )
instance = GreatAI[T](func, version=version) 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(