fix/#652

burr/integrations/pydantic.py

@@ -51,8 +51,10 @@
 
 def model_to_dict(model: pydantic.BaseModel, include: Optional[List[str]] = None) -> dict:
     """Utility function to convert a pydantic model to a dictionary."""
-    keys = model.model_fields.keys()
-    keys = keys if include is None else [item for item in include if item in model.model_fields]
+    keys = type(model).model_fields.keys()
+    keys = (
+        keys if include is None else [item for item in include if item in type(model).model_fields]
+    )
     return {key: getattr(model, key) for key in keys}
 
 

pyproject.toml

@@ -89,7 +89,7 @@ tests = [
   "langchain_core",
   "langchain_community",
   "pandas",
-  "pydantic[email]",
+  "pydantic[email]>=2.0",
   "pyarrow",
   "apache-burr[aiosqlite]",
   "apache-burr[asyncpg]",
@@ -120,7 +120,7 @@ documentation = [
 ]
 
 tracking-client = [
-  "pydantic>1"
+  "pydantic>=2.0"
 ]
 
 tracking-client-s3 = [
@@ -141,7 +141,7 @@ tracking-server = [
   "click",
   "fastapi",
   "uvicorn",
-  "pydantic",
+  "pydantic>=2.0",
   "pydantic-settings",
   "fastapi-pagination",
   "fastapi-utils",
@@ -153,7 +153,7 @@ tracking-server = [
 ]
 
 pydantic = [
-  "pydantic"
+  "pydantic>=2.0"
 ]
 
 haystack = [

tests/integrations/test_burr_pydantic.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import asyncio
+import warnings
 from typing import AsyncGenerator, Generator, List, Optional, Tuple
 
 import pydantic
@@ -134,6 +135,14 @@ class MyModelWithConfig(pydantic.BaseModel):
     assert SubsetModel.model_config == {"arbitrary_types_allowed": True}
 
 
+def test_model_to_dict_no_deprecation_warning():
+    model = OriginalModel(foo=1, bar="bar", nested=NestedModel(nested_field1=1))
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        result = model_to_dict(model)
+    assert "foo" in result
+
+
 def test_merge_to_state():
     model = OriginalModel(
         foo=1,