Span will not finish with langchain_aws.ChatBedrock with streaming=True

[davidol200]

With Datadog LLM Observability and dd-trace-py 2.15.1, the LLM invocation span of the code snippet below will never finish, preventing the associated trace to be sent to the running agent and further to the DD platform.

from langchain_core.prompts.chat import ChatPromptTemplate
import time



chat_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
    ]
)

llm = ChatBedrock(  # type: ignore[call-arg]
    region_name="eu-central-1",
    model_id="anthropic.claude-3-sonnet-20240229-v1:0",
    streaming=True,
    model_kwargs={
        "temperature": 0,
        "max_tokens": 1024,
    },
)

chain = chat_prompt | llm

print(chain.invoke({"input": "What is the capital of France?"}))

time.sleep(15)

Console output:

(env) $ DD_SERVICE=test_OD DD_ENV=dev DD_VERSION=0.1.0 DD_LLMOBS_ENABLED=1 DD_LLMOBS_ML_APP=test_OD_app ddtrace-run python generation.py
content='The capital of France is Paris.' response_metadata={'stop_reason': 'end_turn', 'stop_sequence': None} id='run-a2aeec71-e571-4378-b85e-daf543511977-0' usage_metadata={'input_tokens': 14, 'output_tokens': 10, 'total_tokens': 24}
Shutting down tracer with 1 unfinished spans. Unfinished spans will not be sent to Datadog: trace_id=137129836003462089902841517402087706889 parent_id=3564328065105982738 span_id=17454883787734973218 name=bedrock-runtime.command resource=InvokeModelWithResponseStream started=1730821865.490379 sampling_priority=None

No issue with streaming=False when configuring the model.

I may be missing something.

Thanks in advance.