diff --git a/src/promptflow-tracing/promptflow/tracing/_integrations/_openai_injector.py b/src/promptflow-tracing/promptflow/tracing/_integrations/_openai_injector.py index 5faa7630830..640d36d2749 100644 --- a/src/promptflow-tracing/promptflow/tracing/_integrations/_openai_injector.py +++ b/src/promptflow-tracing/promptflow/tracing/_integrations/_openai_injector.py @@ -21,16 +21,16 @@ IS_LEGACY_OPENAI = version("openai").startswith("0.") -def inject_function_async(args_to_ignore=None, trace_type=TraceType.LLM, name=None): +def inject_function_async(args_to_ignore=None, trace_type=TraceType.LLM): def decorator(func): - return _traced_async(func, args_to_ignore=args_to_ignore, trace_type=trace_type, name=name) + return _traced_async(func, args_to_ignore=args_to_ignore, trace_type=trace_type) return decorator -def inject_function_sync(args_to_ignore=None, trace_type=TraceType.LLM, name=None): +def inject_function_sync(args_to_ignore=None, trace_type=TraceType.LLM): def decorator(func): - return _traced_sync(func, args_to_ignore=args_to_ignore, trace_type=trace_type, name=name) + return _traced_sync(func, args_to_ignore=args_to_ignore, trace_type=trace_type) return decorator @@ -90,67 +90,60 @@ def wrapper(*args, **kwargs): return wrapper -def inject_async(f, trace_type, name): +def inject_async(f, trace_type): wrapper_fun = inject_operation_headers( - (inject_function_async(["api_key", "headers", "extra_headers"], trace_type, name)(f)) + (inject_function_async(["api_key", "headers", "extra_headers"], trace_type)(f)) ) wrapper_fun._original = f return wrapper_fun -def inject_sync(f, trace_type, name): +def inject_sync(f, trace_type): wrapper_fun = inject_operation_headers( - (inject_function_sync(["api_key", "headers", "extra_headers"], trace_type, name)(f)) + (inject_function_sync(["api_key", "headers", "extra_headers"], trace_type)(f)) ) wrapper_fun._original = f return wrapper_fun -def _legacy_openai_apis(): - sync_apis = ( - ("openai", "Completion", "create", TraceType.LLM, "openai_completion_legacy"), - ("openai", "ChatCompletion", "create", TraceType.LLM, "openai_chat_legacy"), - ("openai", "Embedding", "create", TraceType.EMBEDDING, "openai_embedding_legacy"), - ) - async_apis = ( - ("openai", "Completion", "acreate", TraceType.LLM, "openai_completion_legacy"), - ("openai", "ChatCompletion", "acreate", TraceType.LLM, "openai_chat_legacy"), - ("openai", "Embedding", "acreate", TraceType.EMBEDDING, "openai_embedding_legacy"), - ) - return sync_apis, async_apis - - -def _openai_apis(): - sync_apis = ( - ("openai.resources.chat", "Completions", "create", TraceType.LLM, "openai_chat"), - ("openai.resources", "Completions", "create", TraceType.LLM, "openai_completion"), - ("openai.resources", "Embeddings", "create", TraceType.EMBEDDING, "openai_embeddings"), - ) - async_apis = ( - ("openai.resources.chat", "AsyncCompletions", "create", TraceType.LLM, "openai_chat_async"), - ("openai.resources", "AsyncCompletions", "create", TraceType.LLM, "openai_completion_async"), - ("openai.resources", "AsyncEmbeddings", "create", TraceType.EMBEDDING, "openai_embeddings_async"), - ) - return sync_apis, async_apis - - def _openai_api_list(): if IS_LEGACY_OPENAI: - sync_apis, async_apis = _legacy_openai_apis() + sync_apis = ( + ("openai", "Completion", "create", TraceType.LLM), + ("openai", "ChatCompletion", "create", TraceType.LLM), + ("openai", "Embedding", "create", TraceType.EMBEDDING), + ) + + async_apis = ( + ("openai", "Completion", "acreate", TraceType.LLM), + ("openai", "ChatCompletion", "acreate", TraceType.LLM), + ("openai", "Embedding", "acreate", TraceType.EMBEDDING), + ) else: - sync_apis, async_apis = _openai_apis() + sync_apis = ( + ("openai.resources.chat", "Completions", "create", TraceType.LLM), + ("openai.resources", "Completions", "create", TraceType.LLM), + ("openai.resources", "Embeddings", "create", TraceType.EMBEDDING), + ) + + async_apis = ( + ("openai.resources.chat", "AsyncCompletions", "create", TraceType.LLM), + ("openai.resources", "AsyncCompletions", "create", TraceType.LLM), + ("openai.resources", "AsyncEmbeddings", "create", TraceType.EMBEDDING), + ) + yield sync_apis, inject_sync yield async_apis, inject_async def _generate_api_and_injector(apis): for apis, injector in apis: - for module_name, class_name, method_name, trace_type, name in apis: + for module_name, class_name, method_name, trace_type in apis: try: module = importlib.import_module(module_name) api = getattr(module, class_name) if hasattr(api, method_name): - yield api, method_name, trace_type, injector, name + yield api, method_name, trace_type, injector except AttributeError as e: # Log the attribute exception with the missing class information logging.warning( @@ -183,10 +176,10 @@ def inject_openai_api(): 2. Updates the openai api configs from environment variables. """ - for api, method, trace_type, injector, name in available_openai_apis_and_injectors(): + for api, method, trace_type, injector in available_openai_apis_and_injectors(): # Check if the create method of the openai_api class has already been modified if not hasattr(getattr(api, method), "_original"): - setattr(api, method, injector(getattr(api, method), trace_type, name)) + setattr(api, method, injector(getattr(api, method), trace_type)) if IS_LEGACY_OPENAI: # For the openai versions lower than 1.0.0, it reads api configs from environment variables only at @@ -205,6 +198,6 @@ def recover_openai_api(): """This function restores the original create methods of the OpenAI API classes by assigning them back from the _original attributes of the modified methods. """ - for api, method, _, _, _ in available_openai_apis_and_injectors(): + for api, method, _, _ in available_openai_apis_and_injectors(): if hasattr(getattr(api, method), "_original"): setattr(api, method, getattr(getattr(api, method), "_original")) diff --git a/src/promptflow-tracing/promptflow/tracing/_trace.py b/src/promptflow-tracing/promptflow/tracing/_trace.py index 8f3bb92fd31..14b1d164e7e 100644 --- a/src/promptflow-tracing/promptflow/tracing/_trace.py +++ b/src/promptflow-tracing/promptflow/tracing/_trace.py @@ -22,7 +22,7 @@ from ._tracer import Tracer, _create_trace_from_function_call, get_node_name_from_context from ._utils import get_input_names_for_prompt_template, get_prompt_param_name_from_func, serialize from .contracts.generator_proxy import GeneratorProxy -from .contracts.trace import TraceType, Trace +from .contracts.trace import TraceType IS_LEGACY_OPENAI = version("openai").startswith("0.") @@ -85,13 +85,13 @@ def enrich_span_with_context(span): logging.warning(f"Failed to enrich span with context: {e}") -def enrich_span_with_trace(span, trace: Trace): +def enrich_span_with_trace(span, trace): try: span.set_attributes( { "framework": "promptflow", "span_type": trace.type.value, - "function": trace.function, + "function": trace.name, } ) node_name = get_node_name_from_context() @@ -293,11 +293,7 @@ def _traced( def _traced_async( - func: Callable = None, - *, - args_to_ignore: Optional[List[str]] = None, - trace_type=TraceType.FUNCTION, - name: Optional[str] = None, + func: Callable = None, *, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION ) -> Callable: """ Decorator that adds tracing to an asynchronous function. @@ -307,7 +303,6 @@ def _traced_async( args_to_ignore (Optional[List[str]], optional): A list of argument names to be ignored in the trace. Defaults to None. trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION. - name (str, optional): The name of the trace, will set to func name if not provided. Returns: Callable: The traced function. @@ -315,12 +310,7 @@ def _traced_async( def create_trace(func, args, kwargs): return _create_trace_from_function_call( - func, - args=args, - kwargs=kwargs, - args_to_ignore=args_to_ignore, - trace_type=trace_type, - name=name, + func, args=args, kwargs=kwargs, args_to_ignore=args_to_ignore, trace_type=trace_type ) @functools.wraps(func) @@ -353,13 +343,7 @@ async def wrapped(*args, **kwargs): return wrapped -def _traced_sync( - func: Callable = None, - *, - args_to_ignore: Optional[List[str]] = None, - trace_type=TraceType.FUNCTION, - name: Optional[str] = None, -) -> Callable: +def _traced_sync(func: Callable = None, *, args_to_ignore=None, trace_type=TraceType.FUNCTION) -> Callable: """ Decorator that adds tracing to a synchronous function. @@ -368,8 +352,6 @@ def _traced_sync( args_to_ignore (Optional[List[str]], optional): A list of argument names to be ignored in the trace. Defaults to None. trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION. - name (str, optional): The name of the trace, will set to func name if not provided. - Returns: Callable: The traced function. @@ -377,12 +359,7 @@ def _traced_sync( def create_trace(func, args, kwargs): return _create_trace_from_function_call( - func, - args=args, - kwargs=kwargs, - args_to_ignore=args_to_ignore, - trace_type=trace_type, - name=name, + func, args=args, kwargs=kwargs, args_to_ignore=args_to_ignore, trace_type=trace_type ) @functools.wraps(func) diff --git a/src/promptflow-tracing/promptflow/tracing/_tracer.py b/src/promptflow-tracing/promptflow/tracing/_tracer.py index b31d9ee731a..d252fa090d4 100644 --- a/src/promptflow-tracing/promptflow/tracing/_tracer.py +++ b/src/promptflow-tracing/promptflow/tracing/_tracer.py @@ -137,7 +137,7 @@ def _format_error(error: Exception) -> dict: def _create_trace_from_function_call( - f, *, args=None, kwargs=None, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION, name=None, + f, *, args=None, kwargs=None, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION ): """ Creates a trace object from a function call. @@ -149,7 +149,6 @@ def _create_trace_from_function_call( args_to_ignore (Optional[List[str]], optional): A list of argument names to be ignored in the trace. Defaults to None. trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION. - name (str, optional): The name of the trace. Defaults to None. Returns: Trace: The created trace object. @@ -177,17 +176,16 @@ def _create_trace_from_function_call( for key in args_to_ignore: all_kwargs.pop(key, None) - function = f.__qualname__ + name = f.__qualname__ if trace_type in [TraceType.LLM, TraceType.EMBEDDING] and f.__module__: - function = f"{f.__module__}.{function}" + name = f"{f.__module__}.{name}" return Trace( - name=name or function, # Use the function name as the trace name if not provided + name=name, type=trace_type, start_time=datetime.utcnow().timestamp(), inputs=all_kwargs, children=[], - function=function, ) diff --git a/src/promptflow-tracing/promptflow/tracing/contracts/trace.py b/src/promptflow-tracing/promptflow/tracing/contracts/trace.py index 3003967314a..fd697e1da43 100644 --- a/src/promptflow-tracing/promptflow/tracing/contracts/trace.py +++ b/src/promptflow-tracing/promptflow/tracing/contracts/trace.py @@ -52,4 +52,3 @@ class Trace: node_name: Optional[str] = None # The node name of the trace, used for flow level trace parent_id: str = "" # The parent trace id of the trace id: str = "" # The trace id - function: str = "" # The function name of the trace diff --git a/src/promptflow-tracing/tests/e2etests/simple_functions.py b/src/promptflow-tracing/tests/e2etests/simple_functions.py index 85301b8f89d..dd0b8d8458d 100644 --- a/src/promptflow-tracing/tests/e2etests/simple_functions.py +++ b/src/promptflow-tracing/tests/e2etests/simple_functions.py @@ -38,7 +38,7 @@ def greetings(user_id): @trace async def dummy_llm(prompt: str, model: str): - await asyncio.sleep(0.5) + asyncio.sleep(0.5) return "dummy_output" diff --git a/src/promptflow-tracing/tests/e2etests/test_tracing.py b/src/promptflow-tracing/tests/e2etests/test_tracing.py index 05314adbeb8..4ba3f38895c 100644 --- a/src/promptflow-tracing/tests/e2etests/test_tracing.py +++ b/src/promptflow-tracing/tests/e2etests/test_tracing.py @@ -148,18 +148,17 @@ def assert_otel_traces_with_embedding(self, dev_connections, func, inputs, expec self.validate_openai_tokens(span_list) for span in span_list: if span.attributes.get("function", "") in EMBEDDING_FUNCTION_NAMES: - msg = f"Span attributes is not expected: {span.attributes}" - assert "ada" in span.attributes.get("llm.response.model", ""), msg + assert "ada" in span.attributes.get("llm.response.model", "") embeddings = span.attributes.get("embedding.embeddings", "") - assert "embedding.vector" in embeddings, msg - assert "embedding.text" in embeddings, msg + assert "embedding.vector" in embeddings + assert "embedding.text" in embeddings if isinstance(inputs["input"], list): # If the input is a token array, which is list of int, the attribute should contains # the length of the token array ''. - assert "dimensional token" in embeddings, msg + assert "dimensional token" in embeddings else: # If the input is a string, the attribute should contains the original input string. - assert inputs["input"] in embeddings, msg + assert inputs["input"] in embeddings def test_otel_trace_with_multiple_functions(self): execute_function_in_subprocess(self.assert_otel_traces_with_multiple_functions) diff --git a/src/promptflow-tracing/tests/unittests/test_openai_injector.py b/src/promptflow-tracing/tests/unittests/test_openai_injector.py index 723b6158a91..ffbfd0cc0b3 100644 --- a/src/promptflow-tracing/tests/unittests/test_openai_injector.py +++ b/src/promptflow-tracing/tests/unittests/test_openai_injector.py @@ -18,8 +18,6 @@ inject_operation_headers, inject_sync, recover_openai_api, - _legacy_openai_apis, - _openai_apis, ) from promptflow.tracing._operation_context import OperationContext from promptflow.tracing._tracer import Tracer @@ -278,27 +276,43 @@ def generator(): "is_legacy, expected_apis_with_injectors", [ ( - False, + True, [ ( - _openai_apis()[0], + ( + ("openai", "Completion", "create", TraceType.LLM), + ("openai", "ChatCompletion", "create", TraceType.LLM), + ("openai", "Embedding", "create", TraceType.EMBEDDING), + ), inject_sync, ), ( - _openai_apis()[1], + ( + ("openai", "Completion", "acreate", TraceType.LLM), + ("openai", "ChatCompletion", "acreate", TraceType.LLM), + ("openai", "Embedding", "acreate", TraceType.EMBEDDING), + ), inject_async, ), ], ), ( - True, + False, [ ( - _legacy_openai_apis()[0], + ( + ("openai.resources.chat", "Completions", "create", TraceType.LLM), + ("openai.resources", "Completions", "create", TraceType.LLM), + ("openai.resources", "Embeddings", "create", TraceType.EMBEDDING), + ), inject_sync, ), ( - _legacy_openai_apis()[1], + ( + ("openai.resources.chat", "AsyncCompletions", "create", TraceType.LLM), + ("openai.resources", "AsyncCompletions", "create", TraceType.LLM), + ("openai.resources", "AsyncEmbeddings", "create", TraceType.EMBEDDING), + ), inject_async, ), ], @@ -317,13 +331,13 @@ def test_api_list(is_legacy, expected_apis_with_injectors): "apis_with_injectors, expected_output, expected_logs", [ ( - [((("MockModule", "MockAPI", "create", TraceType.LLM, "Mock.create"),), inject_sync)], - [(MockAPI, "create", TraceType.LLM, inject_sync, "Mock.create")], + [((("MockModule", "MockAPI", "create", TraceType.LLM),), inject_sync)], + [(MockAPI, "create", TraceType.LLM, inject_sync)], [], ), ( - [((("MockModule", "MockAPI", "create", TraceType.LLM, "Mock.acreate"),), inject_async)], - [(MockAPI, "create", TraceType.LLM, inject_async, "Mock.acreate")], + [((("MockModule", "MockAPI", "create", TraceType.LLM),), inject_async)], + [(MockAPI, "create", TraceType.LLM, inject_async)], [], ), ], @@ -335,32 +349,21 @@ def test_generate_api_and_injector(apis_with_injectors, expected_output, expecte # Run the generator and collect the output result = list(_generate_api_and_injector(apis_with_injectors)) + # Check if the result matches the expected output + assert result == expected_output + # Check if the logs match the expected logs assert len(caplog.records) == len(expected_logs) for record, expected_message in zip(caplog.records, expected_logs): assert expected_message in record.message - # Check if the result matches the expected output - assert result == expected_output - mock_import_module.assert_called_with("MockModule") def test_generate_api_and_injector_attribute_error_logging(caplog): apis = [ - ((("NonExistentModule", "NonExistentAPI", "create", TraceType.LLM, "create"),), MagicMock()), - ( - ( - ( - "MockModuleMissingMethod", - "MockAPIMissingMethod", - "missing_method", - "missing_trace_type", - "missing_name", - ), - ), - MagicMock(), - ), + ((("NonExistentModule", "NonExistentAPI", "create", TraceType.LLM),), MagicMock()), + ((("MockModuleMissingMethod", "MockAPIMissingMethod", "missing_method", "missing_trace_type"),), MagicMock()), ] # Set up the side effect for the mock @@ -404,7 +407,7 @@ def dummy_api(): pass # Real injector function that adds an _original attribute - def injector(f, trace_type, name): + def injector(f, trace_type): def wrapper_fun(*args, **kwargs): return f(*args, **kwargs) @@ -422,8 +425,8 @@ def wrapper_fun(*args, **kwargs): with patch( "promptflow.tracing._integrations._openai_injector.available_openai_apis_and_injectors", return_value=[ - (FakeAPIWithoutOriginal, "create", TraceType.LLM, injector, "create"), - (FakeAPIWithOriginal, "create", TraceType.LLM, injector, "create"), + (FakeAPIWithoutOriginal, "create", TraceType.LLM, injector), + (FakeAPIWithOriginal, "create", TraceType.LLM, injector), ], ): # Call the function to inject the APIs diff --git a/src/promptflow-tracing/tests/unittests/test_trace.py b/src/promptflow-tracing/tests/unittests/test_trace.py index d7f1548bdf1..b83f88e6067 100644 --- a/src/promptflow-tracing/tests/unittests/test_trace.py +++ b/src/promptflow-tracing/tests/unittests/test_trace.py @@ -70,8 +70,8 @@ def dict(self): class MockTrace: - def __init__(self, function, type): - self.function = function + def __init__(self, name, type): + self.name = name self.type = type