From 612f3ac5e1cc852332f983a6b9b5b9cb57770cba Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Thu, 7 Jul 2022 16:50:10 +0200 Subject: [PATCH] Improve types --- great_ai/deploy/great_ai.py | 85 ++++++++++++++----------- great_ai/tracing/delete_ground_truth.py | 4 +- great_ai/tracing/tracing_context.py | 8 ++- great_ai/views/trace.py | 13 +--- 4 files changed, 58 insertions(+), 52 deletions(-) diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py index 3024545..334d702 100644 --- a/great_ai/deploy/great_ai.py +++ b/great_ai/deploy/great_ai.py @@ -7,6 +7,7 @@ from typing import ( Generic, List, Optional, + Protocol, Sequence, TypeVar, Union, @@ -41,10 +42,26 @@ class GreatAI(Generic[T, V]): __name__: str __doc__: str + class FactoryProtocol(Protocol): + @overload + def __call__( # type: ignore + self, + func: Callable[..., Awaitable[V]], + ) -> "GreatAI[Awaitable[Trace[V]], V]": + + ... + + @overload + def __call__( + self, + func: Callable[..., V], + ) -> "GreatAI[Trace[V], V]": + ... + def __init__( self, func: Callable[..., Union[V, Awaitable[V]]], - version: str, + version: Union[str, int], route_config: RouteConfig, ): func = automatically_decorate_parameters(func) @@ -56,7 +73,7 @@ class GreatAI(Generic[T, V]): wraps(func)(self) self.__doc__ = f"GreatAI wrapper for interacting with the `{self.__name__}` function.\n\n{dedent(self.__doc__ or '')}" - self.version = version + self.version = str(version) flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions) if flat_model_versions: self.version += f"+{flat_model_versions}" @@ -91,44 +108,33 @@ class GreatAI(Generic[T, V]): @overload @staticmethod def create( - func: Optional[Union[Callable[..., V], Callable[..., Awaitable[V]]]] = ..., + func: None = ..., *, - version: str = ..., + version: Union[str, int] = ..., route_config: RouteConfig = ..., - ) -> Callable: + ) -> FactoryProtocol: ... @staticmethod def create( func: Optional[Callable] = None, *, - version: str = "0.0.1", + version: Union[str, int] = "0.0.1", route_config: RouteConfig = RouteConfig(), - ) -> Union[Callable, "GreatAI"]: + ) -> Union[ + FactoryProtocol, "GreatAI[Awaitable[Trace[V]], V]", "GreatAI[Trace[V], V]" + ]: + def factory(_func): # type: ignore + return GreatAI[Trace[V], V]( + _func, + version=version, + route_config=route_config, + ) + if func is None: - - @overload - def inner(func: Awaitable[V]) -> GreatAI[Awaitable[Trace[V]], V]: - ... - - @overload - def inner(func: Callable[..., V]) -> GreatAI[Trace[V], V]: - ... - - def inner(func): # type: ignore - return GreatAI.create( - func, - version=version, - route_config=route_config, - ) - - return inner - - return GreatAI[T, V]( - func, - version=version, - route_config=route_config, - ) + return cast(GreatAI.FactoryProtocol, factory) + else: + return factory(func) def __call__(self, *args: Any, **kwargs: Any) -> T: return self._wrapped_func(*args, **kwargs) @@ -158,29 +164,32 @@ class GreatAI(Generic[T, V]): *args: Any, do_not_persist_traces: bool = False, **kwargs: Any, - ) -> Trace[V]: + ) -> T: with TracingContext[V]( func.__name__, do_not_persist_traces=do_not_persist_traces ) as t: result = func(*args, **kwargs) - return t.finalise(output=result) + return cast(T, t.finalise(output=result)) @alru_cache(maxsize=get_context().prediction_cache_size) async def func_in_tracing_context_async( *args: Any, do_not_persist_traces: bool = False, **kwargs: Any, - ) -> Trace[V]: + ) -> T: with TracingContext[V]( func.__name__, do_not_persist_traces=do_not_persist_traces ) as t: result = await cast(Callable[..., Awaitable], func)(*args, **kwargs) - return t.finalise(output=result) + return cast(T, t.finalise(output=result)) - return ( - func_in_tracing_context_async - if get_function_metadata_store(func).is_asynchronous - else func_in_tracing_context_sync + return cast( + Callable[..., T], + ( + func_in_tracing_context_async + if get_function_metadata_store(func).is_asynchronous + else func_in_tracing_context_sync + ), ) def _bootstrap_rest_api(self, route_config: RouteConfig) -> None: diff --git a/great_ai/tracing/delete_ground_truth.py b/great_ai/tracing/delete_ground_truth.py index e5010bc..ef1620e 100644 --- a/great_ai/tracing/delete_ground_truth.py +++ b/great_ai/tracing/delete_ground_truth.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import List, Optional, Union, cast +from typing import List, Optional, Union from ..context import get_context @@ -19,4 +19,4 @@ def delete_ground_truth( conjunctive_tags=tags, until=until, since=since, has_feedback=True ) - db.delete_batch([cast(str, i.trace_id) for i in items]) + db.delete_batch([i.trace_id for i in items]) diff --git a/great_ai/tracing/tracing_context.py b/great_ai/tracing/tracing_context.py index 701310c..167132b 100644 --- a/great_ai/tracing/tracing_context.py +++ b/great_ai/tracing/tracing_context.py @@ -2,6 +2,7 @@ from contextvars import ContextVar from datetime import datetime from types import TracebackType from typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar +from uuid import uuid4 from ..constants import DEVELOPMENT_TAG_NAME, ONLINE_TAG_NAME, PRODUCTION_TAG_NAME from ..context import get_context @@ -25,11 +26,14 @@ class TracingContext(Generic[T]): def log_model(self, model: Model) -> None: self._models.append(model) - def finalise(self, output: T = None, exception: BaseException = None) -> Trace[T]: + def finalise( + self, output: Optional[T] = None, exception: Optional[BaseException] = None + ) -> Trace[T]: assert self._trace is None, "has been already finalised" delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000 - self._trace = Trace( + self._trace = Trace[T]( + trace_id=str(uuid4()), created=self._start_time.isoformat(), original_execution_time_ms=delta_time, logged_values=self._values, diff --git a/great_ai/views/trace.py b/great_ai/views/trace.py index 9413d85..f952da2 100644 --- a/great_ai/views/trace.py +++ b/great_ai/views/trace.py @@ -1,8 +1,7 @@ from pprint import pformat from typing import Any, Dict, Generic, List, Optional, TypeVar -from uuid import uuid4 -from pydantic import BaseModel, Extra, validator +from pydantic import BaseModel, Extra from ..helper import HashableBaseModel from .model import Model @@ -11,25 +10,19 @@ T = TypeVar("T") class Trace(Generic[T], HashableBaseModel): - trace_id: Optional[str] + trace_id: str created: str original_execution_time_ms: float logged_values: Dict[str, Any] models: List[Model] exception: Optional[str] - output: T + output: Optional[T] feedback: Any = None tags: List[str] class Config: extra = Extra.ignore - @validator("trace_id", always=True) - def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> str: - if v: - return v - return str(uuid4()) - @property def input(self) -> Any: return (