fix bad import

PrefectHQ · Jul 16, 2024 · 1b07f66 · 1b07f66
1 parent 250d6d6
commit 1b07f66
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/prefect/main.py b/src/prefect/main.py
@@ -24,12 +24,13 @@
 import prefect.context
 
 # Perform any forward-ref updates needed for Pydantic models
-import prefect.client.schemas
+from prefect.client.schemas.objects import State
+from prefect.client.schemas.actions import StateCreate
 
 prefect.context.FlowRunContext.model_rebuild()
 prefect.context.TaskRunContext.model_rebuild()
-prefect.client.schemas.State.model_rebuild()
-prefect.client.schemas.StateCreate.model_rebuild()
+State.model_rebuild()
+StateCreate.model_rebuild()
 Transaction.model_rebuild()
 
 # Configure logging