[Draft][Tracing] Modify openai span names (#2424)

# Description

Rename openai spans since openai.resources.chat.Completions.create is
too long.

This pull request primarily introduces changes to the tracing
functionality in the `promptflow` application. The changes allow for
more flexibility and control over the naming of traces, particularly for
asynchronous and synchronous functions. The key changes can be grouped
into two categories: modifications to the function definitions and
adjustments to the function calls.

Modifications to function definitions:

*
[`src/promptflow-tracing/promptflow/tracing/_trace.py`](diffhunk://#diff-580941184737186a94e3e9c06e467e86c06ce846d30ffa42c1c43b45812ddcdcL272-R276):
Both the `_traced_async` and `_traced_sync` functions were modified to
include an optional `name` parameter. This allows for custom naming of
the traces. If no name is provided, the function's name is used as the
trace name.
[[1]](diffhunk://#diff-580941184737186a94e3e9c06e467e86c06ce846d30ffa42c1c43b45812ddcdcL272-R276)
[[2]](diffhunk://#diff-580941184737186a94e3e9c06e467e86c06ce846d30ffa42c1c43b45812ddcdcL322-R334)

Adjustments to function calls:

*
[`src/promptflow-tracing/promptflow/tracing/_integrations/_openai_injector.py`](diffhunk://#diff-01b30cb37ac2a91b223014785c6bc7efe007a66efcb97a0dbb36442a94c2054bL24-R33):
Several function calls were adjusted to include the new `name`
parameter. This includes the `inject_function_async`,
`inject_function_sync`, `inject_async`, and `inject_sync` functions. The
`name` parameter was also added to the tuples in the `_openai_api_list`
function, which are used to generate the API and injector.
[[1]](diffhunk://#diff-01b30cb37ac2a91b223014785c6bc7efe007a66efcb97a0dbb36442a94c2054bL24-R33)
[[2]](diffhunk://#diff-01b30cb37ac2a91b223014785c6bc7efe007a66efcb97a0dbb36442a94c2054bL93-R103)
[[3]](diffhunk://#diff-01b30cb37ac2a91b223014785c6bc7efe007a66efcb97a0dbb36442a94c2054bL112-R132)
[[4]](diffhunk://#diff-01b30cb37ac2a91b223014785c6bc7efe007a66efcb97a0dbb36442a94c2054bL141-R146)
[[5]](diffhunk://#diff-01b30cb37ac2a91b223014785c6bc7efe007a66efcb97a0dbb36442a94c2054bL179-R182)

# All Promptflow Contribution checklist:
- [ ] **The pull request does not introduce [breaking changes].**
- [ ] **CHANGELOG is updated for new features, bug fixes or other
significant changes.**
- [ ] **I have read the [contribution guidelines](../CONTRIBUTING.md).**
- [ ] **Create an issue and link to the pull request to get dedicated
review from promptflow team. Learn more: [suggested
workflow](../CONTRIBUTING.md#suggested-workflow).**

## General Guidelines and Best Practices
- [ ] Title of the pull request is clear and informative.
- [ ] There are a small number of commits, each of which have an
informative message. This means that previously merged commits do not
appear in the history of the PR. For more information on cleaning up the
commits in your PR, [see this
page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).

### Testing Guidelines
- [ ] Pull request includes test coverage for the included changes.

Co-authored-by: Heyi <heta@microsoft.com>

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):
+def inject_function_async(args_to_ignore=None, trace_type=TraceType.LLM, name=None):
     def decorator(func):
-        return _traced_async(func, args_to_ignore=args_to_ignore, trace_type=trace_type)
+        return _traced_async(func, args_to_ignore=args_to_ignore, trace_type=trace_type, name=name)
 
     return decorator
 
 
-def inject_function_sync(args_to_ignore=None, trace_type=TraceType.LLM):
+def inject_function_sync(args_to_ignore=None, trace_type=TraceType.LLM, name=None):
     def decorator(func):
-        return _traced_sync(func, args_to_ignore=args_to_ignore, trace_type=trace_type)
+        return _traced_sync(func, args_to_ignore=args_to_ignore, trace_type=trace_type, name=name)
 
     return decorator
 
@@ -90,60 +90,67 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def inject_async(f, trace_type):
+def inject_async(f, trace_type, name):
     wrapper_fun = inject_operation_headers(
-        (inject_function_async(["api_key", "headers", "extra_headers"], trace_type)(f))
+        (inject_function_async(["api_key", "headers", "extra_headers"], trace_type, name)(f))
     )
     wrapper_fun._original = f
     return wrapper_fun
 
 
-def inject_sync(f, trace_type):
+def inject_sync(f, trace_type, name):
     wrapper_fun = inject_operation_headers(
-        (inject_function_sync(["api_key", "headers", "extra_headers"], trace_type)(f))
+        (inject_function_sync(["api_key", "headers", "extra_headers"], trace_type, name)(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 = (
-            ("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),
-        )
+        sync_apis, async_apis = _legacy_openai_apis()
     else:
-        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),
-        )
-
+        sync_apis, async_apis = _openai_apis()
     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 in apis:
+        for module_name, class_name, method_name, trace_type, name 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
+                    yield api, method_name, trace_type, injector, name
             except AttributeError as e:
                 # Log the attribute exception with the missing class information
                 logging.warning(
@@ -176,10 +183,10 @@ def inject_openai_api():
     2. Updates the openai api configs from environment variables.
     """
 
-    for api, method, trace_type, injector in available_openai_apis_and_injectors():
+    for api, method, trace_type, injector, name 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))
+            setattr(api, method, injector(getattr(api, method), trace_type, name))
 
     if IS_LEGACY_OPENAI:
         # For the openai versions lower than 1.0.0, it reads api configs from environment variables only at
@@ -198,6 +205,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"))

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
+from .contracts.trace import TraceType, Trace
 
 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):
+def enrich_span_with_trace(span, trace: Trace):
     try:
         span.set_attributes(
             {
                 "framework": "promptflow",
                 "span_type": trace.type.value,
-                "function": trace.name,
+                "function": trace.function,
             }
         )
         node_name = get_node_name_from_context()
@@ -293,7 +293,11 @@ def _traced(
 
 
 def _traced_async(
-    func: Callable = None, *, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION
+    func: Callable = None,
+    *,
+    args_to_ignore: Optional[List[str]] = None,
+    trace_type=TraceType.FUNCTION,
+    name: Optional[str] = None,
 ) -> Callable:
     """
     Decorator that adds tracing to an asynchronous function.
@@ -303,14 +307,20 @@ 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.
     """
 
     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
+            func,
+            args=args,
+            kwargs=kwargs,
+            args_to_ignore=args_to_ignore,
+            trace_type=trace_type,
+            name=name,
         )
 
     @functools.wraps(func)
@@ -343,7 +353,13 @@ async def wrapped(*args, **kwargs):
     return wrapped
 
 
-def _traced_sync(func: Callable = None, *, args_to_ignore=None, trace_type=TraceType.FUNCTION) -> Callable:
+def _traced_sync(
+    func: Callable = None,
+    *,
+    args_to_ignore: Optional[List[str]] = None,
+    trace_type=TraceType.FUNCTION,
+    name: Optional[str] = None,
+) -> Callable:
     """
     Decorator that adds tracing to a synchronous function.
 
@@ -352,14 +368,21 @@ def _traced_sync(func: Callable = None, *, args_to_ignore=None, trace_type=Trace
         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.
     """
 
     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
+            func,
+            args=args,
+            kwargs=kwargs,
+            args_to_ignore=args_to_ignore,
+            trace_type=trace_type,
+            name=name,
         )
 
     @functools.wraps(func)

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
+    f, *, args=None, kwargs=None, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION, name=None,
 ):
     """
     Creates a trace object from a function call.
@@ -149,6 +149,7 @@ 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.
@@ -176,16 +177,17 @@ def _create_trace_from_function_call(
     for key in args_to_ignore:
         all_kwargs.pop(key, None)
 
-    name = f.__qualname__
+    function = f.__qualname__
     if trace_type in [TraceType.LLM, TraceType.EMBEDDING] and f.__module__:
-        name = f"{f.__module__}.{name}"
+        function = f"{f.__module__}.{function}"
 
     return Trace(
-        name=name,
+        name=name or function,  # Use the function name as the trace name if not provided
         type=trace_type,
         start_time=datetime.utcnow().timestamp(),
         inputs=all_kwargs,
         children=[],
+        function=function,
     )
 
 

src/promptflow-tracing/promptflow/tracing/contracts/trace.py

@@ -52,3 +52,4 @@ 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

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):
-    asyncio.sleep(0.5)
+    await asyncio.sleep(0.5)
     return "dummy_output"
 
 

src/promptflow-tracing/tests/e2etests/test_tracing.py

@@ -148,17 +148,18 @@ 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:
-                assert "ada" in span.attributes.get("llm.response.model", "")
+                msg = f"Span attributes is not expected: {span.attributes}"
+                assert "ada" in span.attributes.get("llm.response.model", ""), msg
                 embeddings = span.attributes.get("embedding.embeddings", "")
-                assert "embedding.vector" in embeddings
-                assert "embedding.text" in embeddings
+                assert "embedding.vector" in embeddings, msg
+                assert "embedding.text" in embeddings, msg
                 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 '<len(token_array) dimensional token>'.
-                    assert "dimensional token" in embeddings
+                    assert "dimensional token" in embeddings, msg
                 else:
                     # If the input is a string, the attribute should contains the original input string.
-                    assert inputs["input"] in embeddings
+                    assert inputs["input"] in embeddings, msg
 
     def test_otel_trace_with_multiple_functions(self):
         execute_function_in_subprocess(self.assert_otel_traces_with_multiple_functions)