andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Improve types

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-07 16:50:10 +02:00
parent 7a4e75ff63
commit 612f3ac5e1
4 changed files with 58 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

85
great_ai/deploy/great_ai.py
View file

@ -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:

4
great_ai/tracing/delete_ground_truth.py
View file

@ -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])

8
great_ai/tracing/tracing_context.py
View file

@ -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,

13
great_ai/views/trace.py
View file

@ -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 (