update serialization of values to handle broader types

This commit is contained in:
Hardik Shah 2025-02-27 10:56:10 -08:00
parent 200ef29233
commit 0b9b72dd2c
2 changed files with 28 additions and 5 deletions

View file

@ -591,8 +591,17 @@ class ChatAgent(ShieldRunnerMixin):
if event.stop_reason is not None:
stop_reason = event.stop_reason
span.set_attribute("stop_reason", stop_reason)
span.set_attribute("input", [m.model_dump_json() for m in input_messages])
span.set_attribute("output", f"content: {content} tool_calls: {tool_calls}")
span.set_attribute(
"input",
json.dumps([json.loads(m.model_dump_json()) for m in input_messages]),
)
output_attr = json.dumps(
{
"content": content,
"tool_calls": [json.loads(t.model_dump_json()) for t in tool_calls],
}
)
span.set_attribute("output", output_attr)
stop_reason = stop_reason or StopReason.out_of_tokens

View file

@ -6,6 +6,7 @@
import asyncio
import inspect
import json
from functools import wraps
from typing import Any, AsyncGenerator, Callable, Type, TypeVar
@ -17,6 +18,10 @@ T = TypeVar("T")
def serialize_value(value: Any) -> Primitive:
return str(_prepare_for_json(value))
def _prepare_for_json(value: Any) -> str:
"""Serialize a single value into JSON-compatible format."""
if value is None:
return ""
@ -25,8 +30,16 @@ def serialize_value(value: Any) -> Primitive:
elif hasattr(value, "_name_"):
return value._name_
elif isinstance(value, BaseModel):
return value.model_dump_json()
return json.loads(value.model_dump_json())
elif isinstance(value, (list, tuple, set)):
return [_prepare_for_json(item) for item in value]
elif isinstance(value, dict):
return {str(k): _prepare_for_json(v) for k, v in value.items()}
else:
try:
json.dumps(value)
return value
except Exception:
return str(value)
@ -104,7 +117,8 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
result = method(self, *args, **kwargs)
span.set_attribute("output", serialize_value(result))
return result
except Exception as _e:
except Exception as e:
span.set_attribute("error", str(e))
raise
if is_async_gen: