Repository: microsoft/promptflow
File: src/promptflow-tracing/promptflow/tracing/_integrations/_openai_injector.py (Python)

Code before:
    # get promptflow info from operation context
    operation_context = OperationContext.get_instance()
    tracking_info = operation_context._get_tracking_info()

    def is_primitive(value):
        return value is None or isinstance(value, (int, float, str, bool))

    #  Ensure that the telemetry info is primitive
    tracking_info = {k: v for k, v in tracking_info.items() if is_primitive(v)}

    # init headers
    headers = {USER_AGENT_HEADER: operation_context.get_user_agent()}

    # update header with promptflow info
    headers[PROMPTFLOW_HEADER] = json.dumps(tracking_info)

    return headers


def inject_operation_headers(f):
    def inject_headers(kwargs):
        # Inject headers from operation context, overwrite injected header with headers from kwargs.
        injected_headers = get_aoai_telemetry_headers()
        original_headers = kwargs.get("headers" if IS_LEGACY_OPENAI else "extra_headers")
        if original_headers and isinstance(original_headers, dict):
            injected_headers.update(
                {h: " ".join([injected_headers[h], original_headers[h]])
                    if h in injected_headers else original_headers[h] for h in original_headers.keys()})
        kwargs["headers" if IS_LEGACY_OPENAI else "extra_headers"] = injected_headers

    if asyncio.iscoroutinefunction(f):

        @functools.wraps(f)
        async def wrapper(*args, **kwargs):
            inject_headers(kwargs)
            return await f(*args, **kwargs)

    else:

        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            inject_headers(kwargs)
            return f(*args, **kwargs)

    return wrapper


def inject_async(f, trace_type, name):
    wrapper_fun = inject_operation_headers(
        (inject_function_async(["api_key", "headers", "extra_headers"], trace_type, name)(f))
    )

Review comment (suggestion):
The current code is hard to read. 

Here is the version that refined by GPT.

```python
if original_headers and isinstance(original_headers, dict):  
    for key in original_headers.keys():  
        if key in injected_headers:  
            # If the key already exists in injected_headers, concatenate the values with a space  
            injected_headers[key] = " ".join([injected_headers[key], original_headers[key]])  
        else:  
            # If the key does not exist in injected_headers, add it directly  
            injected_headers[key] = original_headers[key]  

```

Code after:
    # get promptflow info from operation context
    operation_context = OperationContext.get_instance()
    tracking_info = operation_context._get_tracking_info()

    def is_primitive(value):
        return value is None or isinstance(value, (int, float, str, bool))

    #  Ensure that the telemetry info is primitive
    tracking_info = {k: v for k, v in tracking_info.items() if is_primitive(v)}

    # init headers
    headers = {USER_AGENT_HEADER: operation_context.get_user_agent()}

    # update header with promptflow info
    headers[PROMPTFLOW_HEADER] = json.dumps(tracking_info)

    return headers


def inject_operation_headers(f):
    def inject_headers(kwargs):
        # Inject headers from operation context, overwrite injected header with headers from kwargs.
        injected_headers = get_aoai_telemetry_headers()
        original_headers = kwargs.get("headers" if IS_LEGACY_OPENAI else "extra_headers")
        if original_headers and isinstance(original_headers, dict):
            for header in original_headers.keys():
                if header in injected_headers:
                    # If the key already exists in injected_headers, concatenate the values with a space
                    injected_headers[header] = " ".join([injected_headers[header], original_headers[header]])
                else:
                    # If the key does not exist in injected_headers, add it directly
                    injected_headers[header] = original_headers[header]

        kwargs["headers" if IS_LEGACY_OPENAI else "extra_headers"] = injected_headers

    if asyncio.iscoroutinefunction(f):

        @functools.wraps(f)
        async def wrapper(*args, **kwargs):
            inject_headers(kwargs)
            return await f(*args, **kwargs)

    else:

        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            inject_headers(kwargs)
            return f(*args, **kwargs)

    return wrapper