Support responses API streaming in langfuse otel (#16153)

* streaming support in langfuse otel

* Added testing for Langfuse Otel tracing in the response API

---------

Co-authored-by: eycjur <eycjur@example.com>
This commit is contained in:
Katsuhiro Muto
2025-11-02 09:36:34 -08:00
committed by GitHub
co-authored by eycjur
parent 3f40613c56
commit 99775fa0f8
5 changed files with 342 additions and 10 deletions
@@ -201,6 +201,10 @@ class MessageAttributes:
"""
The id of the tool call.
"""
MESSAGE_REASONING_SUMMARY = "message.reasoning_summary"
"""
The reasoning summary from the model's chain-of-thought process.
"""
class MessageContentAttributes:
+72 -10
View File
@@ -45,6 +45,7 @@ def set_attributes(span: Span, kwargs, response_obj): # noqa: PLR0915
SpanAttributes,
ToolCallAttributes,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
try:
optional_params = kwargs.get("optional_params", {})
@@ -235,6 +236,7 @@ def set_attributes(span: Span, kwargs, response_obj): # noqa: PLR0915
# Captures response tokens, message, and content.
if hasattr(response_obj, "get"):
# Handle chat completions API (choices field)
for idx, choice in enumerate(response_obj.get("choices", [])):
response_message = choice.get("message", {})
safe_set_attribute(
@@ -256,6 +258,51 @@ def set_attributes(span: Span, kwargs, response_obj): # noqa: PLR0915
response_message.get("content", ""),
)
# Handle responses API (output field)
output_items = response_obj.get("output", [])
if output_items:
for i, item in enumerate(output_items):
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{i}"
if hasattr(item, "type"):
item_type = item.type
# Extract reasoning summary
if item_type == "reasoning" and hasattr(item, "summary"):
for summary in item.summary:
if hasattr(summary, "text"):
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_REASONING_SUMMARY}",
summary.text,
)
# Extract message content
elif item_type == "message" and hasattr(item, "content"):
message_content = ""
content_list = item.content
if content_list and len(content_list) > 0:
first_content = content_list[0]
message_content = getattr(first_content, "text", "")
message_role = getattr(item, "role", "assistant")
safe_set_attribute(
span,
SpanAttributes.OUTPUT_VALUE,
message_content,
)
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
message_content,
)
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
message_role,
)
# Token usage info.
usage = response_obj and response_obj.get("usage")
if usage:
@@ -266,18 +313,33 @@ def set_attributes(span: Span, kwargs, response_obj): # noqa: PLR0915
)
# The number of tokens used in the LLM response (completion).
safe_set_attribute(
span,
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION,
usage.get("completion_tokens"),
)
# Responses API uses "output_tokens", chat completions uses "completion_tokens"
completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens")
if completion_tokens:
safe_set_attribute(
span,
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION,
completion_tokens,
)
# The number of tokens used in the LLM prompt.
safe_set_attribute(
span,
SpanAttributes.LLM_TOKEN_COUNT_PROMPT,
usage.get("prompt_tokens"),
)
# Responses API uses "input_tokens", chat completions uses "prompt_tokens"
prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens")
if prompt_tokens:
safe_set_attribute(
span,
SpanAttributes.LLM_TOKEN_COUNT_PROMPT,
prompt_tokens,
)
# The number of reasoning tokens in the output, if available.
reasoning_tokens = usage.get("output_tokens_details", {}).get("reasoning_tokens")
if reasoning_tokens:
safe_set_attribute(
span,
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING,
reasoning_tokens,
)
except Exception as e:
verbose_logger.error(
@@ -158,6 +158,7 @@ class LangfuseOtelLogger(OpenTelemetry):
# Set observation output (response with tool_calls if present)
if response_obj and hasattr(response_obj, "get"):
# Handle chat completions API (choices field)
choices = response_obj.get("choices", [])
if choices:
# Extract the first choice's message
@@ -212,6 +213,44 @@ class LangfuseOtelLogger(OpenTelemetry):
safe_dumps(output_data),
)
# Handle responses API (output field)
output = response_obj.get("output", [])
if output:
output_data = []
for item in output:
if hasattr(item, "type"):
item_type = item.type
if item_type == "reasoning" and hasattr(item, "summary"):
for summary in item.summary:
if hasattr(summary, "text"):
output_data.append({
"role": "reasoning_summary",
"content": summary.text
})
elif item_type == "message":
output_data.append({
"role": getattr(item, "role", "assistant"),
"content": getattr(getattr(item, "content", [{}])[0], "text", "")
})
elif item_type == "function_call":
arguments_str = getattr(item, "arguments", "{}")
arguments_obj = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str
langfuse_tool_call = {
"id": getattr(item, "id", ""),
"name": getattr(item, "name", ""),
"call_id": getattr(item, "call_id", ""),
"type": "function_call",
"arguments": arguments_obj,
}
output_data.append(langfuse_tool_call)
if output_data:
safe_set_attribute(
span,
LangfuseSpanAttributes.OBSERVATION_OUTPUT.value,
safe_dumps(output_data),
)
@staticmethod
def _get_langfuse_otel_host() -> Optional[str]:
"""
@@ -181,6 +181,107 @@ def test_arize_set_attributes():
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 40)
def test_arize_set_attributes_responses_api():
"""
Test setting attributes for Responses API with mixed output (reasoning + message).
Verifies that multiple output types are correctly handled.
"""
from unittest.mock import MagicMock
from litellm.types.llms.openai import ResponsesAPIResponse, ResponseAPIUsage, OutputTokensDetails
from openai.types.responses import ResponseReasoningItem, ResponseOutputMessage, ResponseOutputText
from openai.types.responses.response_reasoning_item import Summary
span = MagicMock() # Mocked tracing span to test attribute setting
# Construct kwargs to simulate a real LLM request scenario
kwargs = {
"model": "o3-mini",
"messages": [{"role": "user", "content": "What is the answer?"}],
"standard_logging_object": {
"model_parameters": {"user": "test_user", "stream": True},
"metadata": {"key_1": "value_1", "key_2": None},
"call_type": "responses",
},
"optional_params": {
"max_tokens": "100",
"temperature": "1",
"top_p": "5",
"stream": True,
"user": "test_user",
},
"litellm_params": {"custom_llm_provider": "openai"},
}
# Simulate Responses API response with mixed output
response_obj = ResponsesAPIResponse(
id="response-123",
created_at=1625247600,
output=[
ResponseReasoningItem(
id="reasoning-001",
type="reasoning",
summary=[
Summary(
text="First, I need to analyze...",
type="summary_text"
)
]
),
ResponseOutputMessage(
id="msg-001",
type="message",
role="assistant",
status="completed",
content=[
ResponseOutputText(
annotations=[],
text="The answer is 42",
type="output_text",
)
]
)
],
usage=ResponseAPIUsage(
input_tokens=120,
output_tokens=250,
total_tokens=370,
output_tokens_details=OutputTokensDetails(
reasoning_tokens=180
)
)
)
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
# Verify reasoning summary was set (index 0)
span.set_attribute.assert_any_call(
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_REASONING_SUMMARY}",
"First, I need to analyze..."
)
# Verify message content was set (index 1)
span.set_attribute.assert_any_call(
SpanAttributes.OUTPUT_VALUE,
"The answer is 42"
)
span.set_attribute.assert_any_call(
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.1.{MessageAttributes.MESSAGE_CONTENT}",
"The answer is 42"
)
span.set_attribute.assert_any_call(
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.1.{MessageAttributes.MESSAGE_ROLE}",
"assistant"
)
# Verify token counts including reasoning tokens
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 370)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 250)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 120)
span.set_attribute.assert_any_call(
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180
)
class TestArizeLogger(CustomLogger):
"""
Custom logger implementation to capture standard_callback_dynamic_params.
@@ -474,6 +474,132 @@ class TestLangfuseOtelResponsesAPI:
for expected_call in expected_calls:
mock_safe_set_attribute.assert_any_call(*expected_call)
def test_responses_api_with_output(self):
"""Test Langfuse OTEL logger with Responses API output (reasoning + message)."""
from openai.types.responses import ResponseReasoningItem, ResponseOutputMessage, ResponseOutputText
from openai.types.responses.response_reasoning_item import Summary
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
# Create Responses API response with reasoning and message
response_obj = ResponsesAPIResponse(
id="response-456",
created_at=1625247600,
output=[
ResponseReasoningItem(
id="reasoning-001",
type="reasoning",
summary=[
Summary(
text="Let me analyze this problem step by step...",
type="summary_text"
)
]
),
ResponseOutputMessage(
id="msg-001",
type="message",
role="assistant",
status="completed",
content=[
ResponseOutputText(
annotations=[],
text="The weather in San Francisco is sunny, 20°C.",
type="output_text",
)
]
)
]
)
kwargs = {
"call_type": "responses",
"messages": [{"role": "user", "content": "What's the weather in San Francisco?"}],
"model": "gpt-4o",
"optional_params": {},
}
mock_span = MagicMock()
with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj)
# Verify observation output was set
output_calls = [
call for call in mock_safe_set_attribute.call_args_list
if call.args[1] == LangfuseSpanAttributes.OBSERVATION_OUTPUT.value
]
assert len(output_calls) > 0, "observation.output should be set"
output_json = output_calls[0].args[2]
output_data = json.loads(output_json)
# Verify output contains reasoning and message
assert isinstance(output_data, list)
assert len(output_data) == 2
# Verify reasoning summary
assert output_data[0]["role"] == "reasoning_summary"
assert output_data[0]["content"] == "Let me analyze this problem step by step..."
# Verify message
assert output_data[1]["role"] == "assistant"
assert output_data[1]["content"] == "The weather in San Francisco is sunny, 20°C."
def test_responses_api_with_function_calls(self):
"""Test Langfuse OTEL logger with Responses API function_call output."""
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
from openai.types.responses import ResponseFunctionToolCall
# Create Responses API response with function call
response_obj = ResponsesAPIResponse(
id="response-789",
created_at=1625247700,
output=[
ResponseFunctionToolCall(
id="fc-123",
type="function_call",
name="get_weather",
call_id="call-abc",
arguments='{"location": "San Francisco", "unit": "celsius"}',
status="completed"
)
]
)
kwargs = {
"call_type": "responses",
"messages": [{"role": "user", "content": "What's the weather in San Francisco?"}],
"model": "gpt-4o",
"optional_params": {},
}
mock_span = MagicMock()
with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute:
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj)
# Verify observation output was set
output_calls = [
call for call in mock_safe_set_attribute.call_args_list
if call.args[1] == LangfuseSpanAttributes.OBSERVATION_OUTPUT.value
]
assert len(output_calls) > 0, "observation.output should be set"
output_json = output_calls[0].args[2]
output_data = json.loads(output_json)
# Verify output contains function call
assert isinstance(output_data, list)
assert len(output_data) == 1
# Verify function call details
assert output_data[0]["type"] == "function_call"
assert output_data[0]["id"] == "fc-123"
assert output_data[0]["name"] == "get_weather"
assert output_data[0]["call_id"] == "call-abc"
assert output_data[0]["arguments"]["location"] == "San Francisco"
assert output_data[0]["arguments"]["unit"] == "celsius"
if __name__ == "__main__":
pytest.main([__file__])