More experimentation

Signed-off-by: András Schmelczer <andras@schmelczer.dev>

good_ai/src/good_ai/good_ai/models/use_model.py

@@ -1,24 +1,33 @@
-from typing import Literal, Union
+from functools import wraps
+from typing import Any, Callable, Literal, Union
 
-from ..core import function_registry
-from .use_model_plugin import UseModelPlugin
+from ..tracing import Model, TracingContext
+from .load_model import load_model
 
 
 def use_model(
-    key: str, version: Union[int, Literal["latest"]] = None, return_path: bool = False
-):
+    key: str,
+    *,
+    version: Union[int, Literal["latest"]],
+    return_path: bool = False,
+    model_kwarg_name: str = "model"
+) -> Callable[..., Any]:
     assert isinstance(version, int) or version == "latest"
 
-    def inner(f):
-        function_registry.add_plugin(
-            f,
-            UseModelPlugin(
-                f,
-                key=key,
-                version=version if isinstance(version, int) else None,
-                return_path=return_path,
-            ),
-        )
-        return f
+    model, actual_version = load_model(
+        key=key,
+        version=None if version == "latest" else version,
+        return_path=return_path,
+    )
 
-    return inner
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            context = TracingContext.get_current_context()
+            if context:
+                context.log_model(Model(key=key, version=actual_version))
+            return func(*args, **kwargs, **{model_kwarg_name: model})
+
+        return wrapper
+
+    return decorator