Improve batch processing API

great_ai/deploy/great_ai.py

@@ -1,4 +1,4 @@
-from functools import lru_cache, partial, wraps
+from functools import lru_cache, wraps
 from textwrap import dedent
 from typing import (
     Any,
@@ -6,8 +6,10 @@ from typing import (
     Callable,
     Generic,
     List,
+    Literal,
     Optional,
     Sequence,
+    Tuple,
     TypeVar,
     Union,
     cast,
@@ -96,17 +98,66 @@ class GreatAI(Generic[T, V]):
     def __call__(self, *args: Any, **kwargs: Any) -> T:
         return self._wrapped_func(*args, **kwargs)
 
+    @overload
+    def process_batch(
+        self,
+        batch: Sequence[Tuple],
+        *,
+        concurrency: Optional[int] = None,
+        unpack_arguments: Literal[True],
+        do_not_persist_traces: bool = ...,
+    ) -> List[Trace[V]]:
+        ...
+
+    @overload
     def process_batch(
         self,
         batch: Sequence,
+        *,
         concurrency: Optional[int] = None,
-        do_not_persist_traces: Optional[bool] = False,
+        unpack_arguments: Literal[False] = ...,
+        do_not_persist_traces: bool = ...,
     ) -> List[Trace[V]]:
+        ...
+
+    def process_batch(
+        self,
+        batch: Sequence,
+        *,
+        concurrency: Optional[int] = None,
+        unpack_arguments: bool = False,
+        do_not_persist_traces: bool = False,
+    ) -> List[Trace[V]]:
+        wrapped_function = self._wrapped_func
+
+        def inner(value: Any) -> T:
+            return (
+                wrapped_function(*value, do_not_persist_traces=do_not_persist_traces)
+                if unpack_arguments
+                else wrapped_function(
+                    value, do_not_persist_traces=do_not_persist_traces
+                )
+            )
+
+        async def inner_async(value: Any) -> T:
+            return await cast(
+                Awaitable,
+                (
+                    wrapped_function(
+                        *value, do_not_persist_traces=do_not_persist_traces
+                    )
+                    if unpack_arguments
+                    else wrapped_function(
+                        value, do_not_persist_traces=do_not_persist_traces
+                    )
+                ),
+            )
+
         return list(
             parallel_map(
-                partial(
-                    self._wrapped_func, do_not_persist_traces=do_not_persist_traces
-                ),
+                inner_async
+                if get_function_metadata_store(self).is_asynchronous
+                else inner,
                 batch,
                 concurrency=concurrency,
             )

great_ai/tracing/tracing_context.py

@@ -1,7 +1,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 typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar, cast
 from uuid import uuid4
 
 from ..constants import DEVELOPMENT_TAG_NAME, ONLINE_TAG_NAME, PRODUCTION_TAG_NAME
@@ -32,23 +32,26 @@ class TracingContext(Generic[T]):
         assert self._trace is None, "has been already finalised"
 
         delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
-        self._trace = Trace[T](
-            trace_id=str(uuid4()),
-            created=self._start_time.isoformat(),
-            original_execution_time_ms=delta_time,
-            logged_values=self._values,
-            models=self._models,
-            output=output,
-            exception=None
-            if exception is None
-            else f"{type(exception).__name__}: {exception}",
-            tags=[
-                self._name,
-                ONLINE_TAG_NAME,
-                PRODUCTION_TAG_NAME
-                if get_context().is_production
-                else DEVELOPMENT_TAG_NAME,
-            ],
+        self._trace = cast(
+            Trace[T],
+            Trace(  # avoid ValueError: "Trace" object has no field "__orig_class__"
+                trace_id=str(uuid4()),
+                created=self._start_time.isoformat(),
+                original_execution_time_ms=delta_time,
+                logged_values=self._values,
+                models=self._models,
+                output=output,
+                exception=None
+                if exception is None
+                else f"{type(exception).__name__}: {exception}",
+                tags=[
+                    self._name,
+                    ONLINE_TAG_NAME,
+                    PRODUCTION_TAG_NAME
+                    if get_context().is_production
+                    else DEVELOPMENT_TAG_NAME,
+                ],
+            ),
         )
 
         return self._trace