Add ground_truth API and refactor

great_ai/src/great_ai/great_ai/parameters/parameter.py

@@ -1,22 +1,21 @@
 from functools import wraps
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, TypeVar, cast
 
 from ..exceptions import ArgumentValidationError
-from ..helper import (
-    assert_function_is_not_finalised,
-    get_arguments,
-    get_function_metadata_store,
-)
+from ..helper import get_arguments, get_function_metadata_store
+from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
 from ..tracing.tracing_context import TracingContext
 
+F = TypeVar("F", bound=Callable[..., Any])
+
 
 def parameter(
     parameter_name: str,
     *,
     validator: Callable[[Any], bool] = lambda _: True,
     disable_logging: bool = False,
-) -> Callable[..., Any]:
-    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+) -> Callable[[F], F]:
+    def decorator(func: F) -> F:
         get_function_metadata_store(func).input_parameter_names.append(parameter_name)
         assert_function_is_not_finalised(func)
 
@@ -39,7 +38,7 @@ def parameter(
                     f"Argument {parameter_name} in {func.__name__} did not pass validation"
                 )
 
-            context = TracingContext.get_current_context()
+            context = TracingContext.get_current_tracing_context()
             if context and not disable_logging:
                 context.log_value(name=f"{actual_name}:value", value=argument)
                 if isinstance(argument, str):
@@ -47,6 +46,6 @@ def parameter(
 
             return func(*args, **kwargs)
 
-        return wrapper
+        return cast(F, wrapper)
 
     return decorator