From 004a3a00be74b4d13d3149bd08135e509f8cad72 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Wed, 4 Oct 2023 13:42:44 -0400 Subject: [PATCH] Conditional imports to support pydantic>2 --- prefect_openai/completion.py | 8 +++++++- prefect_openai/credentials.py | 7 ++++++- prefect_openai/image.py | 8 +++++++- 3 files changed, 20 insertions(+), 3 deletions(-) diff --git a/prefect_openai/completion.py b/prefect_openai/completion.py index effc0df..69655e3 100644 --- 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 index 52759d1..8f3db54 100644 --- 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 index 783f06b..37d7c46 100644 --- 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