Skip to content

Commit

Permalink
fix: support numpy arrays in span to json encoder
Browse files Browse the repository at this point in the history
  • Loading branch information
RogerHYang committed Mar 14, 2024
1 parent 94b1930 commit 097032a
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 0 deletions.
8 changes: 8 additions & 0 deletions src/phoenix/trace/span_json_encoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
from typing import Any, List
from uuid import UUID

import numpy as np

from phoenix.trace.schemas import (
Span,
SpanContext,
Expand Down Expand Up @@ -45,6 +47,12 @@ def default(self, obj: Any) -> Any:
}
elif isinstance(obj, SpanConversationAttributes):
return {"conversation_id": str(obj.conversation_id)}
elif isinstance(obj, np.ndarray):
return list(obj)
elif isinstance(obj, np.integer):
return int(obj)
elif isinstance(obj, np.floating):
return float(obj)
return super().default(obj)


Expand Down
51 changes: 51 additions & 0 deletions tests/trace/test_span_json_encoder.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
import json
from datetime import datetime, timezone

import numpy as np
from phoenix.trace.schemas import Span, SpanContext, SpanKind, SpanStatusCode
from phoenix.trace.span_json_encoder import span_to_json


def test_span_to_json() -> None:
start_time = datetime.now(timezone.utc)
end_time = datetime.now(timezone.utc)
span = Span(
name="test",
context=SpanContext(
trace_id="1234",
span_id="5678",
),
span_kind=SpanKind.UNKNOWN,
parent_id=None,
start_time=start_time,
end_time=end_time,
status_code=SpanStatusCode.ERROR,
status_message="error",
attributes={
"key": "value",
"integers": np.array([1, 2, 3]),
"floats": np.array([0.1, 0.2, 0.3]),
},
events=[],
conversation=None,
)
assert json.loads(span_to_json(span)) == {
"name": "test",
"context": {
"trace_id": "1234",
"span_id": "5678",
},
"span_kind": "UNKNOWN",
"parent_id": None,
"start_time": start_time.isoformat(),
"end_time": end_time.isoformat(),
"status_code": "ERROR",
"status_message": "error",
"attributes": {
"key": "value",
"integers": [1, 2, 3],
"floats": [0.1, 0.2, 0.3],
},
"events": [],
"conversation": None,
}

0 comments on commit 097032a

Please sign in to comment.