diff --git a/backend/danswer/file_processing/extract_file_text.py b/backend/danswer/file_processing/extract_file_text.py index eeb7b5ae8ba..a9110290e2c 100644 --- a/backend/danswer/file_processing/extract_file_text.py +++ b/backend/danswer/file_processing/extract_file_text.py @@ -283,7 +283,7 @@ def extract_file_text( file_name: str | None, file: IO[Any], break_on_unprocessable: bool = True, -) -> str: +) -> tuple[str, dict]: extension_to_function: dict[str, Callable[[IO[Any]], str]] = { ".pdf": read_pdf_file, ".docx": docx_to_text, diff --git a/backend/danswer/main.py b/backend/danswer/main.py index 6652e5d3c39..8b660f5a9e3 100644 --- a/backend/danswer/main.py +++ b/backend/danswer/main.py @@ -109,6 +109,7 @@ from danswer.tools.built_in_tools import auto_add_search_tool_to_personas from danswer.tools.built_in_tools import load_builtin_tools from danswer.tools.built_in_tools import refresh_built_in_tools_cache +from danswer.utils.gpu_utils import gpu_status_request from danswer.utils.logger import setup_logger from danswer.utils.telemetry import optional_telemetry from danswer.utils.telemetry import RecordType @@ -192,6 +193,34 @@ def setup_postgres(db_session: Session) -> None: auto_add_search_tool_to_personas(db_session) +def update_default_multipass_indexing(db_session: Session) -> None: + docs_exist = check_docs_exist(db_session) + connectors_exist = check_connectors_exist(db_session) + logger.debug(f"Docs exist: {docs_exist}, Connectors exist: {connectors_exist}") + + if not docs_exist and not connectors_exist: + logger.info( + "No existing docs or connectors found. Checking GPU availability for multipass indexing." + ) + gpu_available = gpu_status_request() + logger.info(f"GPU availability: {gpu_available}") + + current_settings = get_current_search_settings(db_session) + + logger.notice(f"Updating multipass indexing setting to: {gpu_available}") + updated_settings = SavedSearchSettings.from_db_model(current_settings) + # Enable multipass indexing if GPU is available or if using a cloud provider + updated_settings.multipass_indexing = ( + gpu_available or current_settings.cloud_provider is not None + ) + update_current_search_settings(db_session, updated_settings) + + else: + logger.debug( + "Existing docs or connectors found. Skipping multipass indexing update." + ) + + def translate_saved_search_settings(db_session: Session) -> None: kv_store = get_dynamic_config_store() @@ -371,6 +400,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator: ), ) + # update multipass indexing setting based on GPU availability + update_default_multipass_indexing(db_session) + optional_telemetry(record_type=RecordType.VERSION, data={"version": __version__}) yield diff --git a/backend/danswer/utils/gpu_utils.py b/backend/danswer/utils/gpu_utils.py new file mode 100644 index 00000000000..70a3dbc2c95 --- /dev/null +++ b/backend/danswer/utils/gpu_utils.py @@ -0,0 +1,30 @@ +import requests +from retry import retry + +from danswer.utils.logger import setup_logger +from shared_configs.configs import INDEXING_MODEL_SERVER_HOST +from shared_configs.configs import INDEXING_MODEL_SERVER_PORT +from shared_configs.configs import MODEL_SERVER_HOST +from shared_configs.configs import MODEL_SERVER_PORT + +logger = setup_logger() + + +@retry(tries=5, delay=5) +def gpu_status_request(indexing: bool = True) -> bool: + if indexing: + model_server_url = f"{INDEXING_MODEL_SERVER_HOST}:{INDEXING_MODEL_SERVER_PORT}" + else: + model_server_url = f"{MODEL_SERVER_HOST}:{MODEL_SERVER_PORT}" + + if "http" not in model_server_url: + model_server_url = f"http://{model_server_url}" + + try: + response = requests.get(f"{model_server_url}/api/gpu-status", timeout=10) + response.raise_for_status() + gpu_status = response.json() + return gpu_status["gpu_available"] + except requests.RequestException as e: + logger.error(f"Error: Unable to fetch GPU status. Error: {str(e)}") + raise # Re-raise exception to trigger a retry diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx index 1c9441d079b..43503a3af14 100644 --- a/web/src/app/chat/ChatPage.tsx +++ b/web/src/app/chat/ChatPage.tsx @@ -2095,6 +2095,7 @@ export function ChatPage({ )} )} + {/* Some padding at the bottom so the search bar has space at the bottom to not cover the last message*/}