From 5b0fde95cea3c7d469ec10106579744958b6eed8 Mon Sep 17 00:00:00 2001 From: Aliaksandr Kuzmik <98702584+alexkuzmik@users.noreply.github.com> Date: Wed, 20 Nov 2024 15:59:23 +0100 Subject: [PATCH] [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 --- .../src/opik/decorator/arguments_helpers.py | 20 ++- .../opik/decorator/base_track_decorator.py | 140 +++++------------- sdks/python/src/opik/decorator/tracker.py | 24 ++- .../anthropic/messages_create_decorator.py | 16 +- .../bedrock/converse_decorator.py | 14 +- .../integrations/openai/openai_decorator.py | 16 +- .../unit/decorator/test_tracker_outputs.py | 66 +++++++++ 7 files changed, 147 insertions(+), 149 deletions(-) diff --git a/sdks/python/src/opik/decorator/arguments_helpers.py b/sdks/python/src/opik/decorator/arguments_helpers.py index fb1cde826b..d570b20934 100644 --- a/sdks/python/src/opik/decorator/arguments_helpers.py +++ b/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] diff --git a/sdks/python/src/opik/decorator/base_track_decorator.py b/sdks/python/src/opik/decorator/base_track_decorator.py index 0d034a4aee..814e240c2c 100644 --- a/sdks/python/src/opik/decorator/base_track_decorator.py +++ b/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,35 +82,34 @@ 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 @@ -116,66 +117,27 @@ def decorator(func: Callable) -> Callable: 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 diff --git a/sdks/python/src/opik/decorator/tracker.py b/sdks/python/src/opik/decorator/tracker.py index 1770f89b1c..711375e0db 100644 --- a/sdks/python/src/opik/decorator/tracker.py +++ b/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 diff --git a/sdks/python/src/opik/integrations/anthropic/messages_create_decorator.py b/sdks/python/src/opik/integrations/anthropic/messages_create_decorator.py index 9a9cc25b6e..7eb101a839 100644 --- a/sdks/python/src/opik/integrations/anthropic/messages_create_decorator.py +++ b/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,19 +23,15 @@ 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 @@ -44,15 +39,14 @@ def _start_span_inputs_preprocessor( 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 diff --git a/sdks/python/src/opik/integrations/bedrock/converse_decorator.py b/sdks/python/src/opik/integrations/bedrock/converse_decorator.py index 8001c25783..abfc3d2032 100644 --- a/sdks/python/src/opik/integrations/bedrock/converse_decorator.py +++ b/sdks/python/src/opik/integrations/bedrock/converse_decorator.py @@ -1,7 +1,6 @@ import logging from typing import List, Any, Dict, Optional, Callable, Tuple, Union, TypedDict, cast from opik import dict_utils -from opik.types import SpanType from opik.decorator import base_track_decorator, arguments_helpers from . import stream_wrappers @@ -33,20 +32,15 @@ class BedrockConverseDecorator(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 BedrockRuntime.Client.converse(**kwargs)" - name = name if name is not None else func.__name__ + name = track_options.name if track_options.name is not None else func.__name__ input, metadata = dict_utils.split_dict_by_keys( kwargs, KWARGS_KEYS_TO_LOG_AS_INPUTS ) @@ -56,10 +50,10 @@ def _start_span_inputs_preprocessor( 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 diff --git a/sdks/python/src/opik/integrations/openai/openai_decorator.py b/sdks/python/src/opik/integrations/openai/openai_decorator.py index a065327ea6..39007600d0 100644 --- a/sdks/python/src/opik/integrations/openai/openai_decorator.py +++ b/sdks/python/src/opik/integrations/openai/openai_decorator.py @@ -2,7 +2,6 @@ from typing import List, Any, Dict, Optional, Callable, Tuple, Union from opik import dict_utils -from opik.types import SpanType from opik.decorator import base_track_decorator, arguments_helpers from . import stream_wrappers @@ -30,20 +29,15 @@ class OpenaiTrackDecorator(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 OpenAI().chat.completion.create(**kwargs)" - name = name if name is not None else func.__name__ - metadata = metadata if metadata is not None else {} + name = track_options.name if track_options.name is not None else func.__name__ + metadata = track_options.metadata if track_options.metadata is not None else {} input, new_metadata = dict_utils.split_dict_by_keys( kwargs, keys=KWARGS_KEYS_TO_LOG_AS_INPUTS @@ -61,10 +55,10 @@ def _start_span_inputs_preprocessor( 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 diff --git a/sdks/python/tests/unit/decorator/test_tracker_outputs.py b/sdks/python/tests/unit/decorator/test_tracker_outputs.py index 102b3d03bc..aa542b1a62 100644 --- a/sdks/python/tests/unit/decorator/test_tracker_outputs.py +++ b/sdks/python/tests/unit/decorator/test_tracker_outputs.py @@ -987,3 +987,69 @@ def f(x): assert len(fake_backend.trace_trees) == 1 assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) + + +def test_tracker__ignore_list_was_passed__ignored_inputs_are_not_logged(fake_backend): + @tracker.track(ignore_arguments=["a", "c", "e", "unknown_argument"]) + def f(a, b, c=3, d=4, e=5): + return {"some-key": "the-output-value"} + + f(1, 2) + tracker.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="f", + input={"b": 2, "d": 4}, + output={"some-key": "the-output-value"}, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + name="f", + input={"b": 2, "d": 4}, + output={"some-key": "the-output-value"}, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + spans=[], + ) + ], + ) + + assert len(fake_backend.trace_trees) == 1 + assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0]) + + +def test_tracker__ignore_list_was_passed__function_does_not_have_any_arguments__input_dicts_are_empty( + fake_backend, +): + @tracker.track(ignore_arguments=["a", "c", "e", "unknown_argument"]) + def f(): + return {"some-key": "the-output-value"} + + f() + tracker.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="f", + input={}, + output={"some-key": "the-output-value"}, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + name="f", + input={}, + output={"some-key": "the-output-value"}, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + spans=[], + ) + ], + ) + + assert len(fake_backend.trace_trees) == 1 + assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])