DRAFT First try (NOW WORKING!) to migrate to Pydantic v2

interlab/queries/json_schema.py

@@ -1,22 +1,24 @@
 import dataclasses
 import json
-from typing import Any
+from typing import Any, Type
 
 import jsonref
 import pydantic
 
 
-def get_pydantic_model(T: type) -> type:
+def get_pydantic_model(T: type) -> pydantic.TypeAdapter | pydantic.BaseModel:
     """
     Convert the type to pydantic BaseModel.
     """
     if not issubclass(T, pydantic.BaseModel):
-        if not dataclasses.is_dataclass(T):
+        if not pydantic.v1.dataclasses.is_builtin_dataclass(
+            T
+        ) and not pydantic.dataclasses.is_pydantic_dataclass(T):
             raise TypeError(
                 "Only pydantic Model, or pydantic or standard dataclasses are accepted"
             )
-        T = pydantic.dataclasses.create_pydantic_model_from_dataclass(T)
-    assert issubclass(T, pydantic.BaseModel)  # In lieu of a test
+        # T = pydantic.v1.dataclasses.create_pydantic_model_from_dataclass(T)
+        return pydantic.TypeAdapter(T)
     return T
 
 
@@ -53,9 +55,9 @@ def _rec(a, seen_ids):
     return d2
 
 
-def get_json_schema(T: type, strip_root=True) -> Any:
+def get_json_schema(pydantic_type: pydantic.TypeAdapter, strip_root=True) -> Any:
     # pydantic schema
-    schema = get_pydantic_model(T).schema()
+    schema = pydantic_type.json_schema()
     if strip_root:
         # remove root title and description (usually misleading to LLMs in dataclasses)
         schema.pop("title", None)

interlab/queries/query_for_json.py

@@ -150,9 +150,9 @@ def query_for_json(
                 d = find_and_parse_json_block(res)
                 # TODO: Is the following conversion/validation working for nested fields as well?
                 # Convert to pydantic type for permissive conversion and validation
-                d = pdT(**d)
                 # Convert back to match expected type (nested types are ok)
-                d = T(**d.dict())
+                d = T(**d)
+                pdT.validate_python(d)
                 assert isinstance(d, T)
                 c.set_result(d)
                 return d