mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-08-12 04:50:39 +00:00
update serialization of values to handle broader types
This commit is contained in:
parent
200ef29233
commit
0b9b72dd2c
2 changed files with 28 additions and 5 deletions
|
@ -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
|
||||
|
||||
|
|
|
@ -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:
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue