Merge pull request #144 from Chainlit/willy/eng-2113-support-pydantic

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,21 +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
-
-    try:
-        from mistralai import UserMessage
 
-        if isinstance(data, UserMessage):
-            return filter_none_values(data.model_dump())
-    except ImportError:
-        pass
+    if isinstance(data, BaseModel):
+        return filter_none_values(data.model_dump())
 
     if isinstance(data, dict):
         return {key: ensure_values_serializable(value) for key, value in data.items()}

literalai/version.py

@@ -1 +1 @@
-__version__ = "0.1.0"
+__version__ = "0.1.1"

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="literalai",
-    version="0.1.0",  # update version in literalai/version.py
+    version="0.1.1",  # update version in literalai/version.py
     description="An SDK for observability in Python applications",
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",