add support for pydantic serialization

Author: willydouhard

literalai/api/gql.py

@@ -1147,6 +1147,8 @@ def steps_variables_builder(steps: List[Union["StepDict", "Step"]]):
                 step["input"] = ensure_values_serializable(step["input"])
             if step.get("output"):
                 step["output"] = ensure_values_serializable(step["output"])
+            if step.get("metadata"):
+                step["metadata"] = ensure_values_serializable(step["metadata"])
             variables.update(serialize_step(step, i))
     return variables
 

literalai/callback/langchain_callback.py

@@ -2,8 +2,6 @@
 from importlib.metadata import version
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict, Union, cast
 
-from pydantic import BaseModel
-
 from literalai.helper import ensure_values_serializable
 from literalai.observability.generation import (
     ChatGeneration,
@@ -150,8 +148,6 @@ def process_content(self, content: Any, root=True):
                     return [self._convert_message(m) for m in content]
             elif self._is_message(content):
                 return self._convert_message(content)
-            elif isinstance(content, BaseModel):
-                return content.model_dump()
             elif isinstance(content, dict):
                 processed_dict = {}
                 for key, value in content.items():

literalai/helper.py

@@ -1,5 +1,7 @@
 from datetime import datetime
 
+from pydantic import BaseModel
+
 
 def filter_none_values(data):
     return {key: value for key, value in data.items() if value is not None}
@@ -9,13 +11,9 @@ def ensure_values_serializable(data):
     """
     Recursively ensures that all values in the input (dict or list) are JSON serializable.
     """
-    try:
-        from openai.types.chat import ChatCompletionMessage
 
-        if isinstance(data, ChatCompletionMessage):
-            return filter_none_values(data.model_dump())
-    except ImportError:
-        pass
+    if isinstance(data, BaseModel):
+        return filter_none_values(data.model_dump())
 
     try:
         from mistralai import UserMessage