ai21[patch]: Update @root_validators for pydantic2 migration (#25401)

Update @root_validators for pydantic 2 migration.
pull/25442/head
Eugene Yurtsev 1 month ago committed by GitHub
parent 6f68c8d6ab
commit a114255b82
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -28,7 +28,7 @@ class AI21Base(BaseModel):
num_retries: Optional[int] = None
"""Maximum number of retries for API requests before giving up."""
@root_validator()
@root_validator(pre=True)
def validate_environment(cls, values: Dict) -> Dict:
api_key = convert_to_secret_str(
values.get("api_key") or os.getenv("AI21_API_KEY") or ""
@ -46,7 +46,13 @@ class AI21Base(BaseModel):
os.getenv("AI21_TIMEOUT_SEC", _DEFAULT_TIMEOUT_SEC)
)
values["timeout_sec"] = timeout_sec
return values
@root_validator(pre=False, skip_on_failure=True)
def post_init(cls, values: Dict) -> Dict:
api_key = values["api_key"]
api_host = values["api_host"]
timeout_sec = values["timeout_sec"]
if values.get("client") is None:
values["client"] = AI21Client(
api_key=api_key.get_secret_value(),

Loading…
Cancel
Save