diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 0d23033e7..5dcb68ca3 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -8,6 +8,7 @@ on: pull_request: branches: - "release/*" + - "patch/*" - "main" jobs: diff --git a/langtest/augmentation/base.py b/langtest/augmentation/base.py index ca6718133..2767df6cb 100644 --- a/langtest/augmentation/base.py +++ b/langtest/augmentation/base.py @@ -19,7 +19,6 @@ from langtest.utils.custom_types.predictions import NERPrediction, SequenceLabel from langtest.utils.custom_types.sample import NERSample from langtest.tasks import TaskManager -from ..utils.lib_manager import try_import_lib from ..errors import Errors @@ -358,6 +357,9 @@ def __init__( # Extend the existing templates list self.__templates.extend(generated_templates[:num_extra_templates]) + except ModuleNotFoundError: + raise ImportError(Errors.E097()) + except Exception as e_msg: raise Errors.E095(e=e_msg) @@ -606,19 +608,19 @@ def __generate_templates( num_extra_templates: int, model_config: Union[OpenAIConfig, AzureOpenAIConfig] = None, ) -> List[str]: - if try_import_lib("openai"): - from langtest.augmentation.utils import ( - generate_templates_azoi, # azoi means Azure OpenAI - generate_templates_openai, - ) + """This method is used to generate extra templates from a given template.""" + from langtest.augmentation.utils import ( + generate_templates_azoi, # azoi means Azure OpenAI + generate_templates_openai, + ) - params = model_config.copy() if model_config else {} + params = model_config.copy() if model_config else {} - if model_config and model_config.get("provider") == "openai": - return generate_templates_openai(template, num_extra_templates, params) + if model_config and model_config.get("provider") == "openai": + return generate_templates_openai(template, num_extra_templates, params) - elif model_config and model_config.get("provider") == "azure": - return generate_templates_azoi(template, num_extra_templates, params) + elif model_config and model_config.get("provider") == "azure": + return generate_templates_azoi(template, num_extra_templates, params) - else: - return generate_templates_openai(template, num_extra_templates) + else: + return generate_templates_openai(template, num_extra_templates) diff --git a/langtest/augmentation/utils.py b/langtest/augmentation/utils.py index a13a8d2e2..ad0051be5 100644 --- a/langtest/augmentation/utils.py +++ b/langtest/augmentation/utils.py @@ -19,15 +19,13 @@ class OpenAIConfig(TypedDict): class AzureOpenAIConfig(TypedDict): """Azure OpenAI Configuration for API Key and Provider.""" - from openai.lib.azure import AzureADTokenProvider - azure_endpoint: str api_version: str api_key: str provider: str azure_deployment: Union[str, None] = None azure_ad_token: Union[str, None] = (None,) - azure_ad_token_provider: Union[AzureADTokenProvider, None] = (None,) + azure_ad_token_provider = (None,) organization: Union[str, None] = (None,) @@ -76,6 +74,7 @@ def generate_templates_azoi( template: str, num_extra_templates: int, model_config: AzureOpenAIConfig ): """Generate new templates based on the provided template using Azure OpenAI API.""" + import openai if "provider" in model_config: @@ -139,6 +138,7 @@ def generate_templates_openai( template: str, num_extra_templates: int, model_config: OpenAIConfig = OpenAIConfig() ): """Generate new templates based on the provided template using OpenAI API.""" + import openai if "provider" in model_config: diff --git a/langtest/errors.py b/langtest/errors.py index d3d7d1bba..c4cb90189 100644 --- a/langtest/errors.py +++ b/langtest/errors.py @@ -275,6 +275,7 @@ class Errors(metaclass=ErrorsWithCodes): E094 = ("Unsupported category: '{category}'. Supported categories: {supported_category}") E095 = ("Failed to make API request: {e}") E096 = ("Failed to generate the templates in Augmentation: {msg}") + E097 = ("Failed to load openai. Please install it using `pip install openai`") class ColumnNameError(Exception):