Expose Span Customization / Event Iteration for handle_ag_ui_request and other AG-UI functions

[Snejj]

Description

When using handle_ag_ui_request() with instrument_all=True, there’s no official way to enrich telemetry spans before they’re emitted (e.g., adding request IDs, inputs/outputs). My attempts today either create duplicate spans or require subclassing Agent despite it being marked as final.

It would be useful if AG-UI exposed a hook or iterator similar to Agent.iter, something like e.g.:

run_input = RunAgentInput.model_validate(await request.json())
request_id = request.headers.get("X-Request-ID")

async def handler(events):
   async for event in events:
        event.span.set_attribute("request.id", request_id)
        event.span.set_attribute("run_input", run_input.json())
        yield event

return await handle_ag_ui_request(
    agent,
    request,
    event_stream_handler=handler,
)

This would make it easy to add metadata and integrate cleanly.

[DouweM]

@Snejj Would it be an option to wrap handle_ag_ui_request in a with logfire.span(..., <those_attributes>):? All the apns created inside handle_ag_ui_request would then have that as the parent span, and this'd prevent duplicating the (potentially large) run_input data on every span.

[Snejj]

I'm currently trying to use LangFuse, and when I try things like this the spans get duplicated. I'm not sure how to start the span outside of the AG-UI functions so that I have access to the span info when I need it

[DouweM]

@Snejj Can you see if you're seeing the same issue with our Pydantic Logfire (https://ai.pydantic.dev/logfire/)? If so, it's something we can look into, if not it's a LangFuse issue.

[github-actions]

This issue is stale, and will be closed in 3 days if no reply is received.

[Snejj]

It seems like logfire sets the inputs and outputs properly automatically, however due to the nature of the contract I still can't update the trace outside of the agui function execution.

[DouweM]

@Snejj Would it be an option to use a SpanProcessor? https://logfire.pydantic.dev/docs/reference/api/logfire/#logfire.configure(additional_span_processors)

[Snejj]

This looks promising as a way to interact with the span before sending out to logfire, thanks.

I can use logfire for personal projects, but can't at work as we're locked into langfuse for now. This seems logfire specific, ideally there would be a way to send a processor on the agui call or at least override the Agent class or something

[DouweM]

@Snejj The logfire package is "just" an OTel SDK that sends to our Pydantic Logfire by default, you can point it at another OTel endpoint like so: https://logfire.pydantic.dev/docs/how-to-guides/alternative-backends/

I agree it would be nice to add span_processors to InstrumentationSettings as well.

[Snejj]

Okay, I'll see how using the logfire sdk but setting the endpoint to our self-hosted langfuse instance goes!

[github-actions]

This issue is stale, and will be closed in 3 days if no reply is received.