fix: Fix Duplicate Trace Calls for Async Functions and Optimize Tracer Implementation

[viniciusdsmello]

Fix: Eliminate Duplicate Trace Calls in Async Functions and Optimize Tracer Implementation

Problem Statement

The Openlayer Python SDK was experiencing duplicate trace calls when using the @trace() decorator on async functions. Instead of generating a single unified trace, async functions created multiple duplicate traces with different data, causing test failures and incorrect tracing behavior.

Root Cause Analysis

The issue was traced to the @trace_async() decorator's incompatibility with async generator functions:

• The decorator used context managers (with create_step()) that interfered with yield statements
• This caused wrapper functions to return coroutine objects instead of async_generator objects
• The context manager became "stuck" waiting for generators to complete before finalizing steps
• Regular async functions worked correctly, but async generators failed

Solution Implementation

1. Async Generator Detection and Handling

• Detection: Added inspect.isasyncgenfunction(func) at decoration time to identify async generators
• Manual Step Management: For async generators, replaced context managers with manual step creation and finalization
• Preserved Context Managers: Regular async functions continue using context managers for optimal performance
• Top-Level Yields: Ensured yield statements remain at the top level of wrapper functions

2. Code Optimization (80% Duplication Reduction)

Identified and eliminated massive code duplication across three wrapper functions:

Before: ~150 lines with 80% duplication
After: ~40 lines with extracted helper functions

New Helper Functions:

• _extract_function_inputs() - Centralized input extraction and cleaning
• _finalize_step_logging() - Unified step timing and logging
• _handle_trace_completion() - Consistent trace completion logic
• _create_step_for_async_generator() - Specialized context manager for async generators

3. Client Authentication Fix

• Issue: _client was initialized at import time before environment variables were set
• Solution: Implemented lazy initialization using _get_client() function
• Result: Client now created when first needed, ensuring API credentials are available

4. Test Refactoring

Created comprehensive nested test scenario with single orchestrator method:

• New Structure: Single intelligent_assistant_main() method with 7 nested operations
• Realistic Workflow: Simulates complex AI assistant with preprocessing, intent extraction, context retrieval, synthesis, streaming, validation, post-processing, and logging
• Single Trace: Generates one comprehensive trace instead of multiple separate tests

[gustavocidornelas]

This function is duplicating code.

For example, the ConfigLlmData is created here and inside the create_step function. Streaming to Openlayer is also happening in both places.

I imagine that the create_step function should use _handle_trace_completion

[gustavocidornelas]

exception is an unused argument. Either use it or remove from function signature

[gustavocidornelas]

step_args and step_kwargs are unused arguments. Either use or remove

[gustavocidornelas]

This comment applies to other function calls as well: some of these new functions, have a lot of arguments. When calling them, it would be better to use the keyword arguments (instead of positional arguments) – e.g., _finalize_step_logging(step=step, inputs=inputs, start_time=step.start_time, ...)).

It's easier to make mistakes when using positional arguments and the code is more readable with kwargs. It would also make it easier to see that some kwargs are not being used (in this case, exception is not being used by _finalize_step_logging)

[gustavocidornelas]

This is a code-style comment: it's fine to separate helper functions with these comments. However, make sure that the sections are consistent and accurate. For example, _create_step_for_async_generator is inside the "Helper functions for tracing" block, but it is not a helper function.

More broadly, we generally aim to declare functions in the order they are used. This helps with readability and makes it easier for others to follow the flow without jumping around the file.

For example, if a main function calls func_1, func_2, and func_3 in that order, we would define them like this:

def main(...):
    func_1()
    func_2()
    func_3()
    
def func_1(...):
    pass
    
def func_2(...):
    pass

def func_3(...):
    pass