fix: Handle iterators not only generators

[omerc7]

Fixes langfuse/langfuse#10797

---

Important

Enhance Langfuse client to handle both iterators and generators, updating relevant functions and tests for comprehensive support.

• Behavior:
  • Replace inspect.isgenerator and inspect.isasyncgen with isinstance checks for Iterator and AsyncIterator in async_wrapper and sync_wrapper functions in observe.py.
  • Update _wrap_sync_generator_result and _wrap_async_generator_result to _wrap_sync_iterator_result and _wrap_async_iterator_result.
  • Modify comments and variable names to reflect iterator handling instead of generator.
• Classes:
  • Rename _ContextPreservedSyncGeneratorWrapper to _ContextPreservedSyncIteratorWrapper.
  • Rename _ContextPreservedAsyncGeneratorWrapper to _ContextPreservedAsyncIteratorWrapper.
• Tests:
  • Add async_iterable_factory fixture in test_decorators.py to test both async generators and iterators.
  • Update test_async_generator_context_preservation to use async_iterable_factory for testing context preservation with both async generators and iterators.

^{This description was created by for 7ecce95. You can customize this summary. It will automatically update as commits are pushed.}

Disclaimer: Experimental PR review

Greptile Overview

Greptile Summary

Changed iterator detection from inspect.isgenerator() and inspect.isasyncgen() to isinstance() checks with Iterator and AsyncIterator from collections.abc. This fixes a bug where custom iterator objects (not generators) were not being properly wrapped for context preservation.

Key changes:

• Replaced inspect.isgenerator(result) with isinstance(result, Iterator)
• Replaced inspect.isasyncgen(result) with isinstance(result, AsyncIterator)
• Renamed methods and variables from "generator" to "iterator" for consistency
• Added parametrized test to verify both generators and custom iterators work correctly

The bug occurred because inspect.isgenerator() only returns True for generator objects created by generator functions, but returns False for custom iterator classes that implement __iter__() and __next__(). Using isinstance() with ABC types correctly identifies all iterators.

Confidence Score: 5/5

• This PR is safe to merge - it correctly fixes a real bug with proper test coverage
• The change correctly addresses the root cause by switching from inspect.isgenerator() (which only detects generator objects) to isinstance() with ABC Iterator types (which detects all iterators). The fix is well-tested with a parametrized test covering both generators and custom iterators. The refactoring from "generator" to "iterator" terminology improves code clarity.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
langfuse/_client/observe.py	5/5	Replaced inspect.isgenerator() and inspect.isasyncgen() with isinstance() checks using Iterator and AsyncIterator from collections.abc to properly handle custom iterators, not just generators
tests/test_decorators.py	5/5	Added parametrized test fixture to verify both generators and custom iterators work correctly with the @observe decorator, ensuring context preservation

Sequence Diagram

sequenceDiagram
    participant User
    participant ObserveDecorator
    participant Function
    participant IteratorWrapper
    participant LangfuseSpan

    User->>ObserveDecorator: Call @observe decorated function
    ObserveDecorator->>LangfuseSpan: Create span/generation
    ObserveDecorator->>Function: Execute function
    Function-->>ObserveDecorator: Return Iterator/AsyncIterator
    
    alt isinstance(result, Iterator) or isinstance(result, AsyncIterator)
        ObserveDecorator->>ObserveDecorator: Detect iterator type (not just generators)
        ObserveDecorator->>IteratorWrapper: Wrap iterator with context preservation
        ObserveDecorator-->>User: Return wrapped iterator
        Note over LangfuseSpan: Span NOT ended yet
        
        User->>IteratorWrapper: Iterate over results
        loop Each iteration
            IteratorWrapper->>Function: Get next item (in preserved context)
            Function-->>IteratorWrapper: Yield item
            IteratorWrapper-->>User: Return item
        end
        
        IteratorWrapper->>LangfuseSpan: Update output and end span
    else Regular return value
        ObserveDecorator->>LangfuseSpan: Update output
        ObserveDecorator->>LangfuseSpan: End span
        ObserveDecorator-->>User: Return result
    end

Loading

[CLAassistant]

All committers have signed the CLA.

[greptile-apps]

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[omerc7]

Hey @hassiebp, did you get a chance to take a look at this PR?

Happy to address any feedback, make changes, or provide additional context if needed.