Skip to content

feat(tracing): enhance tracing for synchronous generator functions #500

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

viniciusdsmello
Copy link
Contributor

Pull Request

Summary

Fixed generator tracing to export concatenated chunks instead of generator objects to Openlayer. Previously, when using @trace or @trace_async decorators on functions that return generators (like OpenAI streaming), Openlayer would receive useless <generator object> instead of the actual concatenated content. This PR implements proper generator tracing that preserves streaming behavior while capturing meaningful trace data.

Changes

  • Added synchronous generator support to @trace decorator

    • Implemented TracedSyncGenerator class-based wrapper that delays trace creation until first iteration
    • Added inspect.isgeneratorfunction() detection to handle generator functions separately from regular functions
    • Created _finalize_sync_generator_step() helper function for proper cleanup and trace completion
  • Enhanced async generator tracing in @trace_async decorator

    • Leveraged existing TracedAsyncGenerator infrastructure (already working correctly)
    • Ensured consistent behavior between sync and async generator tracing
  • Improved chunk concatenation and output handling

    • Updated _join_output_chunks() to properly concatenate generator outputs
    • Ensured traces contain meaningful string outputs instead of generator objects
    • Preserved real-time streaming behavior for applications consuming generators
  • Comprehensive testing and examples

    • Created extensive test suite validating streaming behavior preservation
    • Added comprehensive OpenAI tracing example demonstrating integration with existing trace_openai() function
    • Verified error handling and edge cases (empty generators, exceptions, partial consumption)

Context

Users reported that when using @trace decorators on OpenAI streaming functions, Openlayer was receiving generator objects instead of the actual AI responses. This made the traces useless for monitoring and analysis. The issue affected both sync and async generators, creating duplicate traces with incomplete data.

Before: Openlayer received <generator object at 0x...> as output
After: Openlayer receives concatenated string like "Hello world! How can I help you today?"

Testing

  • Unit tests - Comprehensive test suite covering all generator scenarios
  • Manual testing - Real OpenAI API integration tests with streaming
  • Integration testing - Verified compatibility with existing trace_openai() function
  • Performance testing - Confirmed no impact on streaming latency or behavior
  • Edge case testing - Empty generators, exceptions, partial consumption, multiple instances

Key validation results:

  • ✅ Streaming behavior perfectly preserved (applications receive chunks in real-time)
  • ✅ Single trace per function (no duplicates)
  • ✅ Proper input/output mapping in traces
  • ✅ Error handling with partial results captured
  • ✅ Works with all Python iteration patterns (for loops, comprehensions, etc.)
  • ✅ Compatible with existing OpenAI integrations

Real-world testing:

  • Tested with actual OpenAI streaming API calls
  • Verified hierarchical tracing (user functions + OpenAI client calls)
  • Confirmed proper token counting and latency measurement
  • Validated with various streaming patterns and error scenarios

- Introduced support for tracing synchronous generator functions by implementing a class-based wrapper that initializes tracing upon the first iteration.
- Added a new helper function, `_finalize_sync_generator_step`, to handle the finalization of tracing when the generator is consumed, ensuring proper logging and exception handling.
- Updated the existing tracing decorator to accommodate both regular functions and synchronous generators, improving the overall functionality and maintainability of the tracing system.
- Ensured compliance with coding standards, including comprehensive type annotations and Google-style docstrings for clarity and maintainability.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant