GenAI Utils | Allow passing external Span to the manual LLMInvocation

[Fiery-Fenix]

Description

This PR allows to pass already active Span to TelemetryHandler via LLMInvocation instead of always creating a new Span during start_llm call.
This will allow to control Span lifecycle outside of GenAI utils which might be useful for some cases.

Primary goal for this change - to allow usage of GenAI Utills in botocore instrumentation for bedrock instrumentation. Main feature of botocore instrumentation that it consist of 2 layers:

• main instrumentation with wrappers implementation, which manage Span lifecycle (start/stop)
• "extensions" that are populating Span with useful data

To use instrumentation-genai-util in botocore instrumentation we'll either need to do a big refactoring of the botocore instrumentation to move Span lifecycle control to "extensions" or to allow passing external Span to TelemetryHandler, keeping control on the caller side.

This is also can be useful for instrumenting Streaming GenAI calls when Span lifecycle should be managed on the same level as stream lifecycle control.

Type of change

Please delete options that are not relevant.

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

How Has This Been Tested?

New test was added.

$ tox -e py39-test-util-genai,py310-test-util-genai,py311-test-util-genai,py312-test-util-genai,py313-test-util-genai,py314-test-util-genai
  py39-test-util-genai: OK (12.91=setup[6.41]+cmd[6.51] seconds)
  py310-test-util-genai: OK (16.36=setup[7.09]+cmd[9.27] seconds)
  py311-test-util-genai: OK (18.19=setup[8.96]+cmd[9.23] seconds)
  py312-test-util-genai: OK (15.13=setup[8.66]+cmd[6.47] seconds)
  py313-test-util-genai: OK (18.22=setup[7.69]+cmd[10.52] seconds)
  py314-test-util-genai: OK (18.83=setup[9.22]+cmd[9.61] seconds)

Does This PR Require a Core Repo Change?

• No.

Checklist:

See contributing.md for styleguide, changelog guidelines, and more.

• Followed the style guidelines of this project
• Changelogs have been updated
• Unit tests have been added
• Documentation has been updated

[aabmass]

Can you explain a little more on the use case? Since the span is returned already, why not just use use the returned span?

[Fiery-Fenix]

Sure, main purpose was to make adoption of instrumentation-genai-utils into instrumentation-botocore without big changes on botocore instrumentation side.

In botocore instrumentation Span is started here. Please take a note that this Span is started for any supported "extension" here, not only for Bedrock instrumentation, but also for S3/Lambda/DocumentDB/etc.
At the same time actual LLM Invocation instrumentation is located deeper - in bedrock extension which is actually called few lines below.
Please take a note that bedrock extension can handle both synchronous and stream invokes that bring another complexity for instrumentation-genai-utils.

My idea is to keep top-level instrumentation as is, keeping Span lifecycle management there, at the same time run TelemetryHandler.start_llm/TelemetryHandler.stop_llm exactly in bedrock extension using the provided active Span. This will keep code change more local and will not break code responsibility by introducing dependency on GenAI TelemetryHandler on the botocore instrumentation top level.

[aabmass]

Hmm I see. I don't love it for complicating the API, I think this will pretty much only be used for Bedrock.

It's a bit of an anti pattern to update the span name after starting the span, since the final name isn't available to samplers. But it is what it is. I'll defer to @xrmx who wrote the bedrock code though if he thinks it's worth refactoring.

[Fiery-Fenix]

Actually, current implementation of the TelemetryHandler forcedly overwrites Span name after it's creation, I haven't touched that part.

[aabmass]

Does end_span_on_exit need to be public, or is it only used internally here?

Another option that might be cleaner is to update LLMInvocation to store an ExitStack for cleanup, since you can dynamically push callbacks onto it.

[Fiery-Fenix]

As this field is used by TelemetryHandler to determine if it's required to close the span - I believe yes, it should be public field of LLMInvocation. Usage of this field is very similar to existing LLMInvocation field context_token - it's also seems to be private, but never used by LLMInvocation class itself, only by TelemetryHandler.