Replace use of deprecated dict method with model_dump

replicate · Sep 16, 2024 · 52c1ff0 · 52c1ff0
1 parent 588df97
commit 52c1ff0
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/python/cog/server/webhook.py b/python/cog/server/webhook.py
@@ -8,6 +8,7 @@
 from urllib3.util.retry import Retry
 
 from ..schema import PredictionResponse, Status, WebhookEvent
+from ..types import PYDANTIC_V2
 from .response_throttler import ResponseThrottler
 from .telemetry import current_trace_context
 from .useragent import get_user_agent
@@ -45,7 +46,12 @@ def webhook_caller(webhook: str) -> Callable[[Any], None]:
 
     def caller(response: PredictionResponse) -> None:
         if throttler.should_send_response(response):
-            dict_response = jsonable_encoder(response.dict(exclude_unset=True))
+            if PYDANTIC_V2:
+                dict_response = jsonable_encoder(
+                    response.model_dump(exclude_unset=True)
+                )
+            else:
+                dict_response = jsonable_encoder(response.dict(exclude_unset=True))
             if Status.is_terminal(response.status):
                 # For terminal updates, retry persistently
                 retry_session.post(webhook, json=dict_response)