Conditional imports to support pydantic>2

PrefectHQ · Oct 4, 2023 · 004a3a0 · 004a3a0
1 parent fa78134
commit 004a3a0
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 3 deletions.
diff --git a/prefect_openai/completion.py b/prefect_openai/completion.py
@@ -12,7 +12,13 @@
 from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.tasks import Task
 from prefect.utilities.asyncutils import is_async_fn, sync_compatible
-from pydantic import Field
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field
+else:
+    from pydantic import Field
+
 from typing_extensions import Literal
 
 from prefect_openai import OpenAICredentials

diff --git a/prefect_openai/credentials.py b/prefect_openai/credentials.py
@@ -5,7 +5,12 @@
 
 import openai
 from prefect.blocks.abstract import CredentialsBlock
-from pydantic import Field, SecretStr
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field, SecretStr
+else:
+    from pydantic import Field, SecretStr
 
 
 class OpenAICredentials(CredentialsBlock):

diff --git a/prefect_openai/image.py b/prefect_openai/image.py
@@ -7,7 +7,13 @@
 from prefect.exceptions import MissingContextError
 from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.utilities.asyncutils import sync_compatible
-from pydantic import Field
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field
+else:
+    from pydantic import Field
+
 from typing_extensions import Literal
 
 from prefect_openai import OpenAICredentials