From 136746abc985986760cab5dbc54b88c9bbc09cbf Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 15 May 2024 19:42:39 -0700 Subject: [PATCH] fix test config --- litellm/tests/test_config.py | 24 ++++-------------------- 1 file changed, 4 insertions(+), 20 deletions(-) diff --git a/litellm/tests/test_config.py b/litellm/tests/test_config.py index 96c766919e..e38187e0ea 100644 --- a/litellm/tests/test_config.py +++ b/litellm/tests/test_config.py @@ -5,7 +5,6 @@ import sys, os import traceback from dotenv import load_dotenv -from pydantic import ConfigDict load_dotenv() import os, io @@ -14,36 +13,21 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the, system path import pytest, litellm -from pydantic import BaseModel, VERSION +from pydantic import BaseModel 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 = get_model_config() + + class Config: + protected_namespaces = () @pytest.mark.asyncio