box[major]: use pydantic v2 (#26067)

This commit is contained in:
Bagatur
2024-09-04 14:51:53 -07:00
committed by GitHub
parent 923ce84aa7
commit 567a4ce5aa
4 changed files with 169 additions and 174 deletions

View File

@@ -1,10 +1,11 @@
from typing import Any, Dict, Iterator, List, Optional
from typing import Iterator, List, Optional
from box_sdk_gen import FileBaseTypeField # type: ignore
from langchain_core.document_loaders.base import BaseLoader
from langchain_core.documents import Document
from langchain_core.pydantic_v1 import BaseModel, root_validator
from langchain_core.utils import get_from_dict_or_env
from langchain_core.utils import from_env
from pydantic import BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
from langchain_box.utilities import BoxAuth, _BoxAPIWrapper
@@ -148,7 +149,9 @@ class BoxLoader(BaseLoader, BaseModel):
"""
box_developer_token: Optional[str] = None
box_developer_token: Optional[str] = Field(
default_factory=from_env("BOX_DEVELOPER_TOKEN", default=None)
)
"""String containing the Box Developer Token generated in the developer console"""
box_auth: Optional[BoxAuth] = None
@@ -172,52 +175,47 @@ class BoxLoader(BaseLoader, BaseModel):
_box: Optional[_BoxAPIWrapper]
class Config:
arbitrary_types_allowed = True
extra = "allow"
use_enum_values = True
model_config = ConfigDict(
arbitrary_types_allowed=True,
extra="allow",
use_enum_values=True,
)
@root_validator(allow_reuse=True)
def validate_box_loader_inputs(cls, values: Dict[str, Any]) -> Dict[str, Any]:
@model_validator(mode="after")
def validate_box_loader_inputs(self) -> Self:
_box = None
"""Validate that has either box_file_ids or box_folder_id."""
if not values.get("box_file_ids") and not values.get("box_folder_id"):
if not self.box_file_ids and not self.box_folder_id:
raise ValueError("You must provide box_file_ids or box_folder_id.")
"""Validate that we don't have both box_file_ids and box_folder_id."""
if values.get("box_file_ids") and values.get("box_folder_id"):
if self.box_file_ids and self.box_folder_id:
raise ValueError(
"You must provide either box_file_ids or box_folder_id, not both."
)
"""Validate that we have either a box_developer_token or box_auth."""
if not values.get("box_auth"):
if not get_from_dict_or_env(
values, "box_developer_token", "BOX_DEVELOPER_TOKEN"
):
if not self.box_auth:
if not self.box_developer_token:
raise ValueError(
"you must provide box_developer_token or a box_auth "
"generated with langchain_box.utilities.BoxAuth"
)
else:
token = get_from_dict_or_env(
values, "box_developer_token", "BOX_DEVELOPER_TOKEN"
)
_box = _BoxAPIWrapper( # type: ignore[call-arg]
box_developer_token=token,
character_limit=values.get("character_limit"),
box_developer_token=self.box_developer_token,
character_limit=self.character_limit,
)
else:
_box = _BoxAPIWrapper( # type: ignore[call-arg]
box_auth=values.get("box_auth"),
character_limit=values.get("character_limit"),
box_auth=self.box_auth,
character_limit=self.character_limit,
)
values["_box"] = _box
self._box = _box
return values
return self
def _get_files_from_folder(self, folder_id): # type: ignore[no-untyped-def]
folder_content = self.box.get_folder_items(folder_id)

View File

@@ -1,9 +1,10 @@
from typing import Any, Dict, List, Optional
from typing import List, Optional
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.pydantic_v1 import root_validator
from langchain_core.retrievers import BaseRetriever
from pydantic import ConfigDict, model_validator
from typing_extensions import Self
from langchain_box.utilities import BoxAuth, _BoxAPIWrapper
@@ -129,30 +130,31 @@ class BoxRetriever(BaseRetriever):
_box: Optional[_BoxAPIWrapper]
class Config:
arbitrary_types_allowed = True
extra = "allow"
model_config = ConfigDict(
arbitrary_types_allowed=True,
extra="allow",
)
@root_validator(allow_reuse=True)
def validate_box_loader_inputs(cls, values: Dict[str, Any]) -> Dict[str, Any]:
@model_validator(mode="after")
def validate_box_loader_inputs(self) -> Self:
_box = None
"""Validate that we have either a box_developer_token or box_auth."""
if not values.get("box_auth") and not values.get("box_developer_token"):
if not self.box_auth and not self.box_developer_token:
raise ValueError(
"you must provide box_developer_token or a box_auth "
"generated with langchain_box.utilities.BoxAuth"
)
_box = _BoxAPIWrapper( # type: ignore[call-arg]
box_developer_token=values.get("box_developer_token"),
box_auth=values.get("box_auth"),
character_limit=values.get("character_limit"),
box_developer_token=self.box_developer_token,
box_auth=self.box_auth,
character_limit=self.character_limit,
)
values["_box"] = _box
self._box = _box
return values
return self
def _get_relevant_documents(
self, query: str, *, run_manager: CallbackManagerForRetrieverRun

View File

@@ -6,8 +6,9 @@ from typing import Any, Dict, List, Optional
import box_sdk_gen # type: ignore
import requests
from langchain_core.documents import Document
from langchain_core.pydantic_v1 import BaseModel, root_validator
from langchain_core.utils import get_from_dict_or_env
from langchain_core.utils import from_env
from pydantic import BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
class DocumentFiles(Enum):
@@ -312,17 +313,25 @@ class BoxAuth(BaseModel):
"""``langchain_box.utilities.BoxAuthType``. Enum describing how to
authenticate against Box"""
box_developer_token: Optional[str] = None
box_developer_token: Optional[str] = Field(
default_factory=from_env("BOX_DEVELOPER_TOKEN", default=None)
)
""" If using ``BoxAuthType.TOKEN``, provide your token here"""
box_jwt_path: Optional[str] = None
box_jwt_path: Optional[str] = Field(
default_factory=from_env("BOX_JWT_PATH", default=None)
)
"""If using ``BoxAuthType.JWT``, provide local path to your
JWT configuration file"""
box_client_id: Optional[str] = None
box_client_id: Optional[str] = Field(
default_factory=from_env("BOX_CLIENT_ID", default=None)
)
"""If using ``BoxAuthType.CCG``, provide your app's client ID"""
box_client_secret: Optional[str] = None
box_client_secret: Optional[str] = Field(
default_factory=from_env("BOX_CLIENT_SECRET", default=None)
)
"""If using ``BoxAuthType.CCG``, provide your app's client secret"""
box_enterprise_id: Optional[str] = None
@@ -336,138 +345,123 @@ class BoxAuth(BaseModel):
_box_client: Optional[box_sdk_gen.BoxClient] = None
_custom_header: Dict = dict({"x-box-ai-library": "langchain"})
class Config:
arbitrary_types_allowed = True
use_enum_values = True
extra = "allow"
model_config = ConfigDict(
arbitrary_types_allowed=True,
use_enum_values=True,
extra="allow",
)
@root_validator()
def validate_box_auth_inputs(cls, values: Dict[str, Any]) -> Dict[str, Any]:
@model_validator(mode="after")
def validate_box_auth_inputs(self) -> Self:
"""Validate auth_type is set"""
if not values.get("auth_type"):
if not self.auth_type:
raise ValueError("Auth type must be set.")
"""Validate that TOKEN auth type provides box_developer_token."""
if values.get("auth_type") == "token":
if not get_from_dict_or_env(
values, "box_developer_token", "BOX_DEVELOPER_TOKEN"
):
raise ValueError(
f"{values.get('auth_type')} requires box_developer_token to be set"
)
if self.auth_type == "token" and not self.box_developer_token:
raise ValueError(f"{self.auth_type} requires box_developer_token to be set")
"""Validate that JWT auth type provides box_jwt_path."""
if values.get("auth_type") == "jwt":
if not get_from_dict_or_env(values, "box_jwt_path", "BOX_JWT_PATH"):
raise ValueError(
f"{values.get('auth_type')} requires box_jwt_path to be set"
)
if self.auth_type == "jwt" and not self.box_jwt_path:
raise ValueError(f"{self.auth_type} requires box_jwt_path to be set")
"""Validate that CCG auth type provides box_client_id and
box_client_secret and either box_enterprise_id or box_user_id."""
if values.get("auth_type") == "ccg":
if self.auth_type == "ccg":
if (
not get_from_dict_or_env(values, "box_client_id", "BOX_CLIENT_ID")
or not get_from_dict_or_env(
values, "box_client_secret", "BOX_CLIENT_SECRET"
)
or (
not values.get("box_enterprise_id")
and not values.get("box_user_id")
)
not self.box_client_id
or not self.box_client_secret
or (not self.box_enterprise_id and not self.box_user_id)
):
raise ValueError(
f"{values.get('auth_type')} requires box_client_id, \
box_client_secret, and box_enterprise_id."
f"{self.auth_type} requires box_client_id, \
box_client_secret, and box_enterprise_id/box_user_id."
)
return values
return self
def _authorize(self) -> None:
match self.auth_type:
case "token":
try:
auth = box_sdk_gen.BoxDeveloperTokenAuth(
token=self.box_developer_token
)
self._box_client = box_sdk_gen.BoxClient(
auth=auth
).with_extra_headers(extra_headers=self._custom_header)
except box_sdk_gen.BoxSDKError as bse:
raise RuntimeError(
f"Error getting client from developer token: {bse.message}"
)
except Exception as ex:
raise ValueError(
f"Invalid Box developer token. Please verify your \
token and try again.\n{ex}"
) from ex
case "jwt":
try:
jwt_config = box_sdk_gen.JWTConfig.from_config_file(
config_file_path=self.box_jwt_path
)
auth = box_sdk_gen.BoxJWTAuth(config=jwt_config)
self._box_client = box_sdk_gen.BoxClient(
auth=auth
).with_extra_headers(extra_headers=self._custom_header)
if self.box_user_id is not None:
user_auth = auth.with_user_subject(self.box_user_id)
self._box_client = box_sdk_gen.BoxClient(
auth=user_auth
).with_extra_headers(extra_headers=self._custom_header)
except box_sdk_gen.BoxSDKError as bse:
raise RuntimeError(
f"Error getting client from jwt token: {bse.message}"
)
except Exception as ex:
raise ValueError(
"Error authenticating. Please verify your JWT config \
and try again."
) from ex
case "ccg":
try:
if self.box_user_id is not None:
ccg_config = box_sdk_gen.CCGConfig(
client_id=self.box_client_id,
client_secret=self.box_client_secret,
user_id=self.box_user_id,
)
else:
ccg_config = box_sdk_gen.CCGConfig(
client_id=self.box_client_id,
client_secret=self.box_client_secret,
enterprise_id=self.box_enterprise_id,
)
auth = box_sdk_gen.BoxCCGAuth(config=ccg_config)
self._box_client = box_sdk_gen.BoxClient(
auth=auth
).with_extra_headers(extra_headers=self._custom_header)
except box_sdk_gen.BoxSDKError as bse:
raise RuntimeError(
f"Error getting client from ccg token: {bse.message}"
)
except Exception as ex:
raise ValueError(
"Error authenticating. Please verify you are providing a \
valid client id, secret and either a valid user ID or \
enterprise ID."
) from ex
case _:
raise ValueError(
f"{self.auth_type} is not a valid auth_type. Value must be \
TOKEN, CCG, or JWT."
if self.auth_type == "token":
try:
auth = box_sdk_gen.BoxDeveloperTokenAuth(token=self.box_developer_token)
self._box_client = box_sdk_gen.BoxClient(auth=auth).with_extra_headers(
extra_headers=self._custom_header
)
except box_sdk_gen.BoxSDKError as bse:
raise RuntimeError(
f"Error getting client from developer token: {bse.message}"
)
except Exception as ex:
raise ValueError(
f"Invalid Box developer token. Please verify your \
token and try again.\n{ex}"
) from ex
elif self.auth_type == "jwt":
try:
jwt_config = box_sdk_gen.JWTConfig.from_config_file(
config_file_path=self.box_jwt_path
)
auth = box_sdk_gen.BoxJWTAuth(config=jwt_config)
self._box_client = box_sdk_gen.BoxClient(auth=auth).with_extra_headers(
extra_headers=self._custom_header
)
if self.box_user_id is not None:
user_auth = auth.with_user_subject(self.box_user_id)
self._box_client = box_sdk_gen.BoxClient(
auth=user_auth
).with_extra_headers(extra_headers=self._custom_header)
except box_sdk_gen.BoxSDKError as bse:
raise RuntimeError(
f"Error getting client from jwt token: {bse.message}"
)
except Exception as ex:
raise ValueError(
"Error authenticating. Please verify your JWT config \
and try again."
) from ex
elif self.auth_type == "ccg":
try:
if self.box_user_id is not None:
ccg_config = box_sdk_gen.CCGConfig(
client_id=self.box_client_id,
client_secret=self.box_client_secret,
user_id=self.box_user_id,
)
else:
ccg_config = box_sdk_gen.CCGConfig(
client_id=self.box_client_id,
client_secret=self.box_client_secret,
enterprise_id=self.box_enterprise_id,
)
auth = box_sdk_gen.BoxCCGAuth(config=ccg_config)
self._box_client = box_sdk_gen.BoxClient(auth=auth).with_extra_headers(
extra_headers=self._custom_header
)
except box_sdk_gen.BoxSDKError as bse:
raise RuntimeError(
f"Error getting client from ccg token: {bse.message}"
)
except Exception as ex:
raise ValueError(
"Error authenticating. Please verify you are providing a \
valid client id, secret and either a valid user ID or \
enterprise ID."
) from ex
else:
raise ValueError(
f"{self.auth_type} is not a valid auth_type. Value must be \
TOKEN, CCG, or JWT."
)
def get_client(self) -> box_sdk_gen.BoxClient:
"""Instantiate the Box SDK."""
if self._box_client is None:
@@ -479,7 +473,9 @@ class BoxAuth(BaseModel):
class _BoxAPIWrapper(BaseModel):
"""Wrapper for Box API."""
box_developer_token: Optional[str] = None
box_developer_token: Optional[str] = Field(
default_factory=from_env("BOX_DEVELOPER_TOKEN", default=None)
)
"""String containing the Box Developer Token generated in the developer console"""
box_auth: Optional[BoxAuth] = None
@@ -491,28 +487,27 @@ class _BoxAPIWrapper(BaseModel):
_box: Optional[box_sdk_gen.BoxClient]
class Config:
arbitrary_types_allowed = True
use_enum_values = True
extra = "allow"
model_config = ConfigDict(
arbitrary_types_allowed=True,
use_enum_values=True,
extra="allow",
)
@root_validator(allow_reuse=True)
def validate_box_api_inputs(cls, values: Dict[str, Any]) -> Dict[str, Any]:
values["_box"] = None
@model_validator(mode="after")
def validate_box_api_inputs(self) -> Self:
self._box = None
"""Validate that TOKEN auth type provides box_developer_token."""
if not values.get("box_auth"):
if not get_from_dict_or_env(
values, "box_developer_token", "BOX_DEVELOPER_TOKEN"
):
if not self.box_auth:
if not self.box_developer_token:
raise ValueError(
"You must configure either box_developer_token of box_auth"
)
else:
box_auth = values.get("box_auth")
values["_box"] = box_auth.get_client() # type: ignore[union-attr]
box_auth = self.box_auth
self._box = box_auth.get_client() # type: ignore[union-attr]
return values
return self
def get_box_client(self) -> box_sdk_gen.BoxClient:
box_auth = BoxAuth(

View File

@@ -2,7 +2,7 @@ from unittest.mock import Mock
import pytest
from langchain_core.documents import Document
from pydantic.v1.error_wrappers import ValidationError
from pydantic.error_wrappers import ValidationError
from pytest_mock import MockerFixture
from langchain_box.utilities import BoxAuth, BoxAuthType, _BoxAPIWrapper