Add new features

great_ai/src/great_ai/great_ai/models/use_model.py

@@ -1,7 +1,7 @@
 from functools import wraps
 from typing import Any, Callable, Dict, List, Literal, Union
 
-from ..helper import get_function_metadata_store
+from ..helper import assert_function_is_not_finalised, get_function_metadata_store
 from ..tracing import TracingContext
 from ..views import Model
 from .load_model import load_model
@@ -14,7 +14,9 @@ def use_model(
     return_path: bool = False,
     model_kwarg_name: str = "model",
 ) -> Callable[..., Any]:
-    assert isinstance(version, int) or version == "latest"
+    assert (
+        isinstance(version, int) or version == "latest"
+    ), "Only integers or the string literal `latest` is allowed as version"
 
     model, actual_version = load_model(
         key=key,
@@ -23,17 +25,19 @@ def use_model(
     )
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        assert_function_is_not_finalised(func)
+
         store = get_function_metadata_store(func)
         store.model_parameter_names.append(model_kwarg_name)
-        if store.version:
-            store.version += "|"
-        store.version += f"{key}:{actual_version}"
+        if store.model_versions:
+            store.model_versions += "."
+        store.model_versions += f"{key}-v{actual_version}"
 
         @wraps(func)
         def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
-            context = TracingContext.get_current_context()
-            if context:
-                context.log_model(Model(key=key, version=actual_version))
+            tracing_context = TracingContext.get_current_context()
+            if tracing_context:
+                tracing_context.log_model(Model(key=key, version=actual_version))
             return func(*args, **kwargs, **{model_kwarg_name: model})
 
         return wrapper