Reset LangfuseDecorator context to allow nested traces

[vitalii-diachkov-sumup]

Describe the feature or potential improvement

For our internal use case, we call one LLM "chain" inside other and we want this to be two separate Langfuse traces. Calling once function with @observe() decorator inside another causing second to be created as span.

We want to be able to start a completely new trace within existing langfuse_context

Additional information

We have implemented a hack that patches langfuse.decorators.langfuse_decorator global context variables inside context manager and restores the previous values on manager exit. This invokes nested @observe() inside fresh context, causing Langfuse to create trace instead of span object.

Here's the code:

@contextmanager
def reset_langfuse_context():
    """Reset Langfuse ContextVars to their default values temporarly.
    When existing context, this context manager will restore their default values.
    """
    from langfuse.decorators.langfuse_decorator import (
        _observation_params_context,
        _observation_stack_context,
        _root_trace_id_context,
    )

    # Copy contextvar default factory, so if Langfuse will update it on source
    # we don't have to update this function
    default_factory: Callable[[], dict[str, Any]] = (
        _observation_params_context.get().default_factory
    )

    ctx_vars: dict[ContextVar, Token] = {
        _observation_stack_context: _observation_stack_context.set([]),
        _observation_params_context: _observation_params_context.set(defaultdict(default_factory)),
        _root_trace_id_context: _root_trace_id_context.set(None),
    }

    try:
        yield
    finally:
        for var, prev_state in ctx_vars.items():
            var.reset(prev_state)

The issue I have with this code as we deeply rely on internal implementation of this module and you guys will introduce any other ContextVars or change the way you determine whether to start a new trace or not, our hack will break. I added another test to our test suite that ensures that Langfuse has not introduced any new ContextVar, but that rather looks like a workaround that need to be changed in the future.

I am happy to supply a PR with such changes, introducing new method of LangfuseDecorator class to make it possible to override the global context, but I wanted to first start a discussion here to:

1. ask if Langfuse already supports such use case, but with different API (I know about low-level API and that we can create a trace manually, but this nested "chain" can be called separately and we don't want to change a lot in our code since Langfuse has nice abstractions and the way it works with @observe() decorator fits our needs)
2. ask if this is something you are ready to accept to your codebase

The way it looks in code might be:

with langfuse_context.nested_context():  # This calls `@contextmanager` that resets variables
    call_nested_chain()