Issue: Unexpected behaviour with FastAPI StreamingResponse & AsyncGenerator object in tracing

[fchavat]

Issue you'd like to raise.

Hello,

I am encountering an unexpected result when I try to generate traces for my FastAPI project. I have an endpoint method that returns a StreamingResponse object, I would like to generate a trace where the root span is this method and as a child span I have the result from the async generator. For example, this snippet:

@traceable(name="test_stream_generator")
async def test_stream_generator():
    for i in range(10):
        yield json.dumps({"type": "message_content", "message_fragment": {"content": f"message {i}"}}) + "\n"
        time.sleep(0.2)

@app.post("/test", response_model=None)
@traceable(name="test")
async def test(req: ChatRequest) -> StreamingResponse:
    return StreamingResponse(
            content=test_stream_generator(),
            headers={
                "Content-Type": "application/x-ndjson",
            }
    )

But I get both as separate root span

---

If I explicitly specify the parent as an attribute on langsmith_extra parameter for the async generator:

@app.post("/test", response_model=None)
@traceable(name="test")
async def test(req: ChatRequest) -> StreamingResponse:
    return StreamingResponse(
            content=test_stream_generator(langsmith_extra={"parent": get_current_run_tree()}),
            headers={
                "Content-Type": "application/x-ndjson",
            }
    )

I get a strange result in the LangSmith UI:

The root span appears as it has a child but if you extend the dropdown there's nothing there. If I click the root span I can see the async generator span inside the trace information:

• How could I achieve the child to be shown also in the runs UI?

Thank you!

Suggestion:

No response

[hinthornw]

Thank you for reporting! And thank you for your patience 🙏

What python version are you using? Python's asyncio didn't natively support context copying in tasks until 3.11, which makes implicit context propagation for streaming a challenge.

I'm not sure if this will impact the fastapi inferred endpoint schema, but you can explicitly accept the run tree by adding a named keyword arg

@app.post("/test", response_model=None)
@traceable(name="test")
async def test(req: ChatRequest, run_tree: RunTree) -> StreamingResponse:
    return StreamingResponse(
            content=test_stream_generator(langsmith_extra={"parent": run_tree}),
            headers={
                "Content-Type": "application/x-ndjson",
            }
    )

Or you could potentially use the tracing context manager instead of the decorator:

@app.post("/test", response_model=None)
async def test(req: ChatRequest) -> StreamingResponse:
    with langsmith.trace(name="test") as rt
        return StreamingResponse(
                content=test_stream_generator(langsmith_extra={"parent": rt}),
                headers={
                    "Content-Type": "application/x-ndjson",
                }
        )

though the latter wouldn't give you a pretty output right n ow

[fchavat]

@hinthornw, thank you for your response.

I am currently using Python version 3.10.6 and have also tried version 3.11.5, encountering the same issue. During debugging, I noticed that when trying to access the parent run tree from the async generator method, it appears to have already terminated, likely due to the context manager closing it. This might explain why the child trace isn't visible in the runs UI, although it is accessible when I view the root trace details.

I have implemented your suggestions, unfortunately, without success. Previously, I used get_current_run_tree() as the parent for the async generator method. Is this approach significantly different from passing the run tree as an argument directly?

If there were a way to prevent the parent run tree from closing upon completion, similar to the end_on_exit parameter in OpenTelemetry trace, I believe it could resolve the issue by maintaining the run tree as the parent of the async generator method, thus preserving the visibility of the results in the Smith panel's runs menu.

Thank you!