test: add pydantic v1 compatibility to chat_response_format.py

AI21Labs · Aug 21, 2024 · 9bab854 · 9bab854
1 parent c335204
commit 9bab854
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/examples/studio/chat/chat_response_format.py b/examples/studio/chat/chat_response_format.py
@@ -6,6 +6,8 @@
 from ai21.logger import set_verbose
 from ai21.models.chat import ChatMessage, ResponseFormat
 
+from ai21.models.ai21_base_model import IS_PYDANTIC_V2
+
 set_verbose(True)
 
 
@@ -47,6 +49,9 @@ class ZooTicketsOrder(BaseModel):
 try:
     order = ZooTicketsOrder.model_validate_json(response.choices[0].message.content)
     print("Zoo tickets order details JSON:")
-    print(order.model_dump_json(indent=4))
+    if IS_PYDANTIC_V2:
+        print(order.model_dump_json(indent=4))
+    else:
+        print(order.json(indent=4))
 except ValidationError as exc:
     print(exc)