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: if event.stop_reason is not None:
stop_reason = event.stop_reason stop_reason = event.stop_reason
span.set_attribute("stop_reason", 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(
span.set_attribute("output", f"content: {content} tool_calls: {tool_calls}") "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 stop_reason = stop_reason or StopReason.out_of_tokens

View file

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