diff --git a/backend/ee/onyx/background/celery/tasks/query_history/tasks.py b/backend/ee/onyx/background/celery/tasks/query_history/tasks.py index 50841edfa3..f241a3d78b 100644 --- a/backend/ee/onyx/background/celery/tasks/query_history/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/query_history/tasks.py @@ -4,10 +4,6 @@ from datetime import datetime from celery import Task, shared_task -from ee.onyx.server.query_history.api import ( - ONYX_ANONYMIZED_EMAIL, - fetch_and_process_chat_session_history, -) from ee.onyx.server.query_history.models import QuestionAnswerPairSnapshot from onyx.background.task_utils import construct_query_history_report_name from onyx.configs.app_configs import JOB_TIMEOUT @@ -47,6 +43,11 @@ def export_query_history_task( # Need to include the tenant_id since the TenantAwareTask needs this tenant_id: str, # noqa: ARG001 ) -> None: + from ee.onyx.server.query_history.api import ( + ONYX_ANONYMIZED_EMAIL, + fetch_and_process_chat_session_history, + ) + if not self.request.id: raise RuntimeError("No task id defined for this task; cannot identify it") diff --git a/backend/ee/onyx/background/celery/tasks/sso_domain_revalidation/tasks.py b/backend/ee/onyx/background/celery/tasks/sso_domain_revalidation/tasks.py index b4d97af83a..a6cb254862 100644 --- a/backend/ee/onyx/background/celery/tasks/sso_domain_revalidation/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/sso_domain_revalidation/tasks.py @@ -7,7 +7,6 @@ domain whose DNS TXT proof no longer resolves. from celery import shared_task -from ee.onyx.auth.sso_domain_verification import revalidate_tenant_domains from ee.onyx.db.tenant_sso_domain import reproject_tenant_login_domains from onyx.configs.app_configs import JOB_TIMEOUT from onyx.configs.constants import OnyxCeleryTask @@ -26,6 +25,8 @@ def revalidate_sso_domains_task(*, tenant_id: str) -> None: """Fanned out per tenant by cloud_beat_task_generator. The re-projection is isolated so its failure can't skip the DNS re-check, which is the part that drops routing for a domain whose proof is gone.""" + from ee.onyx.auth.sso_domain_verification import revalidate_tenant_domains + try: reproject_tenant_login_domains(tenant_id) except Exception: diff --git a/backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py b/backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py index 727f84855e..defd756d3c 100644 --- a/backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py @@ -9,12 +9,6 @@ import uuid from celery import Task, shared_task from redis.lock import Lock as RedisLock -from ee.onyx.server.tenants.provisioning import setup_tenant -from ee.onyx.server.tenants.schema_management import ( - create_schema_if_not_exists, - get_current_alembic_version, - run_alembic_migrations, -) from onyx.background.celery.apps.app_base import task_logger from onyx.configs.app_configs import TARGET_AVAILABLE_TENANTS from onyx.configs.constants import ( @@ -135,6 +129,11 @@ def _migrate_stale_pool_tenants() -> None: tenants are always current so that signup doesn't hit schema mismatches (e.g. missing columns added after the tenant was pre-provisioned). """ + from ee.onyx.server.tenants.schema_management import ( + get_current_alembic_version, + run_alembic_migrations, + ) + with get_session_with_shared_schema() as db_session: pool_tenants = db_session.query(AvailableTenant).all() tenant_ids = [t.tenant_id for t in pool_tenants] @@ -178,6 +177,13 @@ def pre_provision_tenant() -> bool: """ # The MULTI_TENANT check is now done at the caller level (check_available_tenants) # rather than inside this function + # Imported here: provisioning reaches every tool implementation (~75 MB) and + # only multi-tenant deployments run this. + from ee.onyx.server.tenants.provisioning import setup_tenant + from ee.onyx.server.tenants.schema_management import ( + create_schema_if_not_exists, + get_current_alembic_version, + ) r = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID) lock_provision: RedisLock = r.lock( diff --git a/backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py b/backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py index e429babec8..4807c2a4f2 100644 --- a/backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py @@ -3,7 +3,6 @@ from uuid import UUID from celery import Task, shared_task -from ee.onyx.server.reporting.usage_export_generation import create_new_usage_report from onyx.configs.app_configs import JOB_TIMEOUT from onyx.configs.constants import OnyxCeleryTask from onyx.db.engine.sql_engine import get_session_with_current_tenant @@ -29,6 +28,8 @@ def generate_usage_report_task( report_id: str | None = None, ) -> None: """User-initiated usage report generation task""" + from ee.onyx.server.reporting.usage_export_generation import create_new_usage_report + # Parse period if provided period = None if period_from and period_to: diff --git a/backend/ee/onyx/external_permissions/sync_params.py b/backend/ee/onyx/external_permissions/sync_params.py index b7466794d7..124c497ce6 100644 --- a/backend/ee/onyx/external_permissions/sync_params.py +++ b/backend/ee/onyx/external_permissions/sync_params.py @@ -1,4 +1,4 @@ -from collections.abc import Generator +from collections.abc import Callable, Generator from typing import TYPE_CHECKING, Optional from pydantic import BaseModel @@ -21,19 +21,6 @@ from ee.onyx.configs.app_configs import ( SLACK_PERMISSION_DOC_SYNC_FREQUENCY, TEAMS_PERMISSION_DOC_SYNC_FREQUENCY, ) -from ee.onyx.external_permissions.box.doc_sync import box_doc_sync -from ee.onyx.external_permissions.box.group_sync import box_group_sync -from ee.onyx.external_permissions.canvas.doc_sync import canvas_doc_sync -from ee.onyx.external_permissions.canvas.group_sync import canvas_group_sync -from ee.onyx.external_permissions.confluence.doc_sync import confluence_doc_sync -from ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync -from ee.onyx.external_permissions.github.doc_sync import github_doc_sync -from ee.onyx.external_permissions.github.group_sync import github_group_sync -from ee.onyx.external_permissions.gmail.doc_sync import gmail_doc_sync -from ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync -from ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync -from ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync -from ee.onyx.external_permissions.jira.group_sync import jira_group_sync from ee.onyx.external_permissions.perm_sync_types import ( CensoringFuncType, DocSyncFuncType, @@ -41,21 +28,157 @@ from ee.onyx.external_permissions.perm_sync_types import ( FetchAllDocumentsIdsFunction, GroupSyncFuncType, ) -from ee.onyx.external_permissions.salesforce.postprocessing import ( - censor_salesforce_chunks, -) -from ee.onyx.external_permissions.sharepoint.doc_sync import sharepoint_doc_sync -from ee.onyx.external_permissions.sharepoint.group_sync import sharepoint_group_sync -from ee.onyx.external_permissions.slack.doc_sync import slack_doc_sync -from ee.onyx.external_permissions.teams.doc_sync import teams_doc_sync from onyx.configs.constants import DocumentSource if TYPE_CHECKING: - from onyx.access.models import DocExternalAccess # noqa + from ee.onyx.db.external_perm import ExternalUserGroup # noqa + from onyx.access.models import DocExternalAccess, ElementExternalAccess # noqa + from onyx.context.search.models import InferenceChunk # noqa from onyx.db.models import ConnectorCredentialPair # noqa from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface # noqa +# Sync implementations import their connector SDKs (box_sdk_gen, office365, +# google, github, ...), which cost ~60 MB. Load them on first call so workers +# that only read this config stay small. +def _lazy_doc_sync(load: Callable[[], DocSyncFuncType]) -> DocSyncFuncType: + def run( + cc_pair: "ConnectorCredentialPair", + fetch_all_docs_fn: FetchAllDocumentsFunction, + fetch_all_docs_ids_fn: FetchAllDocumentsIdsFunction, + callback: Optional["IndexingHeartbeatInterface"], + ) -> Generator["ElementExternalAccess", None, None]: + return load()(cc_pair, fetch_all_docs_fn, fetch_all_docs_ids_fn, callback) + + return run + + +def _lazy_group_sync(load: Callable[[], GroupSyncFuncType]) -> GroupSyncFuncType: + def run( + tenant_id: str, cc_pair: "ConnectorCredentialPair" + ) -> Generator["ExternalUserGroup", None, None]: + return load()(tenant_id, cc_pair) + + return run + + +def _lazy_censoring(load: Callable[[], CensoringFuncType]) -> CensoringFuncType: + def run(chunks: list["InferenceChunk"], user_email: str) -> list["InferenceChunk"]: + return load()(chunks, user_email) + + return run + + +def _load_box_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.box.doc_sync import box_doc_sync + + return box_doc_sync + + +def _load_box_group_sync() -> GroupSyncFuncType: + from ee.onyx.external_permissions.box.group_sync import box_group_sync + + return box_group_sync + + +def _load_canvas_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.canvas.doc_sync import canvas_doc_sync + + return canvas_doc_sync + + +def _load_canvas_group_sync() -> GroupSyncFuncType: + from ee.onyx.external_permissions.canvas.group_sync import canvas_group_sync + + return canvas_group_sync + + +def _load_confluence_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.confluence.doc_sync import confluence_doc_sync + + return confluence_doc_sync + + +def _load_confluence_group_sync() -> GroupSyncFuncType: + from ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync + + return confluence_group_sync + + +def _load_github_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.github.doc_sync import github_doc_sync + + return github_doc_sync + + +def _load_github_group_sync() -> GroupSyncFuncType: + from ee.onyx.external_permissions.github.group_sync import github_group_sync + + return github_group_sync + + +def _load_gmail_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.gmail.doc_sync import gmail_doc_sync + + return gmail_doc_sync + + +def _load_gdrive_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync + + return gdrive_doc_sync + + +def _load_gdrive_group_sync() -> GroupSyncFuncType: + from ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync + + return gdrive_group_sync + + +def _load_jira_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync + + return jira_doc_sync + + +def _load_jira_group_sync() -> GroupSyncFuncType: + from ee.onyx.external_permissions.jira.group_sync import jira_group_sync + + return jira_group_sync + + +def _load_censor_salesforce_chunks() -> CensoringFuncType: + from ee.onyx.external_permissions.salesforce.postprocessing import ( + censor_salesforce_chunks, + ) + + return censor_salesforce_chunks + + +def _load_sharepoint_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.sharepoint.doc_sync import sharepoint_doc_sync + + return sharepoint_doc_sync + + +def _load_sharepoint_group_sync() -> GroupSyncFuncType: + from ee.onyx.external_permissions.sharepoint.group_sync import sharepoint_group_sync + + return sharepoint_group_sync + + +def _load_slack_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.slack.doc_sync import slack_doc_sync + + return slack_doc_sync + + +def _load_teams_doc_sync() -> DocSyncFuncType: + from ee.onyx.external_permissions.teams.doc_sync import teams_doc_sync + + return teams_doc_sync + + class DocSyncConfig(BaseModel): doc_sync_frequency: int doc_sync_func: DocSyncFuncType @@ -96,60 +219,60 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = { DocumentSource.GOOGLE_DRIVE: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=gdrive_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_gdrive_doc_sync), initial_index_should_sync=True, ), group_sync_config=GroupSyncConfig( group_sync_frequency=GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY, - group_sync_func=gdrive_group_sync, + group_sync_func=_lazy_group_sync(_load_gdrive_group_sync), group_sync_is_cc_pair_agnostic=False, ), ), DocumentSource.CONFLUENCE: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=confluence_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_confluence_doc_sync), initial_index_should_sync=False, ), group_sync_config=GroupSyncConfig( group_sync_frequency=CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY, - group_sync_func=confluence_group_sync, + group_sync_func=_lazy_group_sync(_load_confluence_group_sync), group_sync_is_cc_pair_agnostic=True, ), ), DocumentSource.JIRA: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=JIRA_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=jira_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_jira_doc_sync), initial_index_should_sync=True, ), group_sync_config=GroupSyncConfig( group_sync_frequency=JIRA_PERMISSION_GROUP_SYNC_FREQUENCY, - group_sync_func=jira_group_sync, + group_sync_func=_lazy_group_sync(_load_jira_group_sync), group_sync_is_cc_pair_agnostic=True, ), ), DocumentSource.CANVAS: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=CANVAS_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=canvas_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_canvas_doc_sync), initial_index_should_sync=True, ), group_sync_config=GroupSyncConfig( group_sync_frequency=CANVAS_PERMISSION_GROUP_SYNC_FREQUENCY, - group_sync_func=canvas_group_sync, + group_sync_func=_lazy_group_sync(_load_canvas_group_sync), group_sync_is_cc_pair_agnostic=False, ), ), DocumentSource.BOX: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=BOX_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=box_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_box_doc_sync), initial_index_should_sync=True, ), group_sync_config=GroupSyncConfig( group_sync_frequency=BOX_PERMISSION_GROUP_SYNC_FREQUENCY, - group_sync_func=box_group_sync, + group_sync_func=_lazy_group_sync(_load_box_group_sync), group_sync_is_cc_pair_agnostic=False, ), ), @@ -158,32 +281,32 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = { DocumentSource.SLACK: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=SLACK_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=slack_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_slack_doc_sync), initial_index_should_sync=True, ), ), DocumentSource.GMAIL: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=gmail_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_gmail_doc_sync), initial_index_should_sync=False, ), ), DocumentSource.GITHUB: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=GITHUB_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=github_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_github_doc_sync), initial_index_should_sync=True, ), group_sync_config=GroupSyncConfig( group_sync_frequency=GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY, - group_sync_func=github_group_sync, + group_sync_func=_lazy_group_sync(_load_github_group_sync), group_sync_is_cc_pair_agnostic=False, ), ), DocumentSource.SALESFORCE: SyncConfig( censoring_config=CensoringConfig( - chunk_censoring_func=censor_salesforce_chunks, + chunk_censoring_func=_lazy_censoring(_load_censor_salesforce_chunks), ), ), DocumentSource.MOCK_CONNECTOR: SyncConfig( @@ -198,19 +321,19 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = { DocumentSource.TEAMS: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=TEAMS_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=teams_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_teams_doc_sync), initial_index_should_sync=True, ), ), DocumentSource.SHAREPOINT: SyncConfig( doc_sync_config=DocSyncConfig( doc_sync_frequency=SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY, - doc_sync_func=sharepoint_doc_sync, + doc_sync_func=_lazy_doc_sync(_load_sharepoint_doc_sync), initial_index_should_sync=True, ), group_sync_config=GroupSyncConfig( group_sync_frequency=SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY, - group_sync_func=sharepoint_group_sync, + group_sync_func=_lazy_group_sync(_load_sharepoint_group_sync), group_sync_is_cc_pair_agnostic=False, ), ), diff --git a/backend/onyx/background/celery/apps/app_base.py b/backend/onyx/background/celery/apps/app_base.py index 862ee58e3a..465c00cbd7 100644 --- a/backend/onyx/background/celery/apps/app_base.py +++ b/backend/onyx/background/celery/apps/app_base.py @@ -41,7 +41,6 @@ from onyx.configs.app_configs import ( ) from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX, OnyxRedisLocks from onyx.db.engine.sql_engine import get_sqlalchemy_engine -from onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout from onyx.document_index.vespa.shared_utils.utils import wait_for_vespa_with_timeout from onyx.httpx.httpx_pool import HttpxPool from onyx.redis.redis_connector import RedisConnector @@ -689,6 +688,11 @@ def wait_for_document_index_or_shutdown() -> None: raise WorkerShutdown(msg) if ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: + # Imported here: opensearchpy costs ~18 MB and not every worker needs it. + from onyx.document_index.opensearch.client import ( + wait_for_opensearch_with_timeout, + ) + if not wait_for_opensearch_with_timeout(): msg = "[OpenSearch] Readiness probe did not succeed within the timeout. Exiting..." logger.error(msg) diff --git a/backend/onyx/background/celery/apps/light.py b/backend/onyx/background/celery/apps/light.py index 180e43cb15..3f870bc7e7 100644 --- a/backend/onyx/background/celery/apps/light.py +++ b/backend/onyx/background/celery/apps/light.py @@ -164,8 +164,6 @@ celery_app.autodiscover_tasks( "onyx.background.celery.tasks.doc_permission_syncing", "onyx.background.celery.tasks.docprocessing", "onyx.background.celery.tasks.opensearch_migration", - # Sandbox cleanup tasks (build feature) - "onyx.background.celery.tasks.build", ] ) ) diff --git a/backend/onyx/background/celery/tasks/build/tasks.py b/backend/onyx/background/celery/tasks/build/tasks.py index 429dc5a65d..5e49a7c82a 100644 --- a/backend/onyx/background/celery/tasks/build/tasks.py +++ b/backend/onyx/background/celery/tasks/build/tasks.py @@ -19,13 +19,6 @@ from onyx.server.features.build.db.sandbox import ( user_has_stale_active_session, ) from onyx.server.features.build.sandbox.factory import get_sandbox_manager -from onyx.server.features.build.session.locks import get_session_creation_lock -from onyx.server.features.build.session.sandbox_lifecycle import ( - create_session_snapshot_keep_latest, - is_sandbox_idle, - list_snapshotable_session_workspaces, - sleep_sandbox, -) # 100 minutes - snapshotting can take time TIMEOUT_SECONDS = 6000 @@ -54,6 +47,14 @@ def cleanup_idle_sandboxes_task(self: Task, *, tenant_id: str) -> None: # noqa: fail-closed: snapshot failure on a reachable pod keeps the sandbox RUNNING for retry next sweep. """ + from onyx.server.features.build.session.locks import get_session_creation_lock + from onyx.server.features.build.session.sandbox_lifecycle import ( + create_session_snapshot_keep_latest, + is_sandbox_idle, + list_snapshotable_session_workspaces, + sleep_sandbox, + ) + task_logger.info(f"cleanup_idle_sandboxes_task starting for tenant {tenant_id}") redis_client = get_redis_client(tenant_id=tenant_id) diff --git a/backend/onyx/background/celery/tasks/docprocessing/targeted_reindex_task.py b/backend/onyx/background/celery/tasks/docprocessing/targeted_reindex_task.py index 0c7bbb5fdb..a213e656c8 100644 --- a/backend/onyx/background/celery/tasks/docprocessing/targeted_reindex_task.py +++ b/backend/onyx/background/celery/tasks/docprocessing/targeted_reindex_task.py @@ -13,10 +13,6 @@ import logging from celery import Task, shared_task from onyx.background.celery.apps.app_base import task_logger -from onyx.background.indexing.run_targeted_reindex import ( - group_targets_by_cc_pair, - process_targets_for_cc_pair, -) from onyx.configs.constants import OnyxCeleryTask from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import IndexingStatus @@ -39,6 +35,11 @@ def run_targeted_reindex( """Body of the targeted-reindex task. Lifted out of the @shared_task decorator so tests can call it directly without going through celery's binding machinery.""" + from onyx.background.indexing.run_targeted_reindex import ( + group_targets_by_cc_pair, + process_targets_for_cc_pair, + ) + log = logging.LoggerAdapter( task_logger, extra={ diff --git a/backend/onyx/background/celery/tasks/docprocessing/tasks.py b/backend/onyx/background/celery/tasks/docprocessing/tasks.py index 42373bdaa8..86df930345 100644 --- a/backend/onyx/background/celery/tasks/docprocessing/tasks.py +++ b/backend/onyx/background/celery/tasks/docprocessing/tasks.py @@ -109,7 +109,6 @@ from onyx.db.search_settings import ( get_secondary_search_settings, ) from onyx.db.swap_index import check_and_perform_index_swap -from onyx.document_index.factory import get_all_document_indices from onyx.error_handling.exceptions import OnyxError from onyx.file_store.document_batch_storage import ( DocumentBatchStorage, @@ -117,19 +116,6 @@ from onyx.file_store.document_batch_storage import ( ) from onyx.file_store.staging import cleanup_staged_files_for_attempt from onyx.httpx.httpx_pool import HttpxPool -from onyx.indexing.adapters.document_indexing_adapter import ( - DocumentIndexingBatchAdapter, -) -from onyx.indexing.embedder import DefaultIndexingEmbedder -from onyx.indexing.indexing_pipeline import run_indexing_pipeline -from onyx.indexing.persistent_indexing import ( - build_generic_connector_failure, - record_generic_failure, -) -from onyx.natural_language_processing.search_nlp_models import ( - EmbeddingModel, - warm_up_bi_encoder, -) from onyx.redis.redis_connector import RedisConnector from onyx.redis.redis_docprocessing import RedisDocprocessing from onyx.redis.redis_pool import ( @@ -878,6 +864,11 @@ def check_for_indexing(self: Task, *, tenant_id: str) -> int | None: w.r.t previous failed attempt, checkpointing, etc is handled in the docfetching task. """ + from onyx.natural_language_processing.search_nlp_models import ( + EmbeddingModel, + warm_up_bi_encoder, + ) + time_start = time.monotonic() task_logger.warning("check_for_indexing - Starting") @@ -1552,6 +1543,11 @@ def _record_docprocessing_failure_persistent( Every step is wrapped so a follow-on error here does not re-raise out of the Celery task — we have already swallowed the original exception.""" + from onyx.indexing.persistent_indexing import ( + build_generic_connector_failure, + record_generic_failure, + ) + task_logger.info( "PERSISTENT_INDEXING enabled; recording docprocessing failure for " "attempt=%s batch=%s", @@ -1710,6 +1706,14 @@ def _docprocessing_task( cross_batch_db_lock: RedisLock | None = None try: + # Inside the try so a failed first-use import still marks the attempt failed. + from onyx.document_index.factory import get_all_document_indices + from onyx.indexing.adapters.document_indexing_adapter import ( + DocumentIndexingBatchAdapter, + ) + from onyx.indexing.embedder import DefaultIndexingEmbedder + from onyx.indexing.indexing_pipeline import run_indexing_pipeline + # FIX: Monitor memory before loading documents to track problematic batches emit_process_memory( os.getpid(), diff --git a/backend/onyx/background/celery/tasks/evals/tasks.py b/backend/onyx/background/celery/tasks/evals/tasks.py index 1d5dfe7917..e06618ae86 100644 --- a/backend/onyx/background/celery/tasks/evals/tasks.py +++ b/backend/onyx/background/celery/tasks/evals/tasks.py @@ -11,8 +11,6 @@ from onyx.configs.app_configs import ( SCHEDULED_EVAL_PROJECT, ) from onyx.configs.constants import OnyxCeleryTask -from onyx.evals.eval import run_eval -from onyx.evals.models import EvalConfigurationOptions from onyx.utils.logger import setup_logger logger = setup_logger() @@ -31,6 +29,11 @@ def eval_run_task( configuration_dict: dict[str, Any], ) -> None: """Background task to run an evaluation with the given configuration""" + # Imported here: the eval stack pulls in braintrust, every tool and the chat + # pipeline (~90 MB), which the primary worker should not hold at boot. + from onyx.evals.eval import run_eval + from onyx.evals.models import EvalConfigurationOptions + try: configuration = EvalConfigurationOptions.model_validate(configuration_dict) run_eval(configuration, remote_dataset_name=configuration.dataset_name) @@ -90,6 +93,9 @@ def scheduled_eval_task(self: Task, **kwargs: Any) -> None: # noqa: ARG001 dataset_names, ) + from onyx.evals.eval import run_eval + from onyx.evals.models import EvalConfigurationOptions + pipeline_start = datetime.now(timezone.utc) results: list[dict[str, Any]] = [] diff --git a/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py b/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py index 906a88856f..21132f4fdf 100644 --- a/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py +++ b/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py @@ -50,9 +50,6 @@ from onyx.db.scheduled_task import ( insert_run, mark_run_status, ) -from onyx.server.features.build.scheduled_tasks.executor import ( - run_scheduled_task_logic, -) from onyx.server.features.build.timeouts import ( QUEUE_RESIDENCY_SECONDS, SCHEDULED_RUN_HARD_CAP_SECONDS, @@ -233,6 +230,10 @@ def run_scheduled_task(self: Task, *, run_id: str, tenant_id: str) -> None: _ = self # bound only for symmetry with other shared_task wrappers _ = tenant_id try: + from onyx.server.features.build.scheduled_tasks.executor import ( + run_scheduled_task_logic, + ) + run_scheduled_task_logic(UUID(run_id)) except Exception: task_logger.exception( diff --git a/backend/onyx/db/engine/iam_auth.py b/backend/onyx/db/engine/iam_auth.py index c85715e7f7..fface05b9e 100644 --- a/backend/onyx/db/engine/iam_auth.py +++ b/backend/onyx/db/engine/iam_auth.py @@ -3,8 +3,6 @@ import os import ssl from typing import Any -import boto3 - from onyx.configs.app_configs import ( POSTGRES_HOST, POSTGRES_PORT, @@ -20,6 +18,8 @@ def get_iam_auth_token( """ Generate an IAM authentication token using boto3. """ + import boto3 + client = boto3.client("rds", region_name=region) token = client.generate_db_auth_token( DBHostname=host, Port=int(port), DBUsername=user diff --git a/backend/onyx/document_index/opensearch/client.py b/backend/onyx/document_index/opensearch/client.py index a96309b372..377e5e1f46 100644 --- a/backend/onyx/document_index/opensearch/client.py +++ b/backend/onyx/document_index/opensearch/client.py @@ -7,7 +7,6 @@ from contextlib import AbstractContextManager, nullcontext from http import HTTPStatus from typing import Any, Generic, TypeVar -import boto3 from opensearchpy import ( NotFoundError, OpenSearch, @@ -294,6 +293,8 @@ class OpenSearchClient(AbstractContextManager): # SigV4 signing for an AWS managed domain whose FGAC master is an # IAM ARN. Credentials come from the default boto3 chain (env, IRSA, # instance/task role); the signer refreshes them per request. + import boto3 + credentials = boto3.Session().get_credentials() if credentials is None: raise ValueError( diff --git a/backend/onyx/evals/eval_cli.py b/backend/onyx/evals/eval_cli.py index d28d611c83..621b805d49 100644 --- a/backend/onyx/evals/eval_cli.py +++ b/backend/onyx/evals/eval_cli.py @@ -9,7 +9,6 @@ import logging import os from typing import Any -import braintrust import requests from onyx.configs.app_configs import ( @@ -248,6 +247,8 @@ def main() -> None: "--local-only cannot be used with --remote-dataset-name. Use --local-data-path with a local JSON file instead." ) print(f"Loading data from remote dataset: {args.remote_dataset_name}") + import braintrust + dataset = braintrust.init_dataset( project=args.braintrust_project, name=args.remote_dataset_name ) diff --git a/backend/onyx/file_store/file_store.py b/backend/onyx/file_store/file_store.py index a72be20549..2f91246914 100644 --- a/backend/onyx/file_store/file_store.py +++ b/backend/onyx/file_store/file_store.py @@ -5,9 +5,7 @@ from abc import ABC, abstractmethod from io import BytesIO from typing import IO, TYPE_CHECKING, Any, NotRequired, TypedDict, cast -import boto3 import puremagic -from botocore.config import Config from botocore.exceptions import ClientError from sqlalchemy.orm import Session @@ -222,6 +220,10 @@ class S3BackedFileStore(FileStore): """Initialize S3 client if not already done""" if self._s3_client is None: try: + # Imported here: boto3 costs ~16 MB and most workers never build an S3 client. + import boto3 + from botocore.config import Config + client_kwargs: dict[str, Any] = { "service_name": "s3", "region_name": self._aws_region_name, diff --git a/backend/tests/external_dependency_unit/db/test_targeted_reindex_task.py b/backend/tests/external_dependency_unit/db/test_targeted_reindex_task.py index 54dd109079..e51c7f4088 100644 --- a/backend/tests/external_dependency_unit/db/test_targeted_reindex_task.py +++ b/backend/tests/external_dependency_unit/db/test_targeted_reindex_task.py @@ -44,8 +44,7 @@ from tests.external_dependency_unit.indexing_helpers import ( ) _PROCESSOR_PATH = ( - "onyx.background.celery.tasks.docprocessing." - "targeted_reindex_task.process_targets_for_cc_pair" + "onyx.background.indexing.run_targeted_reindex.process_targets_for_cc_pair" ) diff --git a/backend/tests/external_dependency_unit/db/test_tenant_provisioning_rollback.py b/backend/tests/external_dependency_unit/db/test_tenant_provisioning_rollback.py index 717ec91e3a..7f23f22b00 100644 --- a/backend/tests/external_dependency_unit/db/test_tenant_provisioning_rollback.py +++ b/backend/tests/external_dependency_unit/db/test_tenant_provisioning_rollback.py @@ -79,12 +79,12 @@ class TestTenantProvisioningRollback: mock_redis.return_value.lock.return_value = mock_lock with patch( - "ee.onyx.background.celery.tasks.tenant_provisioning.tasks.setup_tenant" + "ee.onyx.server.tenants.provisioning.setup_tenant" ) as mock_setup: mock_setup.side_effect = Exception("Simulated provisioning failure") with patch( - "ee.onyx.background.celery.tasks.tenant_provisioning.tasks.create_schema_if_not_exists", + "ee.onyx.server.tenants.schema_management.create_schema_if_not_exists", side_effect=track_schema_creation, ): # Run pre-provisioning - it should fail and trigger rollback diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_auth.py b/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_auth.py index dd9dd91252..8306476d2e 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_auth.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_auth.py @@ -23,7 +23,7 @@ def test_iam_auth_uses_sigv4_signer() -> None: """ with ( patch("onyx.document_index.opensearch.client.OpenSearch") as mock_os, - patch("onyx.document_index.opensearch.client.boto3.Session") as mock_session, + patch("boto3.Session") as mock_session, ): mock_session.return_value.get_credentials.return_value = MagicMock() OpenSearchClient( @@ -59,7 +59,7 @@ def test_iam_without_region_raises() -> None: def test_iam_without_credentials_raises() -> None: - with patch("onyx.document_index.opensearch.client.boto3.Session") as mock_session: + with patch("boto3.Session") as mock_session: mock_session.return_value.get_credentials.return_value = None with pytest.raises(ValueError, match="no AWS credentials"): OpenSearchClient( diff --git a/tools/ods/internal/lazyimports/lazyimports.go b/tools/ods/internal/lazyimports/lazyimports.go index c93d1d798e..4c65c7ae01 100644 --- a/tools/ods/internal/lazyimports/lazyimports.go +++ b/tools/ods/internal/lazyimports/lazyimports.go @@ -85,6 +85,18 @@ func DefaultLazyImportModules() map[string]LazyImportSettings { "trafilatura": NewLazyImportSettings(), "pypdf": NewLazyImportSettings(), "unstructured_client": NewLazyImportSettings(), + "braintrust": NewLazyImportSettings( + "onyx/evals/providers/braintrust.py", + "onyx/tracing/braintrust_tracing_processor.py", + ), + "exa_py": NewLazyImportSettings( + "onyx/tools/tool_implementations/web_search/clients/exa_client.py", + ), + "playwright": NewLazyImportSettings( + "onyx/connectors/highspot/utils.py", + "onyx/connectors/web/connector.py", + "onyx/utils/playwright_fetch.py", + ), } } diff --git a/tools/ods/internal/lazyimports/lazyimports_test.go b/tools/ods/internal/lazyimports/lazyimports_test.go index 8c6d797340..9daa89461e 100644 --- a/tools/ods/internal/lazyimports/lazyimports_test.go +++ b/tools/ods/internal/lazyimports/lazyimports_test.go @@ -706,6 +706,9 @@ func TestDefaultLazyImportModules(t *testing.T) { "trafilatura", "pypdf", "unstructured_client", + "braintrust", + "exa_py", + "playwright", } for _, mod := range expectedModules { @@ -730,6 +733,71 @@ func TestDefaultLazyImportModules(t *testing.T) { t.Errorf("Expected %s in litellm ignore files", ignore) } } + + // Gateway modules are the only files allowed to import these at module level. + expectedGatewayIgnores := map[string][]string{ + "braintrust": { + "onyx/evals/providers/braintrust.py", + "onyx/tracing/braintrust_tracing_processor.py", + }, + "exa_py": { + "onyx/tools/tool_implementations/web_search/clients/exa_client.py", + }, + "playwright": { + "onyx/connectors/highspot/utils.py", + "onyx/connectors/web/connector.py", + "onyx/utils/playwright_fetch.py", + }, + } + for mod, ignores := range expectedGatewayIgnores { + for _, ignore := range ignores { + if _, ok := modules[mod].IgnoreFiles[ignore]; !ok { + t.Errorf("Expected %s in %s ignore files", ignore, mod) + } + } + } +} + +func TestBraintrustAndExaViolations(t *testing.T) { + // Both packages pull in openai (and braintrust also litellm) at import time. + testContent := ` +import braintrust +from braintrust import Eval +from exa_py import Exa +from exa_py.api import HighlightsContentsOptions +import braintrust_langchain # Different package, not flagged + +def allowed_function(): + from braintrust import init_dataset + from exa_py import Exa + return init_dataset, Exa +` + + testPath := createTempPythonFile(t, testContent) + defer func() { _ = os.Remove(testPath) }() + + patterns := createPatterns([]string{"braintrust", "exa_py"}) + result, err := findEagerImports(testPath, patterns) + if err != nil { + t.Fatalf("findEagerImports failed: %v", err) + } + + lineNumbers := extractLineNumbers(result.ViolationLines) + expectedLines := []int{2, 3, 4, 5} + if len(lineNumbers) != len(expectedLines) { + t.Errorf("Expected %d violations, got %d (lines %v)", len(expectedLines), len(lineNumbers), lineNumbers) + } + for _, expected := range expectedLines { + if !containsLineNum(lineNumbers, expected) { + t.Errorf("Expected line %d in violations", expected) + } + } + + for _, mod := range []string{"braintrust", "exa_py"} { + if _, ok := result.ViolatedModules[mod]; !ok { + t.Errorf("Expected %s in violated modules", mod) + } + } } func TestIsValidPythonFile(t *testing.T) {