From 097032ac8d170635bdd79ffacdd8ed1e428fa21c Mon Sep 17 00:00:00 2001
From: Roger Yang <roger.yang@arize.com>
Date: Thu, 14 Mar 2024 10:08:16 -0700
Subject: [PATCH] fix: support numpy arrays in span to json encoder

---
 src/phoenix/trace/span_json_encoder.py |  8 ++++
 tests/trace/test_span_json_encoder.py  | 51 ++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 tests/trace/test_span_json_encoder.py

diff --git a/src/phoenix/trace/span_json_encoder.py b/src/phoenix/trace/span_json_encoder.py
index ca99897427..b3d89b8722 100644
--- a/src/phoenix/trace/span_json_encoder.py
+++ b/src/phoenix/trace/span_json_encoder.py
@@ -5,6 +5,8 @@
 from typing import Any, List
 from uuid import UUID
 
+import numpy as np
+
 from phoenix.trace.schemas import (
     Span,
     SpanContext,
@@ -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)
 
 
diff --git a/tests/trace/test_span_json_encoder.py b/tests/trace/test_span_json_encoder.py
new file mode 100644
index 0000000000..4fa49095c0
--- /dev/null
+++ b/tests/trace/test_span_json_encoder.py
@@ -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,
+    }