From c0d701a51ea0565c078687af45ee2133d0c618a4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 14 May 2024 17:32:31 -0700 Subject: [PATCH] test(test_config.py): fix linting error --- litellm/tests/test_config.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_config.py b/litellm/tests/test_config.py index d0cd93d1f63..96c766919e3 100644 --- a/litellm/tests/test_config.py +++ b/litellm/tests/test_config.py @@ -14,19 +14,36 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the, system path import pytest, litellm -from pydantic import BaseModel +from pydantic import BaseModel, VERSION from litellm.proxy.proxy_server import ProxyConfig from litellm.proxy.utils import encrypt_value, ProxyLogging, DualCache from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from typing import Literal +# Function to get Pydantic version +def is_pydantic_v2() -> int: + return int(VERSION.split(".")[0]) + + +def get_model_config(arbitrary_types_allowed: bool = False) -> ConfigDict: + # Version-specific configuration + if is_pydantic_v2() >= 2: + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=arbitrary_types_allowed, protected_namespaces=()) # type: ignore + else: + from pydantic import Extra + + model_config = ConfigDict(extra=Extra.allow, arbitrary_types_allowed=arbitrary_types_allowed) # type: ignore + + return model_config + + class DBModel(BaseModel): model_id: str model_name: str model_info: dict litellm_params: dict - model_config = ConfigDict(protected_namespaces=()) + model_config = get_model_config() @pytest.mark.asyncio