diff --git a/python/cog/server/webhook.py b/python/cog/server/webhook.py index b4f1c86a6c..1aca58ae41 100644 --- 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)