diff --git a/source/docq/data_source/googledrive.py b/source/docq/data_source/googledrive.py index 2299d974..96968c67 100644 --- a/source/docq/data_source/googledrive.py +++ b/source/docq/data_source/googledrive.py @@ -23,13 +23,13 @@ def get_config_keys(self: Self) -> List[ConfigKey]: """Get the config keys for google drive.""" return [ ConfigKey( - f"{FileStorageServiceKeys.GOOGLE_DRIVE}-credential", + f"{FileStorageServiceKeys.GOOGLE_DRIVE.name}-credential", "Credential", is_secret=True, ref_link="https://docqai.github.io/docq/user-guide/config-spaces/#data-source-google-drive", ), ConfigKey( - f"{FileStorageServiceKeys.GOOGLE_DRIVE}-root_path", + f"{FileStorageServiceKeys.GOOGLE_DRIVE.name}-root_path", "Select a folder", ref_link="https://docqai.github.io/docq/user-guide/config-spaces/#data-source-google-drive", ), @@ -49,11 +49,11 @@ def lambda_metadata(x: str) -> dict: str(DocumentMetadata.INDEXED_ON.name).lower(): datetime.timestamp(datetime.now().utcnow()), } - root_path = configs[f"{FileStorageServiceKeys.GOOGLE_DRIVE}-root_path"] + root_path = configs[f"{FileStorageServiceKeys.GOOGLE_DRIVE.name}-root_path"] options = { "root": root_path["name"], - "access_token": json.dumps(configs[f"{FileStorageServiceKeys.GOOGLE_DRIVE}-credential"]), + "access_token": json.dumps(configs[f"{FileStorageServiceKeys.GOOGLE_DRIVE.name}-credential"]), } try: @@ -67,7 +67,7 @@ def lambda_metadata(x: str) -> dict: loader = GoogleDriveReader( file_metadata=lambda_metadata, root=root_path["name"], - access_token=configs[f"{FileStorageServiceKeys.GOOGLE_DRIVE}-credential"], + access_token=configs[f"{FileStorageServiceKeys.GOOGLE_DRIVE.name}-credential"], selected_folder_id=root_path["id"] ) diff --git a/source/docq/data_source/main.py b/source/docq/data_source/main.py index b24f8f6f..cdc0922a 100644 --- a/source/docq/data_source/main.py +++ b/source/docq/data_source/main.py @@ -139,21 +139,21 @@ def __init__(self: Self, credential: Callable, root_path: Callable) -> None: """Initialize the handler.""" self.__handlers = {"credential": credential, "root_path": root_path} - def __call__(self: Self, key: str, *args: Any, **kwargs: Any) -> Callable | None: + def __call__(self: Self, operation: str, *args: Any, **kwargs: Any) -> Callable | None: """Call the handler.""" - if key in self.__supported_keys: - return self.__handlers[key](*args, **kwargs) + if operation in self.__supported_keys: + return self.__handlers[operation](*args, **kwargs) class FileStorageServiceHandlers: """File storage service handlers.""" - __handlers: dict[FileStorageServiceKeys, SetFSHandler] = {} + __handlers: dict[str, SetFSHandler] = {} - def __init__(self: Self, handlers: dict[FileStorageServiceKeys, SetFSHandler]) -> None: + def __init__(self: Self, handlers: dict[str, SetFSHandler]) -> None: """Initialize the handlers.""" for key, handler in handlers.items(): - if key in FileStorageServiceKeys: + if key in FileStorageServiceKeys.__members__: self.__handlers[key] = handler def __getitem__(self: Self, key: str) -> Callable | None: diff --git a/web/utils/layout.py b/web/utils/layout.py index 3f56f822..34ca9a8c 100644 --- a/web/utils/layout.py +++ b/web/utils/layout.py @@ -8,9 +8,10 @@ import docq import streamlit as st -from docq import config, services, setup +from docq import services, setup from docq.access_control.main import SpaceAccessType from docq.config import FeatureType, LogType, SpaceType, SystemSettingsKey +from docq.data_source.main import FileStorageServiceHandlers, FileStorageServiceKeys, SetFSHandler from docq.domain import ConfigKey, DocumentListItem, FeatureKey, SpaceKey from docq.model_selection.main import ( ModelUsageSettingsCollection, @@ -799,6 +800,28 @@ def system_settings_ui() -> None: st.divider() +def render_dynamic_options(func: Callable) -> Callable: + """Renders the dynamic options for a config key.""" + + def _drop_down_selector(*args: Any, **kwargs: Any) -> None: + """Streamlit selectbox for rendering dynamic options.""" + options, fmt, disabled = func(*args, **kwargs) + configs: dict = kwargs.get("configs", {}) + configkey: ConfigKey = kwargs.get("configkey", None) + key: str = kwargs.get("key", None) + selected = configs.get(configkey.key) if configs else None + st.selectbox( + f"{configkey.name}{'' if configkey.is_optional else ' *'}", + options=options, + key=key, + format_func=fmt, + help=configkey.ref_link, + disabled=disabled, + index=options.index(selected) if bool(selected and options) else 0, + ) + return _drop_down_selector + + def _get_create_space_config_input_values() -> str: """Get values for space creation from session state.""" space_name = st.session_state.get("create_space_name", "") @@ -812,7 +835,7 @@ def _get_random_key(prefix: str) -> str: return prefix + str(random.randint(0, 1000000)) # noqa E501 -def _render_gdrive_credential_request(configkey: ConfigKey, key: str, configs: Optional[dict]) -> None: +def render_gdrive_credential_request(configkey: ConfigKey, key: str, configs: Optional[dict]) -> None: """Renders the Google Drive credential request input field with an action burron.""" creds = configs.get(configkey.key) if configs else st.session_state.get(key, None) creds = services.google_drive.validate_credentials(creds) @@ -869,7 +892,9 @@ def _render_gdrive_credential_request(configkey: ConfigKey, key: str, configs: O st.stop() -def _list_gdrive_folders(**kwargs: Any) -> tuple[list, Callable, bool]: +@render_dynamic_options +def list_gdrive_folders(**kwargs: Any) -> tuple[list, Callable, bool]: + """List Google Drive folders.""" configs: dict = kwargs.get("configs", {}) configkey: ConfigKey = kwargs.get("configkey", None) saved_settings = configs.get(configkey.key) if configs else None @@ -878,7 +903,7 @@ def _list_gdrive_folders(**kwargs: Any) -> tuple[list, Callable, bool]: return [saved_settings], lambda x: x["name"], True # Make API call to Google Drive to get folders if creating space, i.e No saved settings. - credential_key = config.ConfigKeyHandlers.GET_GDRIVE_CREDENTIAL.name + credential_key = f"{FileStorageServiceKeys.GOOGLE_DRIVE.name}-credential" __creds = configs.get(credential_key) if configs else st.session_state.get(credential_key, None) creds = services.google_drive.validate_credentials(__creds) @@ -886,41 +911,32 @@ def _list_gdrive_folders(**kwargs: Any) -> tuple[list, Callable, bool]: def _render_space_data_source_config_input_fields(data_source: Tuple, prefix: str, configs: Optional[dict] = None) -> None: - config_key_handlers = { - config.ConfigKeyHandlers.GET_GDRIVE_CREDENTIAL.name: _render_gdrive_credential_request, - } - get_config_key_options = { - config.ConfigKeyOptions.GET_GDRIVE_OPTIONS.name: _list_gdrive_folders, - } + fs_handlers = FileStorageServiceHandlers( + { + FileStorageServiceKeys.GOOGLE_DRIVE.name: SetFSHandler( + credential=render_gdrive_credential_request, + root_path=list_gdrive_folders, + ) + } + ) config_key_list: List[ConfigKey] = data_source[2] - for key in config_key_list: - _input_key = prefix + "ds_config_" + key.key - if key.input_element == "credential_request": - config_key_handlers[key.key](key, _input_key, configs) - - elif key.input_element == "selectbox": - options, fmt, disabled = get_config_key_options[key.key](configkey=key, key=_input_key, configs=configs) - selected = configs.get(key.key) if configs else None - st.selectbox( - f"{key.name}{'' if key.is_optional else ' *'}", - options=options, - key=_input_key, - format_func=fmt, - help=key.ref_link, - disabled=disabled, - index=options.index(selected) if bool(selected and options) else 0, - ) + for configkey in config_key_list: + _input_key = prefix + "ds_config_" + configkey.key + service, operation = configkey.key.split("-") + handler = fs_handlers[service] + if handler is not None: + handler(operation, configkey=configkey, key=_input_key, configs=configs) else: - input_type = "password" if key.is_secret else "default" + input_type = "password" if configkey.is_secret else "default" st.text_input( - f"{key.name}{'' if key.is_optional else ' *'}", - value=configs.get(key.key) if configs else "", + f"{configkey.name}{'' if configkey.is_optional else ' *'}", + value=configs.get(configkey.key) if configs else "", key=_input_key, type=input_type, - help=key.ref_link, + help=configkey.ref_link, autocomplete="off", # disable autofill by password manager etc. )