Better Serialize

python/langsmith/client.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import collections
+import dataclasses
 import datetime
 import functools
 import importlib
@@ -150,26 +151,44 @@ def _default_retry_config() -> Retry:
     return Retry(**retry_params)  # type: ignore
 
 
-def _serialize_json(obj: Any) -> Union[str, dict]:
+def _serialize_json(obj: Any) -> Any:
     if isinstance(obj, datetime.datetime):
         return obj.isoformat()
+    if isinstance(obj, uuid.UUID):
+        return str(obj)
+    try:
+        serialization_methods = [
+            ("model_dump_json", True),  # Pydantic V2
+            ("json", True),  # Pydantic V1
+            ("to_json", False),  # dataclass_json
+        ]
 
-    elif hasattr(obj, "model_dump_json") and callable(obj.model_dump_json):
-        # Base models, V2
-        try:
-            return json.loads(obj.model_dump_json(exclude_none=True))
-        except Exception:
-            logger.debug(f"Failed to serialize obj of type {type(obj)} to JSON")
-            return str(obj)
-    elif hasattr(obj, "json") and callable(obj.json):
-        # Base models, V1
+        for attr, exclude_none in serialization_methods:
+            if hasattr(obj, attr) and callable(getattr(obj, attr)):
+                try:
+                    method = getattr(obj, attr)
+                    json_str = (
+                        method(exclude_none=exclude_none) if exclude_none else method()
+                    )
+                    return json.loads(json_str)
+                except Exception as e:
+                    logger.debug(f"Failed to serialize {type(obj)} to JSON: {e}")
+                    return repr(obj)
+
+        if dataclasses.is_dataclass(obj):
+            # Regular dataclass
+            return dataclasses.asdict(obj)
         try:
-            return json.loads(obj.json(exclude_none=True))
-        except Exception:
-            logger.debug(f"Failed to json serialize {type(obj)} to JSON")
+            if hasattr(obj, "__slots__"):
+                return {slot: getattr(obj, slot) for slot in obj.__slots__}
+            else:
+                return vars(obj)
+        except Exception as e:
+            logger.debug(f"Failed to serialize {type(obj)} to JSON using vars: {e}")
             return repr(obj)
-    else:
-        return str(obj)
+    except Exception as e:
+        logger.debug(f"Failed to serialize {type(obj)} to JSON: {e}")
+        return repr(obj)
 
 
 def close_session(session: requests.Session) -> None:

python/tests/unit_tests/test_client.py

@@ -1,17 +1,21 @@
 """Test the LangSmith client."""
 import asyncio
+import dataclasses
 import gc
 import json
 import os
 import time
 import uuid
 import weakref
 from datetime import datetime
+from enum import Enum
 from io import BytesIO
-from typing import Optional
+from typing import NamedTuple, Optional
 from unittest import mock
 from unittest.mock import MagicMock, patch
 
+import attr
+import dataclasses_json
 import pytest
 import requests
 from pydantic import BaseModel
@@ -450,6 +454,99 @@ class MyPydantic(BaseModel):
     assert res2 == {"output": expected}
 
 
+def test_serialize_json() -> None:
+    class MyClass:
+        def __init__(self, x: int) -> None:
+            self.x = x
+            self.y = "y"
+
+    class MyClassWithSlots:
+        __slots__ = ["x", "y"]
+
+        def __init__(self, x: int) -> None:
+            self.x = x
+            self.y = "y"
+
+    class MyPydantic(BaseModel):
+        foo: str
+        bar: int
+
+    @dataclasses.dataclass
+    class MyDataclass:
+        foo: str
+        bar: int
+
+        def something(self) -> None:
+            pass
+
+    class MyEnum(str, Enum):
+        FOO = "foo"
+        BAR = "bar"
+
+    @dataclasses_json.dataclass_json
+    @dataclasses.dataclass
+    class Person:
+        name: str
+
+    @attr.dataclass
+    class AttrDict:
+        foo: str = attr.ib()
+        bar: int
+
+    uid = uuid.uuid4()
+    current_time = datetime.now()
+
+    class NestedClass:
+        __slots__ = ["person"]
+
+        def __init__(self) -> None:
+            self.person = Person(name="foo")
+
+    class MyNamedTuple(NamedTuple):
+        foo: str
+        bar: int
+
+    to_serialize = {
+        "uid": uid,
+        "time": current_time,
+        "my_class": MyClass(1),
+        "my_slotted_class": MyClassWithSlots(1),
+        "my_dataclass": MyDataclass("foo", 1),
+        "my_enum": MyEnum.FOO,
+        "my_pydantic": MyPydantic(foo="foo", bar=1),
+        "person": Person(name="foo"),
+        "a_bool": True,
+        "a_none": None,
+        "a_str": "foo",
+        "an_int": 1,
+        "a_float": 1.1,
+        "nested_class": NestedClass(),
+        "attr_dict": AttrDict(foo="foo", bar=1),
+        "named_tuple": MyNamedTuple(foo="foo", bar=1),
+    }
+
+    res = json.loads(json.dumps(to_serialize, default=_serialize_json))
+    expected = {
+        "uid": str(uid),
+        "time": current_time.isoformat(),
+        "my_class": {"x": 1, "y": "y"},
+        "my_slotted_class": {"x": 1, "y": "y"},
+        "my_dataclass": {"foo": "foo", "bar": 1},
+        "my_enum": "foo",
+        "my_pydantic": {"foo": "foo", "bar": 1},
+        "person": {"name": "foo"},
+        "a_bool": True,
+        "a_none": None,
+        "a_str": "foo",
+        "an_int": 1,
+        "a_float": 1.1,
+        "nested_class": {"person": {"name": "foo"}},
+        "attr_dict": {"foo": "foo", "bar": 1},
+        "named_tuple": ["foo", 1],
+    }
+    assert res == expected
+
+
 def test_host_url() -> None:
     client = Client(api_url="https://api.foobar.com/api", api_key="API_KEY")
     assert client._host_url == "https://api.foobar.com"