Improve types

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:
-
+            return cast(GreatAI.FactoryProtocol, factory)
-            @overload
+        else:
-            def inner(func: Awaitable[V]) -> GreatAI[Awaitable[Trace[V]], V]:
+            return factory(func)
-                ...
-
-            @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,
-        )
 
     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 (
+        return cast(
-            func_in_tracing_context_async
+            Callable[..., T],
-            if get_function_metadata_store(func).is_asynchronous
+            (
-            else func_in_tracing_context_sync
+                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:

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

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,

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 (