[OPIK-285] allow users to specify function parameters to not log when…

… using the track decorator (#677)

* Add a test for ignore argument in track

* Refactor track implementation, all repeatable track arguments that were passed everywhere in BaseTrackDecorator and its inherited classes are moved to a separate dataclass - TrackOptions.

* Remove odd line

* Implement dropping ignored input arguments

* Handle capture_input=True case

* Temporarily remove test

* Revert test back

* Update tests

* Fix lint errors

* Rename ignore -> ignore_arguments

sdks/python/src/opik/decorator/arguments_helpers.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any, Dict, List
+from typing import Optional, Any, Dict, List, Callable
 from ..types import SpanType
 
 import dataclasses
@@ -43,3 +43,21 @@ class StartSpanParameters(BaseArguments):
     metadata: Optional[Dict[str, Any]] = None
     input: Optional[Dict[str, Any]] = None
     project_name: Optional[str] = None
+
+
+@dataclasses.dataclass
+class TrackOptions(BaseArguments):
+    """
+    A storage for all arguments passed to the `track` decorator.
+    """
+
+    name: Optional[str]
+    type: SpanType
+    tags: Optional[List[str]]
+    metadata: Optional[Dict[str, Any]]
+    capture_input: bool
+    ignore_arguments: Optional[List[str]]
+    capture_output: bool
+    generations_aggregator: Optional[Callable[[List[Any]], Any]]
+    flush: bool
+    project_name: Optional[str]

sdks/python/src/opik/decorator/base_track_decorator.py

@@ -47,6 +47,7 @@ def track(
         tags: Optional[List[str]] = None,
         metadata: Optional[Dict[str, Any]] = None,
         capture_input: bool = True,
+        ignore_arguments: Optional[List[str]] = None,
         capture_output: bool = True,
         generations_aggregator: Optional[Callable[[List[Any]], Any]] = None,
         flush: bool = False,
@@ -63,6 +64,7 @@ def track(
             tags: Tags to associate with the span.
             metadata: Metadata to associate with the span.
             capture_input: Whether to capture the input arguments.
+            ignore_arguments: The list of the arguments NOT to include into span/trace inputs.
             capture_output: Whether to capture the output result.
             generations_aggregator: Function to aggregate generation results.
             flush: Whether to flush the client after logging.
@@ -80,102 +82,62 @@ def track(
             and also synchronous and asynchronous generators.
             It automatically detects the function type and applies the appropriate tracking logic.
         """
+        track_options = arguments_helpers.TrackOptions(
+            name=None,
+            type=type,
+            tags=tags,
+            metadata=metadata,
+            capture_input=capture_input,
+            ignore_arguments=ignore_arguments,
+            capture_output=capture_output,
+            generations_aggregator=generations_aggregator,
+            flush=flush,
+            project_name=project_name,
+        )
+
         if callable(name):
             # Decorator was used without '()'. It means that decorated function
             # automatically passed as the first argument of 'track' function - name
             func = name
             return self._decorate(
                 func=func,
-                name=None,
-                type=type,
-                tags=tags,
-                metadata=metadata,
-                capture_input=capture_input,
-                capture_output=capture_output,
-                generations_aggregator=generations_aggregator,
-                flush=flush,
-                project_name=project_name,
+                track_options=track_options,
             )
 
+        track_options.name = name
+
         def decorator(func: Callable) -> Callable:
             return self._decorate(
                 func=func,
-                name=name,
-                type=type,
-                tags=tags,
-                metadata=metadata,
-                capture_input=capture_input,
-                capture_output=capture_output,
-                generations_aggregator=generations_aggregator,
-                flush=flush,
-                project_name=project_name,
+                track_options=track_options,
             )
 
         return decorator
 
     def _decorate(
         self,
         func: Callable,
-        name: Optional[str],
-        type: SpanType,
-        tags: Optional[List[str]],
-        metadata: Optional[Dict[str, Any]],
-        capture_input: bool,
-        capture_output: bool,
-        generations_aggregator: Optional[Callable[[List[Any]], Any]],
-        flush: bool,
-        project_name: Optional[str],
+        track_options: arguments_helpers.TrackOptions,
     ) -> Callable:
         if not inspect_helpers.is_async(func):
             return self._tracked_sync(
                 func=func,
-                name=name,
-                type=type,
-                tags=tags,
-                metadata=metadata,
-                capture_input=capture_input,
-                capture_output=capture_output,
-                generations_aggregator=generations_aggregator,
-                flush=flush,
-                project_name=project_name,
+                track_options=track_options,
             )
 
         return self._tracked_async(
             func=func,
-            name=name,
-            type=type,
-            tags=tags,
-            metadata=metadata,
-            capture_input=capture_input,
-            capture_output=capture_output,
-            generations_aggregator=generations_aggregator,
-            flush=flush,
-            project_name=project_name,
+            track_options=track_options,
         )
 
     def _tracked_sync(
-        self,
-        func: Callable,
-        name: Optional[str],
-        type: SpanType,
-        tags: Optional[List[str]],
-        metadata: Optional[Dict[str, Any]],
-        capture_input: bool,
-        capture_output: bool,
-        generations_aggregator: Optional[Callable[[List[Any]], str]],
-        flush: bool,
-        project_name: Optional[str],
+        self, func: Callable, track_options: arguments_helpers.TrackOptions
     ) -> Callable:
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:  # type: ignore
             self._before_call(
                 func=func,
-                name=name,
-                type=type,
-                tags=tags,
-                metadata=metadata,
-                capture_input=capture_input,
-                project_name=project_name,
+                track_options=track_options,
                 args=args,
                 kwargs=kwargs,
             )
@@ -195,16 +157,16 @@ def wrapper(*args, **kwargs) -> Any:  # type: ignore
             finally:
                 generator_or_generator_container = self._generators_handler(
                     result,
-                    capture_output,
-                    generations_aggregator,
+                    track_options.capture_output,
+                    track_options.generations_aggregator,
                 )
                 if generator_or_generator_container is not None:
                     return generator_or_generator_container
 
                 self._after_call(
                     output=result,
-                    capture_output=capture_output,
-                    flush=flush,
+                    capture_output=track_options.capture_output,
+                    flush=track_options.flush,
                 )
                 if result is not None:
                     return result
@@ -216,26 +178,13 @@ def wrapper(*args, **kwargs) -> Any:  # type: ignore
     def _tracked_async(
         self,
         func: Callable,
-        name: Optional[str],
-        type: SpanType,
-        tags: Optional[List[str]],
-        metadata: Optional[Dict[str, Any]],
-        capture_input: bool,
-        capture_output: bool,
-        generations_aggregator: Optional[Callable[[List[Any]], str]],
-        flush: bool,
-        project_name: Optional[str],
+        track_options: arguments_helpers.TrackOptions,
     ) -> Callable:
         @functools.wraps(func)
         async def wrapper(*args, **kwargs) -> Any:  # type: ignore
             self._before_call(
                 func=func,
-                name=name,
-                type=type,
-                tags=tags,
-                metadata=metadata,
-                capture_input=capture_input,
-                project_name=project_name,
+                track_options=track_options,
                 args=args,
                 kwargs=kwargs,
             )
@@ -254,16 +203,16 @@ async def wrapper(*args, **kwargs) -> Any:  # type: ignore
             finally:
                 generator = self._generators_handler(
                     result,
-                    capture_output,
-                    generations_aggregator,
+                    track_options.capture_output,
+                    track_options.generations_aggregator,
                 )
                 if generator is not None:  # TODO: test this flow for async generators
                     return generator
 
                 self._after_call(
                     output=result,
-                    capture_output=capture_output,
-                    flush=flush,
+                    capture_output=track_options.capture_output,
+                    flush=track_options.flush,
                 )
                 if result is not None:
                     return result
@@ -274,12 +223,7 @@ async def wrapper(*args, **kwargs) -> Any:  # type: ignore
     def _before_call(
         self,
         func: Callable,
-        name: Optional[str],
-        type: SpanType,
-        tags: Optional[List[str]],
-        metadata: Optional[Dict[str, Any]],
-        capture_input: bool,
-        project_name: Optional[str],
+        track_options: arguments_helpers.TrackOptions,
         args: Tuple,
         kwargs: Dict[str, Any],
     ) -> None:
@@ -290,12 +234,7 @@ def _before_call(
 
             start_span_arguments = self._start_span_inputs_preprocessor(
                 func=func,
-                name=name,
-                type=type,
-                tags=tags,
-                metadata=metadata,
-                capture_input=capture_input,
-                project_name=project_name,
+                track_options=track_options,
                 args=args,
                 kwargs=kwargs,
             )
@@ -535,14 +474,9 @@ def _generators_handler(
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
-        name: Optional[str],
-        type: SpanType,
-        tags: Optional[List[str]],
-        metadata: Optional[Dict[str, Any]],
-        capture_input: bool,
+        track_options: arguments_helpers.TrackOptions,
         args: Tuple,
         kwargs: Dict[str, Any],
-        project_name: Optional[str],
     ) -> arguments_helpers.StartSpanParameters:
         """
         Subclasses must override this method to customize generating

sdks/python/src/opik/decorator/tracker.py

@@ -2,7 +2,6 @@
 
 from typing import List, Any, Dict, Optional, Callable, Tuple, Union
 
-from ..types import SpanType
 
 from . import inspect_helpers, arguments_helpers
 from ..api_objects import opik_client
@@ -19,30 +18,29 @@ class OpikTrackDecorator(base_track_decorator.BaseTrackDecorator):
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
-        name: Optional[str],
-        type: SpanType,
-        tags: Optional[List[str]],
-        metadata: Optional[Dict[str, Any]],
-        capture_input: bool,
+        track_options: arguments_helpers.TrackOptions,
         args: Tuple,
         kwargs: Dict[str, Any],
-        project_name: Optional[str],
     ) -> arguments_helpers.StartSpanParameters:
         input = (
             inspect_helpers.extract_inputs(func, args, kwargs)
-            if capture_input
+            if track_options.capture_input
             else None
         )
 
-        name = name if name is not None else func.__name__
+        if input is not None and track_options.ignore_arguments is not None:
+            for argument in track_options.ignore_arguments:
+                input.pop(argument, None)
+
+        name = track_options.name if track_options.name is not None else func.__name__
 
         result = arguments_helpers.StartSpanParameters(
             name=name,
             input=input,
-            type=type,
-            tags=tags,
-            metadata=metadata,
-            project_name=project_name,
+            type=track_options.type,
+            tags=track_options.tags,
+            metadata=track_options.metadata,
+            project_name=track_options.project_name,
         )
 
         return result

sdks/python/src/opik/integrations/anthropic/messages_create_decorator.py

@@ -1,6 +1,5 @@
 import logging
 from typing import List, Any, Dict, Optional, Callable, Tuple, Union
-from opik.types import SpanType
 from opik.decorator import base_track_decorator, arguments_helpers
 from opik import dict_utils
 
@@ -24,35 +23,30 @@ class AnthropicMessagesCreateDecorator(base_track_decorator.BaseTrackDecorator):
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
-        name: Optional[str],
-        type: SpanType,
-        tags: Optional[List[str]],
-        metadata: Optional[Dict[str, Any]],
-        capture_input: bool,
+        track_options: arguments_helpers.TrackOptions,
         args: Optional[Tuple],
         kwargs: Optional[Dict[str, Any]],
-        project_name: Optional[str],
     ) -> arguments_helpers.StartSpanParameters:
         assert (
             kwargs is not None
         ), "Expected kwargs to be not None in Antropic.messages.create(**kwargs)"
-        metadata = metadata if metadata is not None else {}
+        metadata = track_options.metadata if track_options.metadata is not None else {}
+        name = track_options.name if track_options.name is not None else func.__name__
 
         input, metadata_from_kwargs = dict_utils.split_dict_by_keys(
             kwargs, KWARGS_KEYS_TO_LOG_AS_INPUTS
         )
         metadata.update(metadata_from_kwargs)
         metadata["created_from"] = "anthropic"
         tags = ["anthropic"]
-        name = name if name is not None else func.__name__
 
         result = arguments_helpers.StartSpanParameters(
             name=name,
             input=input,
-            type=type,
+            type=track_options.type,
             tags=tags,
             metadata=metadata,
-            project_name=project_name,
+            project_name=track_options.project_name,
         )
 
         return result