fix: initial commit fixing langfuse request/response logging with OTEL

This commit is contained in:
Krrish Dholakia 2025-10-22 14:20:39 -07:00
parent 09fc9deac8
commit a3e4fa5b81
2 changed files with 82 additions and 23 deletions

View file

@ -35,6 +35,47 @@ def safe_set_attribute(span: Span, key: str, value: Any):
span.set_attribute(key, primitive_value)
def set_response_output_messages(span: Span, response_obj):
"""
Sets output message attributes on the span from the LLM response.
Args:
span: The OpenTelemetry span to set attributes on
response_obj: The response object containing choices with messages
"""
from litellm.integrations._types.open_inference import (
MessageAttributes,
SpanAttributes,
)
safe_set_attribute(
span,
"langfuse.observation.output",
response_obj.model_dump_json(),
)
for idx, choice in enumerate(response_obj.get("choices", [])):
response_message = choice.get("message", {})
safe_set_attribute(
span,
SpanAttributes.OUTPUT_VALUE,
response_message.get("content", ""),
)
# This shows up under `output_messages` tab on the span page.
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}"
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
response_message.get("role"),
)
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
response_message.get("content", ""),
)
def set_attributes(span: Span, kwargs, response_obj): # noqa: PLR0915
"""
Populates span with OpenInference-compliant LLM attributes for Arize and Phoenix tracing.
@ -153,6 +194,12 @@ def set_attributes(span: Span, kwargs, response_obj): # noqa: PLR0915
)
messages = kwargs.get("messages")
safe_set_attribute(
span,
"langfuse.observation.input",
json.dumps(messages),
)
# for /chat/completions
# https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
if messages:
@ -235,26 +282,7 @@ def set_attributes(span: Span, kwargs, response_obj): # noqa: PLR0915
# Captures response tokens, message, and content.
if hasattr(response_obj, "get"):
for idx, choice in enumerate(response_obj.get("choices", [])):
response_message = choice.get("message", {})
safe_set_attribute(
span,
SpanAttributes.OUTPUT_VALUE,
response_message.get("content", ""),
)
# This shows up under `output_messages` tab on the span page.
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}"
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
response_message.get("role"),
)
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
response_message.get("content", ""),
)
set_response_output_messages(span, response_obj)
# Token usage info.
usage = response_obj and response_obj.get("usage")

View file

@ -2514,8 +2514,6 @@ def test_completion_azure_key_completion_arg():
pytest.fail(f"Error occurred: {e}")
async def test_re_use_azure_async_client():
try:
print("azure gpt-3.5 ASYNC with clie nttest\n\n")
@ -4357,7 +4355,6 @@ def test_deepseek_reasoning_content_completion():
pytest.skip("Model is timing out")
def test_qwen_text_completion():
# litellm._turn_on_debug()
resp = litellm.completion(
@ -4492,3 +4489,37 @@ def test_completion_gpt_4o_empty_str():
messages=[{"role": "user", "content": ""}],
)
assert resp.choices[0].message.content is not None
def test_edit_note():
litellm.callbacks = ["langfuse_otel"]
response = completion(
model="gpt-4o",
messages=[
{
"role": "system",
"content": "Your only job is to call the edit_note tool with the content specified in the user's message.",
},
{
"role": "user",
"content": "Edit the note with the content: 'This is a test note.'",
},
],
tools=[
{
"type": "function",
"function": {
"name": "edit_note",
"description": "Edit the note with the content specified in the user's message.",
"parameters": {
"type": "object",
"properties": {
"content": {"type": "string"},
},
},
},
},
],
)
return response