Support Numpy

great_ai/deploy/great_ai.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from functools import lru_cache, wraps
 from textwrap import dedent
 from typing import (
@@ -51,7 +53,9 @@ class GreatAI(Generic[T, V]):
 
     IMPORTANT: when a request is served from cache, no new trace is created. Thus, the
     same trace can be returned multiple times. If this is undesirable turn off caching
-    using `configure(prediction_cache_size=0)`.
+    using `configure(prediction_cache_size=0)`. Each caller receives an independent
+    copy. Persisted traces are checked for feedback updates and deletions, including
+    changes made by other workers.
 
     Supports wrapping async and synchronous functions while also maintaining correct
     typing.
@@ -259,36 +263,60 @@ class GreatAI(Generic[T, V]):
         func: Callable[..., Union[V, Awaitable[V]]],
     ) -> Callable[..., T]:
         @lru_cache(maxsize=get_context().prediction_cache_size)
-        def func_in_tracing_context_sync(call: FrozenCall) -> T:
+        def func_in_tracing_context_sync(call: FrozenCall) -> Trace[V]:
             with TracingContext[V](
                 func.__name__, do_not_persist_traces=call.do_not_persist
             ) as t:
                 result = func(*call.args, **call.kwargs)
-                return cast(T, t.finalise(output=result))
+                return t.finalise(output=result)
 
         @alru_cache(maxsize=get_context().prediction_cache_size, cache_exceptions=False)
-        async def func_in_tracing_context_async(call: FrozenCall) -> T:
+        async def func_in_tracing_context_async(call: FrozenCall) -> Trace[V]:
             with TracingContext[V](
                 func.__name__, do_not_persist_traces=call.do_not_persist
             ) as t:
                 result = await cast(Callable[..., Awaitable], func)(
                     *call.args, **call.kwargs
                 )
-                return cast(T, t.finalise(output=result))
+                return t.finalise(output=result)
+
+        def current_trace(trace: Trace[V], call: FrozenCall) -> Optional[Trace[V]]:
+            if call.do_not_persist:
+                return trace.model_copy(deep=True)
+
+            # Consult storage rather than process-local invalidation so that feedback
+            # and deletions remain visible across workers and direct database calls.
+            stored = get_context().tracing_database.get(trace.trace_id)
+            if stored is None:
+                return None
+            result = trace.model_copy(deep=True)
+            result.feedback = stored.feedback
+            return result
 
         @wraps(func)
-        def sync(*args: Any, do_not_persist_traces: bool = False, **kwargs: Any) -> T:
-            return func_in_tracing_context_sync(
-                FrozenCall(args, kwargs, do_not_persist_traces)
-            )
+        def sync(
+            *args: Any, do_not_persist_traces: bool = False, **kwargs: Any
+        ) -> Trace[V]:
+            call = FrozenCall(args, kwargs, do_not_persist_traces)
+            trace = func_in_tracing_context_sync(call)
+            result = current_trace(trace, call)
+            if result is None:
+                func_in_tracing_context_sync.cache_clear()
+                result = func_in_tracing_context_sync(call).model_copy(deep=True)
+            return result
 
         @wraps(func)
         async def asynchronous(
             *args: Any, do_not_persist_traces: bool = False, **kwargs: Any
-        ) -> T:
-            return await func_in_tracing_context_async(
-                FrozenCall(args, kwargs, do_not_persist_traces)
-            )
+        ) -> Trace[V]:
+            call = FrozenCall(args, kwargs, do_not_persist_traces)
+            trace = await func_in_tracing_context_async(call)
+            result = current_trace(trace, call)
+            if result is None:
+                func_in_tracing_context_async.invalidate(call)
+                trace = await func_in_tracing_context_async(call)
+                result = trace.model_copy(deep=True)
+            return result
 
         is_async = get_function_metadata_store(func).is_asynchronous
         wrapped = asynchronous if is_async else sync

great_ai/views/trace.py

@@ -1,12 +1,21 @@
 from pprint import pformat
-from typing import Any, Dict, Generic, List, Optional, TypeVar
+from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, cast
 
-from pydantic import ConfigDict
+from pydantic import (
+    ConfigDict,
+    SerializationInfo,
+    SerializerFunctionWrapHandler,
+    TypeAdapter,
+    field_serializer,
+)
+from pydantic_core import PydanticSerializationError
 
+from ..helper.numpy_to_python import numpy_to_python
 from .hashable_base_model import HashableBaseModel
 from .model import Model
 
 T = TypeVar("T")
+_json_value = TypeAdapter(Any)
 
 
 class Trace(HashableBaseModel, Generic[T]):
@@ -39,8 +48,51 @@ class Trace(HashableBaseModel, Generic[T]):
 
     model_config = ConfigDict(extra="ignore")
 
+    @field_serializer(
+        "output", "logged_values", "feedback", mode="wrap", when_used="json"
+    )
+    def serialize_numpy_values(  # type: ignore[no-untyped-def]
+        self,
+        value: Any,
+        handler: SerializerFunctionWrapHandler,
+        info: SerializationInfo,
+    ):
+        # An explicit return annotation would replace the generic output's JSON
+        # schema. Keep its declared type while providing a fallback for NumPy.
+        try:
+            return handler(value)
+        except PydanticSerializationError:
+            return _json_value.dump_python(
+                value,
+                mode="json",
+                include=cast(Any, info.include),
+                exclude=cast(Any, info.exclude),
+                by_alias=info.by_alias,
+                exclude_unset=info.exclude_unset,
+                exclude_defaults=info.exclude_defaults,
+                exclude_none=info.exclude_none,
+                round_trip=info.round_trip,
+                serialize_as_any=info.serialize_as_any,
+                context=info.context,
+                fallback=numpy_to_python,
+            )
+
     @property
     def input(self) -> Any:
+        """Original input value, or a name-to-value mapping for multiple arguments.
+
+        Online traces omit metrics and derived argument metadata. Arguments marked
+        ``disable_logging=True`` are omitted, so their values cannot be recovered.
+        """
+        if "online" in self.tags and "ground_truth" not in self.tags:
+            arguments = {
+                key[len("arg:") : -len(":value")]: value
+                for key, value in self.logged_values.items()
+                if key.startswith("arg:") and key.endswith(":value")
+            }
+            if len(arguments) == 1:
+                return next(iter(arguments.values()))
+            return arguments
         return (
             self.logged_values["input"]
             if list(self.logged_values.keys()) == ["input"]
@@ -75,11 +127,17 @@ class Trace(HashableBaseModel, Generic[T]):
     def tags_flat(self) -> str:
         return ",\n".join(self.tags)
 
-    def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
+    def to_flat_dict(
+        self,
+        include_original: bool = True,
+        *,
+        mode: Literal["python", "json"] = "python",
+    ) -> Dict[str, Any]:
+        serialized = self.model_dump(mode=mode)
         reserved = set(type(self).model_fields) | {"_id"}
         return {
             **(
-                self.model_dump()
+                serialized
                 if include_original
                 else {
                     "trace_id": self.trace_id,
@@ -93,7 +151,7 @@ class Trace(HashableBaseModel, Generic[T]):
                     if (isinstance(v, float) or isinstance(v, int))
                     else pformat(v, indent=2, compact=True)
                 )
-                for k, v in self.logged_values.items()
+                for k, v in serialized["logged_values"].items()
                 if k not in reserved
             },
             "models_flat": self.models_flat,