diff --git a/libs/partners/box/langchain_box/document_loaders/box.py b/libs/partners/box/langchain_box/document_loaders/box.py index 742f07c361e..32dac91fbbd 100644 --- a/libs/partners/box/langchain_box/document_loaders/box.py +++ b/libs/partners/box/langchain_box/document_loaders/box.py @@ -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) diff --git a/libs/partners/box/langchain_box/retrievers/box.py b/libs/partners/box/langchain_box/retrievers/box.py index de014102769..7715fbd42e2 100644 --- a/libs/partners/box/langchain_box/retrievers/box.py +++ b/libs/partners/box/langchain_box/retrievers/box.py @@ -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 diff --git a/libs/partners/box/langchain_box/utilities/box.py b/libs/partners/box/langchain_box/utilities/box.py index ade944bdc0c..4fef37a41c9 100644 --- a/libs/partners/box/langchain_box/utilities/box.py +++ b/libs/partners/box/langchain_box/utilities/box.py @@ -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( diff --git a/libs/partners/box/tests/unit_tests/utilities/test_box_util.py b/libs/partners/box/tests/unit_tests/utilities/test_box_util.py index 1eabbdf7592..6f6ef907f9a 100644 --- a/libs/partners/box/tests/unit_tests/utilities/test_box_util.py +++ b/libs/partners/box/tests/unit_tests/utilities/test_box_util.py @@ -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