diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 53c931bc1d..7264df9f71 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -70,14 +70,19 @@ repos: args: ["ty", "check"] pass_filenames: true types_or: [python] + # Scope to backend/, matching `[tool.ty.src] include = ["backend"]` in + # pyproject. The full ty run only discovers backend/, but pre-commit + # passes explicit filenames, which bypass that discovery-time include — + # without this, a touched file outside backend/ (examples/, loadtest/, + # tools/, …) would get type-checked against the backend venv it was + # never meant to run in. + files: ^backend/ # sandbox_daemon ships in the sandbox image (own runtime); skills/builtin # are sandbox-only scripts importing deps absent from the backend venv - # (lxml, PIL, office.*). loadtest/ imports locust/gevent, which live in - # the non-default `loadtest` dependency group and so aren't in the venv - # ty runs against. All are excluded from the full ty run via - # [tool.ty.src] in pyproject; repeat here because pre-commit passes - # explicit filenames, which bypass that discovery-time exclude. - exclude: ^(backend/onyx/(server/features/build/sandbox/image/sandbox_daemon/|skills/builtin/)|loadtest/) + # (lxml, PIL, office.*). Both are excluded from the full ty run via + # [tool.ty.src] in pyproject; repeat here because the explicit filenames + # bypass that discovery-time exclude too. + exclude: ^backend/onyx/(server/features/build/sandbox/image/sandbox_daemon/|skills/builtin/) - id: uv-run alias: ruff name: ruff diff --git a/backend/ee/onyx/access/access.py b/backend/ee/onyx/access/access.py index 200478c8f4..804e380ce7 100644 --- a/backend/ee/onyx/access/access.py +++ b/backend/ee/onyx/access/access.py @@ -1,9 +1,13 @@ from sqlalchemy.orm import Session -from ee.onyx.db.external_perm import fetch_external_groups_for_user -from ee.onyx.db.external_perm import fetch_public_external_group_ids -from ee.onyx.db.user_group import fetch_user_groups_for_documents -from ee.onyx.db.user_group import fetch_user_groups_for_user +from ee.onyx.db.external_perm import ( + fetch_external_groups_for_user, + fetch_public_external_group_ids, +) +from ee.onyx.db.user_group import ( + fetch_user_groups_for_documents, + fetch_user_groups_for_user, +) from ee.onyx.external_permissions.sync_params import get_source_perm_sync_config from onyx.access.access import ( _get_access_for_documents as get_access_for_documents_without_groups, @@ -11,12 +15,9 @@ from onyx.access.access import ( from onyx.access.access import _get_acl_for_user as get_acl_for_user_without_groups from onyx.access.access import collect_user_file_access from onyx.access.models import DocumentAccess -from onyx.access.utils import prefix_external_group -from onyx.access.utils import prefix_user_group -from onyx.db.document import get_document_sources -from onyx.db.document import get_documents_by_ids -from onyx.db.models import User -from onyx.db.models import UserFile +from onyx.access.utils import prefix_external_group, prefix_user_group +from onyx.db.document import get_document_sources, get_documents_by_ids +from onyx.db.models import User, UserFile from onyx.db.user_file import fetch_user_files_with_access_relationships from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/auth/users.py b/backend/ee/onyx/auth/users.py index bfa2b9f7f2..32940c21a3 100644 --- a/backend/ee/onyx/auth/users.py +++ b/backend/ee/onyx/auth/users.py @@ -1,15 +1,10 @@ import secrets -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import jwt -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Request -from fastapi import status +from fastapi import Depends, HTTPException, Request, status -from ee.onyx.configs.app_configs import SUPER_CLOUD_API_KEY -from ee.onyx.configs.app_configs import SUPER_USERS +from ee.onyx.configs.app_configs import SUPER_CLOUD_API_KEY, SUPER_USERS from ee.onyx.server.seeding import get_seed_config from onyx.auth.permissions import require_permission from onyx.configs.app_configs import USER_AUTH_SECRET diff --git a/backend/ee/onyx/background/celery/tasks/beat_schedule.py b/backend/ee/onyx/background/celery/tasks/beat_schedule.py index 9078dabf96..378ae98e78 100644 --- a/backend/ee/onyx/background/celery/tasks/beat_schedule.py +++ b/backend/ee/onyx/background/celery/tasks/beat_schedule.py @@ -5,17 +5,17 @@ from ee.onyx.configs.app_configs import CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_H from onyx.background.celery.tasks.beat_schedule import ( beat_cloud_tasks as base_beat_system_tasks, ) -from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT +from onyx.background.celery.tasks.beat_schedule import ( + BEAT_EXPIRES_DEFAULT, + generate_cloud_tasks, +) from onyx.background.celery.tasks.beat_schedule import ( beat_task_templates as base_beat_task_templates, ) -from onyx.background.celery.tasks.beat_schedule import generate_cloud_tasks from onyx.background.celery.tasks.beat_schedule import ( get_tasks_to_schedule as base_get_tasks_to_schedule, ) -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import OnyxCeleryPriority, OnyxCeleryQueues, OnyxCeleryTask from shared_configs.configs import MULTI_TENANT ee_beat_system_tasks: list[dict] = [] diff --git a/backend/ee/onyx/background/celery/tasks/cleanup/tasks.py b/backend/ee/onyx/background/celery/tasks/cleanup/tasks.py index 0e1b5b03bc..92139c86be 100644 --- a/backend/ee/onyx/background/celery/tasks/cleanup/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/cleanup/tasks.py @@ -1,5 +1,4 @@ -from datetime import datetime -from datetime import timedelta +from datetime import datetime, timedelta from celery import shared_task diff --git a/backend/ee/onyx/background/celery/tasks/cloud/tasks.py b/backend/ee/onyx/background/celery/tasks/cloud/tasks.py index bbfa14815b..8a827acf8b 100644 --- a/backend/ee/onyx/background/celery/tasks/cloud/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/cloud/tasks.py @@ -1,26 +1,28 @@ import time -from celery import shared_task -from celery import Task +from celery import shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from redis.lock import Lock as RedisLock from ee.onyx.server.tenants.product_gating import get_gated_tenants from onyx.background.celery.apps.app_base import task_logger from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import ONYX_CLOUD_TENANT_ID -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + ONYX_CLOUD_TENANT_ID, + OnyxCeleryPriority, + OnyxCeleryTask, + OnyxRedisLocks, +) from onyx.db.engine.tenant_utils import get_all_tenant_ids -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import redis_lock_dump -from onyx.redis.redis_tenant_work_gating import cleanup_expired -from onyx.redis.redis_tenant_work_gating import get_active_tenants -from onyx.redis.redis_tenant_work_gating import observe_active_set_size -from onyx.redis.redis_tenant_work_gating import record_full_fanout_cycle -from onyx.redis.redis_tenant_work_gating import record_gate_decision +from onyx.redis.redis_pool import get_redis_client, redis_lock_dump +from onyx.redis.redis_tenant_work_gating import ( + cleanup_expired, + get_active_tenants, + observe_active_set_size, + record_full_fanout_cycle, + record_gate_decision, +) from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.server.runtime.onyx_runtime import OnyxRuntime from shared_configs.configs import IGNORED_SYNCING_TENANT_LIST diff --git a/backend/ee/onyx/background/celery/tasks/doc_permission_syncing/tasks.py b/backend/ee/onyx/background/celery/tasks/doc_permission_syncing/tasks.py index 5e75a05925..f52b7c55cb 100644 --- a/backend/ee/onyx/background/celery/tasks/doc_permission_syncing/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/doc_permission_syncing/tasks.py @@ -1,97 +1,109 @@ import time import traceback -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from time import sleep -from typing import Any -from typing import cast +from typing import Any, cast from uuid import uuid4 -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from pydantic import ValidationError from redis import Redis from redis.exceptions import LockError from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session -from tenacity import retry -from tenacity import retry_if_exception -from tenacity import stop_after_delay -from tenacity import wait_random_exponential +from tenacity import ( + retry, + retry_if_exception, + stop_after_delay, + wait_random_exponential, +) from ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs from ee.onyx.db.document import upsert_document_external_perms from ee.onyx.external_permissions.sync_params import get_source_perm_sync_config -from onyx.access.models import DocExternalAccess -from onyx.access.models import ElementExternalAccess +from onyx.access.models import DocExternalAccess, ElementExternalAccess from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.celery_redis import celery_find_task -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_queue_length -from onyx.background.celery.celery_redis import celery_get_queued_task_ids -from onyx.background.celery.celery_redis import celery_get_unacked_task_ids +from onyx.background.celery.celery_redis import ( + celery_find_task, + celery_get_broker_client, + celery_get_queue_length, + celery_get_queued_task_ids, + celery_get_unacked_task_ids, +) from onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT from onyx.configs.app_configs import JOB_TIMEOUT -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT -from onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import OnyxRedisSignals +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT, + CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT, + DANSWER_REDIS_FUNCTION_LOCK_PREFIX, + DocumentSource, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, + OnyxRedisLocks, + OnyxRedisSignals, +) from onyx.connectors.factory import validate_ccpair_for_user from onyx.db.connector import mark_cc_pair_as_permissions_synced from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.document import get_document_ids_for_connector_credential_pair -from onyx.db.document import get_documents_for_connector_credential_pair_limited_columns -from onyx.db.document import upsert_document_by_connector_credential_pair -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType +from onyx.db.document import ( + get_document_ids_for_connector_credential_pair, + get_documents_for_connector_credential_pair_limited_columns, + upsert_document_by_connector_credential_pair, +) +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_tenant, +) +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + SyncStatus, + SyncType, +) from onyx.db.hierarchy import ( update_hierarchy_node_permissions as db_update_hierarchy_node_permissions, ) from onyx.db.models import ConnectorCredentialPair -from onyx.db.permission_sync_attempt import complete_doc_permission_sync_attempt -from onyx.db.permission_sync_attempt import create_doc_permission_sync_attempt -from onyx.db.permission_sync_attempt import mark_doc_permission_sync_attempt_failed -from onyx.db.permission_sync_attempt import mark_doc_permission_sync_attempt_in_progress -from onyx.db.sync_record import insert_sync_record -from onyx.db.sync_record import update_sync_record_status +from onyx.db.permission_sync_attempt import ( + complete_doc_permission_sync_attempt, + create_doc_permission_sync_attempt, + mark_doc_permission_sync_attempt_failed, + mark_doc_permission_sync_attempt_in_progress, +) +from onyx.db.sync_record import insert_sync_record, update_sync_record_status from onyx.db.users import batch_add_ext_perm_user_if_not_exists -from onyx.db.utils import DocumentRow -from onyx.db.utils import is_retryable_sqlalchemy_error -from onyx.db.utils import SortOrder +from onyx.db.utils import DocumentRow, is_retryable_sqlalchemy_error, SortOrder from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.redis.redis_connector import RedisConnector -from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync -from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSyncPayload -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client -from onyx.redis.redis_pool import redis_lock_dump +from onyx.redis.redis_connector_doc_perm_sync import ( + RedisConnectorPermissionSync, + RedisConnectorPermissionSyncPayload, +) +from onyx.redis.redis_pool import ( + get_redis_client, + get_redis_replica_client, + redis_lock_dump, +) from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active from onyx.redis.tenant_redis_client import TenantRedisClient -from onyx.server.metrics.perm_sync_metrics import inc_doc_perm_sync_docs_processed -from onyx.server.metrics.perm_sync_metrics import inc_doc_perm_sync_errors -from onyx.server.metrics.perm_sync_metrics import observe_doc_perm_sync_duration +from onyx.server.metrics.perm_sync_metrics import ( + inc_doc_perm_sync_docs_processed, + inc_doc_perm_sync_errors, + observe_doc_perm_sync_duration, +) from onyx.server.runtime.onyx_runtime import OnyxRuntime from onyx.server.utils import make_short_id -from onyx.utils.logger import doc_permission_sync_ctx -from onyx.utils.logger import format_error_for_logging -from onyx.utils.logger import LoggerContextVars -from onyx.utils.logger import setup_logger -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType +from onyx.utils.logger import ( + doc_permission_sync_ctx, + format_error_for_logging, + LoggerContextVars, + setup_logger, +) +from onyx.utils.telemetry import optional_telemetry, RecordType from shared_configs.configs import MULTI_TENANT logger = setup_logger() diff --git a/backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py b/backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py index b1e9123a57..dfb404274c 100644 --- a/backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py @@ -1,15 +1,10 @@ import time import traceback -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timedelta, timezone +from typing import Any, cast from uuid import uuid4 -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from pydantic import ValidationError from redis import Redis @@ -18,63 +13,74 @@ from redis.lock import Lock as RedisLock from ee.onyx.background.celery.tasks.external_group_syncing.group_sync_utils import ( mark_all_relevant_cc_pairs_as_external_group_synced, ) -from ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs -from ee.onyx.db.connector_credential_pair import get_cc_pairs_by_source -from ee.onyx.db.external_perm import ExternalUserGroup -from ee.onyx.db.external_perm import mark_old_external_groups_as_stale -from ee.onyx.db.external_perm import remove_stale_external_groups -from ee.onyx.db.external_perm import upsert_external_groups +from ee.onyx.db.connector_credential_pair import ( + get_all_auto_sync_cc_pairs, + get_cc_pairs_by_source, +) +from ee.onyx.db.external_perm import ( + ExternalUserGroup, + mark_old_external_groups_as_stale, + remove_stale_external_groups, + upsert_external_groups, +) from ee.onyx.external_permissions.sync_params import ( get_all_cc_pair_agnostic_group_sync_sources, + get_source_perm_sync_config, ) -from ee.onyx.external_permissions.sync_params import get_source_perm_sync_config from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.celery_redis import celery_find_task -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_unacked_task_ids +from onyx.background.celery.celery_redis import ( + celery_find_task, + celery_get_broker_client, + celery_get_unacked_task_ids, +) from onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT from onyx.background.error_logging import emit_background_error from onyx.configs.app_configs import JOB_TIMEOUT -from onyx.configs.constants import CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import OnyxRedisSignals +from onyx.configs.constants import ( + CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT, + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, + OnyxRedisLocks, + OnyxRedisSignals, +) from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + SyncStatus, + SyncType, +) from onyx.db.models import ConnectorCredentialPair -from onyx.db.permission_sync_attempt import complete_external_group_sync_attempt -from onyx.db.permission_sync_attempt import create_external_group_sync_attempt -from onyx.db.permission_sync_attempt import mark_external_group_sync_attempt_failed -from onyx.db.permission_sync_attempt import mark_external_group_sync_attempt_in_progress -from onyx.db.sync_record import insert_sync_record -from onyx.db.sync_record import update_sync_record_status +from onyx.db.permission_sync_attempt import ( + complete_external_group_sync_attempt, + create_external_group_sync_attempt, + mark_external_group_sync_attempt_failed, + mark_external_group_sync_attempt_in_progress, +) +from onyx.db.sync_record import insert_sync_record, update_sync_record_status from onyx.redis.redis_connector import RedisConnector -from onyx.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync from onyx.redis.redis_connector_ext_group_sync import ( + RedisConnectorExternalGroupSync, RedisConnectorExternalGroupSyncPayload, ) -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client +from onyx.redis.redis_pool import get_redis_client, get_redis_replica_client from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active from onyx.redis.tenant_redis_client import TenantRedisClient -from onyx.server.metrics.perm_sync_metrics import inc_group_sync_errors -from onyx.server.metrics.perm_sync_metrics import inc_group_sync_groups_processed -from onyx.server.metrics.perm_sync_metrics import inc_group_sync_users_processed -from onyx.server.metrics.perm_sync_metrics import observe_group_sync_duration -from onyx.server.metrics.perm_sync_metrics import observe_group_sync_upsert_duration +from onyx.server.metrics.perm_sync_metrics import ( + inc_group_sync_errors, + inc_group_sync_groups_processed, + inc_group_sync_users_processed, + observe_group_sync_duration, + observe_group_sync_upsert_duration, +) from onyx.server.runtime.onyx_runtime import OnyxRuntime from onyx.server.utils import make_short_id -from onyx.utils.logger import format_error_for_logging -from onyx.utils.logger import setup_logger +from onyx.utils.logger import format_error_for_logging, setup_logger from shared_configs.configs import MULTI_TENANT logger = setup_logger() diff --git a/backend/ee/onyx/background/celery/tasks/license_notifications/tasks.py b/backend/ee/onyx/background/celery/tasks/license_notifications/tasks.py index 16d8294e23..ee866c3d48 100644 --- a/backend/ee/onyx/background/celery/tasks/license_notifications/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/license_notifications/tasks.py @@ -2,8 +2,7 @@ from celery import shared_task from ee.onyx.db.license import get_license from ee.onyx.utils.license import verify_license_signature -from ee.onyx.utils.license_expiry import ExpiryWarningStage -from ee.onyx.utils.license_expiry import get_expiry_warning_stage +from ee.onyx.utils.license_expiry import ExpiryWarningStage, get_expiry_warning_stage from ee.onyx.utils.license_notifications import notify_admins_for_stage from onyx.configs.app_configs import JOB_TIMEOUT from onyx.configs.constants import OnyxCeleryTask 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 5a9791be88..545a594587 100644 --- a/backend/ee/onyx/background/celery/tasks/query_history/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/query_history/tasks.py @@ -2,22 +2,27 @@ import csv import io from datetime import datetime -from celery import shared_task -from celery import Task +from celery import shared_task, Task -from ee.onyx.server.query_history.api import fetch_and_process_chat_session_history -from ee.onyx.server.query_history.api import ONYX_ANONYMIZED_EMAIL +from ee.onyx.server.query_history.api import ( + fetch_and_process_chat_session_history, + ONYX_ANONYMIZED_EMAIL, +) 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 -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import FileType -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import QueryHistoryType +from onyx.configs.constants import ( + FileOrigin, + FileType, + OnyxCeleryTask, + QueryHistoryType, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.tasks import delete_task_with_id -from onyx.db.tasks import mark_task_as_finished_with_id -from onyx.db.tasks import mark_task_as_started_with_id +from onyx.db.tasks import ( + delete_task_with_id, + mark_task_as_finished_with_id, + mark_task_as_started_with_id, +) from onyx.file_store.file_store import get_default_file_store from onyx.server.settings.store import load_settings from onyx.utils.csv_utils import sanitize_csv_row 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 af2a6818ad..3e46ae1ef6 100644 --- a/backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py @@ -6,25 +6,27 @@ import asyncio import datetime import uuid -from celery import shared_task -from celery import Task +from celery import shared_task, 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 -from ee.onyx.server.tenants.schema_management import get_current_alembic_version -from ee.onyx.server.tenants.schema_management import run_alembic_migrations +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 ONYX_CLOUD_TENANT_ID -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks +from onyx.configs.constants import ( + ONYX_CLOUD_TENANT_ID, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisLocks, +) from onyx.db.engine.sql_engine import get_session_with_shared_schema from onyx.db.models import AvailableTenant from onyx.redis.redis_pool import get_redis_client -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import TENANT_ID_PREFIX +from shared_configs.configs import MULTI_TENANT, TENANT_ID_PREFIX # Maximum tenants to provision in a single task run. # Each tenant takes ~80s (alembic migrations), so 15 tenants ≈ 20 minutes. diff --git a/backend/ee/onyx/background/celery/tasks/ttl_management/tasks.py b/backend/ee/onyx/background/celery/tasks/ttl_management/tasks.py index f7bf701f95..ebaafce274 100644 --- a/backend/ee/onyx/background/celery/tasks/ttl_management/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/ttl_management/tasks.py @@ -1,18 +1,18 @@ import uuid -from celery import shared_task -from celery import Task +from celery import shared_task, Task from ee.onyx.background.celery_utils import should_perform_chat_ttl_check from onyx.configs.app_configs import JOB_TIMEOUT -from onyx.configs.constants import CELERY_CHAT_TTL_DELETE_TASK_EXPIRES -from onyx.configs.constants import CHAT_TTL_DELETE_BATCH_SIZE -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks -from onyx.db.chat import delete_chat_session -from onyx.db.chat import get_chat_sessions_older_than +from onyx.configs.constants import ( + CELERY_CHAT_TTL_DELETE_TASK_EXPIRES, + CHAT_TTL_DELETE_BATCH_SIZE, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisLocks, +) +from onyx.db.chat import delete_chat_session, get_chat_sessions_older_than from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.redis.redis_pool import get_redis_client from onyx.redis.tenant_redis_client import TenantRedisClient 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 cca3f61ce9..3ae57183cf 100644 --- a/backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py @@ -1,8 +1,7 @@ from datetime import datetime from uuid import UUID -from celery import shared_task -from celery import Task +from celery import shared_task, Task from ee.onyx.server.reporting.usage_export_generation import create_new_usage_report from onyx.configs.app_configs import JOB_TIMEOUT diff --git a/backend/ee/onyx/background/celery/tasks/vespa/tasks.py b/backend/ee/onyx/background/celery/tasks/vespa/tasks.py index c773adc537..cec337dfc1 100644 --- a/backend/ee/onyx/background/celery/tasks/vespa/tasks.py +++ b/backend/ee/onyx/background/celery/tasks/vespa/tasks.py @@ -1,12 +1,13 @@ from sqlalchemy.orm import Session -from ee.onyx.db.user_group import delete_user_group -from ee.onyx.db.user_group import fetch_user_group -from ee.onyx.db.user_group import mark_user_group_as_synced -from ee.onyx.db.user_group import prepare_user_group_for_deletion +from ee.onyx.db.user_group import ( + delete_user_group, + fetch_user_group, + mark_user_group_as_synced, + prepare_user_group_for_deletion, +) from onyx.background.celery.apps.app_base import task_logger -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType +from onyx.db.enums import SyncStatus, SyncType from onyx.db.sync_record import update_sync_record_status from onyx.redis.redis_usergroup import RedisUserGroup from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/ee/onyx/background/celery_utils.py b/backend/ee/onyx/background/celery_utils.py index b401c15854..b9bfee976c 100644 --- a/backend/ee/onyx/background/celery_utils.py +++ b/backend/ee/onyx/background/celery_utils.py @@ -1,8 +1,7 @@ from sqlalchemy.orm import Session from ee.onyx.background.task_name_builders import name_chat_ttl_task -from onyx.db.tasks import check_task_is_live_and_not_timed_out -from onyx.db.tasks import get_latest_task +from onyx.db.tasks import check_task_is_live_and_not_timed_out, get_latest_task from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/db/analytics.py b/backend/ee/onyx/db/analytics.py index 0807dd8f2c..342bae832b 100644 --- a/backend/ee/onyx/db/analytics.py +++ b/backend/ee/onyx/db/analytics.py @@ -2,22 +2,18 @@ import datetime from collections.abc import Sequence from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import case -from sqlalchemy import cast -from sqlalchemy import Date -from sqlalchemy import func -from sqlalchemy import or_ -from sqlalchemy import select +from sqlalchemy import and_, case, cast, Date, func, or_, select from sqlalchemy.orm import Session from onyx.configs.constants import MessageType -from onyx.db.models import ChatMessage -from onyx.db.models import ChatMessageFeedback -from onyx.db.models import ChatSession -from onyx.db.models import Persona -from onyx.db.models import User -from onyx.db.models import UserRole +from onyx.db.models import ( + ChatMessage, + ChatMessageFeedback, + ChatSession, + Persona, + User, + UserRole, +) def fetch_query_analytics( diff --git a/backend/ee/onyx/db/connector_credential_pair.py b/backend/ee/onyx/db/connector_credential_pair.py index 3944150295..e9fa99fb1b 100644 --- a/backend/ee/onyx/db/connector_credential_pair.py +++ b/backend/ee/onyx/db/connector_credential_pair.py @@ -3,11 +3,12 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.db.connector_credential_pair import get_connector_credential_pair -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import UserGroup__ConnectorCredentialPair +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + UserGroup__ConnectorCredentialPair, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/db/document.py b/backend/ee/onyx/db/document.py index ad61cff4fe..d2f5e47f3a 100644 --- a/backend/ee/onyx/db/document.py +++ b/backend/ee/onyx/db/document.py @@ -1,5 +1,4 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from sqlalchemy import select from sqlalchemy.orm import Session diff --git a/backend/ee/onyx/db/document_set.py b/backend/ee/onyx/db/document_set.py index cde439d4c9..1e802dd58c 100644 --- a/backend/ee/onyx/db/document_set.py +++ b/backend/ee/onyx/db/document_set.py @@ -2,13 +2,15 @@ from uuid import UUID from sqlalchemy.orm import Session -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import DocumentSet -from onyx.db.models import DocumentSet__ConnectorCredentialPair -from onyx.db.models import DocumentSet__User -from onyx.db.models import DocumentSet__UserGroup -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import ( + ConnectorCredentialPair, + DocumentSet, + DocumentSet__ConnectorCredentialPair, + DocumentSet__User, + DocumentSet__UserGroup, + User__UserGroup, + UserGroup, +) def make_doc_set_private( diff --git a/backend/ee/onyx/db/external_perm.py b/backend/ee/onyx/db/external_perm.py index b8019fd1c0..c5b78e03ac 100644 --- a/backend/ee/onyx/db/external_perm.py +++ b/backend/ee/onyx/db/external_perm.py @@ -2,19 +2,14 @@ from collections.abc import Sequence from uuid import UUID from pydantic import BaseModel -from sqlalchemy import delete -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import delete, select, update from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session from onyx.access.utils import build_ext_group_name_for_onyx from onyx.configs.constants import DocumentSource -from onyx.db.models import PublicExternalUserGroup -from onyx.db.models import User -from onyx.db.models import User__ExternalUserGroupId -from onyx.db.users import batch_add_ext_perm_user_if_not_exists -from onyx.db.users import get_user_by_email +from onyx.db.models import PublicExternalUserGroup, User, User__ExternalUserGroupId +from onyx.db.users import batch_add_ext_perm_user_if_not_exists, get_user_by_email from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/db/hierarchy.py b/backend/ee/onyx/db/hierarchy.py index d24aa47e10..933aa416b6 100644 --- a/backend/ee/onyx/db/hierarchy.py +++ b/backend/ee/onyx/db/hierarchy.py @@ -4,19 +4,14 @@ This module provides permission-aware hierarchy node access for Enterprise Editi It filters hierarchy nodes based on user email and external group membership. """ -from sqlalchemy import any_ -from sqlalchemy import cast -from sqlalchemy import or_ -from sqlalchemy import select -from sqlalchemy import String +from sqlalchemy import any_, cast, or_, select, String from sqlalchemy.dialects import postgresql from sqlalchemy.orm import Session from sqlalchemy.sql.elements import ColumnElement from onyx.configs.constants import DocumentSource from onyx.db.enums import HierarchyNodeType -from onyx.db.hierarchy import escape_like_pattern -from onyx.db.hierarchy import HIERARCHY_NODE_SEARCH_LIMIT +from onyx.db.hierarchy import escape_like_pattern, HIERARCHY_NODE_SEARCH_LIMIT from onyx.db.models import HierarchyNode diff --git a/backend/ee/onyx/db/license.py b/backend/ee/onyx/db/license.py index 9d69901e75..03fe523d00 100644 --- a/backend/ee/onyx/db/license.py +++ b/backend/ee/onyx/db/license.py @@ -5,20 +5,15 @@ import struct from datetime import datetime from typing import NamedTuple -from sqlalchemy import func -from sqlalchemy import select -from sqlalchemy import text +from sqlalchemy import func, select, text from sqlalchemy.orm import Session -from ee.onyx.server.license.models import LicenseMetadata -from ee.onyx.server.license.models import LicensePayload -from ee.onyx.server.license.models import LicenseSource +from ee.onyx.server.license.models import LicenseMetadata, LicensePayload, LicenseSource from onyx.auth.schemas import UserRole from onyx.cache.factory import get_cache_backend from onyx.configs.constants import ANONYMOUS_USER_EMAIL from onyx.db.enums import AccountType -from onyx.db.models import License -from onyx.db.models import User +from onyx.db.models import License, User from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import get_current_tenant_id @@ -251,8 +246,10 @@ def update_license_cache( The cached LicenseMetadata """ from ee.onyx.utils.license import get_license_status - from ee.onyx.utils.license_expiry import get_expiry_warning_stage - from ee.onyx.utils.license_expiry import get_grace_period_end + from ee.onyx.utils.license_expiry import ( + get_expiry_warning_stage, + get_grace_period_end, + ) tenant = tenant_id or get_current_tenant_id() cache = get_cache_backend(tenant_id=tenant_id) diff --git a/backend/ee/onyx/db/mcp.py b/backend/ee/onyx/db/mcp.py index ec9b2b2330..e0c655f61e 100644 --- a/backend/ee/onyx/db/mcp.py +++ b/backend/ee/onyx/db/mcp.py @@ -2,8 +2,7 @@ from uuid import UUID from sqlalchemy.orm import Session -from onyx.db.models import MCPServer__User -from onyx.db.models import MCPServer__UserGroup +from onyx.db.models import MCPServer__User, MCPServer__UserGroup def make_mcp_server_private( diff --git a/backend/ee/onyx/db/persona.py b/backend/ee/onyx/db/persona.py index 49fb6d1500..d35808bc1c 100644 --- a/backend/ee/onyx/db/persona.py +++ b/backend/ee/onyx/db/persona.py @@ -3,13 +3,13 @@ from uuid import UUID from sqlalchemy.orm import Session from onyx.db.enums import PersonaSharePermission -from onyx.db.models import Persona -from onyx.db.models import Persona__UserGroup -from onyx.db.models import User -from onyx.db.persona import _transfer_persona_ownership -from onyx.db.persona import apply_persona_user_share_diff -from onyx.db.persona import mark_persona_user_files_for_sync -from onyx.db.persona import resolve_desired_user_shares +from onyx.db.models import Persona, Persona__UserGroup, User +from onyx.db.persona import ( + _transfer_persona_ownership, + apply_persona_user_share_diff, + mark_persona_user_files_for_sync, + resolve_desired_user_shares, +) def _resolve_desired_group_shares( diff --git a/backend/ee/onyx/db/query_history.py b/backend/ee/onyx/db/query_history.py index 244787c11c..e4c71d93c3 100644 --- a/backend/ee/onyx/db/query_history.py +++ b/backend/ee/onyx/db/query_history.py @@ -1,26 +1,14 @@ from collections.abc import Sequence from datetime import datetime -from sqlalchemy import asc -from sqlalchemy import BinaryExpression -from sqlalchemy import ColumnElement -from sqlalchemy import desc -from sqlalchemy import distinct -from sqlalchemy.orm import contains_eager -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import Session -from sqlalchemy.sql import case -from sqlalchemy.sql import func -from sqlalchemy.sql import select -from sqlalchemy.sql.expression import literal -from sqlalchemy.sql.expression import UnaryExpression +from sqlalchemy import asc, BinaryExpression, ColumnElement, desc, distinct +from sqlalchemy.orm import contains_eager, joinedload, Session +from sqlalchemy.sql import case, func, select +from sqlalchemy.sql.expression import literal, UnaryExpression from ee.onyx.background.task_name_builders import QUERY_HISTORY_TASK_NAME_PREFIX from onyx.configs.constants import QAFeedbackType -from onyx.db.models import ChatMessage -from onyx.db.models import ChatMessageFeedback -from onyx.db.models import ChatSession -from onyx.db.models import TaskQueueState +from onyx.db.models import ChatMessage, ChatMessageFeedback, ChatSession, TaskQueueState from onyx.db.tasks import get_all_tasks_with_prefix diff --git a/backend/ee/onyx/db/saml.py b/backend/ee/onyx/db/saml.py index 3b9afa1e22..4e14a6c964 100644 --- a/backend/ee/onyx/db/saml.py +++ b/backend/ee/onyx/db/saml.py @@ -2,12 +2,9 @@ import datetime from typing import cast from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import and_, func, select from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS from onyx.db.models import SamlAccount diff --git a/backend/ee/onyx/db/scim.py b/backend/ee/onyx/db/scim.py index 69074298cc..ebeb32b5b9 100644 --- a/backend/ee/onyx/db/scim.py +++ b/backend/ee/onyx/db/scim.py @@ -26,27 +26,23 @@ from __future__ import annotations from uuid import UUID from sqlalchemy import delete as sa_delete -from sqlalchemy import func -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import SQLColumnExpression +from sqlalchemy import func, Select, select, SQLColumnExpression from sqlalchemy.dialects.postgresql import insert as pg_insert -from ee.onyx.server.scim.filtering import ScimFilter -from ee.onyx.server.scim.filtering import ScimFilterOperator +from ee.onyx.server.scim.filtering import ScimFilter, ScimFilterOperator from ee.onyx.server.scim.models import ScimMappingFields from onyx.db.dal import DAL -from onyx.db.enums import AccountType -from onyx.db.enums import GrantSource -from onyx.db.enums import Permission -from onyx.db.models import PermissionGrant -from onyx.db.models import ScimGroupMapping -from onyx.db.models import ScimToken -from onyx.db.models import ScimUserMapping -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup -from onyx.db.models import UserRole +from onyx.db.enums import AccountType, GrantSource, Permission +from onyx.db.models import ( + PermissionGrant, + ScimGroupMapping, + ScimToken, + ScimUserMapping, + User, + User__UserGroup, + UserGroup, + UserRole, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/db/standard_answer.py b/backend/ee/onyx/db/standard_answer.py index f39902f3d4..64ede2e6c9 100644 --- a/backend/ee/onyx/db/standard_answer.py +++ b/backend/ee/onyx/db/standard_answer.py @@ -5,8 +5,7 @@ from collections.abc import Sequence from sqlalchemy import select from sqlalchemy.orm import Session -from onyx.db.models import StandardAnswer -from onyx.db.models import StandardAnswerCategory +from onyx.db.models import StandardAnswer, StandardAnswerCategory from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/db/token_limit.py b/backend/ee/onyx/db/token_limit.py index 52c3bfa7a2..3d687e952b 100644 --- a/backend/ee/onyx/db/token_limit.py +++ b/backend/ee/onyx/db/token_limit.py @@ -1,19 +1,17 @@ from collections.abc import Sequence -from sqlalchemy import exists -from sqlalchemy import Row -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy.orm import aliased -from sqlalchemy.orm import Session +from sqlalchemy import exists, Row, Select, select +from sqlalchemy.orm import aliased, Session from onyx.configs.constants import TokenRateLimitScope -from onyx.db.models import TokenRateLimit -from onyx.db.models import TokenRateLimit__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup -from onyx.db.models import UserRole +from onyx.db.models import ( + TokenRateLimit, + TokenRateLimit__UserGroup, + User, + User__UserGroup, + UserGroup, + UserRole, +) from onyx.server.token_rate_limits.models import TokenRateLimitArgs diff --git a/backend/ee/onyx/db/usage_export.py b/backend/ee/onyx/db/usage_export.py index 985c7a9054..f2bed0c77c 100644 --- a/backend/ee/onyx/db/usage_export.py +++ b/backend/ee/onyx/db/usage_export.py @@ -1,8 +1,7 @@ import uuid from collections.abc import Generator from datetime import datetime -from typing import IO -from typing import Optional +from typing import IO, Optional from fastapi_users_db_sqlalchemy import UUID_ID from sqlalchemy import cast @@ -10,13 +9,13 @@ from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.orm import Session from ee.onyx.db.query_history import fetch_chat_sessions_eagerly_by_time -from ee.onyx.server.reporting.usage_export_models import ChatMessageSkeleton -from ee.onyx.server.reporting.usage_export_models import FlowType -from ee.onyx.server.reporting.usage_export_models import UsageReportMetadata +from ee.onyx.server.reporting.usage_export_models import ( + ChatMessageSkeleton, + FlowType, + UsageReportMetadata, +) from onyx.configs.constants import MessageType -from onyx.db.models import ChatMessage -from onyx.db.models import UsageReport -from onyx.db.models import User +from onyx.db.models import ChatMessage, UsageReport, User from onyx.file_store.file_store import get_default_file_store diff --git a/backend/ee/onyx/db/user_group.py b/backend/ee/onyx/db/user_group.py index f30e025d6c..4f272f5e77 100644 --- a/backend/ee/onyx/db/user_group.py +++ b/backend/ee/onyx/db/user_group.py @@ -3,53 +3,58 @@ from operator import and_ from uuid import UUID from fastapi import HTTPException -from sqlalchemy import delete -from sqlalchemy import func -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import delete, func, Select, select, update from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session -from ee.onyx.server.user_group.models import SetCuratorRequest -from ee.onyx.server.user_group.models import UserGroupCreate -from ee.onyx.server.user_group.models import UserGroupUpdate +from ee.onyx.server.user_group.models import ( + SetCuratorRequest, + UserGroupCreate, + UserGroupUpdate, +) from onyx.configs.app_configs import DISABLE_VECTOR_DB from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import GrantSource -from onyx.db.enums import Permission -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import Credential__UserGroup -from onyx.db.models import Document -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import DocumentSet -from onyx.db.models import DocumentSet__UserGroup -from onyx.db.models import FederatedConnector__DocumentSet -from onyx.db.models import LLMProvider__UserGroup -from onyx.db.models import MCPServer__UserGroup -from onyx.db.models import PermissionGrant -from onyx.db.models import Persona -from onyx.db.models import Persona__User -from onyx.db.models import Persona__UserGroup -from onyx.db.models import TokenRateLimit__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup -from onyx.db.models import UserGroup__ConnectorCredentialPair -from onyx.db.models import UserRole -from onyx.db.permissions import recompute_permissions_for_group__no_commit -from onyx.db.permissions import recompute_user_permissions__no_commit +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + GrantSource, + Permission, +) +from onyx.db.models import ( + ConnectorCredentialPair, + Credential, + Credential__UserGroup, + Document, + DocumentByConnectorCredentialPair, + DocumentSet, + DocumentSet__UserGroup, + FederatedConnector__DocumentSet, + LLMProvider__UserGroup, + MCPServer__UserGroup, + PermissionGrant, + Persona, + Persona__User, + Persona__UserGroup, + TokenRateLimit__UserGroup, + User, + User__UserGroup, + UserGroup, + UserGroup__ConnectorCredentialPair, + UserRole, +) +from onyx.db.permissions import ( + recompute_permissions_for_group__no_commit, + recompute_user_permissions__no_commit, +) from onyx.db.users import fetch_user_by_id from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/external_permissions/box/access.py b/backend/ee/onyx/external_permissions/box/access.py index 8def09287c..dd391d1d01 100644 --- a/backend/ee/onyx/external_permissions/box/access.py +++ b/backend/ee/onyx/external_permissions/box/access.py @@ -1,7 +1,6 @@ from box_sdk_gen import BoxClient from box_sdk_gen.box.errors import BoxAPIError -from box_sdk_gen.schemas.collaboration import Collaboration -from box_sdk_gen.schemas.collaboration import CollaborationStatusField +from box_sdk_gen.schemas.collaboration import Collaboration, CollaborationStatusField from box_sdk_gen.schemas.file import FileSharedLinkField from box_sdk_gen.schemas.file_full import FileFull from box_sdk_gen.schemas.folder import FolderSharedLinkField @@ -9,17 +8,17 @@ from box_sdk_gen.schemas.folder_mini import FolderMini from box_sdk_gen.schemas.group_mini import GroupMini from box_sdk_gen.schemas.user_collaborations import UserCollaborations from box_sdk_gen.schemas.user_mini import UserMini -from box_sdk_gen.schemas.web_link import WebLink -from box_sdk_gen.schemas.web_link import WebLinkSharedLinkField -from pydantic import BaseModel -from pydantic import Field +from box_sdk_gen.schemas.web_link import WebLink, WebLinkSharedLinkField +from pydantic import BaseModel, Field from onyx.access.models import ExternalAccess -from onyx.connectors.box.connector import box_all_enterprise_users_group_id -from onyx.connectors.box.connector import box_api_status_code -from onyx.connectors.box.connector import box_group_id -from onyx.connectors.box.connector import BOX_ROOT_FOLDER_ID -from onyx.connectors.box.connector import normalize_box_login +from onyx.connectors.box.connector import ( + box_all_enterprise_users_group_id, + box_api_status_code, + box_group_id, + BOX_ROOT_FOLDER_ID, + normalize_box_login, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/external_permissions/box/doc_sync.py b/backend/ee/onyx/external_permissions/box/doc_sync.py index 1797c81172..b9a29758eb 100644 --- a/backend/ee/onyx/external_permissions/box/doc_sync.py +++ b/backend/ee/onyx/external_permissions/box/doc_sync.py @@ -1,9 +1,10 @@ from collections.abc import Generator -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction -from ee.onyx.external_permissions.utils import credential_json -from ee.onyx.external_permissions.utils import generic_doc_sync +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) +from ee.onyx.external_permissions.utils import credential_json, generic_doc_sync from onyx.access.models import ElementExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.box.connector import BoxConnector diff --git a/backend/ee/onyx/external_permissions/box/group_sync.py b/backend/ee/onyx/external_permissions/box/group_sync.py index 84501a2917..8feffd9b34 100644 --- a/backend/ee/onyx/external_permissions/box/group_sync.py +++ b/backend/ee/onyx/external_permissions/box/group_sync.py @@ -4,12 +4,14 @@ from box_sdk_gen import BoxClient from ee.onyx.db.external_perm import ExternalUserGroup from ee.onyx.external_permissions.utils import credential_json -from onyx.connectors.box.connector import box_all_enterprise_users_group_id -from onyx.connectors.box.connector import BOX_ENTERPRISE_ID_CREDENTIAL_KEY -from onyx.connectors.box.connector import box_group_id -from onyx.connectors.box.connector import BoxConnector -from onyx.connectors.box.connector import iter_box_enterprise_users -from onyx.connectors.box.connector import normalize_box_login +from onyx.connectors.box.connector import ( + box_all_enterprise_users_group_id, + BOX_ENTERPRISE_ID_CREDENTIAL_KEY, + box_group_id, + BoxConnector, + iter_box_enterprise_users, + normalize_box_login, +) from onyx.db.models import ConnectorCredentialPair from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/external_permissions/canvas/access.py b/backend/ee/onyx/external_permissions/canvas/access.py index fd643b6328..cfcf618a3c 100644 --- a/backend/ee/onyx/external_permissions/canvas/access.py +++ b/backend/ee/onyx/external_permissions/canvas/access.py @@ -1,14 +1,15 @@ from typing import Any -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.access.models import ExternalAccess from onyx.connectors.canvas.client import CanvasApiClient -from onyx.connectors.canvas.connector import canvas_all_users_group_id -from onyx.connectors.canvas.connector import canvas_course_group_id -from onyx.connectors.canvas.connector import canvas_group_group_id -from onyx.connectors.canvas.connector import canvas_section_group_id +from onyx.connectors.canvas.connector import ( + canvas_all_users_group_id, + canvas_course_group_id, + canvas_group_group_id, + canvas_section_group_id, +) from onyx.error_handling.exceptions import OnyxError from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/external_permissions/canvas/doc_sync.py b/backend/ee/onyx/external_permissions/canvas/doc_sync.py index 0c350df750..be8a894ca3 100644 --- a/backend/ee/onyx/external_permissions/canvas/doc_sync.py +++ b/backend/ee/onyx/external_permissions/canvas/doc_sync.py @@ -1,9 +1,10 @@ from collections.abc import Generator -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction -from ee.onyx.external_permissions.utils import credential_json -from ee.onyx.external_permissions.utils import generic_doc_sync +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) +from ee.onyx.external_permissions.utils import credential_json, generic_doc_sync from onyx.access.models import ElementExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.canvas.connector import CanvasConnector diff --git a/backend/ee/onyx/external_permissions/canvas/group_sync.py b/backend/ee/onyx/external_permissions/canvas/group_sync.py index ee223a2be6..b6b078d826 100644 --- a/backend/ee/onyx/external_permissions/canvas/group_sync.py +++ b/backend/ee/onyx/external_permissions/canvas/group_sync.py @@ -3,11 +3,13 @@ from collections.abc import Generator from ee.onyx.db.external_perm import ExternalUserGroup from ee.onyx.external_permissions.canvas.access import build_course_permission_context from ee.onyx.external_permissions.utils import credential_json -from onyx.connectors.canvas.connector import canvas_all_users_group_id -from onyx.connectors.canvas.connector import canvas_course_group_id -from onyx.connectors.canvas.connector import canvas_group_group_id -from onyx.connectors.canvas.connector import canvas_section_group_id -from onyx.connectors.canvas.connector import CanvasConnector +from onyx.connectors.canvas.connector import ( + canvas_all_users_group_id, + canvas_course_group_id, + canvas_group_group_id, + canvas_section_group_id, + CanvasConnector, +) from onyx.db.models import ConnectorCredentialPair from onyx.error_handling.exceptions import OnyxError from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/external_permissions/confluence/doc_sync.py b/backend/ee/onyx/external_permissions/confluence/doc_sync.py index ffa96a73d8..8bba584839 100644 --- a/backend/ee/onyx/external_permissions/confluence/doc_sync.py +++ b/backend/ee/onyx/external_permissions/confluence/doc_sync.py @@ -5,8 +5,10 @@ https://confluence.atlassian.com/conf85/check-who-can-view-a-page-1283360557.htm from collections.abc import Generator -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) from ee.onyx.external_permissions.utils import generic_doc_sync from onyx.access.models import ElementExternalAccess from onyx.configs.constants import DocumentSource diff --git a/backend/ee/onyx/external_permissions/confluence/group_sync.py b/backend/ee/onyx/external_permissions/confluence/group_sync.py index 9236135319..78c57f32df 100644 --- a/backend/ee/onyx/external_permissions/confluence/group_sync.py +++ b/backend/ee/onyx/external_permissions/confluence/group_sync.py @@ -6,8 +6,8 @@ from onyx.background.error_logging import emit_background_error from onyx.configs.app_configs import CONFLUENCE_USE_ONYX_USERS_FOR_GROUP_SYNC from onyx.connectors.confluence.onyx_confluence import ( get_user_email_from_username__server, + OnyxConfluence, ) -from onyx.connectors.confluence.onyx_confluence import OnyxConfluence from onyx.connectors.credentials_provider import OnyxDBCredentialsProvider from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import ConnectorCredentialPair diff --git a/backend/ee/onyx/external_permissions/confluence/page_access.py b/backend/ee/onyx/external_permissions/confluence/page_access.py index 0e7c6203a9..a80fd279bd 100644 --- a/backend/ee/onyx/external_permissions/confluence/page_access.py +++ b/backend/ee/onyx/external_permissions/confluence/page_access.py @@ -6,8 +6,8 @@ from onyx.access.utils import build_ext_group_name_for_onyx from onyx.configs.constants import DocumentSource from onyx.connectors.confluence.onyx_confluence import ( get_user_email_from_username__server, + OnyxConfluence, ) -from onyx.connectors.confluence.onyx_confluence import OnyxConfluence from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/external_permissions/confluence/space_access.py b/backend/ee/onyx/external_permissions/confluence/space_access.py index 75414b22bc..79cc72937f 100644 --- a/backend/ee/onyx/external_permissions/confluence/space_access.py +++ b/backend/ee/onyx/external_permissions/confluence/space_access.py @@ -1,32 +1,22 @@ from ee.onyx.configs.app_configs import CONFLUENCE_ANONYMOUS_ACCESS_IS_PUBLIC -from ee.onyx.external_permissions.confluence.constants import ALL_CONF_EMAILS_GROUP_NAME -from ee.onyx.external_permissions.confluence.constants import REQUEST_PAGINATION_LIMIT from ee.onyx.external_permissions.confluence.constants import ( + ALL_CONF_EMAILS_GROUP_NAME, + REQUEST_PAGINATION_LIMIT, SPACE_PERMISSION_OPERATION_READ, -) -from ee.onyx.external_permissions.confluence.constants import ( SPACE_PERMISSION_SUBJECT_TYPE_GROUP, -) -from ee.onyx.external_permissions.confluence.constants import ( SPACE_PERMISSION_SUBJECT_TYPE_USER, -) -from ee.onyx.external_permissions.confluence.constants import ( SPACE_PERMISSION_TARGET_TYPE_SPACE, + VIEWSPACE_PERMISSION_TYPE, ) -from ee.onyx.external_permissions.confluence.constants import VIEWSPACE_PERMISSION_TYPE from onyx.access.models import ExternalAccess from onyx.access.utils import build_ext_group_name_for_onyx from onyx.configs.constants import DocumentSource from onyx.connectors.confluence.onyx_confluence import ( ConfluenceRestSpacePermissionsNotAvailableError, -) -from onyx.connectors.confluence.onyx_confluence import ( get_user_email_from_userkey__server, -) -from onyx.connectors.confluence.onyx_confluence import ( get_user_email_from_username__server, + OnyxConfluence, ) -from onyx.connectors.confluence.onyx_confluence import OnyxConfluence from onyx.connectors.exceptions import InsufficientPermissionsError from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/external_permissions/github/doc_sync.py b/backend/ee/onyx/external_permissions/github/doc_sync.py index 6e42b9333d..20430ae356 100644 --- a/backend/ee/onyx/external_permissions/github/doc_sync.py +++ b/backend/ee/onyx/external_permissions/github/doc_sync.py @@ -4,26 +4,26 @@ from collections.abc import Generator from github import Github from github.Repository import Repository -from ee.onyx.external_permissions.github.utils import fetch_repository_team_slugs -from ee.onyx.external_permissions.github.utils import form_collaborators_group_id -from ee.onyx.external_permissions.github.utils import form_organization_group_id from ee.onyx.external_permissions.github.utils import ( + fetch_repository_team_slugs, + form_collaborators_group_id, + form_organization_group_id, form_outside_collaborators_group_id, + get_external_access_permission, + get_repository_visibility, + GitHubVisibility, +) +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, ) -from ee.onyx.external_permissions.github.utils import get_external_access_permission -from ee.onyx.external_permissions.github.utils import get_repository_visibility -from ee.onyx.external_permissions.github.utils import GitHubVisibility -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction from ee.onyx.external_permissions.utils import credential_json from onyx.access.models import DocExternalAccess from onyx.access.utils import build_ext_group_name_for_onyx from onyx.configs.constants import DocumentSource -from onyx.connectors.github.connector import DocMetadata -from onyx.connectors.github.connector import GithubConnector +from onyx.connectors.github.connector import DocMetadata, GithubConnector from onyx.db.models import ConnectorCredentialPair -from onyx.db.utils import DocumentRow -from onyx.db.utils import SortOrder +from onyx.db.utils import DocumentRow, SortOrder from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/external_permissions/github/utils.py b/backend/ee/onyx/external_permissions/github/utils.py index 5f6eee1e95..6e326e4ba3 100644 --- a/backend/ee/onyx/external_permissions/github/utils.py +++ b/backend/ee/onyx/external_permissions/github/utils.py @@ -1,12 +1,8 @@ from collections.abc import Callable from enum import Enum -from typing import List -from typing import Optional -from typing import Tuple -from typing import TypeVar +from typing import List, Optional, Tuple, TypeVar -from github import Github -from github import RateLimitExceededException +from github import Github, RateLimitExceededException from github.GithubException import GithubException from github.NamedUser import NamedUser from github.Organization import Organization diff --git a/backend/ee/onyx/external_permissions/gmail/doc_sync.py b/backend/ee/onyx/external_permissions/gmail/doc_sync.py index 7b8253078f..a255f523e1 100644 --- a/backend/ee/onyx/external_permissions/gmail/doc_sync.py +++ b/backend/ee/onyx/external_permissions/gmail/doc_sync.py @@ -1,13 +1,16 @@ from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) from ee.onyx.external_permissions.utils import credential_json -from onyx.access.models import DocExternalAccess -from onyx.access.models import ElementExternalAccess -from onyx.access.models import NodeExternalAccess +from onyx.access.models import ( + DocExternalAccess, + ElementExternalAccess, + NodeExternalAccess, +) from onyx.configs.constants import DocumentSource from onyx.connectors.gmail.connector import GmailConnector from onyx.connectors.interfaces import GenerateSlimDocumentOutput diff --git a/backend/ee/onyx/external_permissions/google_drive/doc_sync.py b/backend/ee/onyx/external_permissions/google_drive/doc_sync.py index 7a1c20289a..a2b4dd40df 100644 --- a/backend/ee/onyx/external_permissions/google_drive/doc_sync.py +++ b/backend/ee/onyx/external_permissions/google_drive/doc_sync.py @@ -1,21 +1,25 @@ from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from ee.onyx.external_permissions.google_drive.models import GoogleDrivePermission -from ee.onyx.external_permissions.google_drive.models import PermissionType +from ee.onyx.external_permissions.google_drive.models import ( + GoogleDrivePermission, + PermissionType, +) from ee.onyx.external_permissions.google_drive.permission_retrieval import ( get_permissions_by_ids, ) -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) from ee.onyx.external_permissions.utils import credential_json -from onyx.access.models import DocExternalAccess -from onyx.access.models import ElementExternalAccess -from onyx.access.models import ExternalAccess -from onyx.access.models import NodeExternalAccess -from onyx.access.utils import build_domain_group_id -from onyx.access.utils import build_ext_group_name_for_onyx +from onyx.access.models import ( + DocExternalAccess, + ElementExternalAccess, + ExternalAccess, + NodeExternalAccess, +) +from onyx.access.utils import build_domain_group_id, build_ext_group_name_for_onyx from onyx.configs.constants import DocumentSource from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.connectors.google_drive.models import GoogleDriveFileType diff --git a/backend/ee/onyx/external_permissions/google_drive/group_sync.py b/backend/ee/onyx/external_permissions/google_drive/group_sync.py index 237a8a930a..9be365ae57 100644 --- a/backend/ee/onyx/external_permissions/google_drive/group_sync.py +++ b/backend/ee/onyx/external_permissions/google_drive/group_sync.py @@ -6,19 +6,21 @@ from pydantic import BaseModel from ee.onyx.db.external_perm import ExternalUserGroup from ee.onyx.external_permissions.google_drive.folder_retrieval import ( get_folder_permissions_by_ids, -) -from ee.onyx.external_permissions.google_drive.folder_retrieval import ( get_modified_folders, ) -from ee.onyx.external_permissions.google_drive.models import GoogleDrivePermission -from ee.onyx.external_permissions.google_drive.models import PermissionType +from ee.onyx.external_permissions.google_drive.models import ( + GoogleDrivePermission, + PermissionType, +) from ee.onyx.external_permissions.utils import credential_json from onyx.access.utils import build_domain_group_id from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.connectors.google_utils.google_utils import execute_paginated_retrieval -from onyx.connectors.google_utils.resources import AdminService -from onyx.connectors.google_utils.resources import get_admin_service -from onyx.connectors.google_utils.resources import get_drive_service +from onyx.connectors.google_utils.resources import ( + AdminService, + get_admin_service, + get_drive_service, +) from onyx.db.models import ConnectorCredentialPair from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/external_permissions/jira/doc_sync.py b/backend/ee/onyx/external_permissions/jira/doc_sync.py index 121d635201..792fc2f3f6 100644 --- a/backend/ee/onyx/external_permissions/jira/doc_sync.py +++ b/backend/ee/onyx/external_permissions/jira/doc_sync.py @@ -1,9 +1,10 @@ from collections.abc import Generator -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction -from ee.onyx.external_permissions.utils import credential_json -from ee.onyx.external_permissions.utils import generic_doc_sync +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) +from ee.onyx.external_permissions.utils import credential_json, generic_doc_sync from onyx.access.models import ElementExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.jira.connector import JiraConnector diff --git a/backend/ee/onyx/external_permissions/jira/models.py b/backend/ee/onyx/external_permissions/jira/models.py index ab30c9a320..87a4a1327f 100644 --- a/backend/ee/onyx/external_permissions/jira/models.py +++ b/backend/ee/onyx/external_permissions/jira/models.py @@ -1,7 +1,6 @@ from typing import Any -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from pydantic.alias_generators import to_camel Holder = dict[str, Any] diff --git a/backend/ee/onyx/external_permissions/jira/page_access.py b/backend/ee/onyx/external_permissions/jira/page_access.py index a3ae65712e..07a401dd94 100644 --- a/backend/ee/onyx/external_permissions/jira/page_access.py +++ b/backend/ee/onyx/external_permissions/jira/page_access.py @@ -5,9 +5,7 @@ from jira import JIRA from jira.resources import PermissionScheme from pydantic import ValidationError -from ee.onyx.external_permissions.jira.models import Holder -from ee.onyx.external_permissions.jira.models import Permission -from ee.onyx.external_permissions.jira.models import User +from ee.onyx.external_permissions.jira.models import Holder, Permission, User from onyx.access.models import ExternalAccess from onyx.access.utils import build_ext_group_name_for_onyx from onyx.configs.constants import DocumentSource diff --git a/backend/ee/onyx/external_permissions/perm_sync_types.py b/backend/ee/onyx/external_permissions/perm_sync_types.py index 26dba539ce..2be14ab0b6 100644 --- a/backend/ee/onyx/external_permissions/perm_sync_types.py +++ b/backend/ee/onyx/external_permissions/perm_sync_types.py @@ -1,16 +1,15 @@ -from collections.abc import Callable -from collections.abc import Generator -from typing import Optional -from typing import Protocol +from collections.abc import Callable, Generator +from typing import Optional, Protocol from ee.onyx.db.external_perm import ExternalUserGroup # noqa -from onyx.access.models import DocExternalAccess # noqa -from onyx.access.models import ElementExternalAccess # noqa -from onyx.access.models import NodeExternalAccess # noqa +from onyx.access.models import ( + DocExternalAccess, # noqa + ElementExternalAccess, # noqa + NodeExternalAccess, # noqa +) from onyx.context.search.models import InferenceChunk from onyx.db.models import ConnectorCredentialPair # noqa -from onyx.db.utils import DocumentRow -from onyx.db.utils import SortOrder +from onyx.db.utils import DocumentRow, SortOrder from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface # noqa diff --git a/backend/ee/onyx/external_permissions/post_query_censoring.py b/backend/ee/onyx/external_permissions/post_query_censoring.py index 74b15c4db8..ba05586c27 100644 --- a/backend/ee/onyx/external_permissions/post_query_censoring.py +++ b/backend/ee/onyx/external_permissions/post_query_censoring.py @@ -1,6 +1,8 @@ from ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs -from ee.onyx.external_permissions.sync_params import get_all_censoring_enabled_sources -from ee.onyx.external_permissions.sync_params import get_source_perm_sync_config +from ee.onyx.external_permissions.sync_params import ( + get_all_censoring_enabled_sources, + get_source_perm_sync_config, +) from onyx.configs.constants import DocumentSource from onyx.context.search.pipeline import InferenceChunk from onyx.db.engine.sql_engine import get_session_with_current_tenant diff --git a/backend/ee/onyx/external_permissions/salesforce/postprocessing.py b/backend/ee/onyx/external_permissions/salesforce/postprocessing.py index 6f02e8f823..55a9f1a0df 100644 --- a/backend/ee/onyx/external_permissions/salesforce/postprocessing.py +++ b/backend/ee/onyx/external_permissions/salesforce/postprocessing.py @@ -3,9 +3,7 @@ import time from ee.onyx.db.external_perm import fetch_external_groups_for_user_email_and_group_ids from ee.onyx.external_permissions.salesforce.utils import ( get_any_salesforce_client_for_doc_id, -) -from ee.onyx.external_permissions.salesforce.utils import get_objects_access_for_user_id -from ee.onyx.external_permissions.salesforce.utils import ( + get_objects_access_for_user_id, get_salesforce_user_id_from_email, ) from onyx.configs.app_configs import BLURB_SIZE diff --git a/backend/ee/onyx/external_permissions/sharepoint/doc_sync.py b/backend/ee/onyx/external_permissions/sharepoint/doc_sync.py index 66ed9809fc..dd24eb6983 100644 --- a/backend/ee/onyx/external_permissions/sharepoint/doc_sync.py +++ b/backend/ee/onyx/external_permissions/sharepoint/doc_sync.py @@ -1,9 +1,10 @@ from collections.abc import Generator -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction -from ee.onyx.external_permissions.utils import credential_json -from ee.onyx.external_permissions.utils import generic_doc_sync +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) +from ee.onyx.external_permissions.utils import credential_json, generic_doc_sync from onyx.access.models import ElementExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.sharepoint.connector import SharepointConnector diff --git a/backend/ee/onyx/external_permissions/sharepoint/group_sync.py b/backend/ee/onyx/external_permissions/sharepoint/group_sync.py index 95f2942c9d..5058542b84 100644 --- a/backend/ee/onyx/external_permissions/sharepoint/group_sync.py +++ b/backend/ee/onyx/external_permissions/sharepoint/group_sync.py @@ -8,8 +8,10 @@ from ee.onyx.external_permissions.sharepoint.permission_utils import ( ) from ee.onyx.external_permissions.utils import credential_json from onyx.configs.app_configs import SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION -from onyx.connectors.sharepoint.connector import acquire_token_for_rest -from onyx.connectors.sharepoint.connector import SharepointConnector +from onyx.connectors.sharepoint.connector import ( + acquire_token_for_rest, + SharepointConnector, +) from onyx.db.models import ConnectorCredentialPair from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/external_permissions/sharepoint/permission_utils.py b/backend/ee/onyx/external_permissions/sharepoint/permission_utils.py index b54ced5303..7bf529b436 100644 --- a/backend/ee/onyx/external_permissions/sharepoint/permission_utils.py +++ b/backend/ee/onyx/external_permissions/sharepoint/permission_utils.py @@ -1,9 +1,7 @@ import re import time from collections import deque -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterable +from collections.abc import Callable, Generator, Iterable from typing import Any from urllib.parse import urlparse @@ -23,10 +21,12 @@ from onyx.access.models import ExternalAccess from onyx.access.utils import build_ext_group_name_for_onyx from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource -from onyx.connectors.sharepoint.connector import GRAPH_API_MAX_RETRIES -from onyx.connectors.sharepoint.connector import GRAPH_API_RETRYABLE_STATUSES -from onyx.connectors.sharepoint.connector import SHARED_DOCUMENTS_MAP_REVERSE -from onyx.connectors.sharepoint.connector import sleep_and_retry +from onyx.connectors.sharepoint.connector import ( + GRAPH_API_MAX_RETRIES, + GRAPH_API_RETRYABLE_STATUSES, + SHARED_DOCUMENTS_MAP_REVERSE, + sleep_and_retry, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_after import parse_retry_after_seconds diff --git a/backend/ee/onyx/external_permissions/slack/channel_access.py b/backend/ee/onyx/external_permissions/slack/channel_access.py index eb8cf9c9a0..4020a61a50 100644 --- a/backend/ee/onyx/external_permissions/slack/channel_access.py +++ b/backend/ee/onyx/external_permissions/slack/channel_access.py @@ -2,10 +2,11 @@ from slack_sdk import WebClient from onyx.access.models import ExternalAccess from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.slack.connector import channel_team_ids -from onyx.connectors.slack.connector import ChannelType -from onyx.connectors.slack.utils import expert_info_from_slack_id -from onyx.connectors.slack.utils import make_paginated_slack_api_call +from onyx.connectors.slack.connector import channel_team_ids, ChannelType +from onyx.connectors.slack.utils import ( + expert_info_from_slack_id, + make_paginated_slack_api_call, +) def get_channel_access( diff --git a/backend/ee/onyx/external_permissions/slack/doc_sync.py b/backend/ee/onyx/external_permissions/slack/doc_sync.py index f3d4ed80d7..c1ddab1b64 100644 --- a/backend/ee/onyx/external_permissions/slack/doc_sync.py +++ b/backend/ee/onyx/external_permissions/slack/doc_sync.py @@ -2,23 +2,28 @@ from collections.abc import Generator from slack_sdk import WebClient -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) from ee.onyx.external_permissions.slack.channel_access import get_channel_access -from ee.onyx.external_permissions.slack.utils import fetch_team_user_emails -from ee.onyx.external_permissions.slack.utils import fetch_user_id_to_email_map +from ee.onyx.external_permissions.slack.utils import ( + fetch_team_user_emails, + fetch_user_id_to_email_map, +) from ee.onyx.external_permissions.utils import credential_json -from onyx.access.models import DocExternalAccess -from onyx.access.models import ExternalAccess +from onyx.access.models import DocExternalAccess, ExternalAccess from onyx.connectors.credentials_provider import OnyxDBCredentialsProvider from onyx.connectors.interfaces import SecondsSinceUnixEpoch from onyx.connectors.models import HierarchyNode -from onyx.connectors.slack.connector import filter_channels -from onyx.connectors.slack.connector import get_channels -from onyx.connectors.slack.connector import get_channels_across_teams -from onyx.connectors.slack.connector import list_grid_team_ids -from onyx.connectors.slack.connector import make_paginated_slack_api_call -from onyx.connectors.slack.connector import SlackConnector +from onyx.connectors.slack.connector import ( + filter_channels, + get_channels, + get_channels_across_teams, + list_grid_team_ids, + make_paginated_slack_api_call, + SlackConnector, +) from onyx.connectors.slack.models import ChannelType from onyx.db.models import ConnectorCredentialPair from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface diff --git a/backend/ee/onyx/external_permissions/sync_params.py b/backend/ee/onyx/external_permissions/sync_params.py index ad59d88f41..0db93e8a45 100644 --- a/backend/ee/onyx/external_permissions/sync_params.py +++ b/backend/ee/onyx/external_permissions/sync_params.py @@ -1,25 +1,26 @@ from collections.abc import Generator -from typing import Optional -from typing import TYPE_CHECKING +from typing import Optional, TYPE_CHECKING from pydantic import BaseModel -from ee.onyx.configs.app_configs import BOX_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import BOX_PERMISSION_GROUP_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import CANVAS_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import CANVAS_PERMISSION_GROUP_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import GITHUB_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import JIRA_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import JIRA_PERMISSION_GROUP_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import SLACK_PERMISSION_DOC_SYNC_FREQUENCY -from ee.onyx.configs.app_configs import TEAMS_PERMISSION_DOC_SYNC_FREQUENCY +from ee.onyx.configs.app_configs import ( + BOX_PERMISSION_DOC_SYNC_FREQUENCY, + BOX_PERMISSION_GROUP_SYNC_FREQUENCY, + CANVAS_PERMISSION_DOC_SYNC_FREQUENCY, + CANVAS_PERMISSION_GROUP_SYNC_FREQUENCY, + CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY, + CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY, + DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY, + GITHUB_PERMISSION_DOC_SYNC_FREQUENCY, + GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY, + GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY, + JIRA_PERMISSION_DOC_SYNC_FREQUENCY, + JIRA_PERMISSION_GROUP_SYNC_FREQUENCY, + SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY, + SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY, + 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 @@ -33,11 +34,13 @@ 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 -from ee.onyx.external_permissions.perm_sync_types import DocSyncFuncType -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction -from ee.onyx.external_permissions.perm_sync_types import GroupSyncFuncType +from ee.onyx.external_permissions.perm_sync_types import ( + CensoringFuncType, + DocSyncFuncType, + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, + GroupSyncFuncType, +) from ee.onyx.external_permissions.salesforce.postprocessing import ( censor_salesforce_chunks, ) diff --git a/backend/ee/onyx/external_permissions/teams/doc_sync.py b/backend/ee/onyx/external_permissions/teams/doc_sync.py index 632b5b8a2a..30884d1469 100644 --- a/backend/ee/onyx/external_permissions/teams/doc_sync.py +++ b/backend/ee/onyx/external_permissions/teams/doc_sync.py @@ -1,9 +1,10 @@ from collections.abc import Generator -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction -from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction -from ee.onyx.external_permissions.utils import credential_json -from ee.onyx.external_permissions.utils import generic_doc_sync +from ee.onyx.external_permissions.perm_sync_types import ( + FetchAllDocumentsFunction, + FetchAllDocumentsIdsFunction, +) +from ee.onyx.external_permissions.utils import credential_json, generic_doc_sync from onyx.access.models import ElementExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.teams.connector import TeamsConnector diff --git a/backend/ee/onyx/external_permissions/utils.py b/backend/ee/onyx/external_permissions/utils.py index 90e5031d6f..bf46c44eee 100644 --- a/backend/ee/onyx/external_permissions/utils.py +++ b/backend/ee/onyx/external_permissions/utils.py @@ -2,13 +2,14 @@ from collections.abc import Generator from typing import Any from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction -from onyx.access.models import DocExternalAccess -from onyx.access.models import ElementExternalAccess -from onyx.access.models import ExternalAccess -from onyx.access.models import NodeExternalAccess +from onyx.access.models import ( + DocExternalAccess, + ElementExternalAccess, + ExternalAccess, + NodeExternalAccess, +) from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync +from onyx.connectors.interfaces import SecondsSinceUnixEpoch, SlimConnectorWithPermSync from onyx.connectors.models import HierarchyNode from onyx.db.models import ConnectorCredentialPair from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface diff --git a/backend/ee/onyx/feature_flags/factory.py b/backend/ee/onyx/feature_flags/factory.py index 47c8ba1634..8b20eca413 100644 --- a/backend/ee/onyx/feature_flags/factory.py +++ b/backend/ee/onyx/feature_flags/factory.py @@ -1,7 +1,6 @@ from ee.onyx.feature_flags.posthog_provider import PostHogFeatureFlagProvider from ee.onyx.utils.posthog_client import posthog -from onyx.feature_flags.interface import FeatureFlagProvider -from onyx.feature_flags.interface import NoOpFeatureFlagProvider +from onyx.feature_flags.interface import FeatureFlagProvider, NoOpFeatureFlagProvider def get_posthog_feature_flag_provider() -> FeatureFlagProvider: diff --git a/backend/ee/onyx/hooks/executor.py b/backend/ee/onyx/hooks/executor.py index c0c3d96731..75bf416176 100644 --- a/backend/ee/onyx/hooks/executor.py +++ b/backend/ee/onyx/hooks/executor.py @@ -48,26 +48,27 @@ The executor uses three sessions: prevent the execution log from being written. This update is best-effort. """ -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar from pydantic import BaseModel from sqlalchemy.orm import Session from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint -from onyx.db.hook import create_hook_execution_log__no_commit -from onyx.db.hook import get_non_deleted_hook_by_hook_point -from onyx.db.hook import update_hook__no_commit +from onyx.db.enums import HookFailStrategy, HookPoint +from onyx.db.hook import ( + create_hook_execution_log__no_commit, + get_non_deleted_hook_by_hook_point, + update_hook__no_commit, +) from onyx.db.models import Hook from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.hooks.executor import HookSkipped -from onyx.hooks.executor import HookSoftFailed -from onyx.utils.external_endpoint import ExternalEndpointConfig -from onyx.utils.external_endpoint import ExternalEndpointOutcome -from onyx.utils.external_endpoint import post_json_to_endpoint +from onyx.hooks.executor import HookSkipped, HookSoftFailed +from onyx.utils.external_endpoint import ( + ExternalEndpointConfig, + ExternalEndpointOutcome, + post_json_to_endpoint, +) from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/ee/onyx/main.py b/backend/ee/onyx/main.py index a1ce40f541..33672e14d2 100644 --- a/backend/ee/onyx/main.py +++ b/backend/ee/onyx/main.py @@ -30,8 +30,7 @@ from ee.onyx.server.query_and_chat.query_backend import basic_router as ee_query from ee.onyx.server.query_and_chat.search_backend import router as search_router from ee.onyx.server.query_history.api import router as query_history_router from ee.onyx.server.reporting.usage_export_api import router as usage_export_router -from ee.onyx.server.scim.api import register_scim_exception_handlers -from ee.onyx.server.scim.api import scim_router +from ee.onyx.server.scim.api import register_scim_exception_handlers, scim_router from ee.onyx.server.seeding import seed_db from ee.onyx.server.tenants.api import router as tenants_router from ee.onyx.server.token_rate_limits.api import ( @@ -39,19 +38,22 @@ from ee.onyx.server.token_rate_limits.api import ( ) from ee.onyx.server.user_group.api import router as user_group_router from ee.onyx.utils.encryption import test_encryption -from onyx.auth.users import auth_backend -from onyx.auth.users import create_onyx_oauth_router -from onyx.configs.app_configs import GOOGLE_LOGIN_BASE_SCOPES -from onyx.configs.app_configs import GOOGLE_OAUTH_SCOPE_OVERRIDE -from onyx.configs.app_configs import OAUTH_CLIENT_ID -from onyx.configs.app_configs import OAUTH_CLIENT_SECRET -from onyx.configs.app_configs import USER_AUTH_SECRET -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.auth.users import auth_backend, create_onyx_oauth_router +from onyx.configs.app_configs import ( + GOOGLE_LOGIN_BASE_SCOPES, + GOOGLE_OAUTH_SCOPE_OVERRIDE, + OAUTH_CLIENT_ID, + OAUTH_CLIENT_SECRET, + USER_AUTH_SECRET, + WEB_DOMAIN, +) from onyx.main import get_application as get_application_base -from onyx.main import include_auth_router_with_prefix -from onyx.main import include_router_with_global_prefix_prepended +from onyx.main import ( + include_auth_router_with_prefix, + include_router_with_global_prefix_prepended, + use_route_function_names_as_operation_ids, +) from onyx.main import lifespan as lifespan_base -from onyx.main import use_route_function_names_as_operation_ids from onyx.server.query_and_chat.query_backend import basic_router as query_router from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import global_version diff --git a/backend/ee/onyx/onyxbot/slack/handlers/handle_standard_answers.py b/backend/ee/onyx/onyxbot/slack/handlers/handle_standard_answers.py index 6fb35eb928..cab62d694a 100644 --- a/backend/ee/onyx/onyxbot/slack/handlers/handle_standard_answers.py +++ b/backend/ee/onyx/onyxbot/slack/handlers/handle_standard_answers.py @@ -1,29 +1,28 @@ from slack_sdk import WebClient -from slack_sdk.models.blocks import ActionsBlock -from slack_sdk.models.blocks import Block -from slack_sdk.models.blocks import ButtonElement -from slack_sdk.models.blocks import SectionBlock +from slack_sdk.models.blocks import ActionsBlock, Block, ButtonElement, SectionBlock from sqlalchemy.orm import Session -from ee.onyx.db.standard_answer import fetch_standard_answer_categories_by_names -from ee.onyx.db.standard_answer import find_matching_standard_answers +from ee.onyx.db.standard_answer import ( + fetch_standard_answer_categories_by_names, + find_matching_standard_answers, +) from onyx.configs.constants import MessageType from onyx.configs.onyxbot_configs import ONYX_BOT_REACT_EMOJI -from onyx.db.chat import create_chat_session -from onyx.db.chat import create_new_chat_message -from onyx.db.chat import get_chat_messages_by_sessions -from onyx.db.chat import get_chat_sessions_by_slack_thread_id -from onyx.db.chat import get_or_create_root_message +from onyx.db.chat import ( + create_chat_session, + create_new_chat_message, + get_chat_messages_by_sessions, + get_chat_sessions_by_slack_thread_id, + get_or_create_root_message, +) from onyx.db.models import SlackChannelConfig from onyx.db.models import StandardAnswer as StandardAnswerModel from onyx.onyxbot.slack.blocks import get_restate_blocks from onyx.onyxbot.slack.constants import GENERATE_ANSWER_BUTTON_ACTION_ID from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.onyxbot.slack.utils import respond_in_thread_or_channel -from onyx.onyxbot.slack.utils import update_emote_react +from onyx.onyxbot.slack.utils import respond_in_thread_or_channel, update_emote_react from onyx.server.manage.models import StandardAnswer as PydanticStandardAnswer -from onyx.utils.logger import OnyxLoggingAdapter -from onyx.utils.logger import setup_logger +from onyx.utils.logger import OnyxLoggingAdapter, setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/search/process_search_query.py b/backend/ee/onyx/search/process_search_query.py index 5c2bb77223..3ef7bdf191 100644 --- a/backend/ee/onyx/search/process_search_query.py +++ b/backend/ee/onyx/search/process_search_query.py @@ -4,18 +4,19 @@ from sqlalchemy.orm import Session from ee.onyx.db.search import create_search_query from ee.onyx.secondary_llm_flows.query_expansion import expand_keywords -from ee.onyx.server.query_and_chat.models import SearchDocWithContent -from ee.onyx.server.query_and_chat.models import SearchFullResponse -from ee.onyx.server.query_and_chat.models import SendSearchQueryRequest -from ee.onyx.server.query_and_chat.streaming_models import LLMSelectedDocsPacket -from ee.onyx.server.query_and_chat.streaming_models import SearchDocsPacket -from ee.onyx.server.query_and_chat.streaming_models import SearchErrorPacket -from ee.onyx.server.query_and_chat.streaming_models import SearchQueriesPacket -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import ChunkSearchRequest -from onyx.context.search.models import InferenceChunk -from onyx.context.search.pipeline import merge_individual_chunks -from onyx.context.search.pipeline import search_pipeline +from ee.onyx.server.query_and_chat.models import ( + SearchDocWithContent, + SearchFullResponse, + SendSearchQueryRequest, +) +from ee.onyx.server.query_and_chat.streaming_models import ( + LLMSelectedDocsPacket, + SearchDocsPacket, + SearchErrorPacket, + SearchQueriesPacket, +) +from onyx.context.search.models import BaseFilters, ChunkSearchRequest, InferenceChunk +from onyx.context.search.pipeline import merge_individual_chunks, search_pipeline from onyx.db.models import User from onyx.db.search_settings import get_current_search_settings from onyx.document_index.factory import get_default_document_index diff --git a/backend/ee/onyx/secondary_llm_flows/query_expansion.py b/backend/ee/onyx/secondary_llm_flows/query_expansion.py index 2bb52b117d..a1a70b91ca 100644 --- a/backend/ee/onyx/secondary_llm_flows/query_expansion.py +++ b/backend/ee/onyx/secondary_llm_flows/query_expansion.py @@ -2,9 +2,7 @@ import re from ee.onyx.prompts.query_expansion import KEYWORD_EXPANSION_PROMPT from onyx.llm.interfaces import LLM -from onyx.llm.models import LanguageModelInput -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import UserMessage +from onyx.llm.models import LanguageModelInput, ReasoningEffort, UserMessage from onyx.llm.utils import llm_response_to_string from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/secondary_llm_flows/search_flow_classification.py b/backend/ee/onyx/secondary_llm_flows/search_flow_classification.py index 9e4db38a88..bb4c685973 100644 --- a/backend/ee/onyx/secondary_llm_flows/search_flow_classification.py +++ b/backend/ee/onyx/secondary_llm_flows/search_flow_classification.py @@ -1,10 +1,10 @@ -from ee.onyx.prompts.search_flow_classification import CHAT_CLASS -from ee.onyx.prompts.search_flow_classification import SEARCH_CHAT_PROMPT -from ee.onyx.prompts.search_flow_classification import SEARCH_CLASS +from ee.onyx.prompts.search_flow_classification import ( + CHAT_CLASS, + SEARCH_CHAT_PROMPT, + SEARCH_CLASS, +) from onyx.llm.interfaces import LLM -from onyx.llm.models import LanguageModelInput -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import UserMessage +from onyx.llm.models import LanguageModelInput, ReasoningEffort, UserMessage from onyx.llm.utils import llm_response_to_string from onyx.utils.logger import setup_logger from onyx.utils.timing import log_function_time diff --git a/backend/ee/onyx/server/analytics/api.py b/backend/ee/onyx/server/analytics/api.py index dd246daf48..93c746e9c0 100644 --- a/backend/ee/onyx/server/analytics/api.py +++ b/backend/ee/onyx/server/analytics/api.py @@ -2,21 +2,21 @@ import datetime from collections import defaultdict from typing import List -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from sqlalchemy.orm import Session -from ee.onyx.db.analytics import fetch_assistant_message_analytics -from ee.onyx.db.analytics import fetch_assistant_unique_users -from ee.onyx.db.analytics import fetch_assistant_unique_users_total -from ee.onyx.db.analytics import fetch_onyxbot_analytics -from ee.onyx.db.analytics import fetch_per_user_query_analytics -from ee.onyx.db.analytics import fetch_persona_message_analytics -from ee.onyx.db.analytics import fetch_persona_unique_users -from ee.onyx.db.analytics import fetch_query_analytics -from ee.onyx.db.analytics import user_can_view_assistant_stats +from ee.onyx.db.analytics import ( + fetch_assistant_message_analytics, + fetch_assistant_unique_users, + fetch_assistant_unique_users_total, + fetch_onyxbot_analytics, + fetch_per_user_query_analytics, + fetch_persona_message_analytics, + fetch_persona_unique_users, + fetch_query_analytics, + user_can_view_assistant_stats, +) from onyx.auth.permissions import require_permission from onyx.configs.constants import PUBLIC_API_TAGS from onyx.db.engine.sql_engine import get_session diff --git a/backend/ee/onyx/server/auth_check.py b/backend/ee/onyx/server/auth_check.py index d8a55a8c70..488379f721 100644 --- a/backend/ee/onyx/server/auth_check.py +++ b/backend/ee/onyx/server/auth_check.py @@ -1,7 +1,6 @@ from fastapi import FastAPI -from onyx.server.auth_check import check_router_auth -from onyx.server.auth_check import PUBLIC_ENDPOINT_SPECS +from onyx.server.auth_check import check_router_auth, PUBLIC_ENDPOINT_SPECS EE_PUBLIC_ENDPOINT_SPECS = PUBLIC_ENDPOINT_SPECS + [ # SCIM 2.0 service discovery — unauthenticated so IdPs can probe diff --git a/backend/ee/onyx/server/billing/api.py b/backend/ee/onyx/server/billing/api.py index 784a27185e..5ba9fd3e3d 100644 --- a/backend/ee/onyx/server/billing/api.py +++ b/backend/ee/onyx/server/billing/api.py @@ -24,28 +24,27 @@ See: https://linear.app/onyx-app/issue/ENG-3533/migrate-tenantsbilling-adminbill import asyncio import httpx -from fastapi import APIRouter -from fastapi import Depends -from pydantic import BaseModel -from pydantic import ValidationError +from fastapi import APIRouter, Depends +from pydantic import BaseModel, ValidationError from redis.exceptions import RedisError from sqlalchemy.orm import Session -from ee.onyx.db.license import get_license -from ee.onyx.db.license import get_used_seats +from ee.onyx.db.license import get_license, get_used_seats from ee.onyx.server.billing.billing_cache import ( invalidate_billing_cache as invalidate_indexing_trial_cache, ) -from ee.onyx.server.billing.models import BillingInformationResponse -from ee.onyx.server.billing.models import CreateCheckoutSessionRequest -from ee.onyx.server.billing.models import CreateCheckoutSessionResponse -from ee.onyx.server.billing.models import CreateCustomerPortalSessionRequest -from ee.onyx.server.billing.models import CreateCustomerPortalSessionResponse -from ee.onyx.server.billing.models import EndTrialResponse -from ee.onyx.server.billing.models import SeatUpdateRequest -from ee.onyx.server.billing.models import SeatUpdateResponse -from ee.onyx.server.billing.models import StripePublishableKeyResponse -from ee.onyx.server.billing.models import SubscriptionStatusResponse +from ee.onyx.server.billing.models import ( + BillingInformationResponse, + CreateCheckoutSessionRequest, + CreateCheckoutSessionResponse, + CreateCustomerPortalSessionRequest, + CreateCustomerPortalSessionResponse, + EndTrialResponse, + SeatUpdateRequest, + SeatUpdateResponse, + StripePublishableKeyResponse, + SubscriptionStatusResponse, +) from ee.onyx.server.billing.service import ( create_checkout_session as create_checkout_service, ) @@ -59,15 +58,16 @@ from ee.onyx.server.billing.service import ( from ee.onyx.server.billing.service import update_seat_count as update_seat_service from onyx.auth.permissions import require_permission from onyx.auth.users import User -from onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_OVERRIDE -from onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_URL -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.app_configs import ( + STRIPE_PUBLISHABLE_KEY_OVERRIDE, + STRIPE_PUBLISHABLE_KEY_URL, + WEB_DOMAIN, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_shared_redis_client +from onyx.redis.redis_pool import get_redis_client, get_shared_redis_client from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/ee/onyx/server/billing/billing_cache.py b/backend/ee/onyx/server/billing/billing_cache.py index bf0957937c..f813a5a3bf 100644 --- a/backend/ee/onyx/server/billing/billing_cache.py +++ b/backend/ee/onyx/server/billing/billing_cache.py @@ -24,8 +24,7 @@ from pydantic import ValidationError from redis.exceptions import RedisError from ee.onyx.server.tenants.billing import fetch_billing_information -from ee.onyx.server.tenants.models import BillingInformation -from ee.onyx.server.tenants.models import SubscriptionStatusResponse +from ee.onyx.server.tenants.models import BillingInformation, SubscriptionStatusResponse from onyx.configs.app_configs import BILLING_CACHE_TTL_SECONDS from onyx.redis.redis_pool import get_shared_redis_client from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/server/billing/service.py b/backend/ee/onyx/server/billing/service.py index 8950ecae43..bb2a31fd4b 100644 --- a/backend/ee/onyx/server/billing/service.py +++ b/backend/ee/onyx/server/billing/service.py @@ -15,13 +15,15 @@ from typing import Literal import httpx from ee.onyx.configs.app_configs import CLOUD_DATA_PLANE_URL -from ee.onyx.server.billing.models import BillingInformationResponse -from ee.onyx.server.billing.models import CreateCheckoutSessionResponse -from ee.onyx.server.billing.models import CreateCustomerPortalSessionResponse -from ee.onyx.server.billing.models import EndTrialResponse -from ee.onyx.server.billing.models import SeatUpdateResponse -from ee.onyx.server.billing.models import StripePortalFlowType -from ee.onyx.server.billing.models import SubscriptionStatusResponse +from ee.onyx.server.billing.models import ( + BillingInformationResponse, + CreateCheckoutSessionResponse, + CreateCustomerPortalSessionResponse, + EndTrialResponse, + SeatUpdateResponse, + StripePortalFlowType, + SubscriptionStatusResponse, +) from ee.onyx.server.tenants.access import generate_data_plane_token from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL from onyx.error_handling.error_codes import OnyxErrorCode diff --git a/backend/ee/onyx/server/documents/cc_pair.py b/backend/ee/onyx/server/documents/cc_pair.py index bfe242358b..c24000e4e8 100644 --- a/backend/ee/onyx/server/documents/cc_pair.py +++ b/backend/ee/onyx/server/documents/cc_pair.py @@ -1,9 +1,7 @@ from datetime import datetime from http import HTTPStatus -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from ee.onyx.background.celery.tasks.doc_permission_syncing.tasks import ( diff --git a/backend/ee/onyx/server/enterprise_settings/api.py b/backend/ee/onyx/server/enterprise_settings/api.py index 6c3d90a736..88f221ae69 100644 --- a/backend/ee/onyx/server/enterprise_settings/api.py +++ b/backend/ee/onyx/server/enterprise_settings/api.py @@ -1,37 +1,38 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import httpx -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Response -from fastapi import status -from fastapi import UploadFile -from pydantic import BaseModel -from pydantic import Field +from fastapi import APIRouter, Depends, HTTPException, Response, status, UploadFile +from pydantic import BaseModel, Field from sqlalchemy.orm import Session from ee.onyx.db.scim import ScimDAL -from ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload -from ee.onyx.server.enterprise_settings.models import EnterpriseSettings -from ee.onyx.server.enterprise_settings.store import get_logo_filename -from ee.onyx.server.enterprise_settings.store import get_logotype_filename -from ee.onyx.server.enterprise_settings.store import load_analytics_script -from ee.onyx.server.enterprise_settings.store import load_settings -from ee.onyx.server.enterprise_settings.store import store_analytics_script -from ee.onyx.server.enterprise_settings.store import store_settings -from ee.onyx.server.enterprise_settings.store import upload_logo +from ee.onyx.server.enterprise_settings.models import ( + AnalyticsScriptUpload, + EnterpriseSettings, +) +from ee.onyx.server.enterprise_settings.store import ( + get_logo_filename, + get_logotype_filename, + load_analytics_script, + load_settings, + store_analytics_script, + store_settings, + upload_logo, +) from ee.onyx.server.scim.auth import generate_scim_token -from ee.onyx.server.scim.models import ScimTokenCreate -from ee.onyx.server.scim.models import ScimTokenCreatedResponse -from ee.onyx.server.scim.models import ScimTokenResponse +from ee.onyx.server.scim.models import ( + ScimTokenCreate, + ScimTokenCreatedResponse, + ScimTokenResponse, +) from ee.onyx.utils.tier import get_tier from onyx.auth.permissions import require_permission -from onyx.auth.users import current_user_with_expired_token -from onyx.auth.users import get_user_manager -from onyx.auth.users import UserManager +from onyx.auth.users import ( + current_user_with_expired_token, + get_user_manager, + UserManager, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User @@ -42,8 +43,7 @@ from onyx.server.settings.models import Tier from onyx.server.settings.tier_order import tier_at_least from onyx.server.utils import BasicAuthenticationError from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA from shared_configs.contextvars import get_current_tenant_id admin_router = APIRouter(prefix="/admin/enterprise-settings") diff --git a/backend/ee/onyx/server/enterprise_settings/models.py b/backend/ee/onyx/server/enterprise_settings/models.py index 5314f7985b..055085d2e3 100644 --- a/backend/ee/onyx/server/enterprise_settings/models.py +++ b/backend/ee/onyx/server/enterprise_settings/models.py @@ -1,11 +1,8 @@ from enum import Enum -from typing import Any -from typing import List +from typing import Any, List from urllib.parse import urlparse -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator +from pydantic import BaseModel, Field, field_validator class NavigationItem(BaseModel): diff --git a/backend/ee/onyx/server/enterprise_settings/store.py b/backend/ee/onyx/server/enterprise_settings/store.py index 2ff446d2e3..2114cd1d0d 100644 --- a/backend/ee/onyx/server/enterprise_settings/store.py +++ b/backend/ee/onyx/server/enterprise_settings/store.py @@ -1,18 +1,19 @@ import os from io import BytesIO -from typing import Any -from typing import cast -from typing import IO +from typing import Any, cast, IO -from fastapi import HTTPException -from fastapi import UploadFile +from fastapi import HTTPException, UploadFile -from ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload -from ee.onyx.server.enterprise_settings.models import EnterpriseSettings -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import KV_CUSTOM_ANALYTICS_SCRIPT_KEY -from onyx.configs.constants import KV_ENTERPRISE_SETTINGS_KEY -from onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME +from ee.onyx.server.enterprise_settings.models import ( + AnalyticsScriptUpload, + EnterpriseSettings, +) +from onyx.configs.constants import ( + FileOrigin, + KV_CUSTOM_ANALYTICS_SCRIPT_KEY, + KV_ENTERPRISE_SETTINGS_KEY, + ONYX_DEFAULT_APPLICATION_NAME, +) from onyx.file_store.file_store import get_default_file_store from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError diff --git a/backend/ee/onyx/server/evals/api.py b/backend/ee/onyx/server/evals/api.py index 1a24b0cd30..3ae0c216bd 100644 --- a/backend/ee/onyx/server/evals/api.py +++ b/backend/ee/onyx/server/evals/api.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from ee.onyx.auth.users import current_cloud_superuser from onyx.background.celery.apps.client import celery_app as client_app diff --git a/backend/ee/onyx/server/features/hooks/api.py b/backend/ee/onyx/server/features/hooks/api.py index a53cb2a46f..956a1e3999 100644 --- a/backend/ee/onyx/server/features/hooks/api.py +++ b/backend/ee/onyx/server/features/hooks/api.py @@ -1,38 +1,36 @@ import httpx -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Query +from fastapi import APIRouter, Depends, Query from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.auth.users import User -from onyx.db.constants import UNSET -from onyx.db.constants import UnsetType -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.db.constants import UNSET, UnsetType +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant from onyx.db.enums import Permission -from onyx.db.hook import create_hook__no_commit -from onyx.db.hook import delete_hook__no_commit -from onyx.db.hook import get_hook_by_id -from onyx.db.hook import get_hook_execution_logs -from onyx.db.hook import get_hooks -from onyx.db.hook import update_hook__no_commit +from onyx.db.hook import ( + create_hook__no_commit, + delete_hook__no_commit, + get_hook_by_id, + get_hook_execution_logs, + get_hooks, + update_hook__no_commit, +) from onyx.db.models import Hook from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.hooks.api_dependencies import require_hook_enabled -from onyx.hooks.models import HookCreateRequest -from onyx.hooks.models import HookExecutionRecord -from onyx.hooks.models import HookPointMetaResponse -from onyx.hooks.models import HookResponse -from onyx.hooks.models import HookUpdateRequest -from onyx.hooks.models import HookValidateResponse -from onyx.hooks.models import HookValidateStatus -from onyx.hooks.registry import get_all_specs -from onyx.hooks.registry import get_hook_point_spec +from onyx.hooks.models import ( + HookCreateRequest, + HookExecutionRecord, + HookPointMetaResponse, + HookResponse, + HookUpdateRequest, + HookValidateResponse, + HookValidateStatus, +) +from onyx.hooks.registry import get_all_specs, get_hook_point_spec from onyx.utils.logger import setup_logger -from onyx.utils.url import SSRFException -from onyx.utils.url import validate_outbound_http_url +from onyx.utils.url import SSRFException, validate_outbound_http_url logger = setup_logger() diff --git a/backend/ee/onyx/server/license/api.py b/backend/ee/onyx/server/license/api.py index 54fbcf814c..63c4630cdf 100644 --- a/backend/ee/onyx/server/license/api.py +++ b/backend/ee/onyx/server/license/api.py @@ -11,25 +11,26 @@ Cloud licensing is managed via the control plane and gated_tenants Redis key. """ import requests -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import UploadFile +from fastapi import APIRouter, Depends, File, UploadFile from sqlalchemy.orm import Session from ee.onyx.configs.app_configs import CLOUD_DATA_PLANE_URL from ee.onyx.db.license import delete_license as db_delete_license -from ee.onyx.db.license import get_license -from ee.onyx.db.license import get_license_metadata -from ee.onyx.db.license import invalidate_license_cache -from ee.onyx.db.license import refresh_license_cache -from ee.onyx.db.license import update_license_cache -from ee.onyx.db.license import upsert_license -from ee.onyx.server.license.models import LicenseResponse -from ee.onyx.server.license.models import LicenseSource -from ee.onyx.server.license.models import LicenseStatusResponse -from ee.onyx.server.license.models import LicenseUploadResponse -from ee.onyx.server.license.models import SeatUsageResponse +from ee.onyx.db.license import ( + get_license, + get_license_metadata, + invalidate_license_cache, + refresh_license_cache, + update_license_cache, + upsert_license, +) +from ee.onyx.server.license.models import ( + LicenseResponse, + LicenseSource, + LicenseStatusResponse, + LicenseUploadResponse, + SeatUsageResponse, +) from ee.onyx.utils.license import verify_license_signature from onyx.auth.permissions import require_permission from onyx.auth.users import User diff --git a/backend/ee/onyx/server/manage/standard_answer.py b/backend/ee/onyx/server/manage/standard_answer.py index d860bef388..78457d9d99 100644 --- a/backend/ee/onyx/server/manage/standard_answer.py +++ b/backend/ee/onyx/server/manage/standard_answer.py @@ -1,25 +1,27 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session -from ee.onyx.db.standard_answer import fetch_standard_answer -from ee.onyx.db.standard_answer import fetch_standard_answer_categories -from ee.onyx.db.standard_answer import fetch_standard_answer_category -from ee.onyx.db.standard_answer import fetch_standard_answers -from ee.onyx.db.standard_answer import insert_standard_answer -from ee.onyx.db.standard_answer import insert_standard_answer_category -from ee.onyx.db.standard_answer import remove_standard_answer -from ee.onyx.db.standard_answer import update_standard_answer -from ee.onyx.db.standard_answer import update_standard_answer_category +from ee.onyx.db.standard_answer import ( + fetch_standard_answer, + fetch_standard_answer_categories, + fetch_standard_answer_category, + fetch_standard_answers, + insert_standard_answer, + insert_standard_answer_category, + remove_standard_answer, + update_standard_answer, + update_standard_answer_category, +) from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.server.manage.models import StandardAnswer -from onyx.server.manage.models import StandardAnswerCategory -from onyx.server.manage.models import StandardAnswerCategoryCreationRequest -from onyx.server.manage.models import StandardAnswerCreationRequest +from onyx.server.manage.models import ( + StandardAnswer, + StandardAnswerCategory, + StandardAnswerCategoryCreationRequest, + StandardAnswerCreationRequest, +) router = APIRouter(prefix="/manage") diff --git a/backend/ee/onyx/server/metrics/license_metrics.py b/backend/ee/onyx/server/metrics/license_metrics.py index 54e90f3349..3837860429 100644 --- a/backend/ee/onyx/server/metrics/license_metrics.py +++ b/backend/ee/onyx/server/metrics/license_metrics.py @@ -14,14 +14,12 @@ so unlicensed instances don't trip "0 seats remaining" alerts. from prometheus_client.core import GaugeMetricFamily from prometheus_client.registry import REGISTRY -from ee.onyx.db.license import get_license_metadata -from ee.onyx.db.license import get_used_seats +from ee.onyx.db.license import get_license_metadata, get_used_seats from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.server.metrics.indexing_pipeline import _CachedCollector from onyx.server.settings.models import ApplicationStatus from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA logger = setup_logger() diff --git a/backend/ee/onyx/server/middleware/license_enforcement.py b/backend/ee/onyx/server/middleware/license_enforcement.py index 6f6d3ea2e7..3ebb374cd9 100644 --- a/backend/ee/onyx/server/middleware/license_enforcement.py +++ b/backend/ee/onyx/server/middleware/license_enforcement.py @@ -55,12 +55,9 @@ For self-hosted deployments: """ import logging -from collections.abc import Awaitable -from collections.abc import Callable +from collections.abc import Awaitable, Callable -from fastapi import FastAPI -from fastapi import Request -from fastapi import Response +from fastapi import FastAPI, Request, Response from fastapi.responses import JSONResponse from sqlalchemy.exc import SQLAlchemyError @@ -68,8 +65,7 @@ from ee.onyx.configs.app_configs import LICENSE_ENFORCEMENT_ENABLED from ee.onyx.configs.license_enforcement_config import ( LICENSE_ENFORCEMENT_ALLOWED_PREFIXES, ) -from ee.onyx.db.license import get_cached_license_metadata -from ee.onyx.db.license import refresh_license_cache +from ee.onyx.db.license import get_cached_license_metadata, refresh_license_cache from onyx.cache.interface import CACHE_TRANSIENT_ERRORS from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.server.settings.models import ApplicationStatus diff --git a/backend/ee/onyx/server/middleware/tenant_tracking.py b/backend/ee/onyx/server/middleware/tenant_tracking.py index 7cacc9f2e4..68038aec57 100644 --- a/backend/ee/onyx/server/middleware/tenant_tracking.py +++ b/backend/ee/onyx/server/middleware/tenant_tracking.py @@ -1,13 +1,9 @@ import asyncio import logging import sys -from collections.abc import Awaitable -from collections.abc import Callable +from collections.abc import Awaitable, Callable -from fastapi import FastAPI -from fastapi import HTTPException -from fastapi import Request -from fastapi import Response +from fastapi import FastAPI, HTTPException, Request, Response from fastapi.responses import JSONResponse from ee.onyx.auth.users import decode_anonymous_user_jwt_token @@ -16,14 +12,14 @@ from ee.onyx.configs.multi_tenant_gating_config import ( ) from ee.onyx.server.tenants.product_gating import is_tenant_gated from onyx.auth.utils import extract_tenant_from_auth_header -from onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME -from onyx.configs.constants import TENANT_ID_COOKIE_NAME +from onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME, TENANT_ID_COOKIE_NAME from onyx.db.engine.sql_engine import is_valid_schema_name from onyx.error_handling.error_codes import OnyxErrorCode -from onyx.redis.redis_pool import retrieve_auth_token_data_from_bearer -from onyx.redis.redis_pool import retrieve_auth_token_data_from_redis -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from onyx.redis.redis_pool import ( + retrieve_auth_token_data_from_bearer, + retrieve_auth_token_data_from_redis, +) +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR # Unauthenticated, non-tenant-scoped endpoints (LB/Prometheus probes, API docs), diff --git a/backend/ee/onyx/server/middleware/tier_gate.py b/backend/ee/onyx/server/middleware/tier_gate.py index 8839e1349b..5a25f3bc27 100644 --- a/backend/ee/onyx/server/middleware/tier_gate.py +++ b/backend/ee/onyx/server/middleware/tier_gate.py @@ -20,18 +20,15 @@ handles GATED_ACCESS, seat limits, and the billing/auth allowlist. import asyncio import logging -from collections.abc import Awaitable -from collections.abc import Callable +from collections.abc import Awaitable, Callable -from fastapi import FastAPI -from fastapi import Request -from fastapi import Response +from fastapi import FastAPI, Request, Response from fastapi.responses import JSONResponse from ee.onyx.configs.license_enforcement_config import ( LICENSE_ENFORCEMENT_ALLOWED_PREFIXES, + PATH_PREFIX_MIN_TIER, ) -from ee.onyx.configs.license_enforcement_config import PATH_PREFIX_MIN_TIER from ee.onyx.utils.tier import get_tier from onyx.error_handling.error_codes import OnyxErrorCode from onyx.server.settings.models import Tier diff --git a/backend/ee/onyx/server/oauth/api.py b/backend/ee/onyx/server/oauth/api.py index 1775e99b61..d71f27018c 100644 --- a/backend/ee/onyx/server/oauth/api.py +++ b/backend/ee/onyx/server/oauth/api.py @@ -1,8 +1,7 @@ import base64 import uuid -from fastapi import Depends -from fastapi import HTTPException +from fastapi import Depends, HTTPException from fastapi.responses import JSONResponse from ee.onyx.server.oauth.api_router import router diff --git a/backend/ee/onyx/server/oauth/confluence_cloud.py b/backend/ee/onyx/server/oauth/confluence_cloud.py index 8ac4c3dd9b..2a377252c7 100644 --- a/backend/ee/onyx/server/oauth/confluence_cloud.py +++ b/backend/ee/onyx/server/oauth/confluence_cloud.py @@ -1,30 +1,29 @@ import base64 import uuid -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timedelta, timezone +from typing import Any, cast import requests -from fastapi import Depends -from fastapi import HTTPException +from fastapi import Depends, HTTPException from fastapi.responses import JSONResponse -from pydantic import BaseModel -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from sqlalchemy.orm import Session from ee.onyx.server.oauth.api_router import router from onyx.auth.permissions import require_permission -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_ID -from onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.app_configs import ( + DEV_MODE, + OAUTH_CONFLUENCE_CLOUD_CLIENT_ID, + OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET, + WEB_DOMAIN, +) from onyx.configs.constants import DocumentSource from onyx.connectors.confluence.utils import CONFLUENCE_OAUTH_TOKEN_URL -from onyx.db.credentials import create_credential -from onyx.db.credentials import fetch_credential_by_id_for_user -from onyx.db.credentials import update_credential_json +from onyx.db.credentials import ( + create_credential, + fetch_credential_by_id_for_user, + update_credential_json, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User diff --git a/backend/ee/onyx/server/oauth/google_drive.py b/backend/ee/onyx/server/oauth/google_drive.py index d822af3886..5dfc360f95 100644 --- a/backend/ee/onyx/server/oauth/google_drive.py +++ b/backend/ee/onyx/server/oauth/google_drive.py @@ -1,33 +1,31 @@ import base64 import json import uuid -from typing import Any -from typing import cast +from typing import Any, cast import requests -from fastapi import Depends -from fastapi import HTTPException +from fastapi import Depends, HTTPException from fastapi.responses import JSONResponse from pydantic import BaseModel from sqlalchemy.orm import Session from ee.onyx.server.oauth.api_router import router from onyx.auth.permissions import require_permission -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_ID -from onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_SECRET -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.app_configs import ( + DEV_MODE, + OAUTH_GOOGLE_DRIVE_CLIENT_ID, + OAUTH_GOOGLE_DRIVE_CLIENT_SECRET, + WEB_DOMAIN, +) from onyx.configs.constants import DocumentSource -from onyx.connectors.google_utils.google_auth import get_google_oauth_creds -from onyx.connectors.google_utils.google_auth import sanitize_oauth_credentials +from onyx.connectors.google_utils.google_auth import ( + get_google_oauth_creds, + sanitize_oauth_credentials, +) from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, -) -from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY -from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_DICT_TOKEN_KEY, DB_CREDENTIALS_PRIMARY_ADMIN_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( GoogleOAuthAuthenticationMethod, ) from onyx.db.credentials import create_credential diff --git a/backend/ee/onyx/server/oauth/slack.py b/backend/ee/onyx/server/oauth/slack.py index 3cee4eafda..db03c3806c 100644 --- a/backend/ee/onyx/server/oauth/slack.py +++ b/backend/ee/onyx/server/oauth/slack.py @@ -3,18 +3,19 @@ import uuid from typing import cast import requests -from fastapi import Depends -from fastapi import HTTPException +from fastapi import Depends, HTTPException from fastapi.responses import JSONResponse from pydantic import BaseModel from sqlalchemy.orm import Session from ee.onyx.server.oauth.api_router import router from onyx.auth.permissions import require_permission -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import OAUTH_SLACK_CLIENT_ID -from onyx.configs.app_configs import OAUTH_SLACK_CLIENT_SECRET -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.app_configs import ( + DEV_MODE, + OAUTH_SLACK_CLIENT_ID, + OAUTH_SLACK_CLIENT_SECRET, + WEB_DOMAIN, +) from onyx.configs.constants import DocumentSource from onyx.db.credentials import create_credential from onyx.db.engine.sql_engine import get_session diff --git a/backend/ee/onyx/server/query_and_chat/models.py b/backend/ee/onyx/server/query_and_chat/models.py index 9dd424d015..45b49759ab 100644 --- a/backend/ee/onyx/server/query_and_chat/models.py +++ b/backend/ee/onyx/server/query_and_chat/models.py @@ -1,12 +1,9 @@ from collections.abc import Sequence from datetime import datetime -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import InferenceSection -from onyx.context.search.models import SearchDoc +from onyx.context.search.models import BaseFilters, InferenceSection, SearchDoc from onyx.server.manage.models import StandardAnswer diff --git a/backend/ee/onyx/server/query_and_chat/query_backend.py b/backend/ee/onyx/server/query_and_chat/query_backend.py index 8b866ac49a..00daf4d822 100644 --- a/backend/ee/onyx/server/query_and_chat/query_backend.py +++ b/backend/ee/onyx/server/query_and_chat/query_backend.py @@ -1,13 +1,13 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from ee.onyx.onyxbot.slack.handlers.handle_standard_answers import ( oneoff_standard_answers, ) -from ee.onyx.server.query_and_chat.models import StandardAnswerRequest -from ee.onyx.server.query_and_chat.models import StandardAnswerResponse +from ee.onyx.server.query_and_chat.models import ( + StandardAnswerRequest, + StandardAnswerResponse, +) from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission diff --git a/backend/ee/onyx/server/query_and_chat/search_backend.py b/backend/ee/onyx/server/query_and_chat/search_backend.py index d3f83bd80c..8f94922ae6 100644 --- a/backend/ee/onyx/server/query_and_chat/search_backend.py +++ b/backend/ee/onyx/server/query_and_chat/search_backend.py @@ -11,29 +11,30 @@ onyx/server/features/search/api.py (POST /api/search). from collections.abc import Generator -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session from ee.onyx.db.search import fetch_search_queries_for_user -from ee.onyx.search.process_search_query import gather_search_stream -from ee.onyx.search.process_search_query import stream_search_query +from ee.onyx.search.process_search_query import ( + gather_search_stream, + stream_search_query, +) from ee.onyx.secondary_llm_flows.search_flow_classification import ( classify_is_search_flow, ) -from ee.onyx.server.query_and_chat.models import SearchFlowClassificationRequest -from ee.onyx.server.query_and_chat.models import SearchFlowClassificationResponse -from ee.onyx.server.query_and_chat.models import SearchFullResponse -from ee.onyx.server.query_and_chat.models import SearchHistoryResponse -from ee.onyx.server.query_and_chat.models import SearchQueryResponse -from ee.onyx.server.query_and_chat.models import SendSearchQueryRequest +from ee.onyx.server.query_and_chat.models import ( + SearchFlowClassificationRequest, + SearchFlowClassificationResponse, + SearchFullResponse, + SearchHistoryResponse, + SearchQueryResponse, + SendSearchQueryRequest, +) from ee.onyx.server.query_and_chat.streaming_models import SearchErrorPacket from onyx.auth.permissions import require_permission from onyx.configs.app_configs import ONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant from onyx.db.enums import Permission from onyx.db.models import User from onyx.llm.factory import get_default_llm diff --git a/backend/ee/onyx/server/query_and_chat/streaming_models.py b/backend/ee/onyx/server/query_and_chat/streaming_models.py index 797dd095ab..9b6de8b73b 100644 --- a/backend/ee/onyx/server/query_and_chat/streaming_models.py +++ b/backend/ee/onyx/server/query_and_chat/streaming_models.py @@ -1,7 +1,6 @@ from typing import Literal -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from ee.onyx.server.query_and_chat.models import SearchDocWithContent diff --git a/backend/ee/onyx/server/query_and_chat/token_limit.py b/backend/ee/onyx/server/query_and_chat/token_limit.py index 2785cae78d..aa9405008d 100644 --- a/backend/ee/onyx/server/query_and_chat/token_limit.py +++ b/backend/ee/onyx/server/query_and_chat/token_limit.py @@ -2,29 +2,30 @@ from collections import defaultdict from collections.abc import Sequence from datetime import datetime from itertools import groupby -from typing import Dict -from typing import List -from typing import Tuple +from typing import Dict, List, Tuple from uuid import UUID from fastapi import HTTPException -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import func, select from sqlalchemy.orm import Session from onyx.db.api_key import is_api_key_email_address from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession -from onyx.db.models import TokenRateLimit -from onyx.db.models import TokenRateLimit__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import ( + ChatMessage, + ChatSession, + TokenRateLimit, + TokenRateLimit__UserGroup, + User, + User__UserGroup, + UserGroup, +) from onyx.db.token_limit import fetch_all_user_token_rate_limits -from onyx.server.query_and_chat.token_limit import _get_cutoff_time -from onyx.server.query_and_chat.token_limit import _is_rate_limited -from onyx.server.query_and_chat.token_limit import _user_is_rate_limited_by_global +from onyx.server.query_and_chat.token_limit import ( + _get_cutoff_time, + _is_rate_limited, + _user_is_rate_limited_by_global, +) from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel diff --git a/backend/ee/onyx/server/query_history/api.py b/backend/ee/onyx/server/query_history/api.py index a968270a84..fb48ba178a 100644 --- a/backend/ee/onyx/server/query_history/api.py +++ b/backend/ee/onyx/server/query_history/api.py @@ -1,56 +1,53 @@ import uuid from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from http import HTTPStatus from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query +from fastapi import APIRouter, Depends, HTTPException, Query from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session from ee.onyx.background.task_name_builders import query_history_task_name -from ee.onyx.db.query_history import get_all_query_history_export_tasks -from ee.onyx.db.query_history import get_page_of_chat_sessions -from ee.onyx.db.query_history import get_total_filtered_chat_sessions_count -from ee.onyx.server.query_history.models import ChatSessionMinimal -from ee.onyx.server.query_history.models import ChatSessionSnapshot -from ee.onyx.server.query_history.models import MessageSnapshot -from ee.onyx.server.query_history.models import QueryHistoryExport +from ee.onyx.db.query_history import ( + get_all_query_history_export_tasks, + get_page_of_chat_sessions, + get_total_filtered_chat_sessions_count, +) +from ee.onyx.server.query_history.models import ( + ChatSessionMinimal, + ChatSessionSnapshot, + MessageSnapshot, + QueryHistoryExport, +) from onyx.auth.permissions import require_permission from onyx.auth.users import get_display_email from onyx.background.celery.versioned_apps.client import app as client_app from onyx.background.task_utils import construct_query_history_report_name from onyx.chat.chat_utils import create_chat_history_chain -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import FileType -from onyx.configs.constants import MessageType -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import PUBLIC_API_TAGS -from onyx.configs.constants import QAFeedbackType -from onyx.configs.constants import QueryHistoryType -from onyx.configs.constants import SessionType -from onyx.db.chat import get_chat_session_by_id -from onyx.db.chat import get_chat_sessions_by_user +from onyx.configs.constants import ( + FileOrigin, + FileType, + MessageType, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + PUBLIC_API_TAGS, + QAFeedbackType, + QueryHistoryType, + SessionType, +) +from onyx.db.chat import get_chat_session_by_id, get_chat_sessions_by_user from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import Permission -from onyx.db.enums import TaskStatus +from onyx.db.enums import Permission, TaskStatus from onyx.db.file_record import get_query_history_export_files -from onyx.db.models import ChatSession -from onyx.db.models import User -from onyx.db.tasks import get_task_with_id -from onyx.db.tasks import register_task +from onyx.db.models import ChatSession, User +from onyx.db.tasks import get_task_with_id, register_task from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import get_default_file_store from onyx.server.documents.models import PaginatedReturn -from onyx.server.query_and_chat.models import ChatSessionDetails -from onyx.server.query_and_chat.models import ChatSessionsResponse +from onyx.server.query_and_chat.models import ChatSessionDetails, ChatSessionsResponse from onyx.server.settings.store import load_settings from onyx.utils.threadpool_concurrency import parallel_yield from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/ee/onyx/server/query_history/models.py b/backend/ee/onyx/server/query_history/models.py index 9540cdc140..27c111d8c0 100644 --- a/backend/ee/onyx/server/query_history/models.py +++ b/backend/ee/onyx/server/query_history/models.py @@ -6,14 +6,9 @@ from pydantic import BaseModel from ee.onyx.background.task_name_builders import QUERY_HISTORY_TASK_NAME_PREFIX from onyx.auth.users import get_display_email from onyx.background.task_utils import extract_task_id_from_query_history_report_name -from onyx.configs.constants import MessageType -from onyx.configs.constants import QAFeedbackType -from onyx.configs.constants import SessionType +from onyx.configs.constants import MessageType, QAFeedbackType, SessionType from onyx.db.enums import TaskStatus -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession -from onyx.db.models import FileRecord -from onyx.db.models import TaskQueueState +from onyx.db.models import ChatMessage, ChatSession, FileRecord, TaskQueueState class AbridgedSearchDoc(BaseModel): diff --git a/backend/ee/onyx/server/reporting/usage_export_api.py b/backend/ee/onyx/server/reporting/usage_export_api.py index 64c6a329d4..918e47daee 100644 --- a/backend/ee/onyx/server/reporting/usage_export_api.py +++ b/backend/ee/onyx/server/reporting/usage_export_api.py @@ -1,17 +1,16 @@ from collections.abc import Generator from datetime import datetime -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Response +from fastapi import APIRouter, Depends, HTTPException, Response from fastapi.responses import StreamingResponse from pydantic import BaseModel from sqlalchemy.orm import Session -from ee.onyx.db.usage_export import get_all_usage_reports -from ee.onyx.db.usage_export import get_usage_report_data -from ee.onyx.db.usage_export import UsageReportMetadata +from ee.onyx.db.usage_export import ( + get_all_usage_reports, + get_usage_report_data, + UsageReportMetadata, +) from onyx.auth.permissions import require_permission from onyx.background.celery.versioned_apps.client import app as client_app from onyx.configs.constants import OnyxCeleryTask diff --git a/backend/ee/onyx/server/reporting/usage_export_generation.py b/backend/ee/onyx/server/reporting/usage_export_generation.py index 6663c0d607..362e12699e 100644 --- a/backend/ee/onyx/server/reporting/usage_export_generation.py +++ b/backend/ee/onyx/server/reporting/usage_export_generation.py @@ -2,25 +2,26 @@ import csv import tempfile import uuid import zipfile -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from fastapi_users_db_sqlalchemy import UUID_ID from sqlalchemy import cast from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.orm import Session -from ee.onyx.db.usage_export import get_all_empty_chat_message_entries -from ee.onyx.db.usage_export import write_usage_report -from ee.onyx.server.reporting.usage_export_models import UsageReportMetadata -from ee.onyx.server.reporting.usage_export_models import UserSkeleton +from ee.onyx.db.usage_export import ( + get_all_empty_chat_message_entries, + write_usage_report, +) +from ee.onyx.server.reporting.usage_export_models import ( + UsageReportMetadata, + UserSkeleton, +) from onyx.configs.constants import FileOrigin from onyx.db.models import User from onyx.db.users import get_all_users from onyx.file_store.constants import MAX_IN_MEMORY_SIZE -from onyx.file_store.file_store import FileStore -from onyx.file_store.file_store import get_default_file_store +from onyx.file_store.file_store import FileStore, get_default_file_store from onyx.utils.csv_utils import sanitize_csv_cell_or_none diff --git a/backend/ee/onyx/server/scim/api.py b/backend/ee/onyx/server/scim/api.py index a42a33abb7..ed82d59b88 100644 --- a/backend/ee/onyx/server/scim/api.py +++ b/backend/ee/onyx/server/scim/api.py @@ -13,57 +13,54 @@ from __future__ import annotations from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import FastAPI -from fastapi import Query -from fastapi import Request -from fastapi import Response +from fastapi import APIRouter, Depends, FastAPI, Query, Request, Response from fastapi.responses import JSONResponse from fastapi_users.password import PasswordHelper from sqlalchemy import func from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session -from ee.onyx.db.license import acquire_seat_lock -from ee.onyx.db.license import check_seat_availability +from ee.onyx.db.license import acquire_seat_lock, check_seat_availability from ee.onyx.db.scim import ScimDAL -from ee.onyx.server.scim.auth import ScimAuthError -from ee.onyx.server.scim.auth import verify_scim_token +from ee.onyx.server.scim.auth import ScimAuthError, verify_scim_token from ee.onyx.server.scim.filtering import parse_scim_filter -from ee.onyx.server.scim.models import SCIM_LIST_RESPONSE_SCHEMA -from ee.onyx.server.scim.models import ScimError -from ee.onyx.server.scim.models import ScimGroupMember -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimListResponse -from ee.onyx.server.scim.models import ScimMappingFields -from ee.onyx.server.scim.models import ScimName -from ee.onyx.server.scim.models import ScimPatchRequest -from ee.onyx.server.scim.models import ScimServiceProviderConfig -from ee.onyx.server.scim.models import ScimUserResource -from ee.onyx.server.scim.patch import apply_group_patch -from ee.onyx.server.scim.patch import apply_user_patch -from ee.onyx.server.scim.patch import ScimPatchError -from ee.onyx.server.scim.providers.base import get_default_provider -from ee.onyx.server.scim.providers.base import ScimProvider -from ee.onyx.server.scim.providers.base import serialize_emails -from ee.onyx.server.scim.schema_definitions import ENTERPRISE_USER_SCHEMA_DEF -from ee.onyx.server.scim.schema_definitions import GROUP_RESOURCE_TYPE -from ee.onyx.server.scim.schema_definitions import GROUP_SCHEMA_DEF -from ee.onyx.server.scim.schema_definitions import SERVICE_PROVIDER_CONFIG -from ee.onyx.server.scim.schema_definitions import USER_RESOURCE_TYPE -from ee.onyx.server.scim.schema_definitions import USER_SCHEMA_DEF +from ee.onyx.server.scim.models import ( + SCIM_LIST_RESPONSE_SCHEMA, + ScimError, + ScimGroupMember, + ScimGroupResource, + ScimListResponse, + ScimMappingFields, + ScimName, + ScimPatchRequest, + ScimServiceProviderConfig, + ScimUserResource, +) +from ee.onyx.server.scim.patch import ( + apply_group_patch, + apply_user_patch, + ScimPatchError, +) +from ee.onyx.server.scim.providers.base import ( + get_default_provider, + ScimProvider, + serialize_emails, +) +from ee.onyx.server.scim.schema_definitions import ( + ENTERPRISE_USER_SCHEMA_DEF, + GROUP_RESOURCE_TYPE, + GROUP_SCHEMA_DEF, + SERVICE_PROVIDER_CONFIG, + USER_RESOURCE_TYPE, + USER_SCHEMA_DEF, +) from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import AccountType -from onyx.db.enums import GrantSource -from onyx.db.enums import Permission -from onyx.db.models import ScimToken -from onyx.db.models import ScimUserMapping -from onyx.db.models import User -from onyx.db.models import UserGroup -from onyx.db.models import UserRole -from onyx.db.permissions import recompute_permissions_for_group__no_commit -from onyx.db.permissions import recompute_user_permissions__no_commit +from onyx.db.enums import AccountType, GrantSource, Permission +from onyx.db.models import ScimToken, ScimUserMapping, User, UserGroup, UserRole +from onyx.db.permissions import ( + recompute_permissions_for_group__no_commit, + recompute_user_permissions__no_commit, +) from onyx.db.users import assign_user_to_default_groups__no_commit from onyx.db.utils import is_unique_violation from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/server/scim/auth.py b/backend/ee/onyx/server/scim/auth.py index e896581505..f904c8c5f1 100644 --- a/backend/ee/onyx/server/scim/auth.py +++ b/backend/ee/onyx/server/scim/auth.py @@ -18,8 +18,7 @@ the admin exactly once at creation time. import hashlib import secrets -from fastapi import Depends -from fastapi import Request +from fastapi import Depends, Request from sqlalchemy.orm import Session from ee.onyx.db.scim import ScimDAL diff --git a/backend/ee/onyx/server/scim/models.py b/backend/ee/onyx/server/scim/models.py index 120da01102..cca34c8994 100644 --- a/backend/ee/onyx/server/scim/models.py +++ b/backend/ee/onyx/server/scim/models.py @@ -11,10 +11,7 @@ from dataclasses import dataclass from datetime import datetime from enum import Enum -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator # --------------------------------------------------------------------------- # SCIM Schema URIs (RFC 7643 §8) diff --git a/backend/ee/onyx/server/scim/patch.py b/backend/ee/onyx/server/scim/patch.py index 4632e18ec0..144aabd742 100644 --- a/backend/ee/onyx/server/scim/patch.py +++ b/backend/ee/onyx/server/scim/patch.py @@ -16,18 +16,19 @@ from __future__ import annotations import logging import re -from dataclasses import dataclass -from dataclasses import field +from dataclasses import dataclass, field from typing import Any -from ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA -from ee.onyx.server.scim.models import ScimGroupMember -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimPatchOperation -from ee.onyx.server.scim.models import ScimPatchOperationType -from ee.onyx.server.scim.models import ScimPatchResourceValue -from ee.onyx.server.scim.models import ScimPatchValue -from ee.onyx.server.scim.models import ScimUserResource +from ee.onyx.server.scim.models import ( + SCIM_ENTERPRISE_USER_SCHEMA, + ScimGroupMember, + ScimGroupResource, + ScimPatchOperation, + ScimPatchOperationType, + ScimPatchResourceValue, + ScimPatchValue, + ScimUserResource, +) logger = logging.getLogger(__name__) diff --git a/backend/ee/onyx/server/scim/providers/base.py b/backend/ee/onyx/server/scim/providers/base.py index a5f847b1a1..a03f7ab1fd 100644 --- a/backend/ee/onyx/server/scim/providers/base.py +++ b/backend/ee/onyx/server/scim/providers/base.py @@ -4,26 +4,26 @@ from __future__ import annotations import json import logging -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from uuid import UUID from pydantic import ValidationError -from ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA -from ee.onyx.server.scim.models import SCIM_USER_SCHEMA -from ee.onyx.server.scim.models import ScimEmail -from ee.onyx.server.scim.models import ScimEnterpriseExtension -from ee.onyx.server.scim.models import ScimGroupMember -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimManagerRef -from ee.onyx.server.scim.models import ScimMappingFields -from ee.onyx.server.scim.models import ScimMeta -from ee.onyx.server.scim.models import ScimName -from ee.onyx.server.scim.models import ScimUserGroupRef -from ee.onyx.server.scim.models import ScimUserResource -from onyx.db.models import User -from onyx.db.models import UserGroup +from ee.onyx.server.scim.models import ( + SCIM_ENTERPRISE_USER_SCHEMA, + SCIM_USER_SCHEMA, + ScimEmail, + ScimEnterpriseExtension, + ScimGroupMember, + ScimGroupResource, + ScimManagerRef, + ScimMappingFields, + ScimMeta, + ScimName, + ScimUserGroupRef, + ScimUserResource, +) +from onyx.db.models import User, UserGroup logger = logging.getLogger(__name__) diff --git a/backend/ee/onyx/server/scim/providers/entra.py b/backend/ee/onyx/server/scim/providers/entra.py index 0cca4cee21..0afea4a1f6 100644 --- a/backend/ee/onyx/server/scim/providers/entra.py +++ b/backend/ee/onyx/server/scim/providers/entra.py @@ -2,10 +2,8 @@ from __future__ import annotations -from ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA -from ee.onyx.server.scim.models import SCIM_USER_SCHEMA -from ee.onyx.server.scim.providers.base import COMMON_IGNORED_PATCH_PATHS -from ee.onyx.server.scim.providers.base import ScimProvider +from ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA, SCIM_USER_SCHEMA +from ee.onyx.server.scim.providers.base import COMMON_IGNORED_PATCH_PATHS, ScimProvider _ENTRA_IGNORED_PATCH_PATHS = COMMON_IGNORED_PATCH_PATHS diff --git a/backend/ee/onyx/server/scim/providers/okta.py b/backend/ee/onyx/server/scim/providers/okta.py index 03bebbbc10..39599fc4dd 100644 --- a/backend/ee/onyx/server/scim/providers/okta.py +++ b/backend/ee/onyx/server/scim/providers/okta.py @@ -2,8 +2,7 @@ from __future__ import annotations -from ee.onyx.server.scim.providers.base import COMMON_IGNORED_PATCH_PATHS -from ee.onyx.server.scim.providers.base import ScimProvider +from ee.onyx.server.scim.providers.base import COMMON_IGNORED_PATCH_PATHS, ScimProvider class OktaProvider(ScimProvider): diff --git a/backend/ee/onyx/server/scim/schema_definitions.py b/backend/ee/onyx/server/scim/schema_definitions.py index c1ad712053..61768d194c 100644 --- a/backend/ee/onyx/server/scim/schema_definitions.py +++ b/backend/ee/onyx/server/scim/schema_definitions.py @@ -4,13 +4,15 @@ Pre-built at import time — these never change at runtime. Separated from api.py to keep the endpoint module focused on request handling. """ -from ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA -from ee.onyx.server.scim.models import SCIM_GROUP_SCHEMA -from ee.onyx.server.scim.models import SCIM_USER_SCHEMA -from ee.onyx.server.scim.models import ScimResourceType -from ee.onyx.server.scim.models import ScimSchemaAttribute -from ee.onyx.server.scim.models import ScimSchemaDefinition -from ee.onyx.server.scim.models import ScimServiceProviderConfig +from ee.onyx.server.scim.models import ( + SCIM_ENTERPRISE_USER_SCHEMA, + SCIM_GROUP_SCHEMA, + SCIM_USER_SCHEMA, + ScimResourceType, + ScimSchemaAttribute, + ScimSchemaDefinition, + ScimServiceProviderConfig, +) SERVICE_PROVIDER_CONFIG = ScimServiceProviderConfig() diff --git a/backend/ee/onyx/server/seeding.py b/backend/ee/onyx/server/seeding.py index 041adf9d48..60fc9c55eb 100644 --- a/backend/ee/onyx/server/seeding.py +++ b/backend/ee/onyx/server/seeding.py @@ -1,28 +1,29 @@ import json import os from copy import deepcopy -from typing import List -from typing import Optional +from typing import List, Optional from pydantic import BaseModel from sqlalchemy.orm import Session from ee.onyx.db.standard_answer import create_initial_default_standard_answer_category -from ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload -from ee.onyx.server.enterprise_settings.models import EnterpriseSettings -from ee.onyx.server.enterprise_settings.models import NavigationItem -from ee.onyx.server.enterprise_settings.store import store_analytics_script +from ee.onyx.server.enterprise_settings.models import ( + AnalyticsScriptUpload, + EnterpriseSettings, + NavigationItem, +) +from ee.onyx.server.enterprise_settings.store import store_analytics_script, upload_logo from ee.onyx.server.enterprise_settings.store import store_settings as store_ee_settings -from ee.onyx.server.enterprise_settings.store import upload_logo from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.llm import fetch_existing_llm_provider_by_name_and_type -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import ( + fetch_existing_llm_provider_by_name_and_type, + update_default_provider, + upsert_llm_provider, +) from onyx.db.models import Tool from onyx.db.persona import upsert_persona from onyx.server.features.persona.models import PersonaUpsertRequest -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView +from onyx.server.manage.llm.models import LLMProviderUpsertRequest, LLMProviderView from onyx.server.settings.models import Settings from onyx.server.settings.store import store_settings as store_base_settings from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/server/settings/api.py b/backend/ee/onyx/server/settings/api.py index e60bccf8f3..27807bff34 100644 --- a/backend/ee/onyx/server/settings/api.py +++ b/backend/ee/onyx/server/settings/api.py @@ -4,16 +4,12 @@ from redis.exceptions import RedisError from sqlalchemy.exc import SQLAlchemyError from ee.onyx.configs.app_configs import LICENSE_ENFORCEMENT_ENABLED -from ee.onyx.db.license import get_cached_license_metadata -from ee.onyx.db.license import refresh_license_cache -from ee.onyx.utils.tier import get_tier -from ee.onyx.utils.tier import tier_from_license_metadata +from ee.onyx.db.license import get_cached_license_metadata, refresh_license_cache +from ee.onyx.utils.tier import get_tier, tier_from_license_metadata from onyx.cache.interface import CACHE_TRANSIENT_ERRORS from onyx.configs.app_configs import ENTERPRISE_EDITION_ENABLED from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.server.settings.models import ApplicationStatus -from onyx.server.settings.models import Settings -from onyx.server.settings.models import Tier +from onyx.server.settings.models import ApplicationStatus, Settings, Tier from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import global_version from shared_configs.configs import MULTI_TENANT diff --git a/backend/ee/onyx/server/tenant_usage_limits.py b/backend/ee/onyx/server/tenant_usage_limits.py index 60eb3ce86f..d5783654b5 100644 --- a/backend/ee/onyx/server/tenant_usage_limits.py +++ b/backend/ee/onyx/server/tenant_usage_limits.py @@ -5,8 +5,7 @@ import time import requests from ee.onyx.server.tenants.access import generate_data_plane_token -from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL -from onyx.configs.app_configs import DEV_MODE +from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL, DEV_MODE from onyx.server.tenant_usage_limits import TenantUsageLimitOverrides from onyx.server.usage_limits import NO_LIMIT from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/server/tenants/access.py b/backend/ee/onyx/server/tenants/access.py index 658df0d916..99765e5547 100644 --- a/backend/ee/onyx/server/tenants/access.py +++ b/backend/ee/onyx/server/tenants/access.py @@ -1,14 +1,9 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import jwt -from fastapi import HTTPException -from fastapi import Request +from fastapi import HTTPException, Request -from onyx.configs.app_configs import DATA_PLANE_SECRET -from onyx.configs.app_configs import EXPECTED_API_KEY -from onyx.configs.app_configs import JWT_ALGORITHM +from onyx.configs.app_configs import DATA_PLANE_SECRET, EXPECTED_API_KEY, JWT_ALGORITHM from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/server/tenants/admin_api.py b/backend/ee/onyx/server/tenants/admin_api.py index afda131f05..19bcd3c9e7 100644 --- a/backend/ee/onyx/server/tenants/admin_api.py +++ b/backend/ee/onyx/server/tenants/admin_api.py @@ -1,22 +1,19 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Response +from fastapi import APIRouter, Depends, HTTPException, Response from fastapi_users import exceptions from ee.onyx.auth.users import current_cloud_superuser from ee.onyx.server.tenants.models import ImpersonateRequest from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email -from onyx.auth.users import auth_backend -from onyx.auth.users import get_redis_strategy -from onyx.auth.users import User +from onyx.auth.users import auth_backend, get_redis_strategy, User from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.users import get_user_by_email -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/ee/onyx/server/tenants/anonymous_users_api.py b/backend/ee/onyx/server/tenants/anonymous_users_api.py index 8a0c552d34..757b52a996 100644 --- a/backend/ee/onyx/server/tenants/anonymous_users_api.py +++ b/backend/ee/onyx/server/tenants/anonymous_users_api.py @@ -1,22 +1,20 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Response +from fastapi import APIRouter, Depends, HTTPException, Response from sqlalchemy.exc import IntegrityError from ee.onyx.auth.users import generate_anonymous_user_jwt_token -from ee.onyx.server.tenants.anonymous_user_path import get_anonymous_user_path from ee.onyx.server.tenants.anonymous_user_path import ( + get_anonymous_user_path, get_tenant_id_for_anonymous_user_path, + modify_anonymous_user_path, + validate_anonymous_user_path, ) -from ee.onyx.server.tenants.anonymous_user_path import modify_anonymous_user_path -from ee.onyx.server.tenants.anonymous_user_path import validate_anonymous_user_path from ee.onyx.server.tenants.models import AnonymousUserPath from onyx.auth.permissions import require_permission -from onyx.auth.users import anonymous_user_enabled -from onyx.auth.users import User -from onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME -from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME +from onyx.auth.users import anonymous_user_enabled, User +from onyx.configs.constants import ( + ANONYMOUS_USER_COOKIE_NAME, + FASTAPI_USERS_AUTH_COOKIE_NAME, +) from onyx.db.engine.sql_engine import get_session_with_shared_schema from onyx.db.enums import Permission from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/server/tenants/billing.py b/backend/ee/onyx/server/tenants/billing.py index 1b9bfe13c1..21c64fe4eb 100644 --- a/backend/ee/onyx/server/tenants/billing.py +++ b/backend/ee/onyx/server/tenants/billing.py @@ -1,6 +1,5 @@ from enum import Enum as PyEnum -from typing import cast -from typing import Literal +from typing import cast, Literal import requests import stripe @@ -9,9 +8,11 @@ from sqlalchemy.orm import Session from ee.onyx.configs.app_configs import STRIPE_SECRET_KEY from ee.onyx.db.license import acquire_seat_lock from ee.onyx.server.tenants.access import generate_data_plane_token -from ee.onyx.server.tenants.models import BillingInformation -from ee.onyx.server.tenants.models import StripeCheckoutSessionResult -from ee.onyx.server.tenants.models import SubscriptionStatusResponse +from ee.onyx.server.tenants.models import ( + BillingInformation, + StripeCheckoutSessionResult, + SubscriptionStatusResponse, +) from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL from onyx.db.engine.sql_engine import get_session_with_shared_schema from onyx.error_handling.error_codes import OnyxErrorCode diff --git a/backend/ee/onyx/server/tenants/billing_api.py b/backend/ee/onyx/server/tenants/billing_api.py index 3b873fd10f..a0eabb29b4 100644 --- a/backend/ee/onyx/server/tenants/billing_api.py +++ b/backend/ee/onyx/server/tenants/billing_api.py @@ -19,43 +19,51 @@ and are NOT part of this migration - they stay here. import asyncio import httpx -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from ee.onyx.server.billing.api import update_seats as _admin_update_seats -from ee.onyx.server.billing.models import SeatUpdateRequest -from ee.onyx.server.billing.models import SeatUpdateResponse +from ee.onyx.server.billing.models import SeatUpdateRequest, SeatUpdateResponse from ee.onyx.server.tenants.access import control_plane_dep -from ee.onyx.server.tenants.billing import fetch_billing_information -from ee.onyx.server.tenants.billing import fetch_customer_portal_session -from ee.onyx.server.tenants.billing import fetch_stripe_checkout_session -from ee.onyx.server.tenants.models import BillingInformation -from ee.onyx.server.tenants.models import CreateCheckoutSessionRequest -from ee.onyx.server.tenants.models import CreateSubscriptionSessionRequest -from ee.onyx.server.tenants.models import ProductGatingFullSyncRequest -from ee.onyx.server.tenants.models import ProductGatingRequest -from ee.onyx.server.tenants.models import ProductGatingResponse -from ee.onyx.server.tenants.models import StripePublishableKeyResponse -from ee.onyx.server.tenants.models import SubscriptionSessionResponse -from ee.onyx.server.tenants.models import SubscriptionStatusResponse -from ee.onyx.server.tenants.models import TierUpdateRequest -from ee.onyx.server.tenants.models import TierUpdateResponse -from ee.onyx.server.tenants.product_gating import overwrite_full_gated_set -from ee.onyx.server.tenants.product_gating import store_product_gating +from ee.onyx.server.tenants.billing import ( + fetch_billing_information, + fetch_customer_portal_session, + fetch_stripe_checkout_session, +) +from ee.onyx.server.tenants.models import ( + BillingInformation, + CreateCheckoutSessionRequest, + CreateSubscriptionSessionRequest, + ProductGatingFullSyncRequest, + ProductGatingRequest, + ProductGatingResponse, + StripePublishableKeyResponse, + SubscriptionSessionResponse, + SubscriptionStatusResponse, + TierUpdateRequest, + TierUpdateResponse, +) +from ee.onyx.server.tenants.product_gating import ( + overwrite_full_gated_set, + store_product_gating, +) from ee.onyx.server.tenants.tier_management import update_tenant_tier from onyx.auth.permissions import require_permission from onyx.auth.users import User -from onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_OVERRIDE -from onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_URL -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.app_configs import ( + STRIPE_PUBLISHABLE_KEY_OVERRIDE, + STRIPE_PUBLISHABLE_KEY_URL, + WEB_DOMAIN, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.utils.logger import setup_logger -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) logger = setup_logger() diff --git a/backend/ee/onyx/server/tenants/product_gating.py b/backend/ee/onyx/server/tenants/product_gating.py index 8fa1e8f9ef..8588e2e190 100644 --- a/backend/ee/onyx/server/tenants/product_gating.py +++ b/backend/ee/onyx/server/tenants/product_gating.py @@ -1,10 +1,8 @@ from ee.onyx.configs.app_configs import GATED_TENANTS_KEY from onyx.configs.constants import ONYX_CLOUD_TENANT_ID -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client +from onyx.redis.redis_pool import get_redis_client, get_redis_replica_client from onyx.server.settings.models import ApplicationStatus -from onyx.server.settings.store import load_settings -from onyx.server.settings.store import store_settings +from onyx.server.settings.store import load_settings, store_settings from onyx.utils.logger import setup_logger from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/ee/onyx/server/tenants/provisioning.py b/backend/ee/onyx/server/tenants/provisioning.py index 1af0145916..7ae23f2279 100644 --- a/backend/ee/onyx/server/tenants/provisioning.py +++ b/backend/ee/onyx/server/tenants/provisioning.py @@ -4,63 +4,82 @@ import uuid import aiohttp # Async HTTP client import httpx import requests -from fastapi import HTTPException -from fastapi import Request +from fastapi import HTTPException, Request from sqlalchemy import select from sqlalchemy.orm import Session from ee.onyx.configs.app_configs import HUBSPOT_TRACKING_URL from ee.onyx.server.tenants.access import generate_data_plane_token -from ee.onyx.server.tenants.models import TenantByDomainResponse -from ee.onyx.server.tenants.models import TenantCreationPayload -from ee.onyx.server.tenants.models import TenantDeletionPayload -from ee.onyx.server.tenants.schema_management import create_schema_if_not_exists -from ee.onyx.server.tenants.schema_management import drop_schema -from ee.onyx.server.tenants.schema_management import run_alembic_migrations -from ee.onyx.server.tenants.user_mapping import add_users_to_tenant -from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email -from ee.onyx.server.tenants.user_mapping import user_owns_a_tenant +from ee.onyx.server.tenants.models import ( + TenantByDomainResponse, + TenantCreationPayload, + TenantDeletionPayload, +) +from ee.onyx.server.tenants.schema_management import ( + create_schema_if_not_exists, + drop_schema, + run_alembic_migrations, +) +from ee.onyx.server.tenants.user_mapping import ( + add_users_to_tenant, + get_tenant_id_for_email, + user_owns_a_tenant, +) from onyx.auth.users import exceptions -from onyx.configs.app_configs import ANTHROPIC_DEFAULT_API_KEY -from onyx.configs.app_configs import AUTO_PROVISION_DEFAULT_LLM_PROVIDERS -from onyx.configs.app_configs import COHERE_DEFAULT_API_KEY -from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import OPENAI_DEFAULT_API_KEY -from onyx.configs.app_configs import OPENROUTER_DEFAULT_API_KEY -from onyx.configs.app_configs import VERTEXAI_DEFAULT_CREDENTIALS -from onyx.configs.app_configs import VERTEXAI_DEFAULT_LOCATION -from onyx.db.engine.sql_engine import get_session_with_shared_schema -from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.configs.app_configs import ( + ANTHROPIC_DEFAULT_API_KEY, + AUTO_PROVISION_DEFAULT_LLM_PROVIDERS, + COHERE_DEFAULT_API_KEY, + CONTROL_PLANE_API_BASE_URL, + DEV_MODE, + OPENAI_DEFAULT_API_KEY, + OPENROUTER_DEFAULT_API_KEY, + VERTEXAI_DEFAULT_CREDENTIALS, + VERTEXAI_DEFAULT_LOCATION, +) +from onyx.db.engine.sql_engine import ( + get_session_with_shared_schema, + get_session_with_tenant, +) from onyx.db.image_generation import create_default_image_gen_config_from_api_key -from onyx.db.llm import fetch_existing_llm_provider_by_name_and_type -from onyx.db.llm import fetch_existing_llm_provider_by_type_nameless -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_cloud_embedding_provider -from onyx.db.llm import upsert_llm_provider -from onyx.db.models import AvailableTenant -from onyx.db.models import IndexModelStatus -from onyx.db.models import SearchSettings -from onyx.db.models import UserTenantMapping +from onyx.db.llm import ( + fetch_existing_llm_provider_by_name_and_type, + fetch_existing_llm_provider_by_type_nameless, + update_default_provider, + upsert_cloud_embedding_provider, + upsert_llm_provider, +) +from onyx.db.models import ( + AvailableTenant, + IndexModelStatus, + SearchSettings, + UserTenantMapping, +) from onyx.llm.well_known_providers.auto_update_models import LLMRecommendations -from onyx.llm.well_known_providers.constants import ANTHROPIC_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import OPENAI_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import OPENROUTER_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import VERTEX_CREDENTIALS_FILE_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_LOCATION_KWARG -from onyx.llm.well_known_providers.constants import VERTEXAI_PROVIDER_NAME -from onyx.llm.well_known_providers.llm_provider_options import get_recommendations +from onyx.llm.well_known_providers.constants import ( + ANTHROPIC_PROVIDER_NAME, + OPENAI_PROVIDER_NAME, + OPENROUTER_PROVIDER_NAME, + VERTEX_CREDENTIALS_FILE_KWARG, + VERTEX_LOCATION_KWARG, + VERTEXAI_PROVIDER_NAME, +) from onyx.llm.well_known_providers.llm_provider_options import ( + get_recommendations, model_configurations_for_provider, ) from onyx.server.manage.embedding.models import CloudEmbeddingProviderCreationRequest -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) from onyx.setup import setup_onyx from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import TENANT_ID_PREFIX +from shared_configs.configs import ( + MULTI_TENANT, + POSTGRES_DEFAULT_SCHEMA, + TENANT_ID_PREFIX, +) from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR from shared_configs.enums import EmbeddingProvider diff --git a/backend/ee/onyx/server/tenants/proxy.py b/backend/ee/onyx/server/tenants/proxy.py index d8a408cbf0..8032c75f06 100644 --- a/backend/ee/onyx/server/tenants/proxy.py +++ b/backend/ee/onyx/server/tenants/proxy.py @@ -22,21 +22,15 @@ Auth levels by endpoint: from typing import Literal import httpx -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Header -from fastapi import HTTPException -from pydantic import BaseModel -from pydantic import ValidationError +from fastapi import APIRouter, Depends, Header, HTTPException +from pydantic import BaseModel, ValidationError from redis.exceptions import RedisError from ee.onyx.configs.app_configs import LICENSE_ENFORCEMENT_ENABLED -from ee.onyx.server.billing.models import SeatUpdateRequest -from ee.onyx.server.billing.models import SeatUpdateResponse +from ee.onyx.server.billing.models import SeatUpdateRequest, SeatUpdateResponse from ee.onyx.server.license.models import LicensePayload from ee.onyx.server.tenants.access import generate_data_plane_token -from ee.onyx.utils.license import is_license_valid -from ee.onyx.utils.license import verify_license_signature +from ee.onyx.utils.license import is_license_valid, verify_license_signature from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL from onyx.redis.redis_pool import get_redis_client from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/ee/onyx/server/tenants/schema_management.py b/backend/ee/onyx/server/tenants/schema_management.py index 60d6ebfcad..27dfdb477b 100644 --- a/backend/ee/onyx/server/tenants/schema_management.py +++ b/backend/ee/onyx/server/tenants/schema_management.py @@ -8,8 +8,7 @@ from sqlalchemy import text from sqlalchemy.orm import Session from sqlalchemy.schema import CreateSchema -from onyx.db.engine.sql_engine import build_connection_string -from onyx.db.engine.sql_engine import get_sqlalchemy_engine +from onyx.db.engine.sql_engine import build_connection_string, get_sqlalchemy_engine from onyx.db.engine.tenant_utils import validate_tenant_id logger = logging.getLogger(__name__) diff --git a/backend/ee/onyx/server/tenants/team_membership_api.py b/backend/ee/onyx/server/tenants/team_membership_api.py index d325087164..e8d89ffe3c 100644 --- a/backend/ee/onyx/server/tenants/team_membership_api.py +++ b/backend/ee/onyx/server/tenants/team_membership_api.py @@ -1,18 +1,17 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from ee.onyx.server.tenants.provisioning import delete_user_from_control_plane -from ee.onyx.server.tenants.user_mapping import remove_all_users_from_tenant -from ee.onyx.server.tenants.user_mapping import remove_users_from_tenant +from ee.onyx.server.tenants.user_mapping import ( + remove_all_users_from_tenant, + remove_users_from_tenant, +) from onyx.auth.permissions import require_permission from onyx.auth.users import User from onyx.db.auth import get_user_count from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.users import delete_user_from_db -from onyx.db.users import get_user_by_email +from onyx.db.users import delete_user_from_db, get_user_by_email from onyx.server.manage.models import UserByEmail from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/ee/onyx/server/tenants/tenant_management_api.py b/backend/ee/onyx/server/tenants/tenant_management_api.py index 0d7cf4faa8..0a18f5c6c4 100644 --- a/backend/ee/onyx/server/tenants/tenant_management_api.py +++ b/backend/ee/onyx/server/tenants/tenant_management_api.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from ee.onyx.server.tenants.models import TenantByDomainResponse from ee.onyx.server.tenants.provisioning import get_tenant_by_domain_from_control_plane diff --git a/backend/ee/onyx/server/tenants/tier_management.py b/backend/ee/onyx/server/tenants/tier_management.py index e84dab33b8..17d56bffb1 100644 --- a/backend/ee/onyx/server/tenants/tier_management.py +++ b/backend/ee/onyx/server/tenants/tier_management.py @@ -6,12 +6,10 @@ Value at TENANT_TIER_KEY is a JSON blob: import json from datetime import datetime -from typing import cast -from typing import NamedTuple +from typing import cast, NamedTuple from ee.onyx.server.license.models import CustomerTier -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client +from onyx.redis.redis_pool import get_redis_client, get_redis_replica_client from onyx.utils.logger import setup_logger # Per-tenant cached CustomerTier; TTL bounds upgrade-visible delay if push is missed. diff --git a/backend/ee/onyx/server/tenants/user_invitations_api.py b/backend/ee/onyx/server/tenants/user_invitations_api.py index f47b007001..41e24e7df4 100644 --- a/backend/ee/onyx/server/tenants/user_invitations_api.py +++ b/backend/ee/onyx/server/tenants/user_invitations_api.py @@ -1,14 +1,16 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException -from ee.onyx.server.tenants.models import ApproveUserRequest -from ee.onyx.server.tenants.models import PendingUserSnapshot -from ee.onyx.server.tenants.models import RequestInviteRequest -from ee.onyx.server.tenants.user_mapping import accept_user_invite -from ee.onyx.server.tenants.user_mapping import approve_user_invite -from ee.onyx.server.tenants.user_mapping import deny_user_invite -from ee.onyx.server.tenants.user_mapping import invite_self_to_tenant +from ee.onyx.server.tenants.models import ( + ApproveUserRequest, + PendingUserSnapshot, + RequestInviteRequest, +) +from ee.onyx.server.tenants.user_mapping import ( + accept_user_invite, + approve_user_invite, + deny_user_invite, + invite_self_to_tenant, +) from onyx.auth.invited_users import get_pending_users from onyx.auth.permissions import require_permission from onyx.auth.users import User diff --git a/backend/ee/onyx/server/tenants/user_mapping.py b/backend/ee/onyx/server/tenants/user_mapping.py index 2b1df08192..11722bb6db 100644 --- a/backend/ee/onyx/server/tenants/user_mapping.py +++ b/backend/ee/onyx/server/tenants/user_mapping.py @@ -1,17 +1,20 @@ from fastapi_users import exceptions from sqlalchemy import select -from onyx.auth.invited_users import get_invited_users -from onyx.auth.invited_users import get_pending_users -from onyx.auth.invited_users import write_invited_users -from onyx.auth.invited_users import write_pending_users -from onyx.db.engine.sql_engine import get_session_with_shared_schema -from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.auth.invited_users import ( + get_invited_users, + get_pending_users, + write_invited_users, + write_pending_users, +) +from onyx.db.engine.sql_engine import ( + get_session_with_shared_schema, + get_session_with_tenant, +) from onyx.db.models import UserTenantMapping from onyx.server.manage.models import TenantSnapshot from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR logger = setup_logger() diff --git a/backend/ee/onyx/server/token_rate_limits/api.py b/backend/ee/onyx/server/token_rate_limits/api.py index fdc7c27aa6..222c3d6feb 100644 --- a/backend/ee/onyx/server/token_rate_limits/api.py +++ b/backend/ee/onyx/server/token_rate_limits/api.py @@ -1,25 +1,30 @@ from collections import defaultdict -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session -from ee.onyx.db.token_limit import fetch_all_user_group_token_rate_limits_by_group -from ee.onyx.db.token_limit import fetch_user_group_token_rate_limits_for_user -from ee.onyx.db.token_limit import insert_user_group_token_rate_limit +from ee.onyx.db.token_limit import ( + fetch_all_user_group_token_rate_limits_by_group, + fetch_user_group_token_rate_limits_for_user, + insert_user_group_token_rate_limit, +) from onyx.auth.permissions import require_permission from onyx.auth.users import current_curator_or_admin_user from onyx.configs.constants import PUBLIC_API_TAGS from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.token_limit import fetch_all_user_token_rate_limits -from onyx.db.token_limit import insert_user_token_rate_limit +from onyx.db.token_limit import ( + fetch_all_user_token_rate_limits, + insert_user_token_rate_limit, +) from onyx.server.query_and_chat.token_limit import ( invalidate_any_rate_limit_exists_cache, ) -from onyx.server.token_rate_limits.models import TokenRateLimitArgs -from onyx.server.token_rate_limits.models import TokenRateLimitDisplay +from onyx.server.token_rate_limits.models import ( + TokenRateLimitArgs, + TokenRateLimitDisplay, +) router = APIRouter(prefix="/admin/token-rate-limits", tags=PUBLIC_API_TAGS) diff --git a/backend/ee/onyx/server/user_group/api.py b/backend/ee/onyx/server/user_group/api.py index c8a01efe62..c359d30b95 100644 --- a/backend/ee/onyx/server/user_group/api.py +++ b/backend/ee/onyx/server/user_group/api.py @@ -1,40 +1,40 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from ee.onyx.db.persona import update_persona_access -from ee.onyx.db.user_group import add_users_to_user_group +from ee.onyx.db.user_group import ( + add_users_to_user_group, + fetch_user_group, + fetch_user_groups, + fetch_user_groups_for_user, + insert_user_group, + prepare_user_group_for_deletion, + rename_user_group, + set_group_permission__no_commit, + update_user_curator_relationship, + update_user_group, +) from ee.onyx.db.user_group import delete_user_group as db_delete_user_group -from ee.onyx.db.user_group import fetch_user_group -from ee.onyx.db.user_group import fetch_user_groups -from ee.onyx.db.user_group import fetch_user_groups_for_user -from ee.onyx.db.user_group import insert_user_group -from ee.onyx.db.user_group import prepare_user_group_for_deletion -from ee.onyx.db.user_group import rename_user_group -from ee.onyx.db.user_group import set_group_permission__no_commit -from ee.onyx.db.user_group import update_user_curator_relationship -from ee.onyx.db.user_group import update_user_group -from ee.onyx.server.user_group.models import AddUsersToUserGroupRequest -from ee.onyx.server.user_group.models import MinimalUserGroupSnapshot -from ee.onyx.server.user_group.models import SetCuratorRequest -from ee.onyx.server.user_group.models import SetPermissionRequest -from ee.onyx.server.user_group.models import SetPermissionResponse -from ee.onyx.server.user_group.models import UpdateGroupAgentsRequest -from ee.onyx.server.user_group.models import UserGroup -from ee.onyx.server.user_group.models import UserGroupCreate -from ee.onyx.server.user_group.models import UserGroupRename -from ee.onyx.server.user_group.models import UserGroupUpdate -from onyx.auth.permissions import NON_TOGGLEABLE_PERMISSIONS -from onyx.auth.permissions import require_permission +from ee.onyx.server.user_group.models import ( + AddUsersToUserGroupRequest, + MinimalUserGroupSnapshot, + SetCuratorRequest, + SetPermissionRequest, + SetPermissionResponse, + UpdateGroupAgentsRequest, + UserGroup, + UserGroupCreate, + UserGroupRename, + UserGroupUpdate, +) +from onyx.auth.permissions import NON_TOGGLEABLE_PERMISSIONS, require_permission from onyx.auth.users import current_curator_or_admin_user from onyx.configs.app_configs import DISABLE_VECTOR_DB from onyx.configs.constants import PUBLIC_API_TAGS from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.models import User -from onyx.db.models import UserRole +from onyx.db.models import User, UserRole from onyx.db.persona import get_persona_by_id from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/ee/onyx/server/user_group/models.py b/backend/ee/onyx/server/user_group/models.py index 3dd878c48f..1efd6c3de8 100644 --- a/backend/ee/onyx/server/user_group/models.py +++ b/backend/ee/onyx/server/user_group/models.py @@ -4,13 +4,14 @@ from pydantic import BaseModel from onyx.auth.permissions import Permission from onyx.db.models import UserGroup as UserGroupModel -from onyx.server.documents.models import ConnectorCredentialPairDescriptor -from onyx.server.documents.models import ConnectorSnapshot -from onyx.server.documents.models import CredentialSnapshot +from onyx.server.documents.models import ( + ConnectorCredentialPairDescriptor, + ConnectorSnapshot, + CredentialSnapshot, +) from onyx.server.features.document_set.models import DocumentSet from onyx.server.features.persona.models import PersonaSnapshot -from onyx.server.manage.models import UserInfo -from onyx.server.manage.models import UserPreferences +from onyx.server.manage.models import UserInfo, UserPreferences class UserGroup(BaseModel): diff --git a/backend/ee/onyx/utils/encryption.py b/backend/ee/onyx/utils/encryption.py index b1be2154f0..37ee226b48 100644 --- a/backend/ee/onyx/utils/encryption.py +++ b/backend/ee/onyx/utils/encryption.py @@ -3,9 +3,7 @@ from os import urandom from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import padding -from cryptography.hazmat.primitives.ciphers import algorithms -from cryptography.hazmat.primitives.ciphers import Cipher -from cryptography.hazmat.primitives.ciphers import modes +from cryptography.hazmat.primitives.ciphers import algorithms, Cipher, modes from onyx.configs.app_configs import ENCRYPTION_KEY_SECRET from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/utils/license.py b/backend/ee/onyx/utils/license.py index 138f0e3b1a..d404c009f6 100644 --- a/backend/ee/onyx/utils/license.py +++ b/backend/ee/onyx/utils/license.py @@ -3,18 +3,15 @@ import base64 import json import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from pathlib import Path from cryptography.exceptions import InvalidSignature -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import padding from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey -from ee.onyx.server.license.models import LicenseData -from ee.onyx.server.license.models import LicensePayload +from ee.onyx.server.license.models import LicenseData, LicensePayload from onyx.server.settings.models import ApplicationStatus from onyx.utils.logger import setup_logger diff --git a/backend/ee/onyx/utils/license_expiry.py b/backend/ee/onyx/utils/license_expiry.py index 3740d614c8..3f5086b4fd 100644 --- a/backend/ee/onyx/utils/license_expiry.py +++ b/backend/ee/onyx/utils/license_expiry.py @@ -11,9 +11,7 @@ Stages: GRACE — license already expired, within 14-day grace window """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from enum import Enum LICENSE_GRACE_PERIOD_DAYS = 14 diff --git a/backend/ee/onyx/utils/license_notifications.py b/backend/ee/onyx/utils/license_notifications.py index fd053b4487..5dba68ce19 100644 --- a/backend/ee/onyx/utils/license_notifications.py +++ b/backend/ee/onyx/utils/license_notifications.py @@ -7,24 +7,22 @@ admins for a given (stage, expires_at[, sent_date]) tuple are skipped — only freshly-notified admins receive an email. """ -from datetime import date -from datetime import datetime -from datetime import timezone +from datetime import date, datetime, timezone from typing import Any from sqlalchemy.orm import Session from ee.onyx.db.license import get_license from ee.onyx.utils.license import verify_license_signature -from ee.onyx.utils.license_expiry import ExpiryWarningStage -from ee.onyx.utils.license_expiry import get_expiry_warning_stage -from ee.onyx.utils.license_expiry import get_grace_days_remaining -from onyx.auth.email_utils import build_html_email -from onyx.auth.email_utils import send_email +from ee.onyx.utils.license_expiry import ( + ExpiryWarningStage, + get_expiry_warning_stage, + get_grace_days_remaining, +) +from onyx.auth.email_utils import build_html_email, send_email from onyx.auth.schemas import UserRole from onyx.configs.app_configs import EMAIL_CONFIGURED -from onyx.configs.constants import NotificationType -from onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME +from onyx.configs.constants import NotificationType, ONYX_DEFAULT_APPLICATION_NAME from onyx.db.models import User from onyx.db.notification import batch_create_notifications from onyx.db.users import get_active_admin_users diff --git a/backend/ee/onyx/utils/posthog_client.py b/backend/ee/onyx/utils/posthog_client.py index e87c74212a..3a62093411 100644 --- a/backend/ee/onyx/utils/posthog_client.py +++ b/backend/ee/onyx/utils/posthog_client.py @@ -4,10 +4,12 @@ from urllib.parse import unquote from posthog import Posthog -from ee.onyx.configs.app_configs import MARKETING_POSTHOG_API_KEY -from ee.onyx.configs.app_configs import POSTHOG_API_KEY -from ee.onyx.configs.app_configs import POSTHOG_DEBUG_LOGS_ENABLED -from ee.onyx.configs.app_configs import POSTHOG_HOST +from ee.onyx.configs.app_configs import ( + MARKETING_POSTHOG_API_KEY, + POSTHOG_API_KEY, + POSTHOG_DEBUG_LOGS_ENABLED, + POSTHOG_HOST, +) from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/ee/onyx/utils/tier.py b/backend/ee/onyx/utils/tier.py index 8b6d48c7a0..ae3f6e04c3 100644 --- a/backend/ee/onyx/utils/tier.py +++ b/backend/ee/onyx/utils/tier.py @@ -15,29 +15,23 @@ naturally resolves back to BUSINESS without waiting on a webhook. from __future__ import annotations -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import requests from redis.exceptions import RedisError -from sqlalchemy.exc import ProgrammingError -from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.exc import ProgrammingError, SQLAlchemyError from ee.onyx.configs.app_configs import LICENSE_ENFORCEMENT_ENABLED -from ee.onyx.db.license import get_cached_license_metadata -from ee.onyx.db.license import refresh_license_cache +from ee.onyx.db.license import get_cached_license_metadata, refresh_license_cache from ee.onyx.server.license.models import CustomerTier from ee.onyx.server.tenants.billing import fetch_billing_information -from ee.onyx.server.tenants.models import BillingInformation -from ee.onyx.server.tenants.models import SubscriptionStatusResponse -from ee.onyx.server.tenants.tier_management import get_cached_tier -from ee.onyx.server.tenants.tier_management import update_tenant_tier +from ee.onyx.server.tenants.models import BillingInformation, SubscriptionStatusResponse +from ee.onyx.server.tenants.tier_management import get_cached_tier, update_tenant_tier from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import AccessType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.settings.models import ApplicationStatus -from onyx.server.settings.models import Tier +from onyx.server.settings.models import ApplicationStatus, Tier from onyx.server.settings.tier_order import tier_at_least from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import global_version diff --git a/backend/model_server/encoders.py b/backend/model_server/encoders.py index 528ab60753..6f77bdb5f3 100644 --- a/backend/model_server/encoders.py +++ b/backend/model_server/encoders.py @@ -1,18 +1,13 @@ import asyncio import time -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING -from fastapi import APIRouter -from fastapi import HTTPException -from fastapi import Request +from fastapi import APIRouter, HTTPException, Request from model_server.utils import simple_log_function_time from onyx.utils.logger import setup_logger from shared_configs.enums import EmbedTextType -from shared_configs.model_server_models import Embedding -from shared_configs.model_server_models import EmbedRequest -from shared_configs.model_server_models import EmbedResponse +from shared_configs.model_server_models import Embedding, EmbedRequest, EmbedResponse if TYPE_CHECKING: from sentence_transformers import SentenceTransformer diff --git a/backend/model_server/main.py b/backend/model_server/main.py index 49ca3c3376..333fb98ee1 100644 --- a/backend/model_server/main.py +++ b/backend/model_server/main.py @@ -16,18 +16,20 @@ from transformers import logging as transformer_logging from model_server.ca_certs import configure_trusted_ca_bundle from model_server.encoders import router as encoders_router from model_server.management_endpoints import router as management_router -from model_server.utils import get_cgroup_cpu_limit -from model_server.utils import get_gpu_type +from model_server.utils import get_cgroup_cpu_limit, get_gpu_type from onyx import __version__ -from onyx.utils.logger import setup_logger -from onyx.utils.logger import setup_uvicorn_logger -from onyx.utils.middleware import add_onyx_request_id_middleware -from onyx.utils.middleware import add_onyx_tenant_id_middleware -from shared_configs.configs import INDEXING_ONLY -from shared_configs.configs import MIN_THREADS_ML_MODELS -from shared_configs.configs import MODEL_SERVER_PORT -from shared_configs.configs import SENTRY_DSN -from shared_configs.configs import SENTRY_TRACES_SAMPLE_RATE +from onyx.utils.logger import setup_logger, setup_uvicorn_logger +from onyx.utils.middleware import ( + add_onyx_request_id_middleware, + add_onyx_tenant_id_middleware, +) +from shared_configs.configs import ( + INDEXING_ONLY, + MIN_THREADS_ML_MODELS, + MODEL_SERVER_PORT, + SENTRY_DSN, + SENTRY_TRACES_SAMPLE_RATE, +) os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" diff --git a/backend/model_server/management_endpoints.py b/backend/model_server/management_endpoints.py index a722911ae6..877e764411 100644 --- a/backend/model_server/management_endpoints.py +++ b/backend/model_server/management_endpoints.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Response +from fastapi import APIRouter, Response from model_server.constants import GPUStatus from model_server.utils import get_gpu_type diff --git a/backend/model_server/utils.py b/backend/model_server/utils.py index 072dd8b94f..9a3336e286 100644 --- a/backend/model_server/utils.py +++ b/backend/model_server/utils.py @@ -1,13 +1,9 @@ import asyncio import time -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Callable, Generator, Iterator from functools import wraps from pathlib import Path -from typing import Any -from typing import cast -from typing import TypeVar +from typing import Any, cast, TypeVar import torch diff --git a/backend/onyx/access/access.py b/backend/onyx/access/access.py index 3507732dbd..0519c1110d 100644 --- a/backend/onyx/access/access.py +++ b/backend/onyx/access/access.py @@ -2,33 +2,33 @@ from collections.abc import Callable from typing import cast from sqlalchemy import cast as sa_cast -from sqlalchemy import or_ -from sqlalchemy import select +from sqlalchemy import or_, select from sqlalchemy.dialects import postgresql from sqlalchemy.orm import Session from onyx.access.models import DocumentAccess from onyx.access.utils import prefix_user_email -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import PUBLIC_DOC_PAT -from onyx.db.document import get_access_info_for_document -from onyx.db.document import get_access_info_for_documents -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession -from onyx.db.models import ChatSessionSharedStatus -from onyx.db.models import Connector -from onyx.db.models import Document -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import FileRecord -from onyx.db.models import Persona -from onyx.db.models import Persona__User -from onyx.db.models import Persona__UserFile -from onyx.db.models import User -from onyx.db.models import UserFile +from onyx.configs.constants import DocumentSource, FileOrigin, PUBLIC_DOC_PAT +from onyx.db.document import get_access_info_for_document, get_access_info_for_documents +from onyx.db.models import ( + ChatMessage, + ChatSession, + ChatSessionSharedStatus, + Connector, + Document, + DocumentByConnectorCredentialPair, + FileRecord, + Persona, + Persona__User, + Persona__UserFile, + User, + UserFile, +) from onyx.db.user_file import fetch_user_files_with_access_relationships -from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop -from onyx.utils.variable_functionality import fetch_versioned_implementation +from onyx.utils.variable_functionality import ( + fetch_ee_implementation_or_noop, + fetch_versioned_implementation, +) def _get_access_for_document( diff --git a/backend/onyx/access/models.py b/backend/onyx/access/models.py index fbf2550e1e..0ddf041ddf 100644 --- a/backend/onyx/access/models.py +++ b/backend/onyx/access/models.py @@ -1,8 +1,10 @@ from dataclasses import dataclass -from onyx.access.utils import prefix_external_group -from onyx.access.utils import prefix_user_email -from onyx.access.utils import prefix_user_group +from onyx.access.utils import ( + prefix_external_group, + prefix_user_email, + prefix_user_group, +) from onyx.configs.constants import PUBLIC_DOC_PAT diff --git a/backend/onyx/auth/anonymous_user.py b/backend/onyx/auth/anonymous_user.py index b8aa809d3d..52aaf3f9d0 100644 --- a/backend/onyx/auth/anonymous_user.py +++ b/backend/onyx/auth/anonymous_user.py @@ -1,17 +1,15 @@ from collections.abc import Mapping -from typing import Any -from typing import cast +from typing import Any, cast from onyx.auth.schemas import UserRole -from onyx.configs.constants import ANONYMOUS_USER_EMAIL -from onyx.configs.constants import ANONYMOUS_USER_INFO_ID -from onyx.configs.constants import KV_ANONYMOUS_USER_PERSONALIZATION_KEY -from onyx.configs.constants import KV_ANONYMOUS_USER_PREFERENCES_KEY -from onyx.key_value_store.store import KeyValueStore -from onyx.key_value_store.store import KvKeyNotFoundError -from onyx.server.manage.models import UserInfo -from onyx.server.manage.models import UserPersonalization -from onyx.server.manage.models import UserPreferences +from onyx.configs.constants import ( + ANONYMOUS_USER_EMAIL, + ANONYMOUS_USER_INFO_ID, + KV_ANONYMOUS_USER_PERSONALIZATION_KEY, + KV_ANONYMOUS_USER_PREFERENCES_KEY, +) +from onyx.key_value_store.store import KeyValueStore, KvKeyNotFoundError +from onyx.server.manage.models import UserInfo, UserPersonalization, UserPreferences def set_anonymous_user_preferences( diff --git a/backend/onyx/auth/api_key.py b/backend/onyx/auth/api_key.py index 6b28942bd0..bff3a346fe 100644 --- a/backend/onyx/auth/api_key.py +++ b/backend/onyx/auth/api_key.py @@ -7,9 +7,11 @@ from fastapi import Request from passlib.hash import sha256_crypt from pydantic import BaseModel -from onyx.auth.constants import API_KEY_LENGTH -from onyx.auth.constants import API_KEY_PREFIX -from onyx.auth.constants import DEPRECATED_API_KEY_PREFIX +from onyx.auth.constants import ( + API_KEY_LENGTH, + API_KEY_PREFIX, + DEPRECATED_API_KEY_PREFIX, +) from onyx.auth.schemas import UserRole from onyx.auth.utils import get_hashed_bearer_token_from_request from onyx.configs.app_configs import API_KEY_HASH_ROUNDS diff --git a/backend/onyx/auth/captcha.py b/backend/onyx/auth/captcha.py index f6a5ff7fad..398c6ac567 100644 --- a/backend/onyx/auth/captcha.py +++ b/backend/onyx/auth/captcha.py @@ -17,22 +17,22 @@ the OAuth cookie. import hashlib import hmac import time -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from enum import StrEnum import httpx -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field -from onyx.configs.app_configs import CAPTCHA_COOKIE_TTL_SECONDS -from onyx.configs.app_configs import CAPTCHA_ENABLED -from onyx.configs.app_configs import RECAPTCHA_ENTERPRISE_API_KEY -from onyx.configs.app_configs import RECAPTCHA_ENTERPRISE_PROJECT_ID -from onyx.configs.app_configs import RECAPTCHA_HOSTNAME_ALLOWLIST -from onyx.configs.app_configs import RECAPTCHA_SCORE_THRESHOLD -from onyx.configs.app_configs import RECAPTCHA_SITE_KEY -from onyx.configs.app_configs import USER_AUTH_SECRET +from onyx.configs.app_configs import ( + CAPTCHA_COOKIE_TTL_SECONDS, + CAPTCHA_ENABLED, + RECAPTCHA_ENTERPRISE_API_KEY, + RECAPTCHA_ENTERPRISE_PROJECT_ID, + RECAPTCHA_HOSTNAME_ALLOWLIST, + RECAPTCHA_SCORE_THRESHOLD, + RECAPTCHA_SITE_KEY, + USER_AUTH_SECRET, +) from onyx.redis.redis_pool import get_async_redis_connection from onyx.utils.logger import setup_logger diff --git a/backend/onyx/auth/email_utils.py b/backend/onyx/auth/email_utils.py index 6dfced3086..ce9f0a6e82 100644 --- a/backend/onyx/auth/email_utils.py +++ b/backend/onyx/auth/email_utils.py @@ -4,33 +4,35 @@ from datetime import datetime from email.mime.image import MIMEImage from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText -from email.utils import formatdate -from email.utils import make_msgid +from email.utils import formatdate, make_msgid import sendgrid -from sendgrid.helpers.mail import Attachment -from sendgrid.helpers.mail import Content -from sendgrid.helpers.mail import ContentId -from sendgrid.helpers.mail import Disposition -from sendgrid.helpers.mail import Email -from sendgrid.helpers.mail import FileContent -from sendgrid.helpers.mail import FileName -from sendgrid.helpers.mail import FileType -from sendgrid.helpers.mail import Mail -from sendgrid.helpers.mail import To +from sendgrid.helpers.mail import ( + Attachment, + Content, + ContentId, + Disposition, + Email, + FileContent, + FileName, + FileType, + Mail, + To, +) -from onyx.configs.app_configs import EMAIL_ARCHIVE_BCC_ADDRESSES -from onyx.configs.app_configs import EMAIL_CONFIGURED -from onyx.configs.app_configs import EMAIL_FROM -from onyx.configs.app_configs import SENDGRID_API_KEY -from onyx.configs.app_configs import SMTP_PASS -from onyx.configs.app_configs import SMTP_PORT -from onyx.configs.app_configs import SMTP_SERVER -from onyx.configs.app_configs import SMTP_STARTTLS -from onyx.configs.app_configs import SMTP_USER -from onyx.configs.app_configs import WEB_DOMAIN -from onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME -from onyx.configs.constants import ONYX_DISCORD_URL +from onyx.configs.app_configs import ( + EMAIL_ARCHIVE_BCC_ADDRESSES, + EMAIL_CONFIGURED, + EMAIL_FROM, + SENDGRID_API_KEY, + SMTP_PASS, + SMTP_PORT, + SMTP_SERVER, + SMTP_STARTTLS, + SMTP_USER, + WEB_DOMAIN, +) +from onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME, ONYX_DISCORD_URL from onyx.db.models import User from onyx.server.runtime.onyx_runtime import OnyxRuntime from onyx.utils.logger import setup_logger diff --git a/backend/onyx/auth/invited_users.py b/backend/onyx/auth/invited_users.py index 60838d3fe8..bf68e800ed 100644 --- a/backend/onyx/auth/invited_users.py +++ b/backend/onyx/auth/invited_users.py @@ -1,7 +1,6 @@ from typing import cast -from onyx.configs.constants import KV_PENDING_USERS_KEY -from onyx.configs.constants import KV_USER_STORE_KEY +from onyx.configs.constants import KV_PENDING_USERS_KEY, KV_USER_STORE_KEY from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError from onyx.utils.special_types import JSON_ro diff --git a/backend/onyx/auth/jwt.py b/backend/onyx/auth/jwt.py index 4622dd5682..a06bc000c0 100644 --- a/backend/onyx/auth/jwt.py +++ b/backend/onyx/auth/jwt.py @@ -1,15 +1,13 @@ import json from enum import Enum from functools import lru_cache -from typing import Any -from typing import cast +from typing import Any, cast import jwt import requests from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey from jwt import decode as jwt_decode -from jwt import InvalidTokenError -from jwt import PyJWTError +from jwt import InvalidTokenError, PyJWTError from jwt.algorithms import RSAAlgorithm # ty: ignore[possibly-missing-import] from onyx.configs.app_configs import JWT_PUBLIC_KEY_URL diff --git a/backend/onyx/auth/mobile_sso/code_store.py b/backend/onyx/auth/mobile_sso/code_store.py index 7009fd7d1d..4bd85e1d85 100644 --- a/backend/onyx/auth/mobile_sso/code_store.py +++ b/backend/onyx/auth/mobile_sso/code_store.py @@ -23,8 +23,7 @@ import json import secrets from onyx.auth.pkce import compute_s256_challenge -from onyx.configs.app_configs import MOBILE_SSO_CODE_PREFIX -from onyx.configs.app_configs import MOBILE_SSO_CODE_TTL_SECONDS +from onyx.configs.app_configs import MOBILE_SSO_CODE_PREFIX, MOBILE_SSO_CODE_TTL_SECONDS from onyx.redis.redis_pool import get_async_redis_connection from onyx.utils.logger import setup_logger diff --git a/backend/onyx/auth/oauth_claims_capture.py b/backend/onyx/auth/oauth_claims_capture.py index b6fa0ddbef..ffd4d789d9 100644 --- a/backend/onyx/auth/oauth_claims_capture.py +++ b/backend/onyx/auth/oauth_claims_capture.py @@ -23,15 +23,16 @@ claims and token metadata (key names, scope, expiry). import asyncio import json -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import httpx import jwt -from onyx.configs.app_configs import IDP_PROFILE_CLAIM_MAP -from onyx.configs.app_configs import IDP_PROFILE_ENRICHMENT_ENABLED +from onyx.configs.app_configs import ( + IDP_PROFILE_CLAIM_MAP, + IDP_PROFILE_ENRICHMENT_ENABLED, +) from onyx.redis.redis_pool import get_async_redis_connection from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/auth/oauth_refresher.py b/backend/onyx/auth/oauth_refresher.py index 1ba945fbb1..9364690fa2 100644 --- a/backend/onyx/auth/oauth_refresher.py +++ b/backend/onyx/auth/oauth_refresher.py @@ -2,23 +2,19 @@ import asyncio import os import time import uuid -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast -from typing import Dict -from typing import List -from typing import Optional +from datetime import datetime, timezone +from typing import Any, cast, Dict, List, Optional import httpx from fastapi_users.manager import BaseUserManager from sqlalchemy.ext.asyncio import AsyncSession -from onyx.configs.app_configs import OAUTH_CLIENT_ID -from onyx.configs.app_configs import OAUTH_CLIENT_SECRET -from onyx.configs.app_configs import OPENID_CONFIG_URL -from onyx.db.models import OAuthAccount -from onyx.db.models import User +from onyx.configs.app_configs import ( + OAUTH_CLIENT_ID, + OAUTH_CLIENT_SECRET, + OPENID_CONFIG_URL, +) +from onyx.db.models import OAuthAccount, User from onyx.server.security.store import get_security_settings from onyx.utils.logger import setup_logger diff --git a/backend/onyx/auth/oauth_token_manager.py b/backend/onyx/auth/oauth_token_manager.py index 0dc6cce917..9604e79f66 100644 --- a/backend/onyx/auth/oauth_token_manager.py +++ b/backend/onyx/auth/oauth_token_manager.py @@ -7,10 +7,8 @@ import requests from pydantic import BaseModel from sqlalchemy.orm import Session -from onyx.db.models import OAuthConfig -from onyx.db.models import OAuthUserToken -from onyx.db.oauth_config import get_user_oauth_token -from onyx.db.oauth_config import upsert_user_oauth_token +from onyx.db.models import OAuthConfig, OAuthUserToken +from onyx.db.oauth_config import get_user_oauth_token, upsert_user_oauth_token from onyx.server.security.models import outbound_ssrf_params from onyx.server.security.store import get_security_settings from onyx.utils.logger import setup_logger diff --git a/backend/onyx/auth/oidc_client.py b/backend/onyx/auth/oidc_client.py index 194d9563d2..651b4d14b1 100644 --- a/backend/onyx/auth/oidc_client.py +++ b/backend/onyx/auth/oidc_client.py @@ -4,11 +4,9 @@ verified, and the discovery document's issuer must own the configured discovery URL, so one provider's tokens cannot be replayed against another provider's callback (OIDC mix-up defense).""" -from urllib.parse import unquote -from urllib.parse import urlsplit +from urllib.parse import unquote, urlsplit -from httpx_oauth.clients.openid import BASE_SCOPES -from httpx_oauth.clients.openid import OpenID +from httpx_oauth.clients.openid import BASE_SCOPES, OpenID from httpx_oauth.exceptions import GetIdEmailError diff --git a/backend/onyx/auth/pat.py b/backend/onyx/auth/pat.py index 57c0f98c96..c9f4606a56 100644 --- a/backend/onyx/auth/pat.py +++ b/backend/onyx/auth/pat.py @@ -2,15 +2,12 @@ import hashlib import secrets -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from urllib.parse import quote from fastapi import Request -from onyx.auth.constants import PAT_LENGTH -from onyx.auth.constants import PAT_PREFIX +from onyx.auth.constants import PAT_LENGTH, PAT_PREFIX from onyx.auth.utils import get_hashed_bearer_token_from_request from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/auth/permissions.py b/backend/onyx/auth/permissions.py index 0b8b05b1db..080419e01d 100644 --- a/backend/onyx/auth/permissions.py +++ b/backend/onyx/auth/permissions.py @@ -6,15 +6,12 @@ loaded for free with every auth query. Implied permissions are expanded at read time — only directly granted permissions are persisted. """ -from collections.abc import Callable -from collections.abc import Coroutine +from collections.abc import Callable, Coroutine from typing import Any -from fastapi import Depends -from fastapi import Request +from fastapi import Depends, Request -from onyx.auth.users import current_chat_accessible_user -from onyx.auth.users import current_user +from onyx.auth.users import current_chat_accessible_user, current_user from onyx.db.enums import Permission from onyx.db.models import User from onyx.db.permissions import parse_permission_values diff --git a/backend/onyx/auth/session_tokens.py b/backend/onyx/auth/session_tokens.py index 0794b19e04..d6af546279 100644 --- a/backend/onyx/auth/session_tokens.py +++ b/backend/onyx/auth/session_tokens.py @@ -16,18 +16,12 @@ only where authentication has conclusively failed for the request. from contextvars import ContextVar from dataclasses import dataclass -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from enum import Enum -from pydantic import AwareDatetime -from pydantic import BaseModel -from pydantic import ValidationError +from pydantic import AwareDatetime, BaseModel, ValidationError -from onyx.auth.constants import API_KEY_PREFIX -from onyx.auth.constants import DEPRECATED_API_KEY_PREFIX -from onyx.auth.constants import PAT_PREFIX +from onyx.auth.constants import API_KEY_PREFIX, DEPRECATED_API_KEY_PREFIX, PAT_PREFIX from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.utils.logger import setup_logger diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index e9155dfa23..d8a4028cf6 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -5,72 +5,58 @@ import random import secrets import string import uuid -from collections.abc import AsyncGenerator -from collections.abc import Sequence -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from collections.abc import AsyncGenerator, Sequence +from datetime import datetime, timedelta, timezone from functools import partial -from typing import Any -from typing import cast -from typing import Dict -from typing import List -from typing import Literal -from typing import Optional -from typing import Protocol -from typing import Tuple -from typing import TypeVar +from typing import Any, cast, Dict, List, Literal, Optional, Protocol, Tuple, TypeVar from urllib.parse import urlparse import jwt -from email_validator import EmailNotValidError -from email_validator import EmailUndeliverableError -from email_validator import validate_email -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query -from fastapi import Request -from fastapi import Response -from fastapi import status -from fastapi import WebSocket -from fastapi.responses import JSONResponse -from fastapi.responses import RedirectResponse +from email_validator import EmailNotValidError, EmailUndeliverableError, validate_email +from fastapi import ( + APIRouter, + Depends, + HTTPException, + Query, + Request, + Response, + status, + WebSocket, +) +from fastapi.responses import JSONResponse, RedirectResponse from fastapi.routing import APIRoute from fastapi.security import OAuth2PasswordRequestForm -from fastapi_users import BaseUserManager -from fastapi_users import exceptions -from fastapi_users import FastAPIUsers -from fastapi_users import models -from fastapi_users import schemas -from fastapi_users import UUIDIDMixin -from fastapi_users.authentication import AuthenticationBackend -from fastapi_users.authentication import BearerTransport -from fastapi_users.authentication import CookieTransport -from fastapi_users.authentication import JWTStrategy +from fastapi_users import ( + BaseUserManager, + exceptions, + FastAPIUsers, + models, + schemas, + UUIDIDMixin, +) from fastapi_users.authentication import ( + AuthenticationBackend, + BearerTransport, + CookieTransport, + JWTStrategy, RedisStrategy, # ty: ignore[possibly-missing-import] + Strategy, +) +from fastapi_users.authentication.strategy.db import ( + AccessTokenDatabase, + DatabaseStrategy, ) -from fastapi_users.authentication import Strategy -from fastapi_users.authentication.strategy.db import AccessTokenDatabase -from fastapi_users.authentication.strategy.db import DatabaseStrategy from fastapi_users.exceptions import UserAlreadyExists -from fastapi_users.jwt import decode_jwt -from fastapi_users.jwt import generate_jwt -from fastapi_users.jwt import SecretType +from fastapi_users.jwt import decode_jwt, generate_jwt, SecretType from fastapi_users.manager import UserManagerDependency from fastapi_users.openapi import OpenAPIResponseType -from fastapi_users.router.common import ErrorCode -from fastapi_users.router.common import ErrorModel +from fastapi_users.router.common import ErrorCode, ErrorModel from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase from httpx_oauth.exceptions import GetIdEmailError from httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallback -from httpx_oauth.oauth2 import BaseOAuth2 -from httpx_oauth.oauth2 import GetAccessTokenError -from httpx_oauth.oauth2 import OAuth2Token +from httpx_oauth.oauth2 import BaseOAuth2, GetAccessTokenError, OAuth2Token from pydantic import BaseModel -from sqlalchemy import nulls_last -from sqlalchemy import select +from sqlalchemy import nulls_last, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session @@ -78,98 +64,112 @@ from starlette.routing import BaseRoute from onyx.auth.api_key import get_hashed_api_key_from_request from onyx.auth.disposable_email_validator import is_disposable_email -from onyx.auth.email_utils import send_forgot_password_email -from onyx.auth.email_utils import send_user_verification_email -from onyx.auth.invited_users import get_invited_users -from onyx.auth.invited_users import remove_user_from_invited_users +from onyx.auth.email_utils import ( + send_forgot_password_email, + send_user_verification_email, +) +from onyx.auth.invited_users import get_invited_users, remove_user_from_invited_users from onyx.auth.jwt import verify_jwt_token -from onyx.auth.mobile_sso.sso_completion import apply_mobile_state -from onyx.auth.mobile_sso.sso_completion import complete_mobile_sso -from onyx.auth.mobile_sso.sso_completion import is_mobile_sso +from onyx.auth.mobile_sso.sso_completion import ( + apply_mobile_state, + complete_mobile_sso, + is_mobile_sso, +) from onyx.auth.oauth_claims_capture import capture_oauth_login_claims from onyx.auth.pat import get_hashed_pat_from_request -from onyx.auth.schemas import AuthBackend -from onyx.auth.schemas import UserCreate -from onyx.auth.schemas import UserRole -from onyx.auth.session_tokens import build_session_rejection_error -from onyx.auth.session_tokens import build_session_token_value -from onyx.auth.session_tokens import build_session_tombstone_value -from onyx.auth.session_tokens import classify_session_token_value -from onyx.auth.session_tokens import compute_session_expires_at -from onyx.auth.session_tokens import may_be_session_token -from onyx.auth.session_tokens import physical_session_ttl_seconds -from onyx.auth.session_tokens import record_session_rejection -from onyx.auth.session_tokens import SESSION_TOKEN_GRACE_PERIOD_SECONDS -from onyx.auth.session_tokens import SessionRejection +from onyx.auth.schemas import AuthBackend, UserCreate, UserRole +from onyx.auth.session_tokens import ( + build_session_rejection_error, + build_session_token_value, + build_session_tombstone_value, + classify_session_token_value, + compute_session_expires_at, + may_be_session_token, + physical_session_ttl_seconds, + record_session_rejection, + SESSION_TOKEN_GRACE_PERIOD_SECONDS, + SessionRejection, +) from onyx.auth.signup_rate_limit import enforce_signup_rate_limit -from onyx.configs.app_configs import AUTH_BACKEND -from onyx.configs.app_configs import AUTH_COOKIE_EXPIRE_TIME_SECONDS -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import EMAIL_CONFIGURED -from onyx.configs.app_configs import INTEGRATION_TESTS_MODE -from onyx.configs.app_configs import JWT_PUBLIC_KEY_URL -from onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX -from onyx.configs.app_configs import REQUIRE_EMAIL_VERIFICATION -from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS -from onyx.configs.app_configs import USER_AUTH_SECRET -from onyx.configs.app_configs import WEB_DOMAIN -from onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME -from onyx.configs.constants import ANONYMOUS_USER_EMAIL -from onyx.configs.constants import ANONYMOUS_USER_UUID -from onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN -from onyx.configs.constants import DANSWER_API_KEY_PREFIX -from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME -from onyx.configs.constants import MilestoneRecordType -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import PASSWORD_SPECIAL_CHARS -from onyx.configs.constants import UNNAMED_KEY_PLACEHOLDER +from onyx.configs.app_configs import ( + AUTH_BACKEND, + AUTH_COOKIE_EXPIRE_TIME_SECONDS, + DEV_MODE, + EMAIL_CONFIGURED, + INTEGRATION_TESTS_MODE, + JWT_PUBLIC_KEY_URL, + REDIS_AUTH_KEY_PREFIX, + REQUIRE_EMAIL_VERIFICATION, + SESSION_EXPIRE_TIME_SECONDS, + USER_AUTH_SECRET, + WEB_DOMAIN, +) +from onyx.configs.constants import ( + ANONYMOUS_USER_COOKIE_NAME, + ANONYMOUS_USER_EMAIL, + ANONYMOUS_USER_UUID, + DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN, + DANSWER_API_KEY_PREFIX, + FASTAPI_USERS_AUTH_COOKIE_NAME, + MilestoneRecordType, + OnyxRedisLocks, + PASSWORD_SPECIAL_CHARS, + UNNAMED_KEY_PLACEHOLDER, +) from onyx.db.api_key import fetch_user_for_api_key -from onyx.db.auth import get_access_token_db -from onyx.db.auth import get_default_admin_user_emails -from onyx.db.auth import get_user_count -from onyx.db.auth import get_user_db -from onyx.db.auth import SQLAlchemyUserAdminDB -from onyx.db.engine.async_sql_engine import get_async_session -from onyx.db.engine.async_sql_engine import get_async_session_context_manager -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.enums import AccountType -from onyx.db.enums import Permission -from onyx.db.models import AccessToken -from onyx.db.models import OAuthAccount -from onyx.db.models import Persona -from onyx.db.models import User +from onyx.db.auth import ( + get_access_token_db, + get_default_admin_user_emails, + get_user_count, + get_user_db, + SQLAlchemyUserAdminDB, +) +from onyx.db.engine.async_sql_engine import ( + get_async_session, + get_async_session_context_manager, +) +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_tenant, +) +from onyx.db.enums import AccountType, Permission +from onyx.db.models import AccessToken, OAuthAccount, Persona, User from onyx.db.pat import resolve_pat -from onyx.db.users import assign_user_to_default_groups__no_commit -from onyx.db.users import get_user_by_email -from onyx.db.users import is_limited_user +from onyx.db.users import ( + assign_user_to_default_groups__no_commit, + get_user_by_email, + is_limited_user, +) from onyx.error_handling.error_codes import OnyxErrorCode -from onyx.error_handling.exceptions import log_onyx_error -from onyx.error_handling.exceptions import onyx_error_to_json_response -from onyx.error_handling.exceptions import OnyxError -from onyx.redis.redis_pool import get_async_redis_connection -from onyx.redis.redis_pool import retrieve_ws_token_data +from onyx.error_handling.exceptions import ( + log_onyx_error, + onyx_error_to_json_response, + OnyxError, +) +from onyx.redis.redis_pool import get_async_redis_connection, retrieve_ws_token_data from onyx.server.security.store import get_security_settings from onyx.server.settings.store import load_settings from onyx.server.utils import BasicAuthenticationError -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditActor -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import AuditAction, AuditActor, AuditOutcome, emit_audit_event from onyx.utils.logger import setup_logger -from onyx.utils.telemetry import mt_cloud_identify_user -from onyx.utils.telemetry import mt_cloud_telemetry -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType +from onyx.utils.telemetry import ( + mt_cloud_identify_user, + mt_cloud_telemetry, + optional_telemetry, + RecordType, +) from onyx.utils.timing import log_function_time -from onyx.utils.url import add_url_params -from onyx.utils.url import sanitize_next_url +from onyx.utils.url import add_url_params, sanitize_next_url from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop -from shared_configs.configs import async_return_default_schema -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.configs import ( + async_return_default_schema, + MULTI_TENANT, + POSTGRES_DEFAULT_SCHEMA, +) +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) logger = setup_logger() @@ -590,10 +590,12 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): await enforce_signup_rate_limit(request) # Verify captcha if enabled (for cloud signup protection) - from onyx.auth.captcha import CaptchaAction - from onyx.auth.captcha import CaptchaVerificationError - from onyx.auth.captcha import is_captcha_enabled - from onyx.auth.captcha import verify_captcha_token + from onyx.auth.captcha import ( + CaptchaAction, + CaptchaVerificationError, + is_captcha_enabled, + verify_captcha_token, + ) if is_captcha_enabled() and request is not None: # Get captcha token from request body or headers @@ -1395,8 +1397,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): ) if not verified: # Raise some HTTPException (or your custom exception) if old password is invalid: - from fastapi import HTTPException - from fastapi import status + from fastapi import HTTPException, status raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, diff --git a/backend/onyx/auth/utils.py b/backend/onyx/auth/utils.py index dba698c345..d730c3da52 100644 --- a/backend/onyx/auth/utils.py +++ b/backend/onyx/auth/utils.py @@ -5,12 +5,14 @@ from urllib.parse import unquote from fastapi import Request -from onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME -from onyx.auth.constants import API_KEY_HEADER_NAME -from onyx.auth.constants import API_KEY_PREFIX -from onyx.auth.constants import BEARER_PREFIX -from onyx.auth.constants import DEPRECATED_API_KEY_PREFIX -from onyx.auth.constants import PAT_PREFIX +from onyx.auth.constants import ( + API_KEY_HEADER_ALTERNATIVE_NAME, + API_KEY_HEADER_NAME, + API_KEY_PREFIX, + BEARER_PREFIX, + DEPRECATED_API_KEY_PREFIX, + PAT_PREFIX, +) def get_hashed_bearer_token_from_request( diff --git a/backend/onyx/background/celery/apps/app_base.py b/backend/onyx/background/celery/apps/app_base.py index c2a7399ac5..d81e934a12 100644 --- a/backend/onyx/background/celery/apps/app_base.py +++ b/backend/onyx/background/celery/apps/app_base.py @@ -3,16 +3,15 @@ import multiprocessing import os import sys import time -from typing import Any -from typing import cast +from typing import Any, cast -from celery import bootsteps # ty: ignore[unresolved-import] -from celery import Task +from celery import ( + bootsteps, # ty: ignore[unresolved-import] + Task, +) from celery.app import trace # ty: ignore[unresolved-import] from celery.exceptions import WorkerShutdown -from celery.signals import before_task_publish -from celery.signals import task_postrun -from celery.signals import task_prerun +from celery.signals import before_task_publish, task_postrun, task_prerun from celery.states import READY_STATES from celery.utils.log import get_task_logger from celery.worker import strategy # ty: ignore[unresolved-import] @@ -22,18 +21,25 @@ from sentry_sdk.integrations.celery import CeleryIntegration from sqlalchemy import text from sqlalchemy.orm import Session -from onyx.background.celery.apps.task_formatters import CeleryTaskColoredFormatter -from onyx.background.celery.apps.task_formatters import CeleryTaskJsonFormatter -from onyx.background.celery.apps.task_formatters import CeleryTaskPlainFormatter -from onyx.background.celery.celery_utils import celery_is_worker_primary -from onyx.background.celery.celery_utils import make_probe_path -from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_PREFIX -from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_TASKSET_KEY -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX -from onyx.configs.app_configs import ONYX_DISABLE_VESPA -from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX -from onyx.configs.constants import OnyxRedisLocks +from onyx.background.celery.apps.task_formatters import ( + CeleryTaskColoredFormatter, + CeleryTaskJsonFormatter, + CeleryTaskPlainFormatter, +) +from onyx.background.celery.celery_utils import ( + celery_is_worker_primary, + make_probe_path, +) +from onyx.background.celery.tasks.vespa.document_sync import ( + DOCUMENT_SYNC_PREFIX, + DOCUMENT_SYNC_TASKSET_KEY, +) +from onyx.configs.app_configs import ( + DISABLE_VECTOR_DB, + ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + ONYX_DISABLE_VESPA, +) +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 @@ -47,19 +53,23 @@ from onyx.redis.redis_document_set import RedisDocumentSet from onyx.redis.redis_pool import get_redis_client from onyx.redis.redis_usergroup import RedisUserGroup from onyx.tracing.setup import setup_tracing -from onyx.utils.logger import ColoredFormatter -from onyx.utils.logger import get_json_formatter -from onyx.utils.logger import get_log_level_from_str -from onyx.utils.logger import LoggerContextVars -from onyx.utils.logger import PlainFormatter -from onyx.utils.logger import setup_logger -from shared_configs.configs import DEV_LOGGING_ENABLED -from shared_configs.configs import JSON_LOGGING -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import SENTRY_CELERY_TRACES_SAMPLE_RATE -from shared_configs.configs import SENTRY_DSN -from shared_configs.configs import TENANT_ID_PREFIX +from onyx.utils.logger import ( + ColoredFormatter, + get_json_formatter, + get_log_level_from_str, + LoggerContextVars, + PlainFormatter, + setup_logger, +) +from shared_configs.configs import ( + DEV_LOGGING_ENABLED, + JSON_LOGGING, + MULTI_TENANT, + POSTGRES_DEFAULT_SCHEMA, + SENTRY_CELERY_TRACES_SAMPLE_RATE, + SENTRY_DSN, + TENANT_ID_PREFIX, +) from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR logger = setup_logger() diff --git a/backend/onyx/background/celery/apps/beat.py b/backend/onyx/background/celery/apps/beat.py index 8b858974db..8cd60fe38f 100644 --- a/backend/onyx/background/celery/apps/beat.py +++ b/backend/onyx/background/celery/apps/beat.py @@ -1,8 +1,7 @@ from datetime import timedelta from typing import Any -from celery import Celery -from celery import signals +from celery import Celery, signals from celery.beat import PersistentScheduler # ty: ignore[unresolved-import] from celery.signals import beat_init from celery.utils.log import get_task_logger @@ -15,8 +14,7 @@ from onyx.db.engine.sql_engine import SqlEngine from onyx.db.engine.tenant_utils import get_all_tenant_ids from onyx.server.runtime.onyx_runtime import OnyxRuntime from onyx.utils.variable_functionality import fetch_versioned_implementation -from shared_configs.configs import IGNORED_SYNCING_TENANT_LIST -from shared_configs.configs import MULTI_TENANT +from shared_configs.configs import IGNORED_SYNCING_TENANT_LIST, MULTI_TENANT task_logger = get_task_logger(__name__) diff --git a/backend/onyx/background/celery/apps/docfetching.py b/backend/onyx/background/celery/apps/docfetching.py index c28f36a953..148c4c23f9 100644 --- a/backend/onyx/background/celery/apps/docfetching.py +++ b/backend/onyx/background/celery/apps/docfetching.py @@ -1,15 +1,14 @@ -from typing import Any -from typing import cast +from typing import Any, cast -from celery import Celery -from celery import signals -from celery import Task +from celery import Celery, signals, Task from celery.apps.worker import Worker -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown -from celery.signals import worker_shutting_down +from celery.signals import ( + celeryd_init, + worker_init, + worker_ready, + worker_shutdown, + worker_shutting_down, +) import onyx.background.celery.apps.app_base as app_base from onyx.background.celery.tasks.docfetching.worker_shutdown import ( @@ -17,13 +16,17 @@ from onyx.background.celery.tasks.docfetching.worker_shutdown import ( ) from onyx.configs.constants import POSTGRES_CELERY_WORKER_DOCFETCHING_APP_NAME from onyx.db.engine.sql_engine import SqlEngine -from onyx.server.metrics.celery_task_metrics import on_celery_task_postrun -from onyx.server.metrics.celery_task_metrics import on_celery_task_prerun -from onyx.server.metrics.celery_task_metrics import on_celery_task_rejected -from onyx.server.metrics.celery_task_metrics import on_celery_task_retry -from onyx.server.metrics.celery_task_metrics import on_celery_task_revoked -from onyx.server.metrics.indexing_task_metrics import on_indexing_task_postrun -from onyx.server.metrics.indexing_task_metrics import on_indexing_task_prerun +from onyx.server.metrics.celery_task_metrics import ( + on_celery_task_postrun, + on_celery_task_prerun, + on_celery_task_rejected, + on_celery_task_retry, + on_celery_task_revoked, +) +from onyx.server.metrics.indexing_task_metrics import ( + on_indexing_task_postrun, + on_indexing_task_prerun, +) from onyx.server.metrics.metrics_server import start_metrics_server from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/background/celery/apps/docprocessing.py b/backend/onyx/background/celery/apps/docprocessing.py index c41aae4991..1a032d7c1e 100644 --- a/backend/onyx/background/celery/apps/docprocessing.py +++ b/backend/onyx/background/celery/apps/docprocessing.py @@ -1,32 +1,33 @@ -from typing import Any -from typing import cast +from typing import Any, cast -from celery import Celery -from celery import signals -from celery import Task +from celery import Celery, signals, Task from celery.apps.worker import Worker -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_process_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown +from celery.signals import ( + celeryd_init, + worker_init, + worker_process_init, + worker_ready, + worker_shutdown, +) import onyx.background.celery.apps.app_base as app_base from onyx.background.celery.tasks.docprocessing.batch_counters import ( on_docprocessing_task_postrun, -) -from onyx.background.celery.tasks.docprocessing.batch_counters import ( on_docprocessing_task_prerun, ) from onyx.configs.constants import POSTGRES_CELERY_WORKER_DOCPROCESSING_APP_NAME from onyx.db.engine.sql_engine import SqlEngine -from onyx.server.metrics.celery_task_metrics import on_celery_task_postrun -from onyx.server.metrics.celery_task_metrics import on_celery_task_prerun -from onyx.server.metrics.celery_task_metrics import on_celery_task_rejected -from onyx.server.metrics.celery_task_metrics import on_celery_task_retry -from onyx.server.metrics.celery_task_metrics import on_celery_task_revoked -from onyx.server.metrics.indexing_task_metrics import on_indexing_task_postrun -from onyx.server.metrics.indexing_task_metrics import on_indexing_task_prerun +from onyx.server.metrics.celery_task_metrics import ( + on_celery_task_postrun, + on_celery_task_prerun, + on_celery_task_rejected, + on_celery_task_retry, + on_celery_task_revoked, +) +from onyx.server.metrics.indexing_task_metrics import ( + on_indexing_task_postrun, + on_indexing_task_prerun, +) from onyx.server.metrics.metrics_server import start_metrics_server from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/background/celery/apps/heavy.py b/backend/onyx/background/celery/apps/heavy.py index d86c369b93..ba3011b3e6 100644 --- a/backend/onyx/background/celery/apps/heavy.py +++ b/backend/onyx/background/celery/apps/heavy.py @@ -1,23 +1,19 @@ -from typing import Any -from typing import cast +from typing import Any, cast -from celery import Celery -from celery import signals -from celery import Task +from celery import Celery, signals, Task from celery.apps.worker import Worker -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown +from celery.signals import celeryd_init, worker_init, worker_ready, worker_shutdown import onyx.background.celery.apps.app_base as app_base from onyx.configs.constants import POSTGRES_CELERY_WORKER_HEAVY_APP_NAME from onyx.db.engine.sql_engine import SqlEngine -from onyx.server.metrics.celery_task_metrics import on_celery_task_postrun -from onyx.server.metrics.celery_task_metrics import on_celery_task_prerun -from onyx.server.metrics.celery_task_metrics import on_celery_task_rejected -from onyx.server.metrics.celery_task_metrics import on_celery_task_retry -from onyx.server.metrics.celery_task_metrics import on_celery_task_revoked +from onyx.server.metrics.celery_task_metrics import ( + on_celery_task_postrun, + on_celery_task_prerun, + on_celery_task_rejected, + on_celery_task_retry, + on_celery_task_revoked, +) from onyx.server.metrics.metrics_server import start_metrics_server from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/background/celery/apps/light.py b/backend/onyx/background/celery/apps/light.py index 9bd040cc89..5c08167e5f 100644 --- a/backend/onyx/background/celery/apps/light.py +++ b/backend/onyx/background/celery/apps/light.py @@ -1,26 +1,25 @@ from typing import Any -from celery import Celery -from celery import signals -from celery import Task +from celery import Celery, signals, Task from celery.apps.worker import Worker -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown +from celery.signals import celeryd_init, worker_init, worker_ready, worker_shutdown import onyx.background.celery.apps.app_base as app_base from onyx.background.celery.celery_utils import httpx_init_vespa_pool -from onyx.configs.app_configs import MANAGED_VESPA -from onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH -from onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH +from onyx.configs.app_configs import ( + MANAGED_VESPA, + VESPA_CLOUD_CERT_PATH, + VESPA_CLOUD_KEY_PATH, +) from onyx.configs.constants import POSTGRES_CELERY_WORKER_LIGHT_APP_NAME from onyx.db.engine.sql_engine import SqlEngine -from onyx.server.metrics.celery_task_metrics import on_celery_task_postrun -from onyx.server.metrics.celery_task_metrics import on_celery_task_prerun -from onyx.server.metrics.celery_task_metrics import on_celery_task_rejected -from onyx.server.metrics.celery_task_metrics import on_celery_task_retry -from onyx.server.metrics.celery_task_metrics import on_celery_task_revoked +from onyx.server.metrics.celery_task_metrics import ( + on_celery_task_postrun, + on_celery_task_prerun, + on_celery_task_rejected, + on_celery_task_retry, + on_celery_task_revoked, +) from onyx.server.metrics.metrics_server import start_metrics_server from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/background/celery/apps/monitoring.py b/backend/onyx/background/celery/apps/monitoring.py index ae852fdd3d..cef01ae9b6 100644 --- a/backend/onyx/background/celery/apps/monitoring.py +++ b/backend/onyx/background/celery/apps/monitoring.py @@ -1,13 +1,8 @@ import multiprocessing from typing import Any -from celery import Celery -from celery import signals -from celery import Task -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown +from celery import Celery, signals, Task +from celery.signals import celeryd_init, worker_init, worker_ready, worker_shutdown import onyx.background.celery.apps.app_base as app_base from onyx.configs.constants import POSTGRES_CELERY_WORKER_MONITORING_APP_NAME diff --git a/backend/onyx/background/celery/apps/primary.py b/backend/onyx/background/celery/apps/primary.py index 1bf50e896d..460903020d 100644 --- a/backend/onyx/background/celery/apps/primary.py +++ b/backend/onyx/background/celery/apps/primary.py @@ -1,19 +1,17 @@ import logging import os -from typing import Any -from typing import cast +from typing import Any, cast -from celery import bootsteps # ty: ignore[unresolved-import] -from celery import Celery -from celery import signals -from celery import Task +from celery import ( + bootsteps, # ty: ignore[unresolved-import] + Celery, + signals, + Task, +) from celery.apps.worker import Worker from celery.exceptions import WorkerShutdown from celery.result import AsyncResult -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown +from celery.signals import celeryd_init, worker_init, worker_ready, worker_shutdown from redis.lock import Lock as RedisLock import onyx.background.celery.apps.app_base as app_base @@ -21,14 +19,14 @@ from onyx.background.celery.apps.app_base import task_logger from onyx.background.celery.celery_utils import celery_is_worker_primary from onyx.background.celery.tasks.vespa.document_sync import reset_document_sync from onyx.configs.app_configs import CELERY_WORKER_PRIMARY_POOL_OVERFLOW -from onyx.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT -from onyx.configs.constants import OnyxRedisConstants -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine -from onyx.db.index_attempt import get_index_attempt -from onyx.db.index_attempt import mark_attempt_canceled +from onyx.configs.constants import ( + CELERY_PRIMARY_WORKER_LOCK_TIMEOUT, + OnyxRedisConstants, + OnyxRedisLocks, + POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME, +) +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine +from onyx.db.index_attempt import get_index_attempt, mark_attempt_canceled from onyx.db.indexing_coordination import IndexingCoordination from onyx.redis.redis_connector_delete import RedisConnectorDelete from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync @@ -38,15 +36,16 @@ from onyx.redis.redis_connector_stop import RedisConnectorStop from onyx.redis.redis_document_set import RedisDocumentSet from onyx.redis.redis_pool import get_redis_client from onyx.redis.redis_usergroup import RedisUserGroup -from onyx.server.metrics.celery_task_metrics import on_celery_task_postrun -from onyx.server.metrics.celery_task_metrics import on_celery_task_prerun -from onyx.server.metrics.celery_task_metrics import on_celery_task_rejected -from onyx.server.metrics.celery_task_metrics import on_celery_task_retry -from onyx.server.metrics.celery_task_metrics import on_celery_task_revoked +from onyx.server.metrics.celery_task_metrics import ( + on_celery_task_postrun, + on_celery_task_prerun, + on_celery_task_rejected, + on_celery_task_retry, + on_celery_task_revoked, +) from onyx.server.metrics.metrics_server import start_metrics_server from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA logger = setup_logger() diff --git a/backend/onyx/background/celery/apps/scheduled_tasks.py b/backend/onyx/background/celery/apps/scheduled_tasks.py index 2da3706518..c32fb7a012 100644 --- a/backend/onyx/background/celery/apps/scheduled_tasks.py +++ b/backend/onyx/background/celery/apps/scheduled_tasks.py @@ -1,23 +1,19 @@ -from typing import Any -from typing import cast +from typing import Any, cast -from celery import Celery -from celery import signals -from celery import Task +from celery import Celery, signals, Task from celery.apps.worker import Worker -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown +from celery.signals import celeryd_init, worker_init, worker_ready, worker_shutdown import onyx.background.celery.apps.app_base as app_base from onyx.configs.constants import POSTGRES_CELERY_WORKER_SCHEDULED_TASKS_APP_NAME from onyx.db.engine.sql_engine import SqlEngine -from onyx.server.metrics.celery_task_metrics import on_celery_task_postrun -from onyx.server.metrics.celery_task_metrics import on_celery_task_prerun -from onyx.server.metrics.celery_task_metrics import on_celery_task_rejected -from onyx.server.metrics.celery_task_metrics import on_celery_task_retry -from onyx.server.metrics.celery_task_metrics import on_celery_task_revoked +from onyx.server.metrics.celery_task_metrics import ( + on_celery_task_postrun, + on_celery_task_prerun, + on_celery_task_rejected, + on_celery_task_retry, + on_celery_task_revoked, +) from onyx.server.metrics.metrics_server import start_metrics_server from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/background/celery/apps/task_formatters.py b/backend/onyx/background/celery/apps/task_formatters.py index f303e3dd57..a013205a76 100644 --- a/backend/onyx/background/celery/apps/task_formatters.py +++ b/backend/onyx/background/celery/apps/task_formatters.py @@ -2,9 +2,7 @@ import logging from celery import current_task -from onyx.utils.logger import ColoredFormatter -from onyx.utils.logger import get_json_formatter -from onyx.utils.logger import PlainFormatter +from onyx.utils.logger import ColoredFormatter, get_json_formatter, PlainFormatter class CeleryTaskJsonFormatter(logging.Formatter): diff --git a/backend/onyx/background/celery/apps/user_file_processing.py b/backend/onyx/background/celery/apps/user_file_processing.py index 5157574707..61a473b1c8 100644 --- a/backend/onyx/background/celery/apps/user_file_processing.py +++ b/backend/onyx/background/celery/apps/user_file_processing.py @@ -1,15 +1,14 @@ -from typing import Any -from typing import cast +from typing import Any, cast -from celery import Celery -from celery import signals -from celery import Task +from celery import Celery, signals, Task from celery.apps.worker import Worker -from celery.signals import celeryd_init -from celery.signals import worker_init -from celery.signals import worker_process_init -from celery.signals import worker_ready -from celery.signals import worker_shutdown +from celery.signals import ( + celeryd_init, + worker_init, + worker_process_init, + worker_ready, + worker_shutdown, +) import onyx.background.celery.apps.app_base as app_base from onyx.configs.constants import POSTGRES_CELERY_WORKER_USER_FILE_PROCESSING_APP_NAME diff --git a/backend/onyx/background/celery/celery_redis.py b/backend/onyx/background/celery/celery_redis.py index 0ccab40415..480e960a7d 100644 --- a/backend/onyx/background/celery/celery_redis.py +++ b/backend/onyx/background/celery/celery_redis.py @@ -1,16 +1,14 @@ # These are helper objects for tracking the keys we need to write in redis import json import threading -from typing import Any -from typing import cast +from typing import Any, cast from celery import Celery from redis import Redis from onyx.background.celery.configs.base import CELERY_SEPARATOR from onyx.configs.app_configs import REDIS_HEALTH_CHECK_INTERVAL -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS +from onyx.configs.constants import OnyxCeleryPriority, REDIS_SOCKET_KEEPALIVE_OPTIONS _broker_client: Redis | None = None _broker_url: str | None = None diff --git a/backend/onyx/background/celery/celery_utils.py b/backend/onyx/background/celery/celery_utils.py index f07780c720..bbbeca9215 100644 --- a/backend/onyx/background/celery/celery_utils.py +++ b/backend/onyx/background/celery/celery_utils.py @@ -1,38 +1,43 @@ import time -from collections.abc import Generator -from collections.abc import Iterator -from collections.abc import Sequence -from datetime import datetime -from datetime import timezone +from collections.abc import Generator, Iterator, Sequence +from datetime import datetime, timezone from pathlib import Path -from typing import Any -from typing import cast -from typing import TypeVar +from typing import Any, cast, TypeVar import httpx from pydantic import BaseModel -from onyx.configs.app_configs import MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE -from onyx.configs.app_configs import VESPA_REQUEST_TIMEOUT +from onyx.configs.app_configs import ( + MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE, + VESPA_REQUEST_TIMEOUT, +) from onyx.connectors.connector_runner import CheckpointOutputWrapper from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.interfaces import BaseConnector -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import ConnectorCheckpoint -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.file_store.staging import build_tracking_raw_file_callback -from onyx.file_store.staging import delete_files_best_effort +from onyx.connectors.interfaces import ( + BaseConnector, + CheckpointedConnector, + ConnectorCheckpoint, + LoadConnector, + PollConnector, + SlimConnector, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + ConnectorFailure, + Document, + HierarchyNode, + SlimDocument, +) +from onyx.file_store.staging import ( + build_tracking_raw_file_callback, + delete_files_best_effort, +) from onyx.httpx.httpx_pool import HttpxPool from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface -from onyx.server.metrics.pruning_metrics import inc_pruning_rate_limit_error -from onyx.server.metrics.pruning_metrics import observe_pruning_enumeration_duration +from onyx.server.metrics.pruning_metrics import ( + inc_pruning_rate_limit_error, + observe_pruning_enumeration_duration, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/background/celery/configs/base.py b/backend/onyx/background/celery/configs/base.py index 7690f76617..70cb1cedd8 100644 --- a/backend/onyx/background/celery/configs/base.py +++ b/backend/onyx/background/celery/configs/base.py @@ -1,25 +1,26 @@ # docs: https://docs.celeryq.dev/en/stable/userguide/configuration.html import urllib.parse -from onyx.configs.app_configs import CELERY_BROKER_POOL_LIMIT -from onyx.configs.app_configs import CELERY_RESULT_EXPIRES -from onyx.configs.app_configs import REDIS_DB_NUMBER_CELERY -from onyx.configs.app_configs import REDIS_DB_NUMBER_CELERY_RESULT_BACKEND -from onyx.configs.app_configs import REDIS_HEALTH_CHECK_INTERVAL -from onyx.configs.app_configs import REDIS_HOST -from onyx.configs.app_configs import REDIS_PASSWORD -from onyx.configs.app_configs import REDIS_PORT -from onyx.configs.app_configs import REDIS_SENTINEL_HOSTS -from onyx.configs.app_configs import REDIS_SENTINEL_MASTER_NAME -from onyx.configs.app_configs import REDIS_SENTINEL_PASSWORD -from onyx.configs.app_configs import REDIS_SSL -from onyx.configs.app_configs import REDIS_SSL_CA_CERTS -from onyx.configs.app_configs import REDIS_SSL_CERT_REQS -from onyx.configs.app_configs import REDIS_SSL_CERTFILE -from onyx.configs.app_configs import REDIS_SSL_KEYFILE -from onyx.configs.app_configs import USE_REDIS_IAM_AUTH -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS +from onyx.configs.app_configs import ( + CELERY_BROKER_POOL_LIMIT, + CELERY_RESULT_EXPIRES, + REDIS_DB_NUMBER_CELERY, + REDIS_DB_NUMBER_CELERY_RESULT_BACKEND, + REDIS_HEALTH_CHECK_INTERVAL, + REDIS_HOST, + REDIS_PASSWORD, + REDIS_PORT, + REDIS_SENTINEL_HOSTS, + REDIS_SENTINEL_MASTER_NAME, + REDIS_SENTINEL_PASSWORD, + REDIS_SSL, + REDIS_SSL_CA_CERTS, + REDIS_SSL_CERT_REQS, + REDIS_SSL_CERTFILE, + REDIS_SSL_KEYFILE, + USE_REDIS_IAM_AUTH, +) +from onyx.configs.constants import OnyxCeleryPriority, REDIS_SOCKET_KEEPALIVE_OPTIONS CELERY_SEPARATOR = ":" diff --git a/backend/onyx/background/celery/configs/light.py b/backend/onyx/background/celery/configs/light.py index 5eaa091796..9a3f8e34b1 100644 --- a/backend/onyx/background/celery/configs/light.py +++ b/backend/onyx/background/celery/configs/light.py @@ -1,6 +1,8 @@ import onyx.background.celery.configs.base as shared_config -from onyx.configs.app_configs import CELERY_WORKER_LIGHT_CONCURRENCY -from onyx.configs.app_configs import CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER +from onyx.configs.app_configs import ( + CELERY_WORKER_LIGHT_CONCURRENCY, + CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER, +) broker_url = shared_config.broker_url broker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup diff --git a/backend/onyx/background/celery/memory_monitoring.py b/backend/onyx/background/celery/memory_monitoring.py index 046bca4f27..6a33e68ffa 100644 --- a/backend/onyx/background/celery/memory_monitoring.py +++ b/backend/onyx/background/celery/memory_monitoring.py @@ -7,8 +7,10 @@ from logging.handlers import RotatingFileHandler import psutil -from onyx.configs.app_configs import INDEXING_WORKER_MEMORY_LIMIT_MB -from onyx.configs.app_configs import INDEXING_WORKER_TRACEMALLOC +from onyx.configs.app_configs import ( + INDEXING_WORKER_MEMORY_LIMIT_MB, + INDEXING_WORKER_TRACEMALLOC, +) from onyx.utils.logger import setup_logger from onyx.utils.platform_utils import is_running_in_container diff --git a/backend/onyx/background/celery/tasks/beat_schedule.py b/backend/onyx/background/celery/tasks/beat_schedule.py index ce389a32c9..042e64d3c7 100644 --- a/backend/onyx/background/celery/tasks/beat_schedule.py +++ b/backend/onyx/background/celery/tasks/beat_schedule.py @@ -4,18 +4,22 @@ from typing import Any from celery.schedules import crontab -from onyx.configs.app_configs import AUTO_LLM_CONFIG_URL -from onyx.configs.app_configs import AUTO_LLM_UPDATE_INTERVAL_SECONDS -from onyx.configs.app_configs import DISABLE_OPENSEARCH_MIGRATION_TASK -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX -from onyx.configs.app_configs import ENTERPRISE_EDITION_ENABLED -from onyx.configs.app_configs import ONYX_DISABLE_VESPA -from onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES -from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.app_configs import ( + AUTO_LLM_CONFIG_URL, + AUTO_LLM_UPDATE_INTERVAL_SECONDS, + DISABLE_OPENSEARCH_MIGRATION_TASK, + DISABLE_VECTOR_DB, + ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + ENTERPRISE_EDITION_ENABLED, + ONYX_DISABLE_VESPA, + SCHEDULED_EVAL_DATASET_NAMES, +) +from onyx.configs.constants import ( + ONYX_CLOUD_CELERY_TASK_PREFIX, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, +) from onyx.server.features.build.configs import SANDBOX_IDLE_CLEANUP_INTERVAL_SECONDS from onyx.utils.variable_functionality import _LICENSE_ENFORCEMENT_ENABLED from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/background/celery/tasks/build/tasks.py b/backend/onyx/background/celery/tasks/build/tasks.py index 0042f5ec73..a7089ed982 100644 --- a/backend/onyx/background/celery/tasks/build/tasks.py +++ b/backend/onyx/background/celery/tasks/build/tasks.py @@ -3,27 +3,27 @@ import datetime import time -from celery import shared_task -from celery import Task +from celery import shared_task, Task from redis.lock import Lock as RedisLock from onyx.background.celery.apps.app_base import task_logger -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks +from onyx.configs.constants import OnyxCeleryTask, OnyxRedisLocks from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import Sandbox from onyx.redis.redis_pool import get_redis_client from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active from onyx.server.features.build.configs import SANDBOX_IDLE_TIMEOUT_SECONDS -from onyx.server.features.build.db.sandbox import get_latest_snapshot_for_session -from onyx.server.features.build.db.sandbox import get_running_sandboxes -from onyx.server.features.build.db.sandbox import user_has_stale_active_session +from onyx.server.features.build.db.sandbox import ( + get_latest_snapshot_for_session, + get_running_sandboxes, + user_has_stale_active_session, +) from onyx.server.features.build.sandbox.factory import get_sandbox_manager from onyx.server.features.build.session.sandbox_lifecycle import ( create_session_snapshot_keep_latest, + is_sandbox_idle, + sleep_sandbox, ) -from onyx.server.features.build.session.sandbox_lifecycle import is_sandbox_idle -from onyx.server.features.build.session.sandbox_lifecycle import sleep_sandbox # 100 minutes - snapshotting can take time TIMEOUT_SECONDS = 6000 diff --git a/backend/onyx/background/celery/tasks/connector_deletion/tasks.py b/backend/onyx/background/celery/tasks/connector_deletion/tasks.py index 9b2680fb14..eb4e51e5bf 100644 --- a/backend/onyx/background/celery/tasks/connector_deletion/tasks.py +++ b/backend/onyx/background/celery/tasks/connector_deletion/tasks.py @@ -1,12 +1,8 @@ import traceback -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timezone +from typing import Any, cast -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from pydantic import ValidationError from redis import Redis @@ -14,64 +10,71 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_queue_length -from onyx.background.celery.celery_redis import celery_get_queued_task_ids +from onyx.background.celery.celery_redis import ( + celery_get_broker_client, + celery_get_queue_length, + celery_get_queued_task_ids, +) from onyx.configs.app_configs import JOB_TIMEOUT -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import OnyxRedisSignals +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, + OnyxRedisLocks, + OnyxRedisSignals, +) from onyx.db.connector import fetch_connector_by_id -from onyx.db.connector_credential_pair import add_deletion_failure_message from onyx.db.connector_credential_pair import ( + add_deletion_failure_message, delete_connector_credential_pair__no_commit, + get_connector_credential_pair_from_id, + get_connector_credential_pairs, ) -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.connector_credential_pair import get_connector_credential_pairs from onyx.db.document import ( delete_all_documents_by_connector_credential_pair__no_commit, + get_document_ids_for_connector_credential_pair, ) -from onyx.db.document import get_document_ids_for_connector_credential_pair from onyx.db.document_set import delete_document_set_cc_pair_relationship__no_commit from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingStatus -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType -from onyx.db.index_attempt import delete_index_attempts -from onyx.db.index_attempt import get_recent_attempts_for_cc_pair +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + IndexingStatus, + SyncStatus, + SyncType, +) +from onyx.db.index_attempt import delete_index_attempts, get_recent_attempts_for_cc_pair from onyx.db.permission_sync_attempt import ( delete_doc_permission_sync_attempts__no_commit, -) -from onyx.db.permission_sync_attempt import ( delete_external_group_permission_sync_attempts__no_commit, ) -from onyx.db.port_attempt import get_active_port_attempt -from onyx.db.port_attempt import request_port_cancel +from onyx.db.port_attempt import get_active_port_attempt, request_port_cancel from onyx.db.search_settings import get_all_search_settings -from onyx.db.sync_record import cleanup_sync_records -from onyx.db.sync_record import insert_sync_record -from onyx.db.sync_record import update_sync_record_status +from onyx.db.sync_record import ( + cleanup_sync_records, + insert_sync_record, + update_sync_record_status, +) from onyx.db.tag import delete_orphan_tags__no_commit from onyx.redis.redis_connector import RedisConnector -from onyx.redis.redis_connector_delete import RedisConnectorDelete -from onyx.redis.redis_connector_delete import RedisConnectorDeletePayload -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client +from onyx.redis.redis_connector_delete import ( + RedisConnectorDelete, + RedisConnectorDeletePayload, +) +from onyx.redis.redis_pool import get_redis_client, get_redis_replica_client from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active from onyx.redis.tenant_redis_client import TenantRedisClient -from onyx.server.metrics.deletion_metrics import inc_deletion_blocked -from onyx.server.metrics.deletion_metrics import inc_deletion_completed -from onyx.server.metrics.deletion_metrics import inc_deletion_fence_reset -from onyx.server.metrics.deletion_metrics import inc_deletion_started -from onyx.server.metrics.deletion_metrics import observe_deletion_taskset_duration +from onyx.server.metrics.deletion_metrics import ( + inc_deletion_blocked, + inc_deletion_completed, + inc_deletion_fence_reset, + inc_deletion_started, + observe_deletion_taskset_duration, +) from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, + noop_fallback, ) -from onyx.utils.variable_functionality import noop_fallback class TaskDependencyError(RuntimeError): diff --git a/backend/onyx/background/celery/tasks/docfetching/task_creation_utils.py b/backend/onyx/background/celery/tasks/docfetching/task_creation_utils.py index 5626442299..e248408db8 100644 --- a/backend/onyx/background/celery/tasks/docfetching/task_creation_utils.py +++ b/backend/onyx/background/celery/tasks/docfetching/task_creation_utils.py @@ -5,15 +5,16 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.background.celery.apps.app_base import task_logger -from onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import ( + DANSWER_REDIS_FUNCTION_LOCK_PREFIX, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, +) from onyx.db.enums import ConnectorCredentialPairStatus from onyx.db.index_attempt import mark_attempt_failed from onyx.db.indexing_coordination import IndexingCoordination -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import SearchSettings +from onyx.db.models import ConnectorCredentialPair, SearchSettings from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/onyx/background/celery/tasks/docfetching/tasks.py b/backend/onyx/background/celery/tasks/docfetching/tasks.py index ceb30e3278..cd6e6319c9 100644 --- a/backend/onyx/background/celery/tasks/docfetching/tasks.py +++ b/backend/onyx/background/celery/tasks/docfetching/tasks.py @@ -5,47 +5,56 @@ import traceback from time import sleep import psutil -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.memory_monitoring import emit_process_memory -from onyx.background.celery.memory_monitoring import start_memory_observer -from onyx.background.celery.memory_monitoring import stop_memory_observer +from onyx.background.celery.memory_monitoring import ( + emit_process_memory, + start_memory_observer, + stop_memory_observer, +) from onyx.background.celery.tasks.docfetching.worker_shutdown import ( is_worker_shutting_down, ) -from onyx.background.celery.tasks.docprocessing.heartbeat import start_heartbeat -from onyx.background.celery.tasks.docprocessing.heartbeat import stop_heartbeat +from onyx.background.celery.tasks.docprocessing.heartbeat import ( + start_heartbeat, + stop_heartbeat, +) from onyx.background.celery.tasks.docprocessing.tasks import ConnectorIndexingLogBuilder from onyx.background.celery.tasks.docprocessing.utils import IndexingCallback -from onyx.background.celery.tasks.models import DocProcessingContext -from onyx.background.celery.tasks.models import IndexingWatchdogTerminalStatus -from onyx.background.celery.tasks.models import SimpleJobResult -from onyx.background.indexing.job_client import SimpleJob -from onyx.background.indexing.job_client import SimpleJobClient -from onyx.background.indexing.job_client import SimpleJobException +from onyx.background.celery.tasks.models import ( + DocProcessingContext, + IndexingWatchdogTerminalStatus, + SimpleJobResult, +) +from onyx.background.indexing.job_client import ( + SimpleJob, + SimpleJobClient, + SimpleJobException, +) from onyx.background.indexing.run_docfetching import run_docfetching_entrypoint from onyx.configs.app_configs import INDEXING_WORKER_MEMORY_LIMIT_MB -from onyx.configs.constants import CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT -from onyx.configs.constants import CELERY_INDEXING_WATCHDOG_SIGTERM_GRACE_SECONDS -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import ( + CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT, + CELERY_INDEXING_WATCHDOG_SIGTERM_GRACE_SECONDS, + OnyxCeleryTask, +) from onyx.connectors.exceptions import ConnectorValidationError from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import IndexingStatus -from onyx.db.index_attempt import get_index_attempt -from onyx.db.index_attempt import mark_attempt_canceled -from onyx.db.index_attempt import mark_attempt_failed -from onyx.db.index_attempt import mark_attempt_interrupted +from onyx.db.index_attempt import ( + get_index_attempt, + mark_attempt_canceled, + mark_attempt_failed, + mark_attempt_interrupted, +) from onyx.db.indexing_coordination import IndexingCoordination from onyx.redis.redis_connector import RedisConnector from onyx.server.metrics.connector_health_metrics import on_index_attempt_status_change from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import global_version -from shared_configs.configs import SENTRY_CELERY_TRACES_SAMPLE_RATE -from shared_configs.configs import SENTRY_DSN +from shared_configs.configs import SENTRY_CELERY_TRACES_SAMPLE_RATE, SENTRY_DSN logger = setup_logger() 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 6cd13bdec2..f406f5105f 100644 --- a/backend/onyx/background/celery/tasks/docprocessing/targeted_reindex_task.py +++ b/backend/onyx/background/celery/tasks/docprocessing/targeted_reindex_task.py @@ -10,19 +10,22 @@ connector invocation + indexing pipeline plumbing live in import datetime import logging -from celery import shared_task -from celery import Task +from celery import shared_task, Task from onyx.background.celery.apps.app_base import task_logger -from onyx.background.indexing.run_targeted_reindex import group_targets_by_cc_pair -from onyx.background.indexing.run_targeted_reindex import process_targets_for_cc_pair +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 -from onyx.db.targeted_reindex import get_index_attempts_for_targeted_reindex_job -from onyx.db.targeted_reindex import get_targeted_reindex_job -from onyx.db.targeted_reindex import get_targets_for_job -from onyx.db.targeted_reindex import resolve_failure_derived_targets +from onyx.db.targeted_reindex import ( + get_index_attempts_for_targeted_reindex_job, + get_targeted_reindex_job, + get_targets_for_job, + resolve_failure_derived_targets, +) from shared_configs.contextvars import get_current_tenant_id _TARGETED_REINDEX_SOFT_TIME_LIMIT = 60 * 30 # 30 minutes diff --git a/backend/onyx/background/celery/tasks/docprocessing/tasks.py b/backend/onyx/background/celery/tasks/docprocessing/tasks.py index e8dc197c1b..6c33c543e4 100644 --- a/backend/onyx/background/celery/tasks/docprocessing/tasks.py +++ b/backend/onyx/background/celery/tasks/docprocessing/tasks.py @@ -4,104 +4,118 @@ import time import traceback from collections import defaultdict from dataclasses import dataclass -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any -from celery import Celery -from celery import current_app -from celery import shared_task -from celery import Task +from celery import Celery, current_app, shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from pydantic import BaseModel from redis.lock import Lock as RedisLock -from sqlalchemy import exists -from sqlalchemy import select +from sqlalchemy import exists, select from sqlalchemy.orm import Session from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_queued_task_ids -from onyx.background.celery.celery_redis import celery_get_unacked_task_ids +from onyx.background.celery.celery_redis import ( + celery_get_broker_client, + celery_get_queued_task_ids, + celery_get_unacked_task_ids, +) from onyx.background.celery.celery_utils import httpx_init_vespa_pool from onyx.background.celery.memory_monitoring import emit_process_memory from onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT from onyx.background.celery.tasks.docfetching.task_creation_utils import ( try_creating_docfetching_task, ) -from onyx.background.celery.tasks.docprocessing.heartbeat import start_heartbeat -from onyx.background.celery.tasks.docprocessing.heartbeat import stop_heartbeat +from onyx.background.celery.tasks.docprocessing.heartbeat import ( + start_heartbeat, + stop_heartbeat, +) from onyx.background.celery.tasks.docprocessing.targeted_reindex_task import ( # noqa: F401 # registers @shared_task with celery targeted_reindex_task, ) -from onyx.background.celery.tasks.docprocessing.utils import IndexingCallback -from onyx.background.celery.tasks.docprocessing.utils import is_in_repeated_error_state -from onyx.background.celery.tasks.docprocessing.utils import should_index +from onyx.background.celery.tasks.docprocessing.utils import ( + IndexingCallback, + is_in_repeated_error_state, + should_index, +) from onyx.background.celery.tasks.models import DocProcessingContext -from onyx.background.indexing.checkpointing_utils import cleanup_checkpoint from onyx.background.indexing.checkpointing_utils import ( + cleanup_checkpoint, get_index_attempts_with_old_checkpoints, ) -from onyx.background.indexing.index_attempt_utils import cleanup_index_attempts -from onyx.background.indexing.index_attempt_utils import get_old_index_attempt_ids -from onyx.configs.app_configs import MANAGED_VESPA -from onyx.configs.app_configs import PERSISTENT_INDEXING -from onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH -from onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_INDEXING_LOCK_TIMEOUT -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MilestoneRecordType -from onyx.configs.constants import NotificationType -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import OnyxRedisSignals -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import IndexAttemptMetadata +from onyx.background.indexing.index_attempt_utils import ( + cleanup_index_attempts, + get_old_index_attempt_ids, +) +from onyx.configs.app_configs import ( + MANAGED_VESPA, + PERSISTENT_INDEXING, + VESPA_CLOUD_CERT_PATH, + VESPA_CLOUD_KEY_PATH, +) +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + CELERY_INDEXING_LOCK_TIMEOUT, + DocumentSource, + MilestoneRecordType, + NotificationType, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, + OnyxRedisLocks, + OnyxRedisSignals, +) +from onyx.connectors.models import ConnectorFailure, Document, IndexAttemptMetadata from onyx.db.connector import mark_ccpair_with_indexing_trigger from onyx.db.connector_credential_pair import ( fetch_indexable_standard_connector_credential_pair_ids, + get_connector_credential_pair_from_id, + set_cc_pair_repeated_error_state, + update_connector_credential_pair_from_id, ) -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.connector_credential_pair import set_cc_pair_repeated_error_state -from onyx.db.connector_credential_pair import update_connector_credential_pair_from_id from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.engine.time_utils import get_db_current_time -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingMode -from onyx.db.enums import IndexingStatus -from onyx.db.enums import SwitchoverType -from onyx.db.index_attempt import create_index_attempt_error -from onyx.db.index_attempt import get_index_attempt -from onyx.db.index_attempt import get_index_attempt_errors_for_cc_pair -from onyx.db.index_attempt import get_stale_not_started_index_attempts -from onyx.db.index_attempt import IndexAttemptError -from onyx.db.index_attempt import mark_attempt_canceled -from onyx.db.index_attempt import mark_attempt_failed -from onyx.db.index_attempt import mark_attempt_partially_succeeded -from onyx.db.index_attempt import mark_attempt_succeeded -from onyx.db.index_attempt_metrics import IndexAttemptStage -from onyx.db.index_attempt_metrics import safe_record_single_event -from onyx.db.index_attempt_metrics import time_stage -from onyx.db.indexing_coordination import CoordinationStatus -from onyx.db.indexing_coordination import IndexingCoordination -from onyx.db.models import IndexAttempt -from onyx.db.models import SearchSettings -from onyx.db.notification import batch_create_notifications -from onyx.db.notification import delete_notifications_by_additional_data -from onyx.db.search_settings import get_current_search_settings -from onyx.db.search_settings import get_secondary_search_settings +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + IndexingMode, + IndexingStatus, + SwitchoverType, +) +from onyx.db.index_attempt import ( + create_index_attempt_error, + get_index_attempt, + get_index_attempt_errors_for_cc_pair, + get_stale_not_started_index_attempts, + IndexAttemptError, + mark_attempt_canceled, + mark_attempt_failed, + mark_attempt_partially_succeeded, + mark_attempt_succeeded, +) +from onyx.db.index_attempt_metrics import ( + IndexAttemptStage, + safe_record_single_event, + time_stage, +) +from onyx.db.indexing_coordination import CoordinationStatus, IndexingCoordination +from onyx.db.models import IndexAttempt, SearchSettings +from onyx.db.notification import ( + batch_create_notifications, + delete_notifications_by_additional_data, +) +from onyx.db.search_settings import ( + get_current_search_settings, + get_secondary_search_settings, +) from onyx.db.swap_index import check_and_perform_index_swap from onyx.db.users import get_active_admin_users 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 -from onyx.file_store.document_batch_storage import get_document_batch_storage +from onyx.file_store.document_batch_storage import ( + DocumentBatchStorage, + get_document_batch_storage, +) 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 ( @@ -109,34 +123,44 @@ from onyx.indexing.adapters.document_indexing_adapter import ( ) 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 -from onyx.indexing.persistent_indexing import record_generic_failure -from onyx.natural_language_processing.search_nlp_models import EmbeddingModel -from onyx.natural_language_processing.search_nlp_models import warm_up_bi_encoder +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 get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client -from onyx.redis.redis_pool import redis_lock_dump -from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT +from onyx.redis.redis_pool import ( + get_redis_client, + get_redis_replica_client, + redis_lock_dump, + SCAN_ITER_COUNT_DEFAULT, +) from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active from onyx.redis.redis_utils import is_fence from onyx.redis.tenant_redis_client import TenantRedisClient -from onyx.server.metrics.connector_health_metrics import on_connector_error_state_change -from onyx.server.metrics.connector_health_metrics import on_connector_indexing_success -from onyx.server.metrics.connector_health_metrics import on_index_attempt_status_change +from onyx.server.metrics.connector_health_metrics import ( + on_connector_error_state_change, + on_connector_indexing_success, + on_index_attempt_status_change, +) from onyx.server.runtime.onyx_runtime import OnyxRuntime from onyx.utils.logger import setup_logger from onyx.utils.middleware import make_randomized_onyx_request_id -from onyx.utils.telemetry import mt_cloud_telemetry -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType -from shared_configs.configs import INDEXING_MODEL_SERVER_HOST -from shared_configs.configs import INDEXING_MODEL_SERVER_PORT -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import USAGE_LIMITS_ENABLED -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR +from onyx.utils.telemetry import mt_cloud_telemetry, optional_telemetry, RecordType +from shared_configs.configs import ( + INDEXING_MODEL_SERVER_HOST, + INDEXING_MODEL_SERVER_PORT, + MULTI_TENANT, + USAGE_LIMITS_ENABLED, +) +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + INDEX_ATTEMPT_INFO_CONTEXTVAR, +) # Bound for waiting on the shared cross-batch DB lock (see usage below). CROSS_BATCH_DB_LOCK_ACQUIRE_TIMEOUT_S = 300 @@ -1847,8 +1871,7 @@ def _docprocessing_task( # Track chunk indexing usage for cloud usage limits if USAGE_LIMITS_ENABLED and index_pipeline_result.total_chunks > 0: try: - from onyx.db.usage import increment_usage - from onyx.db.usage import UsageType + from onyx.db.usage import increment_usage, UsageType with get_session_with_current_tenant() as usage_db_session: increment_usage( diff --git a/backend/onyx/background/celery/tasks/docprocessing/utils.py b/backend/onyx/background/celery/tasks/docprocessing/utils.py index 3ebfc4bc17..fb32b7627f 100644 --- a/backend/onyx/background/celery/tasks/docprocessing/utils.py +++ b/backend/onyx/background/celery/tasks/docprocessing/utils.py @@ -1,22 +1,23 @@ import time -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from redis.exceptions import LockError from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import DocumentSource +from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT, DocumentSource from onyx.db.engine.time_utils import get_db_current_time -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.index_attempt import get_last_attempt_for_cc_pair -from onyx.db.index_attempt import get_recent_attempts_for_cc_pair -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import SearchSettings +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + IndexingStatus, + IndexModelStatus, +) +from onyx.db.index_attempt import ( + get_last_attempt_for_cc_pair, + get_recent_attempts_for_cc_pair, +) +from onyx.db.models import ConnectorCredentialPair, SearchSettings from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.redis.redis_connector import RedisConnector from onyx.redis.redis_pool import redis_lock_dump diff --git a/backend/onyx/background/celery/tasks/evals/tasks.py b/backend/onyx/background/celery/tasks/evals/tasks.py index bce3132eb7..9f9045a75e 100644 --- a/backend/onyx/background/celery/tasks/evals/tasks.py +++ b/backend/onyx/background/celery/tasks/evals/tasks.py @@ -1,15 +1,15 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any -from celery import shared_task -from celery import Task +from celery import shared_task, Task -from onyx.configs.app_configs import BRAINTRUST_API_KEY -from onyx.configs.app_configs import JOB_TIMEOUT -from onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES -from onyx.configs.app_configs import SCHEDULED_EVAL_PERMISSIONS_EMAIL -from onyx.configs.app_configs import SCHEDULED_EVAL_PROJECT +from onyx.configs.app_configs import ( + BRAINTRUST_API_KEY, + JOB_TIMEOUT, + SCHEDULED_EVAL_DATASET_NAMES, + SCHEDULED_EVAL_PERMISSIONS_EMAIL, + SCHEDULED_EVAL_PROJECT, +) from onyx.configs.constants import OnyxCeleryTask from onyx.evals.eval import run_eval from onyx.evals.models import EvalConfigurationOptions diff --git a/backend/onyx/background/celery/tasks/hierarchyfetching/tasks.py b/backend/onyx/background/celery/tasks/hierarchyfetching/tasks.py index d9ea664d00..ba92648d6b 100644 --- a/backend/onyx/background/celery/tasks/hierarchyfetching/tasks.py +++ b/backend/onyx/background/celery/tasks/hierarchyfetching/tasks.py @@ -9,45 +9,48 @@ structural information from the connector source. """ import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from uuid import uuid4 -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.background.celery.apps.app_base import task_logger from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks -from onyx.connectors.factory import ConnectorMissingException -from onyx.connectors.factory import identify_connector_class -from onyx.connectors.factory import instantiate_connector +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + DANSWER_REDIS_FUNCTION_LOCK_PREFIX, + DocumentSource, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisLocks, +) +from onyx.connectors.factory import ( + ConnectorMissingException, + identify_connector_class, + instantiate_connector, +) from onyx.connectors.interfaces import HierarchyConnector from onyx.connectors.models import HierarchyNode as PydanticHierarchyNode from onyx.db.connector import mark_cc_pair_as_hierarchy_fetched from onyx.db.connector_credential_pair import ( fetch_indexable_standard_connector_credential_pair_ids, + get_connector_credential_pair_from_id, ) -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries -from onyx.db.hierarchy import upsert_hierarchy_nodes_batch +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus +from onyx.db.hierarchy import ( + upsert_hierarchy_node_cc_pair_entries, + upsert_hierarchy_nodes_batch, +) from onyx.db.models import ConnectorCredentialPair -from onyx.redis.redis_hierarchy import cache_hierarchy_nodes_batch -from onyx.redis.redis_hierarchy import ensure_source_node_exists -from onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry +from onyx.redis.redis_hierarchy import ( + cache_hierarchy_nodes_batch, + ensure_source_node_exists, + HierarchyNodeCacheEntry, +) from onyx.redis.redis_pool import get_redis_client from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.utils.logger import setup_logger diff --git a/backend/onyx/background/celery/tasks/llm_model_update/tasks.py b/backend/onyx/background/celery/tasks/llm_model_update/tasks.py index c334c1d088..1f3e8321b4 100644 --- a/backend/onyx/background/celery/tasks/llm_model_update/tasks.py +++ b/backend/onyx/background/celery/tasks/llm_model_update/tasks.py @@ -1,5 +1,4 @@ -from celery import shared_task -from celery import Task +from celery import shared_task, Task from onyx.background.celery.apps.app_base import task_logger from onyx.configs.app_configs import AUTO_LLM_CONFIG_URL diff --git a/backend/onyx/background/celery/tasks/monitoring/tasks.py b/backend/onyx/background/celery/tasks/monitoring/tasks.py index 6bd581f3cb..a838bcb80f 100644 --- a/backend/onyx/background/celery/tasks/monitoring/tasks.py +++ b/backend/onyx/background/celery/tasks/monitoring/tasks.py @@ -2,53 +2,54 @@ import json import time from datetime import timedelta from itertools import islice -from typing import Any -from typing import cast -from typing import Literal +from typing import Any, cast, Literal import psutil -from celery import shared_task -from celery import Task +from celery import shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from pydantic import BaseModel from redis import Redis from redis.lock import Lock as RedisLock -from sqlalchemy import select -from sqlalchemy import text +from sqlalchemy import select, text from sqlalchemy.orm import Session from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_queue_length -from onyx.background.celery.celery_redis import celery_get_unacked_task_ids +from onyx.background.celery.celery_redis import ( + celery_get_broker_client, + celery_get_queue_length, + celery_get_unacked_task_ids, +) from onyx.background.celery.memory_monitoring import emit_process_memory -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import ONYX_CLOUD_TENANT_ID -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import get_session_with_shared_schema -from onyx.db.engine.tenant_utils import get_all_tenant_ids -from onyx.db.engine.tenant_utils import validate_tenant_id +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + ONYX_CLOUD_TENANT_ID, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisLocks, +) +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_shared_schema, +) +from onyx.db.engine.tenant_utils import get_all_tenant_ids, validate_tenant_id from onyx.db.engine.time_utils import get_db_current_time -from onyx.db.enums import IndexingStatus -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import DocumentSet -from onyx.db.models import IndexAttempt -from onyx.db.models import SyncRecord -from onyx.db.models import UserGroup +from onyx.db.enums import IndexingStatus, SyncStatus, SyncType +from onyx.db.models import ( + ConnectorCredentialPair, + DocumentSet, + IndexAttempt, + SyncRecord, + UserGroup, +) from onyx.db.search_settings import get_active_search_settings_list -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_shared_redis_client -from onyx.redis.redis_pool import redis_lock_dump +from onyx.redis.redis_pool import ( + get_redis_client, + get_shared_redis_client, + redis_lock_dump, +) from onyx.redis.tenant_redis_client import TenantRedisClient -from onyx.utils.platform_utils import is_running_in_container -from onyx.utils.platform_utils import is_running_in_kubernetes -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType +from onyx.utils.platform_utils import is_running_in_container, is_running_in_kubernetes +from onyx.utils.telemetry import optional_telemetry, RecordType from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/onyx/background/celery/tasks/opensearch_migration/tasks.py b/backend/onyx/background/celery/tasks/opensearch_migration/tasks.py index 8b50f33352..60c161d256 100644 --- a/backend/onyx/background/celery/tasks/opensearch_migration/tasks.py +++ b/backend/onyx/background/celery/tasks/opensearch_migration/tasks.py @@ -3,47 +3,35 @@ import time import traceback -from celery import shared_task -from celery import Task +from celery import shared_task, Task from redis.lock import Lock as RedisLock from onyx.background.celery.apps.app_base import task_logger from onyx.background.celery.tasks.opensearch_migration.constants import ( FINISHED_VISITING_SLICE_CONTINUATION_TOKEN, -) -from onyx.background.celery.tasks.opensearch_migration.constants import ( GET_VESPA_CHUNKS_PAGE_SIZE, -) -from onyx.background.celery.tasks.opensearch_migration.constants import ( MIGRATION_TASK_LOCK_BLOCKING_TIMEOUT_S, -) -from onyx.background.celery.tasks.opensearch_migration.constants import ( MIGRATION_TASK_LOCK_TIMEOUT_S, -) -from onyx.background.celery.tasks.opensearch_migration.constants import ( MIGRATION_TASK_SOFT_TIME_LIMIT_S, -) -from onyx.background.celery.tasks.opensearch_migration.constants import ( MIGRATION_TASK_TIME_LIMIT_S, ) from onyx.background.celery.tasks.opensearch_migration.transformer import ( transform_vespa_chunks_to_opensearch_chunks, ) -from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX -from onyx.configs.app_configs import VESPA_MIGRATION_REQUEST_TIMEOUT_S -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks +from onyx.configs.app_configs import ( + ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + VESPA_MIGRATION_REQUEST_TIMEOUT_S, +) +from onyx.configs.constants import OnyxCeleryTask, OnyxRedisLocks from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.opensearch_migration import build_sanitized_to_original_doc_id_mapping -from onyx.db.opensearch_migration import get_vespa_visit_state -from onyx.db.opensearch_migration import is_migration_completed from onyx.db.opensearch_migration import ( + build_sanitized_to_original_doc_id_mapping, + get_vespa_visit_state, + is_migration_completed, mark_migration_completed_time_if_not_set_with_commit, -) -from onyx.db.opensearch_migration import ( try_insert_opensearch_tenant_migration_record_with_commit, + update_vespa_visit_progress_with_commit, ) -from onyx.db.opensearch_migration import update_vespa_visit_progress_with_commit from onyx.db.search_settings import get_current_search_settings from onyx.document_index.interfaces_new import TenantState from onyx.document_index.opensearch.opensearch_document_index import ( diff --git a/backend/onyx/background/celery/tasks/opensearch_migration/transformer.py b/backend/onyx/background/celery/tasks/opensearch_migration/transformer.py index e395b3ee43..fdfdd8dbb3 100644 --- a/backend/onyx/background/celery/tasks/opensearch_migration/transformer.py +++ b/backend/onyx/background/celery/tasks/opensearch_migration/transformer.py @@ -1,37 +1,38 @@ import traceback -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from onyx.configs.constants import PUBLIC_DOC_PAT from onyx.document_index.interfaces_new import TenantState from onyx.document_index.opensearch.schema import DocumentChunk -from onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST -from onyx.document_index.vespa_constants import BLURB -from onyx.document_index.vespa_constants import BOOST -from onyx.document_index.vespa_constants import CHUNK_CONTEXT -from onyx.document_index.vespa_constants import CHUNK_ID -from onyx.document_index.vespa_constants import CONTENT -from onyx.document_index.vespa_constants import DOC_SUMMARY -from onyx.document_index.vespa_constants import DOC_UPDATED_AT -from onyx.document_index.vespa_constants import DOCUMENT_ID -from onyx.document_index.vespa_constants import DOCUMENT_SETS -from onyx.document_index.vespa_constants import EMBEDDINGS -from onyx.document_index.vespa_constants import FULL_CHUNK_EMBEDDING_KEY -from onyx.document_index.vespa_constants import HIDDEN -from onyx.document_index.vespa_constants import IMAGE_FILE_NAME -from onyx.document_index.vespa_constants import METADATA_LIST -from onyx.document_index.vespa_constants import METADATA_SUFFIX -from onyx.document_index.vespa_constants import PERSONAS -from onyx.document_index.vespa_constants import PRIMARY_OWNERS -from onyx.document_index.vespa_constants import SECONDARY_OWNERS -from onyx.document_index.vespa_constants import SEMANTIC_IDENTIFIER -from onyx.document_index.vespa_constants import SOURCE_LINKS -from onyx.document_index.vespa_constants import SOURCE_TYPE -from onyx.document_index.vespa_constants import TENANT_ID -from onyx.document_index.vespa_constants import TITLE -from onyx.document_index.vespa_constants import TITLE_EMBEDDING -from onyx.document_index.vespa_constants import USER_PROJECT +from onyx.document_index.vespa_constants import ( + ACCESS_CONTROL_LIST, + BLURB, + BOOST, + CHUNK_CONTEXT, + CHUNK_ID, + CONTENT, + DOC_SUMMARY, + DOC_UPDATED_AT, + DOCUMENT_ID, + DOCUMENT_SETS, + EMBEDDINGS, + FULL_CHUNK_EMBEDDING_KEY, + HIDDEN, + IMAGE_FILE_NAME, + METADATA_LIST, + METADATA_SUFFIX, + PERSONAS, + PRIMARY_OWNERS, + SECONDARY_OWNERS, + SEMANTIC_IDENTIFIER, + SOURCE_LINKS, + SOURCE_TYPE, + TENANT_ID, + TITLE, + TITLE_EMBEDDING, + USER_PROJECT, +) from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/background/celery/tasks/port/tasks.py b/backend/onyx/background/celery/tasks/port/tasks.py index bcd9af7bba..92065ce922 100644 --- a/backend/onyx/background/celery/tasks/port/tasks.py +++ b/backend/onyx/background/celery/tasks/port/tasks.py @@ -12,58 +12,61 @@ SUCCESS/CANCELED are left alone). import logging import time -from collections.abc import Callable -from collections.abc import MutableMapping -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from collections.abc import Callable, MutableMapping +from datetime import datetime, timedelta, timezone from typing import Any -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.background.celery.apps.app_base import task_logger from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import MAX_CONCURRENT_PORT_ATTEMPTS -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks +from onyx.configs.app_configs import INDEX_BATCH_SIZE, MAX_CONCURRENT_PORT_ATTEMPTS +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisLocks, +) from onyx.db.connector_credential_pair import ( fetch_indexable_standard_connector_credential_pair_ids, + get_connector_credential_pair_from_id, +) +from onyx.db.document import ( + filter_existing_cc_pair_document_ids, + get_document_ids_for_cc_pair_batch, + get_max_document_id_for_cc_pair, ) -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.document import filter_existing_cc_pair_document_ids -from onyx.db.document import get_document_ids_for_cc_pair_batch -from onyx.db.document import get_max_document_id_for_cc_pair from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import PortAttemptStatus -from onyx.db.enums import SwitchoverType -from onyx.db.models import PortAttempt -from onyx.db.models import SearchSettings -from onyx.db.port_attempt import commit_port_cursor -from onyx.db.port_attempt import count_active_port_attempts -from onyx.db.port_attempt import count_consecutive_failed_port_attempts_no_progress -from onyx.db.port_attempt import create_port_attempt -from onyx.db.port_attempt import get_active_port_attempt -from onyx.db.port_attempt import get_latest_port_attempt -from onyx.db.port_attempt import get_port_attempt -from onyx.db.port_attempt import get_stale_in_progress_port_attempts -from onyx.db.port_attempt import mark_port_canceled -from onyx.db.port_attempt import mark_port_failed -from onyx.db.port_attempt import mark_port_in_progress -from onyx.db.port_attempt import mark_port_succeeded -from onyx.db.port_attempt import port_backfill_has_pending_work -from onyx.db.port_attempt import touch_port_progress -from onyx.db.search_settings import get_current_search_settings -from onyx.db.search_settings import get_search_settings_by_id -from onyx.db.search_settings import get_secondary_search_settings +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + PortAttemptStatus, + SwitchoverType, +) +from onyx.db.models import PortAttempt, SearchSettings +from onyx.db.port_attempt import ( + commit_port_cursor, + count_active_port_attempts, + count_consecutive_failed_port_attempts_no_progress, + create_port_attempt, + get_active_port_attempt, + get_latest_port_attempt, + get_port_attempt, + get_stale_in_progress_port_attempts, + mark_port_canceled, + mark_port_failed, + mark_port_in_progress, + mark_port_succeeded, + port_backfill_has_pending_work, + touch_port_progress, +) +from onyx.db.search_settings import ( + get_current_search_settings, + get_search_settings_by_id, + get_secondary_search_settings, +) from onyx.document_index.opensearch.port_copy import PortCopier from onyx.redis.redis_pool import get_redis_client diff --git a/backend/onyx/background/celery/tasks/pruning/tasks.py b/backend/onyx/background/celery/tasks/pruning/tasks.py index 5fc7105c9b..70482be57f 100644 --- a/backend/onyx/background/celery/tasks/pruning/tasks.py +++ b/backend/onyx/background/celery/tasks/pruning/tasks.py @@ -1,14 +1,9 @@ import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timedelta, timezone +from typing import Any, cast from uuid import uuid4 -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from pydantic import ValidationError from redis import Redis @@ -16,72 +11,87 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.celery_redis import celery_find_task -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_queue_length -from onyx.background.celery.celery_redis import celery_get_queued_task_ids -from onyx.background.celery.celery_redis import celery_get_unacked_task_ids +from onyx.background.celery.celery_redis import ( + celery_find_task, + celery_get_broker_client, + celery_get_queue_length, + celery_get_queued_task_ids, + celery_get_unacked_task_ids, +) from onyx.background.celery.celery_utils import extract_ids_from_runnable_connector from onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT from onyx.background.celery.tasks.docprocessing.utils import IndexingCallbackBase -from onyx.configs.app_configs import ALLOW_SIMULTANEOUS_PRUNING -from onyx.configs.app_configs import JOB_TIMEOUT -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_PRUNING_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT -from onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import OnyxRedisSignals +from onyx.configs.app_configs import ALLOW_SIMULTANEOUS_PRUNING, JOB_TIMEOUT +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + CELERY_PRUNING_LOCK_TIMEOUT, + CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT, + DANSWER_REDIS_FUNCTION_LOCK_PREFIX, + DocumentSource, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, + OnyxRedisLocks, + OnyxRedisSignals, +) from onyx.connectors.factory import instantiate_connector from onyx.connectors.interfaces import BaseConnector from onyx.connectors.models import InputType from onyx.db.connector import mark_ccpair_as_pruned -from onyx.db.connector_credential_pair import get_connector_credential_pair -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.connector_credential_pair import get_connector_credential_pairs -from onyx.db.document import backfill_docs_created_at__no_commit -from onyx.db.document import get_documents_for_connector_credential_pair +from onyx.db.connector_credential_pair import ( + get_connector_credential_pair, + get_connector_credential_pair_from_id, + get_connector_credential_pairs, +) +from onyx.db.document import ( + backfill_docs_created_at__no_commit, + get_documents_for_connector_credential_pair, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType -from onyx.db.hierarchy import delete_orphaned_hierarchy_nodes -from onyx.db.hierarchy import remove_stale_hierarchy_node_cc_pair_entries -from onyx.db.hierarchy import reparent_orphaned_hierarchy_nodes -from onyx.db.hierarchy import update_document_parent_hierarchy_nodes -from onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries -from onyx.db.hierarchy import upsert_hierarchy_nodes_batch +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + SyncStatus, + SyncType, +) +from onyx.db.hierarchy import ( + delete_orphaned_hierarchy_nodes, + remove_stale_hierarchy_node_cc_pair_entries, + reparent_orphaned_hierarchy_nodes, + update_document_parent_hierarchy_nodes, + upsert_hierarchy_node_cc_pair_entries, + upsert_hierarchy_nodes_batch, +) from onyx.db.models import ConnectorCredentialPair from onyx.db.models import HierarchyNode as DBHierarchyNode -from onyx.db.sync_record import insert_sync_record -from onyx.db.sync_record import update_sync_record_status +from onyx.db.sync_record import insert_sync_record, update_sync_record_status from onyx.db.tag import delete_orphan_tags_batched from onyx.redis.redis_connector import RedisConnector -from onyx.redis.redis_connector_prune import RedisConnectorPrune -from onyx.redis.redis_connector_prune import RedisConnectorPrunePayload -from onyx.redis.redis_hierarchy import cache_hierarchy_nodes_batch -from onyx.redis.redis_hierarchy import ensure_source_node_exists -from onyx.redis.redis_hierarchy import evict_hierarchy_nodes_from_cache -from onyx.redis.redis_hierarchy import get_node_id_from_raw_id -from onyx.redis.redis_hierarchy import get_source_node_id_from_cache -from onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client +from onyx.redis.redis_connector_prune import ( + RedisConnectorPrune, + RedisConnectorPrunePayload, +) +from onyx.redis.redis_hierarchy import ( + cache_hierarchy_nodes_batch, + ensure_source_node_exists, + evict_hierarchy_nodes_from_cache, + get_node_id_from_raw_id, + get_source_node_id_from_cache, + HierarchyNodeCacheEntry, +) +from onyx.redis.redis_pool import get_redis_client, get_redis_replica_client from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.server.metrics.pruning_metrics import observe_pruning_diff_duration from onyx.server.runtime.onyx_runtime import OnyxRuntime from onyx.server.utils import make_short_id -from onyx.utils.logger import format_error_for_logging -from onyx.utils.logger import LoggerContextVars -from onyx.utils.logger import pruning_ctx -from onyx.utils.logger import setup_logger +from onyx.utils.logger import ( + format_error_for_logging, + LoggerContextVars, + pruning_ctx, + setup_logger, +) from shared_configs.configs import MULTI_TENANT logger = setup_logger() diff --git a/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py b/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py index d8586174ac..400acac270 100644 --- a/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py +++ b/backend/onyx/background/celery/tasks/scheduled_tasks/tasks.py @@ -28,33 +28,32 @@ Per CLAUDE.md: from __future__ import annotations import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from uuid import UUID -from celery import shared_task -from celery import Task +from celery import shared_task, Task from onyx.background.celery.apps.app_base import task_logger -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import OnyxCeleryPriority, OnyxCeleryQueues, OnyxCeleryTask from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import ScheduledTaskErrorClass -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import ScheduledTaskSkipReason -from onyx.db.enums import ScheduledTaskTriggerSource -from onyx.db.scheduled_task import advance_next_run_at -from onyx.db.scheduled_task import claim_due_scheduled_tasks -from onyx.db.scheduled_task import find_stuck_runs -from onyx.db.scheduled_task import has_in_flight_run_for_task -from onyx.db.scheduled_task import insert_run -from onyx.db.scheduled_task import mark_run_status +from onyx.db.enums import ( + ScheduledTaskErrorClass, + ScheduledTaskRunStatus, + ScheduledTaskSkipReason, + ScheduledTaskTriggerSource, +) +from onyx.db.scheduled_task import ( + advance_next_run_at, + claim_due_scheduled_tasks, + find_stuck_runs, + has_in_flight_run_for_task, + insert_run, + mark_run_status, +) from onyx.server.features.build.scheduled_tasks.executor import ( DEFAULT_EXECUTOR_BUDGET_SECONDS, + run_scheduled_task_logic, ) -from onyx.server.features.build.scheduled_tasks.executor import run_scheduled_task_logic from onyx.server.features.build.utils import is_craft_enabled_for_user from onyx.server.settings.store import load_settings diff --git a/backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py b/backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py index db6e42ef21..1c618af19c 100644 --- a/backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py +++ b/backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py @@ -1,11 +1,12 @@ import httpx -from tenacity import retry -from tenacity import retry_if_exception_type -from tenacity import stop_after_delay -from tenacity import wait_random_exponential +from tenacity import ( + retry, + retry_if_exception_type, + stop_after_delay, + wait_random_exponential, +) -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import MetadataUpdateRequest +from onyx.document_index.interfaces_new import DocumentIndex, MetadataUpdateRequest class RetryDocumentIndex: diff --git a/backend/onyx/background/celery/tasks/shared/tasks.py b/backend/onyx/background/celery/tasks/shared/tasks.py index 485e5b956d..89f0970e5e 100644 --- a/backend/onyx/background/celery/tasks/shared/tasks.py +++ b/backend/onyx/background/celery/tasks/shared/tasks.py @@ -3,23 +3,23 @@ from enum import Enum from http import HTTPStatus import httpx -from celery import shared_task -from celery import Task +from celery import shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from tenacity import RetryError from onyx.access.access import get_access_for_document from onyx.background.celery.apps.app_base import task_logger from onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex -from onyx.configs.constants import ONYX_CELERY_BEAT_HEARTBEAT_KEY -from onyx.configs.constants import OnyxCeleryTask -from onyx.db.document import delete_document_by_connector_credential_pair__no_commit -from onyx.db.document import delete_documents_complete -from onyx.db.document import fetch_chunk_count_for_document -from onyx.db.document import get_document -from onyx.db.document import get_document_connector_count -from onyx.db.document import mark_document_as_modified -from onyx.db.document import mark_document_as_synced +from onyx.configs.constants import ONYX_CELERY_BEAT_HEARTBEAT_KEY, OnyxCeleryTask +from onyx.db.document import ( + delete_document_by_connector_credential_pair__no_commit, + delete_documents_complete, + fetch_chunk_count_for_document, + get_document, + get_document_connector_count, + mark_document_as_modified, + mark_document_as_synced, +) from onyx.db.document_set import fetch_document_sets_for_document from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.relationships import delete_document_references_from_kg diff --git a/backend/onyx/background/celery/tasks/user_file_processing/tasks.py b/backend/onyx/background/celery/tasks/user_file_processing/tasks.py index 9e155a9f10..1e2372f452 100644 --- a/backend/onyx/background/celery/tasks/user_file_processing/tasks.py +++ b/backend/onyx/background/celery/tasks/user_file_processing/tasks.py @@ -3,53 +3,62 @@ import time from uuid import UUID import sqlalchemy as sa -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from redis.lock import Lock as RedisLock from sqlalchemy import select from onyx.access.access import build_access_for_user_files from onyx.access.models import DocumentAccess from onyx.background.celery.apps.app_base import task_logger -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_queue_length +from onyx.background.celery.celery_redis import ( + celery_get_broker_client, + celery_get_queue_length, +) from onyx.background.celery.celery_utils import httpx_init_vespa_pool from onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import MANAGED_VESPA -from onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH -from onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_USER_FILE_DELETE_TASK_EXPIRES -from onyx.configs.constants import CELERY_USER_FILE_PROCESSING_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_USER_FILE_PROCESSING_TASK_EXPIRES -from onyx.configs.constants import CELERY_USER_FILE_PROJECT_SYNC_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import USER_FILE_DELETE_MAX_QUEUE_DEPTH -from onyx.configs.constants import USER_FILE_PROCESSING_MAX_QUEUE_DEPTH -from onyx.configs.constants import USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH +from onyx.configs.app_configs import ( + DISABLE_VECTOR_DB, + MANAGED_VESPA, + VESPA_CLOUD_CERT_PATH, + VESPA_CLOUD_KEY_PATH, +) +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + CELERY_USER_FILE_DELETE_TASK_EXPIRES, + CELERY_USER_FILE_PROCESSING_LOCK_TIMEOUT, + CELERY_USER_FILE_PROCESSING_TASK_EXPIRES, + CELERY_USER_FILE_PROJECT_SYNC_LOCK_TIMEOUT, + CELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES, + DocumentSource, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisLocks, + USER_FILE_DELETE_MAX_QUEUE_DEPTH, + USER_FILE_PROCESSING_MAX_QUEUE_DEPTH, + USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH, +) from onyx.connectors.file.connector import LocalFileConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import UserFileStatus from onyx.db.models import UserFile -from onyx.db.search_settings import get_active_search_settings -from onyx.db.search_settings import get_active_search_settings_list +from onyx.db.search_settings import ( + get_active_search_settings, + get_active_search_settings_list, +) from onyx.db.user_file import fetch_user_files_with_access_relationships from onyx.document_index.factory import get_all_document_indices from onyx.document_index.interfaces_new import MetadataUpdateRequest from onyx.file_store.file_store import get_default_file_store -from onyx.file_store.staging import build_tracking_raw_file_callback -from onyx.file_store.staging import delete_files_best_effort -from onyx.file_store.utils import store_user_file_plaintext -from onyx.file_store.utils import user_file_id_to_plaintext_file_name +from onyx.file_store.staging import ( + build_tracking_raw_file_callback, + delete_files_best_effort, +) +from onyx.file_store.utils import ( + store_user_file_plaintext, + user_file_id_to_plaintext_file_name, +) from onyx.httpx.httpx_pool import HttpxPool from onyx.indexing.adapters.user_file_indexing_adapter import UserFileIndexingAdapter from onyx.indexing.embedder import DefaultIndexingEmbedder @@ -268,8 +277,7 @@ def _process_user_file_without_vector_db( Opens its own short DB session only for the final status write, so the caller does not need to hold a session open during the text/token work. """ - from onyx.llm.factory import get_default_llm - from onyx.llm.factory import get_llm_tokenizer_encode_func + from onyx.llm.factory import get_default_llm, get_llm_tokenizer_encode_func user_file_uuid = _as_uuid(user_file_id) diff --git a/backend/onyx/background/celery/tasks/vespa/document_sync.py b/backend/onyx/background/celery/tasks/vespa/document_sync.py index b08899f66a..9279bbbc01 100644 --- a/backend/onyx/background/celery/tasks/vespa/document_sync.py +++ b/backend/onyx/background/celery/tasks/vespa/document_sync.py @@ -7,17 +7,19 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.configs.app_configs import DB_YIELD_PER_DEFAULT -from onyx.configs.constants import CELERY_DOCUMENT_SYNC_TASK_EXPIRES -from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants +from onyx.configs.constants import ( + CELERY_DOCUMENT_SYNC_TASK_EXPIRES, + CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, +) from onyx.db.document import ( construct_document_id_select_by_needs_sync_or_secondary_pending, + count_documents_by_needs_sync_or_secondary_pending, + count_secondary_only_sync_pending_documents, ) -from onyx.db.document import count_documents_by_needs_sync_or_secondary_pending -from onyx.db.document import count_secondary_only_sync_pending_documents from onyx.db.port_attempt import any_future_port_in_progress from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/onyx/background/celery/tasks/vespa/tasks.py b/backend/onyx/background/celery/tasks/vespa/tasks.py index a7f99b3557..2cacc55ee2 100644 --- a/backend/onyx/background/celery/tasks/vespa/tasks.py +++ b/backend/onyx/background/celery/tasks/vespa/tasks.py @@ -1,13 +1,10 @@ import time from collections.abc import Callable from http import HTTPStatus -from typing import Any -from typing import cast +from typing import Any, cast import httpx -from celery import Celery -from celery import shared_task -from celery import Task +from celery import Celery, shared_task, Task from celery.exceptions import SoftTimeLimitExceeded from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session @@ -16,58 +13,69 @@ from tenacity import RetryError from onyx.access.access import get_access_for_document from onyx.background.celery.apps.app_base import task_logger from onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex -from onyx.background.celery.tasks.shared.tasks import LIGHT_SOFT_TIME_LIMIT -from onyx.background.celery.tasks.shared.tasks import LIGHT_TIME_LIMIT -from onyx.background.celery.tasks.shared.tasks import OnyxCeleryTaskCompletionStatus -from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_FENCE_KEY -from onyx.background.celery.tasks.vespa.document_sync import get_document_sync_payload -from onyx.background.celery.tasks.vespa.document_sync import get_document_sync_remaining -from onyx.background.celery.tasks.vespa.document_sync import reset_document_sync +from onyx.background.celery.tasks.shared.tasks import ( + LIGHT_SOFT_TIME_LIMIT, + LIGHT_TIME_LIMIT, + OnyxCeleryTaskCompletionStatus, +) from onyx.background.celery.tasks.vespa.document_sync import ( + DOCUMENT_SYNC_FENCE_KEY, + get_document_sync_payload, + get_document_sync_remaining, + reset_document_sync, try_generate_stale_document_sync_tasks, ) -from onyx.configs.app_configs import JOB_TIMEOUT -from onyx.configs.app_configs import VESPA_SYNC_MAX_TASKS -from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants -from onyx.configs.constants import OnyxRedisLocks -from onyx.db.document import document_has_indexable_cc_pair -from onyx.db.document import get_document -from onyx.db.document import mark_document_as_synced -from onyx.db.document import mark_document_synced_secondary_pending -from onyx.db.document_set import delete_document_set -from onyx.db.document_set import fetch_document_sets -from onyx.db.document_set import fetch_document_sets_for_document -from onyx.db.document_set import get_document_set_by_id -from onyx.db.document_set import mark_document_set_as_synced +from onyx.configs.app_configs import JOB_TIMEOUT, VESPA_SYNC_MAX_TASKS +from onyx.configs.constants import ( + CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT, + OnyxCeleryTask, + OnyxRedisConstants, + OnyxRedisLocks, +) +from onyx.db.document import ( + document_has_indexable_cc_pair, + get_document, + mark_document_as_synced, + mark_document_synced_secondary_pending, +) +from onyx.db.document_set import ( + delete_document_set, + fetch_document_sets, + fetch_document_sets_for_document, + get_document_set_by_id, + mark_document_set_as_synced, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType -from onyx.db.models import DocumentSet -from onyx.db.models import UserGroup +from onyx.db.enums import SyncStatus, SyncType +from onyx.db.models import DocumentSet, UserGroup from onyx.db.port_attempt import port_backfill_has_pending_work from onyx.db.search_settings import get_active_search_settings -from onyx.db.sync_record import cleanup_sync_records -from onyx.db.sync_record import insert_sync_record -from onyx.db.sync_record import update_sync_record_status +from onyx.db.sync_record import ( + cleanup_sync_records, + insert_sync_record, + update_sync_record_status, +) from onyx.document_index.factory import get_all_document_indices -from onyx.document_index.interfaces_new import MetadataUpdateRequest -from onyx.document_index.interfaces_new import SecondaryIndexDocumentMissingError +from onyx.document_index.interfaces_new import ( + MetadataUpdateRequest, + SecondaryIndexDocumentMissingError, +) from onyx.httpx.httpx_pool import HttpxPool from onyx.redis.redis_document_set import RedisDocumentSet -from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_pool import get_redis_replica_client -from onyx.redis.redis_pool import redis_lock_dump +from onyx.redis.redis_pool import ( + get_redis_client, + get_redis_replica_client, + redis_lock_dump, +) from onyx.redis.redis_usergroup import RedisUserGroup from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.utils.logger import setup_logger -from onyx.utils.variable_functionality import fetch_versioned_implementation from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, fetch_versioned_implementation_with_fallback, + global_version, + noop_fallback, ) -from onyx.utils.variable_functionality import global_version -from onyx.utils.variable_functionality import noop_fallback logger = setup_logger() diff --git a/backend/onyx/background/celery/versioned_apps/heavy.py b/backend/onyx/background/celery/versioned_apps/heavy.py index 645be23d86..6a61de2210 100644 --- a/backend/onyx/background/celery/versioned_apps/heavy.py +++ b/backend/onyx/background/celery/versioned_apps/heavy.py @@ -4,8 +4,10 @@ fetch. Port over the code in those files if we add an EE version of this worker. from celery import Celery -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + set_is_ee_based_on_env_variable, +) set_is_ee_based_on_env_variable() app: Celery = fetch_versioned_implementation( diff --git a/backend/onyx/background/celery/versioned_apps/light.py b/backend/onyx/background/celery/versioned_apps/light.py index 20fff856c6..f1980d869a 100644 --- a/backend/onyx/background/celery/versioned_apps/light.py +++ b/backend/onyx/background/celery/versioned_apps/light.py @@ -4,8 +4,10 @@ fetch. Port over the code in those files if we add an EE version of this worker. from celery import Celery -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + set_is_ee_based_on_env_variable, +) set_is_ee_based_on_env_variable() app: Celery = fetch_versioned_implementation( diff --git a/backend/onyx/background/celery/versioned_apps/monitoring.py b/backend/onyx/background/celery/versioned_apps/monitoring.py index d9e7898573..ca81f2deca 100644 --- a/backend/onyx/background/celery/versioned_apps/monitoring.py +++ b/backend/onyx/background/celery/versioned_apps/monitoring.py @@ -2,8 +2,10 @@ from celery import Celery -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + set_is_ee_based_on_env_variable, +) set_is_ee_based_on_env_variable() app: Celery = fetch_versioned_implementation( diff --git a/backend/onyx/background/celery/versioned_apps/primary.py b/backend/onyx/background/celery/versioned_apps/primary.py index f52a054fe3..b411646744 100644 --- a/backend/onyx/background/celery/versioned_apps/primary.py +++ b/backend/onyx/background/celery/versioned_apps/primary.py @@ -2,8 +2,10 @@ from celery import Celery -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + set_is_ee_based_on_env_variable, +) set_is_ee_based_on_env_variable() app: Celery = fetch_versioned_implementation( diff --git a/backend/onyx/background/indexing/checkpointing_utils.py b/backend/onyx/background/indexing/checkpointing_utils.py index b4903f2dad..f4a590332b 100644 --- a/backend/onyx/background/indexing/checkpointing_utils.py +++ b/backend/onyx/background/indexing/checkpointing_utils.py @@ -1,18 +1,17 @@ -from datetime import datetime -from datetime import timedelta +from datetime import datetime, timedelta from io import BytesIO from sqlalchemy import and_ from sqlalchemy.orm import Session -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import NUM_DAYS_TO_KEEP_CHECKPOINTS -from onyx.connectors.interfaces import BaseConnector -from onyx.connectors.interfaces import CheckpointedConnector +from onyx.configs.constants import FileOrigin, NUM_DAYS_TO_KEEP_CHECKPOINTS +from onyx.connectors.interfaces import BaseConnector, CheckpointedConnector from onyx.connectors.models import ConnectorCheckpoint from onyx.db.engine.time_utils import get_db_current_time -from onyx.db.index_attempt import get_index_attempt -from onyx.db.index_attempt import get_recent_completed_attempts_for_cc_pair +from onyx.db.index_attempt import ( + get_index_attempt, + get_recent_completed_attempts_for_cc_pair, +) from onyx.db.models import IndexAttempt from onyx.file_store.file_store import get_default_file_store from onyx.utils.logger import setup_logger diff --git a/backend/onyx/background/indexing/index_attempt_utils.py b/backend/onyx/background/indexing/index_attempt_utils.py index 63c21b54d5..1bd873ba0e 100644 --- a/backend/onyx/background/indexing/index_attempt_utils.py +++ b/backend/onyx/background/indexing/index_attempt_utils.py @@ -5,8 +5,7 @@ from sqlalchemy.orm import Session from onyx.configs.constants import NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS from onyx.db.engine.time_utils import get_db_current_time -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError +from onyx.db.models import IndexAttempt, IndexAttemptError # Always retain at least this many attempts per connector/search settings pair NUM_RECENT_INDEX_ATTEMPTS_TO_KEEP = 10 diff --git a/backend/onyx/background/indexing/job_client.py b/backend/onyx/background/indexing/job_client.py index eb06055c7b..47969d9652 100644 --- a/backend/onyx/background/indexing/job_client.py +++ b/backend/onyx/background/indexing/job_client.py @@ -11,15 +11,12 @@ import traceback from collections.abc import Callable from dataclasses import dataclass from multiprocessing.context import SpawnProcess -from typing import Any -from typing import Literal -from typing import Optional +from typing import Any, Literal, Optional from onyx.configs.constants import POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME from onyx.db.engine.sql_engine import SqlEngine from onyx.utils.logger import setup_logger -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import TENANT_ID_PREFIX +from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA, TENANT_ID_PREFIX from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR logger = setup_logger() diff --git a/backend/onyx/background/indexing/run_docfetching.py b/backend/onyx/background/indexing/run_docfetching.py index 7a4a118718..30b101bb47 100644 --- a/backend/onyx/background/indexing/run_docfetching.py +++ b/backend/onyx/background/indexing/run_docfetching.py @@ -1,11 +1,8 @@ import sys import time import traceback -from collections.abc import Generator -from collections.abc import Iterable -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from collections.abc import Generator, Iterable +from datetime import datetime, timedelta, timezone from typing import TypeVar import sentry_sdk @@ -13,72 +10,94 @@ from celery import Celery from sqlalchemy.orm import Session from onyx.access.access import source_should_fetch_permissions_during_indexing -from onyx.background.indexing.checkpointing_utils import check_checkpoint_size -from onyx.background.indexing.checkpointing_utils import get_latest_valid_checkpoint -from onyx.background.indexing.checkpointing_utils import save_checkpoint +from onyx.background.indexing.checkpointing_utils import ( + check_checkpoint_size, + get_latest_valid_checkpoint, + save_checkpoint, +) from onyx.background.indexing.memory_tracer import MemoryTracer -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import INDEXING_SIZE_WARNING_THRESHOLD -from onyx.configs.app_configs import INDEXING_TRACER_INTERVAL -from onyx.configs.app_configs import INTEGRATION_TESTS_MODE -from onyx.configs.app_configs import LEAVE_CONNECTOR_ACTIVE_ON_INITIALIZATION_FAILURE -from onyx.configs.app_configs import MAX_FILE_SIZE_BYTES -from onyx.configs.app_configs import PERSISTENT_INDEXING -from onyx.configs.app_configs import POLL_CONNECTOR_OFFSET -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.app_configs import ( + INDEX_BATCH_SIZE, + INDEXING_SIZE_WARNING_THRESHOLD, + INDEXING_TRACER_INTERVAL, + INTEGRATION_TESTS_MODE, + LEAVE_CONNECTOR_ACTIVE_ON_INITIALIZATION_FAILURE, + MAX_FILE_SIZE_BYTES, + PERSISTENT_INDEXING, + POLL_CONNECTOR_OFFSET, +) +from onyx.configs.constants import OnyxCeleryPriority, OnyxCeleryQueues, OnyxCeleryTask from onyx.connectors.connector_runner import ConnectorRunner -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import UnexpectedValidationError +from onyx.connectors.exceptions import ( + ConnectorValidationError, + UnexpectedValidationError, +) from onyx.connectors.factory import instantiate_connector from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorStopSignal -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + ConnectorFailure, + ConnectorStopSignal, + Document, + HierarchyNode, + TextSection, +) from onyx.db.connector import mark_ccpair_with_indexing_trigger -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.connector_credential_pair import get_last_successful_attempt_poll_range_end -from onyx.db.connector_credential_pair import update_connector_credential_pair +from onyx.db.connector_credential_pair import ( + get_connector_credential_pair_from_id, + get_last_successful_attempt_poll_range_end, + update_connector_credential_pair, +) from onyx.db.constants import CONNECTOR_VALIDATION_ERROR_MESSAGE_PREFIX from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries -from onyx.db.hierarchy import upsert_hierarchy_nodes_batch -from onyx.db.index_attempt import create_index_attempt_error -from onyx.db.index_attempt import get_index_attempt -from onyx.db.index_attempt import get_recent_completed_attempts_for_cc_pair -from onyx.db.index_attempt import mark_attempt_canceled -from onyx.db.index_attempt import mark_attempt_failed -from onyx.db.index_attempt import transition_attempt_to_in_progress -from onyx.db.index_attempt_metrics import IndexAttemptStage -from onyx.db.index_attempt_metrics import StageEventBuffer -from onyx.db.index_attempt_metrics import time_stage +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + IndexingStatus, + IndexModelStatus, +) +from onyx.db.hierarchy import ( + upsert_hierarchy_node_cc_pair_entries, + upsert_hierarchy_nodes_batch, +) +from onyx.db.index_attempt import ( + create_index_attempt_error, + get_index_attempt, + get_recent_completed_attempts_for_cc_pair, + mark_attempt_canceled, + mark_attempt_failed, + transition_attempt_to_in_progress, +) +from onyx.db.index_attempt_metrics import ( + IndexAttemptStage, + StageEventBuffer, + time_stage, +) from onyx.db.indexing_coordination import IndexingCoordination -from onyx.db.models import Connector -from onyx.db.models import Credential -from onyx.db.models import IndexAttempt -from onyx.file_store.document_batch_storage import DocumentBatchStorage -from onyx.file_store.document_batch_storage import get_document_batch_storage -from onyx.file_store.staging import build_raw_file_callback -from onyx.file_store.staging import RawFileCallback -from onyx.file_store.staging import reap_prior_attempt_staged_files +from onyx.db.models import Connector, Credential, IndexAttempt +from onyx.file_store.document_batch_storage import ( + DocumentBatchStorage, + get_document_batch_storage, +) +from onyx.file_store.staging import ( + build_raw_file_callback, + RawFileCallback, + reap_prior_attempt_staged_files, +) from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.redis.redis_docprocessing import RedisDocprocessing -from onyx.redis.redis_hierarchy import cache_hierarchy_nodes_batch -from onyx.redis.redis_hierarchy import ensure_source_node_exists -from onyx.redis.redis_hierarchy import get_node_id_from_raw_id -from onyx.redis.redis_hierarchy import get_source_node_id_from_cache -from onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry +from onyx.redis.redis_hierarchy import ( + cache_hierarchy_nodes_batch, + ensure_source_node_exists, + get_node_id_from_raw_id, + get_source_node_id_from_cache, + HierarchyNodeCacheEntry, +) from onyx.redis.redis_pool import get_redis_client from onyx.utils.logger import setup_logger -from onyx.utils.postgres_sanitization import sanitize_document_for_postgres -from onyx.utils.postgres_sanitization import sanitize_hierarchy_nodes_for_postgres +from onyx.utils.postgres_sanitization import ( + sanitize_document_for_postgres, + sanitize_hierarchy_nodes_for_postgres, +) from onyx.utils.variable_functionality import global_version from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR diff --git a/backend/onyx/background/indexing/run_targeted_reindex.py b/backend/onyx/background/indexing/run_targeted_reindex.py index a5740ffeb7..3cc31fcd87 100644 --- a/backend/onyx/background/indexing/run_targeted_reindex.py +++ b/backend/onyx/background/indexing/run_targeted_reindex.py @@ -17,8 +17,7 @@ reindex isn't supported yet for that source. import datetime from collections import defaultdict -from collections.abc import Iterable -from collections.abc import Sequence +from collections.abc import Iterable, Sequence from more_itertools import chunked from sqlalchemy.orm import Session @@ -26,29 +25,39 @@ from sqlalchemy.orm import Session from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.connectors.factory import instantiate_connector from onyx.connectors.interfaces import Resolver -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import IndexAttemptMetadata +from onyx.connectors.models import ( + ConnectorFailure, + Document, + HierarchyNode, + IndexAttemptMetadata, +) from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.db.enums import AccessType -from onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries -from onyx.db.hierarchy import upsert_hierarchy_nodes_batch -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import TargetedReindexJobTarget +from onyx.db.hierarchy import ( + upsert_hierarchy_node_cc_pair_entries, + upsert_hierarchy_nodes_batch, +) +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + TargetedReindexJobTarget, +) from onyx.db.targeted_reindex import targets_to_connector_failures from onyx.document_index.factory import get_all_document_indices -from onyx.file_store.staging import build_tracking_raw_file_callback -from onyx.file_store.staging import delete_files_best_effort +from onyx.file_store.staging import ( + build_tracking_raw_file_callback, + delete_files_best_effort, +) 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.redis.redis_hierarchy import cache_hierarchy_nodes_batch -from onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry +from onyx.redis.redis_hierarchy import ( + cache_hierarchy_nodes_batch, + HierarchyNodeCacheEntry, +) from onyx.redis.redis_pool import get_redis_client from onyx.utils.logger import setup_logger from onyx.utils.middleware import make_randomized_onyx_request_id diff --git a/backend/onyx/background/periodic_poller.py b/backend/onyx/background/periodic_poller.py index 80414d827e..9178abc250 100644 --- a/backend/onyx/background/periodic_poller.py +++ b/backend/onyx/background/periodic_poller.py @@ -15,8 +15,7 @@ that runs from the API server process. Two responsibilities: import threading import time from collections.abc import Callable -from dataclasses import dataclass -from dataclasses import field +from dataclasses import dataclass, field from onyx.utils.logger import setup_logger @@ -66,10 +65,12 @@ def _run_cache_cleanup() -> None: def _run_scheduled_eval() -> None: - from onyx.configs.app_configs import BRAINTRUST_API_KEY - from onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES - from onyx.configs.app_configs import SCHEDULED_EVAL_PERMISSIONS_EMAIL - from onyx.configs.app_configs import SCHEDULED_EVAL_PROJECT + from onyx.configs.app_configs import ( + BRAINTRUST_API_KEY, + SCHEDULED_EVAL_DATASET_NAMES, + SCHEDULED_EVAL_PERMISSIONS_EMAIL, + SCHEDULED_EVAL_PROJECT, + ) if not all( [ @@ -81,8 +82,7 @@ def _run_scheduled_eval() -> None: ): return - from datetime import datetime - from datetime import timezone + from datetime import datetime, timezone from onyx.evals.eval import run_eval from onyx.evals.models import EvalConfigurationOptions @@ -111,10 +111,12 @@ _CACHE_CLEANUP_INTERVAL_SECONDS = 300 def _build_periodic_tasks() -> list[_PeriodicTaskDef]: from onyx.cache.interface import CacheBackendType - from onyx.configs.app_configs import AUTO_LLM_CONFIG_URL - from onyx.configs.app_configs import AUTO_LLM_UPDATE_INTERVAL_SECONDS - from onyx.configs.app_configs import CACHE_BACKEND - from onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES + from onyx.configs.app_configs import ( + AUTO_LLM_CONFIG_URL, + AUTO_LLM_UPDATE_INTERVAL_SECONDS, + CACHE_BACKEND, + SCHEDULED_EVAL_DATASET_NAMES, + ) tasks: list[_PeriodicTaskDef] = [] if CACHE_BACKEND == CacheBackendType.POSTGRES: @@ -159,8 +161,7 @@ def _try_claim_task(task_def: _PeriodicTaskDef) -> bool: ``KVStore`` timestamp for cross-instance dedup. The DB session is held only for this brief claim transaction, not during task execution. """ - from datetime import datetime - from datetime import timezone + from datetime import datetime, timezone from sqlalchemy import text @@ -218,9 +219,11 @@ def _try_run_periodic_task(task_def: _PeriodicTaskDef) -> None: def _run_drain_loops(tenant_id: str) -> None: - from onyx.background.task_utils import drain_delete_loop - from onyx.background.task_utils import drain_processing_loop - from onyx.background.task_utils import drain_project_sync_loop + from onyx.background.task_utils import ( + drain_delete_loop, + drain_processing_loop, + drain_project_sync_loop, + ) drain_processing_loop(tenant_id) drain_delete_loop(tenant_id) diff --git a/backend/onyx/cache/factory.py b/backend/onyx/cache/factory.py index 7c194c650a..9ecda8339f 100644 --- a/backend/onyx/cache/factory.py +++ b/backend/onyx/cache/factory.py @@ -1,7 +1,6 @@ from collections.abc import Callable -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheBackendType +from onyx.cache.interface import CacheBackend, CacheBackendType from onyx.configs.app_configs import CACHE_BACKEND diff --git a/backend/onyx/cache/locks.py b/backend/onyx/cache/locks.py index dd9065bb76..7feef01768 100644 --- a/backend/onyx/cache/locks.py +++ b/backend/onyx/cache/locks.py @@ -1,8 +1,7 @@ import time from collections.abc import Generator from contextlib import contextmanager -from logging import Logger -from logging import LoggerAdapter +from logging import Logger, LoggerAdapter from onyx.cache.factory import get_shared_cache_backend from onyx.cache.interface import CacheLockAcquisitionError diff --git a/backend/onyx/cache/postgres_backend.py b/backend/onyx/cache/postgres_backend.py index 73bd94522a..482cad01ee 100644 --- a/backend/onyx/cache/postgres_backend.py +++ b/backend/onyx/cache/postgres_backend.py @@ -9,22 +9,18 @@ import struct import time import uuid from contextlib import AbstractContextManager -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone -from sqlalchemy import delete -from sqlalchemy import func -from sqlalchemy import or_ -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import delete, func, or_, select, update from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheLock -from onyx.cache.interface import TTL_KEY_NOT_FOUND -from onyx.cache.interface import TTL_NO_EXPIRY +from onyx.cache.interface import ( + CacheBackend, + CacheLock, + TTL_KEY_NOT_FOUND, + TTL_NO_EXPIRY, +) from onyx.db.models import CacheStore _LIST_KEY_PREFIX = "_q:" diff --git a/backend/onyx/cache/redis_backend.py b/backend/onyx/cache/redis_backend.py index 7735b5f9df..d3d663b958 100644 --- a/backend/onyx/cache/redis_backend.py +++ b/backend/onyx/cache/redis_backend.py @@ -3,9 +3,7 @@ import math from redis.exceptions import LockNotOwnedError from redis.lock import Lock as RedisLock -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheLock -from onyx.cache.interface import CacheLockLostError +from onyx.cache.interface import CacheBackend, CacheLock, CacheLockLostError from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/onyx/chat/chat_state.py b/backend/onyx/chat/chat_state.py index 0d4b096701..0265e54121 100644 --- a/backend/onyx/chat/chat_state.py +++ b/backend/onyx/chat/chat_state.py @@ -7,22 +7,20 @@ from pydantic import BaseModel from onyx.cache.interface import CacheBackend from onyx.chat.citation_processor import CitationMapping -from onyx.chat.models import ChatLoadedFile -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import ExtractedContextFiles -from onyx.chat.models import FileToolMetadata -from onyx.chat.models import SearchParams +from onyx.chat.models import ( + ChatLoadedFile, + ChatMessageSimple, + ExtractedContextFiles, + FileToolMetadata, + SearchParams, +) from onyx.context.search.models import SearchDoc from onyx.db.memory import UserMemoryContext -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession -from onyx.db.models import Persona -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMUserIdentity +from onyx.db.models import ChatMessage, ChatSession, Persona +from onyx.llm.interfaces import LLM, LLMUserIdentity from onyx.onyxbot.slack.models import SlackContext from onyx.server.query_and_chat.models import SendMessageRequest -from onyx.tools.models import ChatFile -from onyx.tools.models import ToolCallInfo +from onyx.tools.models import ChatFile, ToolCallInfo # Type alias for search doc deduplication key # Simple key: just document_id (str) diff --git a/backend/onyx/chat/chat_utils.py b/backend/onyx/chat/chat_utils.py index c872c668b4..c5eb550403 100644 --- a/backend/onyx/chat/chat_utils.py +++ b/backend/onyx/chat/chat_utils.py @@ -8,48 +8,51 @@ from fastapi.datastructures import Headers from pydantic import BaseModel from sqlalchemy.orm import Session -from onyx.chat.models import ChatHistoryResult -from onyx.chat.models import ChatLoadedFile -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import FileToolMetadata -from onyx.chat.models import ToolCallSimple -from onyx.configs.constants import DEFAULT_PERSONA_ID -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import MessageType -from onyx.configs.constants import TMP_DRALPHA_PERSONA_NAME +from onyx.chat.models import ( + ChatHistoryResult, + ChatLoadedFile, + ChatMessageSimple, + FileToolMetadata, + ToolCallSimple, +) +from onyx.configs.constants import ( + DEFAULT_PERSONA_ID, + FileOrigin, + MessageType, + TMP_DRALPHA_PERSONA_NAME, +) from onyx.context.search.models import SearchDoc from onyx.context.search.utils import sandbox_filename_for_document -from onyx.db.chat import create_chat_session -from onyx.db.chat import get_chat_messages_by_session -from onyx.db.chat import get_or_create_root_message +from onyx.db.chat import ( + create_chat_session, + get_chat_messages_by_session, + get_or_create_root_message, +) from onyx.db.file_record import FileRecordNotFoundError -from onyx.db.kg_config import get_kg_config_settings -from onyx.db.kg_config import is_kg_config_settings_enabled_valid -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession -from onyx.db.models import Persona +from onyx.db.kg_config import ( + get_kg_config_settings, + is_kg_config_settings_enabled_valid, +) +from onyx.db.models import ChatMessage, ChatSession, Persona, User, UserFile from onyx.db.models import SearchDoc as DbSearchDoc -from onyx.db.models import User -from onyx.db.models import UserFile from onyx.db.persona import user_can_access_persona from onyx.db.projects import check_project_ownership from onyx.file_processing.extract_file_text import extract_file_text from onyx.file_store.file_store import get_default_file_store -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import FileDescriptor -from onyx.file_store.utils import plaintext_file_name_for_id -from onyx.file_store.utils import store_plaintext +from onyx.file_store.models import ChatFileType, FileDescriptor +from onyx.file_store.utils import plaintext_file_name_for_id, store_plaintext from onyx.kg.models import KGException from onyx.kg.setup.kg_default_entity_definitions import ( populate_missing_default_entity_types__commit, ) -from onyx.prompts.chat_prompts import ADDITIONAL_CONTEXT_PROMPT -from onyx.prompts.chat_prompts import TOOL_CALL_RESPONSE_CROSS_MESSAGE +from onyx.prompts.chat_prompts import ( + ADDITIONAL_CONTEXT_PROMPT, + TOOL_CALL_RESPONSE_CROSS_MESSAGE, +) from onyx.prompts.tool_prompts import TOOL_CALL_FAILURE_PROMPT from onyx.server.query_and_chat.models import ChatSessionCreationRequest from onyx.server.query_and_chat.streaming_models import CitationInfo -from onyx.tools.models import ChatFile -from onyx.tools.models import ToolCallKickoff +from onyx.tools.models import ChatFile, ToolCallKickoff from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel from onyx.utils.timing import log_function_time diff --git a/backend/onyx/chat/citation_utils.py b/backend/onyx/chat/citation_utils.py index 229087bdd2..026d4fff88 100644 --- a/backend/onyx/chat/citation_utils.py +++ b/backend/onyx/chat/citation_utils.py @@ -1,7 +1,6 @@ import re -from onyx.chat.citation_processor import CitationMapping -from onyx.chat.citation_processor import DynamicCitationProcessor +from onyx.chat.citation_processor import CitationMapping, DynamicCitationProcessor from onyx.context.search.models import SearchDocsResponse from onyx.tools.built_in_tools import CITEABLE_TOOLS_NAMES from onyx.tools.models import ToolResponse diff --git a/backend/onyx/chat/compression.py b/backend/onyx/chat/compression.py index c29583c024..66db27ba0b 100644 --- a/backend/onyx/chat/compression.py +++ b/backend/onyx/chat/compression.py @@ -19,20 +19,23 @@ from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import ChatMessage from onyx.db.tools import get_tools from onyx.llm.interfaces import LLM -from onyx.llm.models import AssistantMessage -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import SystemMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ( + AssistantMessage, + ChatCompletionMessage, + SystemMessage, + UserMessage, +) from onyx.natural_language_processing.utils import get_tokenizer -from onyx.prompts.compression_prompts import PROGRESSIVE_SUMMARY_SYSTEM_PROMPT_BLOCK -from onyx.prompts.compression_prompts import PROGRESSIVE_USER_REMINDER -from onyx.prompts.compression_prompts import SUMMARIZATION_CUTOFF_MARKER -from onyx.prompts.compression_prompts import SUMMARIZATION_PROMPT -from onyx.prompts.compression_prompts import USER_REMINDER +from onyx.prompts.compression_prompts import ( + PROGRESSIVE_SUMMARY_SYSTEM_PROMPT_BLOCK, + PROGRESSIVE_USER_REMINDER, + SUMMARIZATION_CUTOFF_MARKER, + SUMMARIZATION_PROMPT, + USER_REMINDER, +) from onyx.tracing.flows import LLMFlow from onyx.tracing.framework.create import ensure_trace -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/chat/llm_loop.py b/backend/onyx/chat/llm_loop.py index 676c1f8c6a..6ca47d5014 100644 --- a/backend/onyx/chat/llm_loop.py +++ b/backend/onyx/chat/llm_loop.py @@ -1,65 +1,70 @@ import json import time from collections.abc import Callable -from typing import Any -from typing import Literal +from typing import Any, Literal from onyx.chat.chat_state import ChatStateContainer -from onyx.chat.chat_utils import build_python_chat_files_from_search_docs -from onyx.chat.chat_utils import create_tool_call_failure_messages -from onyx.chat.citation_processor import CitationMapping -from onyx.chat.citation_processor import CitationMode -from onyx.chat.citation_processor import DynamicCitationProcessor +from onyx.chat.chat_utils import ( + build_python_chat_files_from_search_docs, + create_tool_call_failure_messages, +) +from onyx.chat.citation_processor import ( + CitationMapping, + CitationMode, + DynamicCitationProcessor, +) from onyx.chat.citation_utils import update_citation_processor_from_tool_response from onyx.chat.emitter import Emitter -from onyx.chat.llm_step import _looks_like_xml_tool_call_payload -from onyx.chat.llm_step import extract_tool_calls_from_response_text -from onyx.chat.llm_step import run_llm_step -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import ContextFileMetadata -from onyx.chat.models import ExtractedContextFiles -from onyx.chat.models import FileToolMetadata -from onyx.chat.models import LlmStepResult -from onyx.chat.models import ToolCallSimple -from onyx.chat.prompt_utils import build_reminder_message -from onyx.chat.prompt_utils import build_system_prompt -from onyx.chat.prompt_utils import get_default_base_system_prompt -from onyx.chat.prompt_utils import process_prompt_template +from onyx.chat.llm_step import ( + _looks_like_xml_tool_call_payload, + extract_tool_calls_from_response_text, + run_llm_step, +) +from onyx.chat.models import ( + ChatMessageSimple, + ContextFileMetadata, + ExtractedContextFiles, + FileToolMetadata, + LlmStepResult, + ToolCallSimple, +) +from onyx.chat.prompt_utils import ( + build_reminder_message, + build_system_prompt, + get_default_base_system_prompt, + process_prompt_template, +) from onyx.configs.app_configs import INTEGRATION_TESTS_MODE from onyx.configs.chat_configs import MAX_LLM_CYCLES -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MessageType +from onyx.configs.constants import DocumentSource, MessageType from onyx.configs.model_configs import GEN_AI_INPUT_TOKEN_SAFETY_MARGIN -from onyx.context.search.models import SearchDoc -from onyx.context.search.models import SearchDocsResponse +from onyx.context.search.models import SearchDoc, SearchDocsResponse from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.memory import add_memory -from onyx.db.memory import update_memory_at_index -from onyx.db.memory import UserMemoryContext +from onyx.db.memory import add_memory, update_memory_at_index, UserMemoryContext from onyx.db.models import Persona from onyx.llm.constants import LlmProviderNames -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMUserIdentity -from onyx.llm.interfaces import ToolChoiceOptions +from onyx.llm.interfaces import LLM, LLMUserIdentity, ToolChoiceOptions from onyx.llm.model_capabilities import is_true_openai_model -from onyx.prompts.chat_prompts import IMAGE_GEN_REMINDER -from onyx.prompts.chat_prompts import OPEN_URL_REMINDER +from onyx.prompts.chat_prompts import IMAGE_GEN_REMINDER, OPEN_URL_REMINDER from onyx.prompts.prompt_utils import substitute_user_placeholders from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ToolCallDebug -from onyx.server.query_and_chat.streaming_models import TopLevelBranching -from onyx.tools.built_in_tools import CITEABLE_TOOLS_NAMES -from onyx.tools.built_in_tools import STOPPING_TOOLS_NAMES +from onyx.server.query_and_chat.streaming_models import ( + OverallStop, + Packet, + ToolCallDebug, + TopLevelBranching, +) +from onyx.tools.built_in_tools import CITEABLE_TOOLS_NAMES, STOPPING_TOOLS_NAMES from onyx.tools.interface import Tool -from onyx.tools.models import ChatFile -from onyx.tools.models import CustomToolCallSummary -from onyx.tools.models import MemoryToolResponseSnapshot -from onyx.tools.models import PythonToolRichResponse -from onyx.tools.models import ToolCallInfo -from onyx.tools.models import ToolCallKickoff -from onyx.tools.models import ToolResponse +from onyx.tools.models import ( + ChatFile, + CustomToolCallSummary, + MemoryToolResponseSnapshot, + PythonToolRichResponse, + ToolCallInfo, + ToolCallKickoff, + ToolResponse, +) from onyx.tools.tool_implementations.images.models import FinalImageGenerationResponse from onyx.tools.tool_implementations.memory.models import MemoryToolResponse from onyx.tools.tool_implementations.open_url.open_url_tool import OpenURLTool diff --git a/backend/onyx/chat/llm_step.py b/backend/onyx/chat/llm_step.py index 1b1b022985..fa5eec1f60 100644 --- a/backend/onyx/chat/llm_step.py +++ b/backend/onyx/chat/llm_step.py @@ -2,58 +2,59 @@ import json import re import time import uuid -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Mapping -from collections.abc import Sequence +from collections.abc import Callable, Generator, Mapping, Sequence from html import unescape -from typing import Any -from typing import cast +from typing import Any, cast from onyx.chat.chat_state import ChatStateContainer from onyx.chat.citation_processor import DynamicCitationProcessor from onyx.chat.emitter import Emitter -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import LlmStepResult +from onyx.chat.models import ChatMessageSimple, LlmStepResult from onyx.chat.tool_call_args_streaming import maybe_emit_argument_delta -from onyx.configs.app_configs import ENABLE_AZURE_IMAGE_CAP -from onyx.configs.app_configs import LOG_ONYX_MODEL_INTERACTIONS -from onyx.configs.app_configs import PROMPT_CACHE_CHAT_HISTORY +from onyx.configs.app_configs import ( + ENABLE_AZURE_IMAGE_CAP, + LOG_ONYX_MODEL_INTERACTIONS, + PROMPT_CACHE_CHAT_HISTORY, +) from onyx.configs.constants import MessageType from onyx.context.search.models import SearchDoc from onyx.file_store.models import ChatFileType from onyx.llm.constants import LlmProviderNames -from onyx.llm.interfaces import LanguageModelInput -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMConfig -from onyx.llm.interfaces import LLMUserIdentity -from onyx.llm.interfaces import ToolChoiceOptions +from onyx.llm.interfaces import ( + LanguageModelInput, + LLM, + LLMConfig, + LLMUserIdentity, + ToolChoiceOptions, +) from onyx.llm.model_response import Delta -from onyx.llm.models import AssistantMessage -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import FunctionCall -from onyx.llm.models import ImageContentPart -from onyx.llm.models import ImageUrlDetail -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import SystemMessage -from onyx.llm.models import TextContentPart -from onyx.llm.models import ToolCall -from onyx.llm.models import ToolMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ( + AssistantMessage, + ChatCompletionMessage, + FunctionCall, + ImageContentPart, + ImageUrlDetail, + ReasoningEffort, + SystemMessage, + TextContentPart, + ToolCall, + ToolMessage, + UserMessage, +) from onyx.llm.prompt_cache.processor import process_with_prompt_cache from onyx.llm.utils import model_needs_formatting_reenabled -from onyx.prompts.chat_prompts import CODE_BLOCK_MARKDOWN -from onyx.prompts.chat_prompts import IMAGE_DROP_REMINDER -from onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE -from onyx.prompts.constants import SYSTEM_REMINDER_TAG_OPEN +from onyx.prompts.chat_prompts import CODE_BLOCK_MARKDOWN, IMAGE_DROP_REMINDER +from onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE, SYSTEM_REMINDER_TAG_OPEN from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import CitationInfo -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ReasoningDelta -from onyx.server.query_and_chat.streaming_models import ReasoningDone -from onyx.server.query_and_chat.streaming_models import ReasoningStart +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, + CitationInfo, + Packet, + ReasoningDelta, + ReasoningDone, + ReasoningStart, +) from onyx.tools.models import ToolCallKickoff from onyx.tools.tool_name import sanitize_tool_name from onyx.tracing.flows import LLMFlow diff --git a/backend/onyx/chat/models.py b/backend/onyx/chat/models.py index 29215c4727..ed6d033e05 100644 --- a/backend/onyx/chat/models.py +++ b/backend/onyx/chat/models.py @@ -1,21 +1,22 @@ from collections.abc import Iterator -from typing import Any -from typing import Callable +from typing import Any, Callable from uuid import UUID from pydantic import BaseModel from onyx.configs.constants import MessageType from onyx.context.search.models import SearchDoc -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import InMemoryChatFile -from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.models import MultiModelMessageResponseIDInfo -from onyx.server.query_and_chat.streaming_models import CitationInfo -from onyx.server.query_and_chat.streaming_models import GeneratedImage -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.tools.models import SearchToolUsage -from onyx.tools.models import ToolCallKickoff +from onyx.file_store.models import ChatFileType, InMemoryChatFile +from onyx.server.query_and_chat.models import ( + MessageResponseIDInfo, + MultiModelMessageResponseIDInfo, +) +from onyx.server.query_and_chat.streaming_models import ( + CitationInfo, + GeneratedImage, + Packet, +) +from onyx.tools.models import SearchToolUsage, ToolCallKickoff from onyx.tools.tool_implementations.custom.base_tool_types import ToolResultType diff --git a/backend/onyx/chat/process_message.py b/backend/onyx/chat/process_message.py index fb6e509076..48a80888db 100644 --- a/backend/onyx/chat/process_message.py +++ b/backend/onyx/chat/process_message.py @@ -11,127 +11,128 @@ import re import threading import time import traceback -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from concurrent.futures import ThreadPoolExecutor from contextvars import Token from enum import Enum -from typing import cast -from typing import Final +from typing import cast, Final from uuid import UUID from sqlalchemy.orm import Session from onyx.cache.factory import get_cache_backend from onyx.chat.chat_processing_checker import set_processing_status -from onyx.chat.chat_state import AvailableFiles -from onyx.chat.chat_state import ChatStateContainer -from onyx.chat.chat_state import ChatTurnSetup -from onyx.chat.chat_utils import build_file_context -from onyx.chat.chat_utils import convert_chat_history -from onyx.chat.chat_utils import create_chat_history_chain -from onyx.chat.chat_utils import create_chat_session_from_request -from onyx.chat.chat_utils import get_custom_agent_prompt -from onyx.chat.chat_utils import is_last_assistant_message_clarification -from onyx.chat.chat_utils import load_all_chat_files -from onyx.chat.compression import calculate_total_history_tokens -from onyx.chat.compression import compress_chat_history -from onyx.chat.compression import find_summary_for_branch -from onyx.chat.compression import get_compression_params +from onyx.chat.chat_state import AvailableFiles, ChatStateContainer, ChatTurnSetup +from onyx.chat.chat_utils import ( + build_file_context, + convert_chat_history, + create_chat_history_chain, + create_chat_session_from_request, + get_custom_agent_prompt, + is_last_assistant_message_clarification, + load_all_chat_files, +) +from onyx.chat.compression import ( + calculate_total_history_tokens, + compress_chat_history, + find_summary_for_branch, + get_compression_params, +) from onyx.chat.emitter import Emitter -from onyx.chat.llm_loop import EmptyLLMResponseError -from onyx.chat.llm_loop import run_llm_loop -from onyx.chat.models import AnswerStream -from onyx.chat.models import AnswerStreamPart -from onyx.chat.models import ChatBasicResponse -from onyx.chat.models import ChatFullResponse -from onyx.chat.models import ChatLoadedFile -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import ContextFileMetadata -from onyx.chat.models import CreateChatSessionID -from onyx.chat.models import ExtractedContextFiles -from onyx.chat.models import FileToolMetadata -from onyx.chat.models import SearchParams -from onyx.chat.models import StreamingError -from onyx.chat.models import ToolCallResponse +from onyx.chat.llm_loop import EmptyLLMResponseError, run_llm_loop +from onyx.chat.models import ( + AnswerStream, + AnswerStreamPart, + ChatBasicResponse, + ChatFullResponse, + ChatLoadedFile, + ChatMessageSimple, + ContextFileMetadata, + CreateChatSessionID, + ExtractedContextFiles, + FileToolMetadata, + SearchParams, + StreamingError, + ToolCallResponse, +) from onyx.chat.prompt_utils import calculate_reserved_tokens from onyx.chat.save_chat import save_chat_turn from onyx.chat.stop_signal_checker import is_connected as check_stop_signal from onyx.chat.stop_signal_checker import reset_cancel_status from onyx.chat.stream_buffer import StreamBufferWriter -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import INTEGRATION_TESTS_MODE +from onyx.configs.app_configs import DISABLE_VECTOR_DB, INTEGRATION_TESTS_MODE from onyx.configs.chat_configs import CHAT_HEARTBEAT_INTERVAL_S -from onyx.configs.constants import DEFAULT_PERSONA_ID -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MessageType -from onyx.configs.constants import MilestoneRecordType -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import SearchDoc -from onyx.db.chat import create_new_chat_message -from onyx.db.chat import get_chat_session_by_id -from onyx.db.chat import get_or_create_root_message -from onyx.db.chat import reserve_message_id -from onyx.db.chat import reserve_multi_model_message_ids +from onyx.configs.constants import ( + DEFAULT_PERSONA_ID, + DocumentSource, + MessageType, + MilestoneRecordType, +) +from onyx.context.search.models import BaseFilters, SearchDoc +from onyx.db.chat import ( + create_new_chat_message, + get_chat_session_by_id, + get_or_create_root_message, + reserve_message_id, + reserve_multi_model_message_ids, +) from onyx.db.document_set import filter_document_set_names_by_user_access from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import HookPoint from onyx.db.memory import get_memories -from onyx.db.models import ChatMessage -from onyx.db.models import Persona -from onyx.db.models import User -from onyx.db.models import UserFile +from onyx.db.models import ChatMessage, Persona, User, UserFile from onyx.db.projects import get_user_files_from_project from onyx.db.tools import get_tools from onyx.deep_research.dr_loop import run_deep_research_llm_loop from onyx.error_handling.error_codes import OnyxErrorCode -from onyx.error_handling.exceptions import log_onyx_error -from onyx.error_handling.exceptions import OnyxError +from onyx.error_handling.exceptions import log_onyx_error, OnyxError from onyx.file_processing.extract_file_text import extract_file_text -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import InMemoryChatFile -from onyx.file_store.utils import get_default_file_store -from onyx.file_store.utils import load_in_memory_chat_files -from onyx.file_store.utils import verify_user_files -from onyx.hooks.executor import execute_hook -from onyx.hooks.executor import HookSkipped -from onyx.hooks.executor import HookSoftFailed -from onyx.hooks.points.query_processing import QueryProcessingPayload -from onyx.hooks.points.query_processing import QueryProcessingResponse -from onyx.llm.factory import get_llm_for_persona -from onyx.llm.factory import get_llm_token_counter -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMUserIdentity +from onyx.file_store.models import ChatFileType, InMemoryChatFile +from onyx.file_store.utils import ( + get_default_file_store, + load_in_memory_chat_files, + verify_user_files, +) +from onyx.hooks.executor import execute_hook, HookSkipped, HookSoftFailed +from onyx.hooks.points.query_processing import ( + QueryProcessingPayload, + QueryProcessingResponse, +) +from onyx.llm.factory import get_llm_for_persona, get_llm_token_counter +from onyx.llm.interfaces import LLM, LLMUserIdentity from onyx.llm.override_models import LLMOverride -from onyx.llm.request_context import reset_llm_mock_response -from onyx.llm.request_context import set_llm_mock_response +from onyx.llm.request_context import reset_llm_mock_response, set_llm_mock_response from onyx.llm.utils import litellm_exception_to_error_msg from onyx.onyxbot.slack.models import SlackContext from onyx.prompts.prompt_utils import substitute_user_placeholders from onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type -from onyx.server.query_and_chat.models import AUTO_PLACE_AFTER_LATEST_MESSAGE -from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.models import ModelResponseSlot -from onyx.server.query_and_chat.models import MultiModelMessageResponseIDInfo -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.server.query_and_chat.models import ( + AUTO_PLACE_AFTER_LATEST_MESSAGE, + MessageResponseIDInfo, + ModelResponseSlot, + MultiModelMessageResponseIDInfo, + SendMessageRequest, +) from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import CitationInfo -from onyx.server.query_and_chat.streaming_models import heartbeat_packet -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, + CitationInfo, + heartbeat_packet, + OverallStop, + Packet, +) from onyx.server.settings.store import load_settings from onyx.server.usage_limits import check_llm_cost_limit_for_provider from onyx.server.utils import get_json_line -from onyx.tools.constants import FILE_READER_TOOL_ID -from onyx.tools.constants import SEARCH_TOOL_ID -from onyx.tools.models import ChatFile -from onyx.tools.models import SearchToolUsage -from onyx.tools.tool_constructor import construct_tools -from onyx.tools.tool_constructor import CustomToolConfig -from onyx.tools.tool_constructor import FileReaderToolConfig -from onyx.tools.tool_constructor import SearchToolConfig +from onyx.tools.constants import FILE_READER_TOOL_ID, SEARCH_TOOL_ID +from onyx.tools.models import ChatFile, SearchToolUsage +from onyx.tools.tool_constructor import ( + construct_tools, + CustomToolConfig, + FileReaderToolConfig, + SearchToolConfig, +) from onyx.utils.logger import setup_logger from onyx.utils.telemetry import mt_cloud_telemetry from onyx.utils.timing import log_function_time diff --git a/backend/onyx/chat/prompt_utils.py b/backend/onyx/chat/prompt_utils.py index 58954fa180..1e8995826f 100644 --- a/backend/onyx/chat/prompt_utils.py +++ b/backend/onyx/chat/prompt_utils.py @@ -1,5 +1,4 @@ -from collections.abc import Callable -from collections.abc import Sequence +from collections.abc import Callable, Sequence from uuid import UUID from sqlalchemy.orm import Session @@ -8,29 +7,34 @@ from onyx.db.memory import UserMemoryContext from onyx.db.persona import get_default_behavior_persona from onyx.db.user_file import calculate_user_files_token_count from onyx.file_store.models import FileDescriptor -from onyx.prompts.chat_prompts import CITATION_REMINDER -from onyx.prompts.chat_prompts import DEFAULT_SYSTEM_PROMPT -from onyx.prompts.chat_prompts import FILE_REMINDER -from onyx.prompts.chat_prompts import LAST_CYCLE_CITATION_REMINDER -from onyx.prompts.chat_prompts import REQUIRE_CITATION_GUIDANCE -from onyx.prompts.prompt_utils import apply_prompt_placeholders -from onyx.prompts.prompt_utils import get_company_context -from onyx.prompts.tool_prompts import GENERATE_IMAGE_GUIDANCE -from onyx.prompts.tool_prompts import INTERNAL_SEARCH_GUIDANCE -from onyx.prompts.tool_prompts import MEMORY_GUIDANCE -from onyx.prompts.tool_prompts import OPEN_URLS_GUIDANCE -from onyx.prompts.tool_prompts import PYTHON_TOOL_GUIDANCE -from onyx.prompts.tool_prompts import TOOL_DESCRIPTION_SEARCH_GUIDANCE -from onyx.prompts.tool_prompts import TOOL_SECTION_HEADER -from onyx.prompts.tool_prompts import WEB_SEARCH_GUIDANCE -from onyx.prompts.tool_prompts import WEB_SEARCH_SITE_DISABLED_GUIDANCE -from onyx.prompts.user_info import BASIC_INFORMATION_PROMPT -from onyx.prompts.user_info import ORGANIZATION_PROFILE_PROMPT -from onyx.prompts.user_info import TEAM_INFORMATION_PROMPT -from onyx.prompts.user_info import USER_INFORMATION_HEADER -from onyx.prompts.user_info import USER_MEMORIES_PROMPT -from onyx.prompts.user_info import USER_PREFERENCES_PROMPT -from onyx.prompts.user_info import USER_ROLE_PROMPT +from onyx.prompts.chat_prompts import ( + CITATION_REMINDER, + DEFAULT_SYSTEM_PROMPT, + FILE_REMINDER, + LAST_CYCLE_CITATION_REMINDER, + REQUIRE_CITATION_GUIDANCE, +) +from onyx.prompts.prompt_utils import apply_prompt_placeholders, get_company_context +from onyx.prompts.tool_prompts import ( + GENERATE_IMAGE_GUIDANCE, + INTERNAL_SEARCH_GUIDANCE, + MEMORY_GUIDANCE, + OPEN_URLS_GUIDANCE, + PYTHON_TOOL_GUIDANCE, + TOOL_DESCRIPTION_SEARCH_GUIDANCE, + TOOL_SECTION_HEADER, + WEB_SEARCH_GUIDANCE, + WEB_SEARCH_SITE_DISABLED_GUIDANCE, +) +from onyx.prompts.user_info import ( + BASIC_INFORMATION_PROMPT, + ORGANIZATION_PROFILE_PROMPT, + TEAM_INFORMATION_PROMPT, + USER_INFORMATION_HEADER, + USER_MEMORIES_PROMPT, + USER_PREFERENCES_PROMPT, + USER_ROLE_PROMPT, +) from onyx.tools.interface import Tool from onyx.tools.tool_implementations.images.image_generation_tool import ( ImageGenerationTool, diff --git a/backend/onyx/chat/save_chat.py b/backend/onyx/chat/save_chat.py index 5a6393152e..fedeb36bed 100644 --- a/backend/onyx/chat/save_chat.py +++ b/backend/onyx/chat/save_chat.py @@ -3,19 +3,18 @@ import mimetypes from sqlalchemy.orm import Session -from onyx.chat.chat_state import ChatStateContainer -from onyx.chat.chat_state import SearchDocKey +from onyx.chat.chat_state import ChatStateContainer, SearchDocKey from onyx.configs.constants import DocumentSource from onyx.context.search.models import SearchDoc -from onyx.db.chat import add_search_docs_to_chat_message -from onyx.db.chat import add_search_docs_to_tool_call -from onyx.db.chat import create_db_search_doc -from onyx.db.models import ChatMessage -from onyx.db.models import ToolCall +from onyx.db.chat import ( + add_search_docs_to_chat_message, + add_search_docs_to_tool_call, + create_db_search_doc, +) +from onyx.db.models import ChatMessage, ToolCall from onyx.db.tools import create_tool_call_no_commit from onyx.file_store.models import FileDescriptor -from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.natural_language_processing.utils import get_tokenizer +from onyx.natural_language_processing.utils import BaseTokenizer, get_tokenizer from onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type from onyx.tools.models import ToolCallInfo from onyx.utils.logger import setup_logger diff --git a/backend/onyx/chat/stream_buffer.py b/backend/onyx/chat/stream_buffer.py index 968da1859e..9016a80bc8 100644 --- a/backend/onyx/chat/stream_buffer.py +++ b/backend/onyx/chat/stream_buffer.py @@ -13,13 +13,14 @@ to the DB-rendered message rather than replaying a broken sequence. import zlib from uuid import UUID -from pydantic import BaseModel -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from onyx.cache.interface import CacheBackend -from onyx.configs.chat_configs import CHAT_STREAM_BUFFER_DONE_TTL_S -from onyx.configs.chat_configs import CHAT_STREAM_BUFFER_MAX_BYTES -from onyx.configs.chat_configs import CHAT_STREAM_BUFFER_TTL_S +from onyx.configs.chat_configs import ( + CHAT_STREAM_BUFFER_DONE_TTL_S, + CHAT_STREAM_BUFFER_MAX_BYTES, + CHAT_STREAM_BUFFER_TTL_S, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/chat/tool_call_args_streaming.py b/backend/onyx/chat/tool_call_args_streaming.py index 2520042bed..b28bb12754 100644 --- a/backend/onyx/chat/tool_call_args_streaming.py +++ b/backend/onyx/chat/tool_call_args_streaming.py @@ -1,12 +1,9 @@ -from collections.abc import Generator -from collections.abc import Mapping -from typing import Any -from typing import Type +from collections.abc import Generator, Mapping +from typing import Any, Type from onyx.llm.model_response import ChatCompletionDeltaToolCall from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta +from onyx.server.query_and_chat.streaming_models import Packet, ToolCallArgumentDelta from onyx.tools.built_in_tools import TOOL_NAME_TO_CLASS from onyx.tools.interface import Tool from onyx.utils.jsonriver import Parser diff --git a/backend/onyx/configs/app_configs.py b/backend/onyx/configs/app_configs.py index 8df41fd9d1..4bacd7322a 100644 --- a/backend/onyx/configs/app_configs.py +++ b/backend/onyx/configs/app_configs.py @@ -1,8 +1,7 @@ import json import os import urllib.parse -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import cast from onyx.auth.schemas import AuthBackend @@ -10,8 +9,10 @@ from onyx.cache.interface import CacheBackendType from onyx.configs.constants import QueryHistoryType from onyx.document_index.opensearch.constants import OpenSearchAuthMethod from onyx.file_processing.enums import HtmlBasedConnectorTransformLinksStrategy -from onyx.prompts.image_analysis import DEFAULT_IMAGE_SUMMARIZATION_SYSTEM_PROMPT -from onyx.prompts.image_analysis import DEFAULT_IMAGE_SUMMARIZATION_USER_PROMPT +from onyx.prompts.image_analysis import ( + DEFAULT_IMAGE_SUMMARIZATION_SYSTEM_PROMPT, + DEFAULT_IMAGE_SUMMARIZATION_USER_PROMPT, +) from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/configs/constants.py b/backend/onyx/configs/constants.py index 6e61a6f9f2..756f55f858 100644 --- a/backend/onyx/configs/constants.py +++ b/backend/onyx/configs/constants.py @@ -2,8 +2,7 @@ import os import platform import re import socket -from enum import auto -from enum import Enum +from enum import auto, Enum ONYX_DEFAULT_APPLICATION_NAME = "Onyx" ONYX_DISCORD_URL = "https://discord.gg/4NA5SbzrWb" diff --git a/backend/onyx/configs/sentry.py b/backend/onyx/configs/sentry.py index 1a14d59003..daba2ee4ca 100644 --- a/backend/onyx/configs/sentry.py +++ b/backend/onyx/configs/sentry.py @@ -1,10 +1,8 @@ from __future__ import annotations -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING -from sentry_sdk.scrubber import DEFAULT_DENYLIST -from sentry_sdk.scrubber import EventScrubber +from sentry_sdk.scrubber import DEFAULT_DENYLIST, EventScrubber from sentry_sdk.types import Event from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/airtable/airtable_connector.py b/backend/onyx/connectors/airtable/airtable_connector.py index 8dfeb9f835..8054176d5b 100644 --- a/backend/onyx/connectors/airtable/airtable_connector.py +++ b/backend/onyx/connectors/airtable/airtable_connector.py @@ -1,30 +1,21 @@ import contextvars import re -from concurrent.futures import as_completed -from concurrent.futures import Future -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, Future, ThreadPoolExecutor from io import BytesIO -from typing import Any -from typing import cast +from typing import Any, cast import requests from pyairtable import Api as AirtableApi from pyairtable.api.types import RecordDict from pyairtable.models.schema import TableSchema -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import extract_file_text -from onyx.file_processing.extract_file_text import get_file_ext +from onyx.connectors.interfaces import GenerateDocumentsOutput, LoadConnector +from onyx.connectors.models import Document, HierarchyNode, ImageSection, TextSection +from onyx.file_processing.extract_file_text import extract_file_text, get_file_ext from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/asana/connector.py b/backend/onyx/connectors/asana/connector.py index 097006cc1e..1d2715f190 100755 --- a/backend/onyx/connectors/asana/connector.py +++ b/backend/onyx/connectors/asana/connector.py @@ -1,17 +1,16 @@ import datetime from typing import Any -from onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE -from onyx.configs.app_configs import INDEX_BATCH_SIZE +from onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE, INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.asana import asana_api -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import Document, HierarchyNode, TextSection from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/axero/connector.py b/backend/onyx/connectors/axero/connector.py index 00073926ab..78e7aa0656 100644 --- a/backend/onyx/connectors/axero/connector.py +++ b/backend/onyx/connectors/axero/connector.py @@ -1,23 +1,27 @@ import time -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import requests from pydantic import BaseModel -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource -from onyx.connectors.cross_connector_utils.miscellaneous_utils import process_in_batches -from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc +from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( + process_in_batches, + time_str_to_utc, +) from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + TextSection, +) from onyx.file_processing.html_utils import parse_html_page_basic from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/bitbucket/connector.py b/backend/onyx/connectors/bitbucket/connector.py index add7daab9d..1fbebc34b4 100644 --- a/backend/onyx/connectors/bitbucket/connector.py +++ b/backend/onyx/connectors/bitbucket/connector.py @@ -1,38 +1,42 @@ from __future__ import annotations import copy -from collections.abc import Callable -from collections.abc import Iterator -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import TYPE_CHECKING +from collections.abc import Callable, Iterator +from datetime import datetime, timezone +from typing import Any, TYPE_CHECKING from typing_extensions import override -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource -from onyx.connectors.bitbucket.utils import build_auth_client -from onyx.connectors.bitbucket.utils import list_repositories -from onyx.connectors.bitbucket.utils import map_pr_to_document -from onyx.connectors.bitbucket.utils import paginate -from onyx.connectors.bitbucket.utils import parse_bitbucket_datetime -from onyx.connectors.bitbucket.utils import PR_LIST_RESPONSE_FIELDS -from onyx.connectors.bitbucket.utils import SLIM_PR_LIST_RESPONSE_FIELDS -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument +from onyx.connectors.bitbucket.utils import ( + build_auth_client, + list_repositories, + map_pr_to_document, + paginate, + parse_bitbucket_datetime, + PR_LIST_RESPONSE_FIELDS, + SLIM_PR_LIST_RESPONSE_FIELDS, +) +from onyx.connectors.exceptions import ( + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + DocumentFailure, + HierarchyNode, + SlimDocument, +) from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/bitbucket/utils.py b/backend/onyx/connectors/bitbucket/utils.py index a883dfbd8a..4068cfad51 100644 --- a/backend/onyx/connectors/bitbucket/utils.py +++ b/backend/onyx/connectors/bitbucket/utils.py @@ -1,8 +1,7 @@ from __future__ import annotations import time -from collections.abc import Callable -from collections.abc import Iterator +from collections.abc import Callable, Iterator from datetime import datetime from typing import Any @@ -11,10 +10,7 @@ import httpx from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import Document -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import BasicExpertInfo, Document, ImageSection, TextSection from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger from onyx.utils.retry_after import parse_retry_after_seconds diff --git a/backend/onyx/connectors/blob/connector.py b/backend/onyx/connectors/blob/connector.py index 667c849dcc..0b7a0251f6 100644 --- a/backend/onyx/connectors/blob/connector.py +++ b/backend/onyx/connectors/blob/connector.py @@ -2,51 +2,48 @@ import mimetypes import os import time from collections.abc import Mapping -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from io import BytesIO from numbers import Integral -from typing import Any -from typing import Optional -from typing import TYPE_CHECKING +from typing import Any, Optional, TYPE_CHECKING from urllib.parse import quote import boto3 from botocore.client import Config from botocore.credentials import RefreshableCredentials -from botocore.exceptions import ClientError -from botocore.exceptions import NoCredentialsError -from botocore.exceptions import PartialCredentialsError +from botocore.exceptions import ClientError, NoCredentialsError, PartialCredentialsError from botocore.session import get_session -from onyx.configs.app_configs import BLOB_STORAGE_SIZE_THRESHOLD -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.constants import BlobType -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.app_configs import BLOB_STORAGE_SIZE_THRESHOLD, INDEX_BATCH_SIZE +from onyx.configs.constants import BlobType, DocumentSource, FileOrigin from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( process_onyx_metadata, ) from onyx.connectors.cross_connector_utils.tabular_section_utils import ( extract_and_stage_tabular_file, + is_tabular_file, ) -from onyx.connectors.cross_connector_utils.tabular_section_utils import is_tabular_file -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import extract_text_and_images -from onyx.file_processing.extract_file_text import get_file_ext +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + ImageSection, + TabularSection, + TextSection, +) +from onyx.file_processing.extract_file_text import extract_text_and_images, get_file_ext from onyx.file_processing.file_types import OnyxFileExtensions from onyx.file_processing.image_utils import store_image_and_create_section from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/bookstack/connector.py b/backend/onyx/connectors/bookstack/connector.py index ac86fc3c20..ba31f96952 100644 --- a/backend/onyx/connectors/bookstack/connector.py +++ b/backend/onyx/connectors/bookstack/connector.py @@ -1,26 +1,33 @@ import html import time from collections.abc import Callable -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.bookstack.client import BookStackApiClient -from onyx.connectors.bookstack.client import BookStackClientRequestFailedError +from onyx.connectors.bookstack.client import ( + BookStackApiClient, + BookStackClientRequestFailedError, +) from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.file_processing.html_utils import parse_html_page_basic diff --git a/backend/onyx/connectors/box/access.py b/backend/onyx/connectors/box/access.py index 6a8760b43e..78e26f6999 100644 --- a/backend/onyx/connectors/box/access.py +++ b/backend/onyx/connectors/box/access.py @@ -9,8 +9,8 @@ from box_sdk_gen.schemas.web_link import WebLink from onyx.access.models import ExternalAccess from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, + noop_fallback, ) -from onyx.utils.variable_functionality import noop_fallback def resolve_box_ancestor_access( diff --git a/backend/onyx/connectors/box/connector.py b/backend/onyx/connectors/box/connector.py index ea74472800..48d5a70818 100644 --- a/backend/onyx/connectors/box/connector.py +++ b/backend/onyx/connectors/box/connector.py @@ -1,15 +1,12 @@ from collections import deque from collections.abc import Generator from copy import deepcopy -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from io import BytesIO from typing import Any from urllib.parse import urlparse -from box_sdk_gen import BoxCCGAuth -from box_sdk_gen import BoxClient -from box_sdk_gen import CCGConfig +from box_sdk_gen import BoxCCGAuth, BoxClient, CCGConfig from box_sdk_gen.box.errors import BoxAPIError from box_sdk_gen.schemas.file_full import FileFull from box_sdk_gen.schemas.folder_mini import FolderMini @@ -17,41 +14,45 @@ from box_sdk_gen.schemas.user_full import UserFull from box_sdk_gen.schemas.web_link import WebLink from onyx.access.models import ExternalAccess -from onyx.configs.app_configs import BOX_CONNECTOR_SIZE_THRESHOLD -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin -from onyx.connectors.box.access import resolve_box_ancestor_access -from onyx.connectors.box.access import resolve_box_file_access -from onyx.connectors.box.access import resolve_box_folder_access -from onyx.connectors.box.access import resolve_box_web_link_access -from onyx.connectors.box.models import BoxConnectorCheckpoint -from onyx.connectors.box.models import BoxFolderFrontierEntry +from onyx.configs.app_configs import BOX_CONNECTOR_SIZE_THRESHOLD, INDEX_BATCH_SIZE +from onyx.configs.constants import DocumentSource, FileOrigin +from onyx.connectors.box.access import ( + resolve_box_ancestor_access, + resolve_box_file_access, + resolve_box_folder_access, + resolve_box_web_link_access, +) +from onyx.connectors.box.models import BoxConnectorCheckpoint, BoxFolderFrontierEntry from onyx.connectors.cross_connector_utils.miscellaneous_utils import datetime_to_utc -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import NormalizationResult -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + GenerateSlimDocumentOutput, + NormalizationResult, + SecondsSinceUnixEpoch, + SlimConnector, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + EntityFailure, + HierarchyNode, + ImageSection, + SlimDocument, + TextSection, +) from onyx.db.enums import HierarchyNodeType -from onyx.file_processing.extract_file_text import extract_text_and_images -from onyx.file_processing.extract_file_text import get_file_ext +from onyx.file_processing.extract_file_text import extract_text_and_images, get_file_ext from onyx.file_processing.file_types import OnyxFileExtensions from onyx.file_processing.image_utils import store_image_and_create_section from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface diff --git a/backend/onyx/connectors/box/models.py b/backend/onyx/connectors/box/models.py index 7241f69265..3bb5747971 100644 --- a/backend/onyx/connectors/box/models.py +++ b/backend/onyx/connectors/box/models.py @@ -1,5 +1,4 @@ -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.access.models import ExternalAccess from onyx.connectors.models import ConnectorCheckpoint diff --git a/backend/onyx/connectors/braintrust/connector.py b/backend/onyx/connectors/braintrust/connector.py index 02cea0a352..7cbf324b7e 100644 --- a/backend/onyx/connectors/braintrust/connector.py +++ b/backend/onyx/connectors/braintrust/connector.py @@ -2,11 +2,8 @@ import copy import csv import io import json -from collections.abc import Callable -from collections.abc import Mapping -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from collections.abc import Callable, Mapping +from datetime import datetime, timedelta, timezone from enum import Enum from typing import Any @@ -14,22 +11,25 @@ import requests from pydantic import BaseModel from onyx.configs.constants import DocumentSource -from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder from onyx.connectors.cross_connector_utils.rate_limit_wrapper import ( + rate_limit_builder, wrap_request_to_handle_ratelimiting, ) -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialInvalidError -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ConnectorValidationError, CredentialInvalidError +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + EntityFailure, + TabularSection, + TextSection, +) from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/canvas/access.py b/backend/onyx/connectors/canvas/access.py index e335277236..2e0b9a3bd6 100644 --- a/backend/onyx/connectors/canvas/access.py +++ b/backend/onyx/connectors/canvas/access.py @@ -1,11 +1,12 @@ from collections.abc import Callable -from typing import Any -from typing import cast +from typing import Any, cast from onyx.access.models import ExternalAccess from onyx.connectors.canvas.client import CanvasApiClient -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) def get_course_permissions( diff --git a/backend/onyx/connectors/canvas/connector.py b/backend/onyx/connectors/canvas/connector.py index fddeb761fe..b2a9840d45 100644 --- a/backend/onyx/connectors/canvas/connector.py +++ b/backend/onyx/connectors/canvas/connector.py @@ -1,41 +1,45 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from enum import StrEnum -from typing import Any -from typing import cast -from typing import NoReturn +from typing import Any, cast, NoReturn -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from typing_extensions import override from onyx.access.models import ExternalAccess from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.canvas.access import build_course_permission_context -from onyx.connectors.canvas.access import get_announcement_permissions -from onyx.connectors.canvas.access import get_assignment_permissions -from onyx.connectors.canvas.access import get_page_permissions +from onyx.connectors.canvas.access import ( + build_course_permission_context, + get_announcement_permissions, + get_assignment_permissions, + get_page_permissions, +) from onyx.connectors.canvas.client import CanvasApiClient -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + EntityFailure, + HierarchyNode, + ImageSection, + SlimDocument, + TextSection, +) from onyx.error_handling.exceptions import OnyxError from onyx.file_processing.html_utils import parse_html_page_basic from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface diff --git a/backend/onyx/connectors/clickup/connector.py b/backend/onyx/connectors/clickup/connector.py index 42ca0578a8..937785e926 100644 --- a/backend/onyx/connectors/clickup/connector.py +++ b/backend/onyx/connectors/clickup/connector.py @@ -1,23 +1,24 @@ -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import Optional +from datetime import datetime, timezone +from typing import Any, Optional import requests -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.utils.retry_wrapper import retry_builder CLICKUP_API_BASE_URL = "https://api.clickup.com/api/v2" diff --git a/backend/onyx/connectors/coda/connector.py b/backend/onyx/connectors/coda/connector.py index f35305ce49..75e5d3eb0d 100644 --- a/backend/onyx/connectors/coda/connector.py +++ b/backend/onyx/connectors/coda/connector.py @@ -1,28 +1,30 @@ import os from collections.abc import Generator from datetime import datetime -from typing import Any -from typing import cast -from typing import Dict -from typing import List -from typing import Optional +from typing import Any, cast, Dict, List, Optional from pydantic import BaseModel from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rl_requests -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + ImageSection, + TextSection, +) from onyx.utils.batching import batch_generator from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/confluence/access.py b/backend/onyx/connectors/confluence/access.py index 8c9fe6678c..51dbcb7590 100644 --- a/backend/onyx/connectors/confluence/access.py +++ b/backend/onyx/connectors/confluence/access.py @@ -1,11 +1,12 @@ from collections.abc import Callable -from typing import Any -from typing import cast +from typing import Any, cast from onyx.access.models import ExternalAccess from onyx.connectors.confluence.onyx_confluence import OnyxConfluence -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) def get_page_restrictions( diff --git a/backend/onyx/connectors/confluence/connector.py b/backend/onyx/connectors/confluence/connector.py index 667b576e8d..e79d927a04 100644 --- a/backend/onyx/connectors/confluence/connector.py +++ b/backend/onyx/connectors/confluence/connector.py @@ -1,9 +1,7 @@ import copy import re from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from urllib.parse import quote @@ -12,51 +10,63 @@ from requests.exceptions import HTTPError from typing_extensions import override from onyx.access.models import ExternalAccess -from onyx.configs.app_configs import CONFLUENCE_CONNECTOR_LABELS_TO_SKIP -from onyx.configs.app_configs import CONFLUENCE_TIMEZONE_OFFSET -from onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE -from onyx.configs.app_configs import INDEX_BATCH_SIZE +from onyx.configs.app_configs import ( + CONFLUENCE_CONNECTOR_LABELS_TO_SKIP, + CONFLUENCE_TIMEZONE_OFFSET, + CONTINUE_ON_CONNECTOR_FAILURE, + INDEX_BATCH_SIZE, +) from onyx.configs.constants import DocumentSource -from onyx.connectors.confluence.access import get_all_space_permissions -from onyx.connectors.confluence.access import get_page_restrictions from onyx.connectors.confluence.access import ( + get_all_space_permissions, + get_page_restrictions, get_page_restrictions_with_per_ancestor_fetch, ) -from onyx.connectors.confluence.onyx_confluence import Confcloud77618Error -from onyx.connectors.confluence.onyx_confluence import extract_text_from_confluence_html -from onyx.connectors.confluence.onyx_confluence import OnyxConfluence -from onyx.connectors.confluence.utils import build_confluence_document_id -from onyx.connectors.confluence.utils import convert_attachment_to_content -from onyx.connectors.confluence.utils import datetime_from_string -from onyx.connectors.confluence.utils import update_param_in_path -from onyx.connectors.confluence.utils import validate_attachment_filetype +from onyx.connectors.confluence.onyx_confluence import ( + Confcloud77618Error, + extract_text_from_confluence_html, + OnyxConfluence, +) +from onyx.connectors.confluence.utils import ( + build_confluence_document_id, + convert_attachment_to_content, + datetime_from_string, + update_param_in_path, + validate_attachment_filetype, +) from onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( is_atlassian_date_error, ) -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import ConnectorCheckpoint -from onyx.connectors.interfaces import ConnectorFailure -from onyx.connectors.interfaces import CredentialsConnector -from onyx.connectors.interfaces import CredentialsProviderInterface -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import Resolver -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + ConnectorCheckpoint, + ConnectorFailure, + CredentialsConnector, + CredentialsProviderInterface, + GenerateSlimDocumentOutput, + Resolver, + SecondsSinceUnixEpoch, + SlimConnector, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + HierarchyNode, + ImageSection, + SlimDocument, + TextSection, +) from onyx.db.enums import HierarchyNodeType from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/confluence/onyx_confluence.py b/backend/onyx/connectors/confluence/onyx_confluence.py index 18a73a65e9..e147037cca 100644 --- a/backend/onyx/connectors/confluence/onyx_confluence.py +++ b/backend/onyx/connectors/confluence/onyx_confluence.py @@ -14,15 +14,9 @@ We haven't explored all of the cloud APIs' pagination strategies. @raunakab take import json import time -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterator -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import cast -from typing import TypeVar +from collections.abc import Callable, Generator, Iterator +from datetime import datetime, timedelta, timezone +from typing import Any, cast, TypeVar from urllib.parse import quote import bs4 @@ -30,20 +24,26 @@ import requests from atlassian import Confluence from requests import HTTPError -from onyx.configs.app_configs import CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE -from onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_ID -from onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET +from onyx.configs.app_configs import ( + CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE, + OAUTH_CONFLUENCE_CLOUD_CLIENT_ID, + OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET, +) from onyx.connectors.confluence.models import ConfluenceUser from onyx.connectors.confluence.user_profile_override import ( process_confluence_user_profiles_override, ) -from onyx.connectors.confluence.utils import _handle_http_error -from onyx.connectors.confluence.utils import confluence_refresh_tokens -from onyx.connectors.confluence.utils import get_start_param_from_url -from onyx.connectors.confluence.utils import update_param_in_path +from onyx.connectors.confluence.utils import ( + _handle_http_error, + confluence_refresh_tokens, + get_start_param_from_url, + update_param_in_path, +) from onyx.connectors.cross_connector_utils.miscellaneous_utils import scoped_url -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import InsufficientPermissionsError +from onyx.connectors.exceptions import ( + ConnectorValidationError, + InsufficientPermissionsError, +) from onyx.connectors.interfaces import CredentialsProviderInterface from onyx.file_processing.html_utils import format_document_soup from onyx.redis.redis_pool import get_redis_client diff --git a/backend/onyx/connectors/confluence/utils.py b/backend/onyx/connectors/confluence/utils.py index beab4c9e6f..7024d4b080 100644 --- a/backend/onyx/connectors/confluence/utils.py +++ b/backend/onyx/connectors/confluence/utils.py @@ -1,33 +1,23 @@ import math import time from collections.abc import Callable -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from io import BytesIO from pathlib import Path -from typing import Any -from typing import cast -from typing import TYPE_CHECKING -from typing import TypeVar -from urllib.parse import parse_qs -from urllib.parse import quote -from urllib.parse import urljoin -from urllib.parse import urlparse +from typing import Any, cast, TYPE_CHECKING, TypeVar +from urllib.parse import parse_qs, quote, urljoin, urlparse import requests from pydantic import BaseModel from onyx.configs.app_configs import ( CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD, + CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD, + REQUEST_TIMEOUT_SECONDS, ) -from onyx.configs.app_configs import CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import FileOrigin -from onyx.file_processing.extract_file_text import extract_file_text -from onyx.file_processing.extract_file_text import get_file_ext -from onyx.file_processing.file_types import OnyxFileExtensions -from onyx.file_processing.file_types import OnyxMimeTypes +from onyx.file_processing.extract_file_text import extract_file_text, get_file_ext +from onyx.file_processing.file_types import OnyxFileExtensions, OnyxMimeTypes from onyx.file_processing.image_utils import store_image_and_create_section from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/connector_runner.py b/backend/onyx/connectors/connector_runner.py index 2f7e924f4e..b06d73741c 100644 --- a/backend/onyx/connectors/connector_runner.py +++ b/backend/onyx/connectors/connector_runner.py @@ -2,19 +2,22 @@ import sys import time from collections.abc import Generator from datetime import datetime -from typing import Generic -from typing import TypeVar +from typing import Generic, TypeVar -from onyx.connectors.interfaces import BaseConnector -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.interfaces import ( + BaseConnector, + CheckpointedConnector, + CheckpointedConnectorWithPermSync, + CheckpointOutput, + LoadConnector, + PollConnector, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + Document, + HierarchyNode, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/cross_connector_utils/miscellaneous_utils.py b/backend/onyx/connectors/cross_connector_utils/miscellaneous_utils.py index 563d0b5858..3d3246c6cf 100644 --- a/backend/onyx/connectors/cross_connector_utils/miscellaneous_utils.py +++ b/backend/onyx/connectors/cross_connector_utils/miscellaneous_utils.py @@ -1,22 +1,16 @@ import re -from collections.abc import Callable -from collections.abc import Iterator -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Iterator +from datetime import datetime, timezone from email.utils import parsedate_to_datetime -from typing import Any -from typing import TypeVar -from urllib.parse import urljoin -from urllib.parse import urlparse +from typing import Any, TypeVar +from urllib.parse import urljoin, urlparse import requests from dateutil.parser import parse from onyx.configs.app_configs import CONNECTOR_LOCALHOST_OVERRIDE -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import IGNORE_FOR_QA -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import OnyxMetadata +from onyx.configs.constants import DocumentSource, IGNORE_FOR_QA +from onyx.connectors.models import BasicExpertInfo, OnyxMetadata from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger from onyx.utils.text_processing import is_valid_email diff --git a/backend/onyx/connectors/cross_connector_utils/rate_limit_wrapper.py b/backend/onyx/connectors/cross_connector_utils/rate_limit_wrapper.py index acf8b793b5..6a63505b90 100644 --- a/backend/onyx/connectors/cross_connector_utils/rate_limit_wrapper.py +++ b/backend/onyx/connectors/cross_connector_utils/rate_limit_wrapper.py @@ -1,9 +1,7 @@ import time from collections.abc import Callable from functools import wraps -from typing import Any -from typing import cast -from typing import TypeVar +from typing import Any, cast, TypeVar import requests diff --git a/backend/onyx/connectors/cross_connector_utils/section_utils.py b/backend/onyx/connectors/cross_connector_utils/section_utils.py index 3c9499ab52..398128874c 100644 --- a/backend/onyx/connectors/cross_connector_utils/section_utils.py +++ b/backend/onyx/connectors/cross_connector_utils/section_utils.py @@ -1,7 +1,5 @@ from onyx.configs.app_configs import CONNECTOR_MAX_EXTRACTED_TEXT_CHARS -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import ImageSection, TabularSection, TextSection from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/cross_connector_utils/tabular_section_utils.py b/backend/onyx/connectors/cross_connector_utils/tabular_section_utils.py index 66da80a0d9..d4687812dc 100644 --- a/backend/onyx/connectors/cross_connector_utils/tabular_section_utils.py +++ b/backend/onyx/connectors/cross_connector_utils/tabular_section_utils.py @@ -5,8 +5,7 @@ from typing import IO from pydantic import BaseModel from onyx.connectors.models import TabularSection -from onyx.file_processing.extract_file_text import file_io_to_text -from onyx.file_processing.extract_file_text import stage_xlsx_sheets +from onyx.file_processing.extract_file_text import file_io_to_text, stage_xlsx_sheets from onyx.file_processing.file_types import OnyxFileExtensions from onyx.file_store.staging import RawFileCallback from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/discord/connector.py b/backend/onyx/connectors/discord/connector.py index 7bf0f48aed..c6578a0bcd 100644 --- a/backend/onyx/connectors/discord/connector.py +++ b/backend/onyx/connectors/discord/connector.py @@ -1,15 +1,10 @@ import asyncio -from collections.abc import AsyncGenerator -from collections.abc import AsyncIterable -from collections.abc import Iterable -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast +from collections.abc import AsyncGenerator, AsyncIterable, Iterable +from datetime import datetime, timezone +from typing import Any, cast from discord import Client -from discord.channel import TextChannel -from discord.channel import Thread +from discord.channel import TextChannel, Thread from discord.enums import MessageType from discord.errors import LoginFailure from discord.flags import Intents @@ -18,15 +13,19 @@ from discord.message import Message as DiscordMessage from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.exceptions import CredentialInvalidError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + ImageSection, + TextSection, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/discourse/connector.py b/backend/onyx/connectors/discourse/connector.py index 44edc1a7d2..de453cf693 100644 --- a/backend/onyx/connectors/discourse/connector.py +++ b/backend/onyx/connectors/discourse/connector.py @@ -1,28 +1,29 @@ import time import urllib.parse -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timezone +from typing import Any, cast import requests from pydantic import BaseModel from requests import Response -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + ImageSection, + TextSection, +) from onyx.file_processing.html_utils import parse_html_page_basic from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/document360/connector.py b/backend/onyx/connectors/document360/connector.py index 92004e6ae3..cff00c630d 100644 --- a/backend/onyx/connectors/document360/connector.py +++ b/backend/onyx/connectors/document360/connector.py @@ -1,25 +1,25 @@ -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import List -from typing import Optional +from datetime import datetime, timezone +from typing import Any, List, Optional import requests -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder from onyx.connectors.document360.utils import flatten_child_categories -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.file_processing.html_utils import parse_html_page_basic from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/dropbox/connector.py b/backend/onyx/connectors/dropbox/connector.py index 65e43ecf75..aa9795c13d 100644 --- a/backend/onyx/connectors/dropbox/connector.py +++ b/backend/onyx/connectors/dropbox/connector.py @@ -2,24 +2,28 @@ from io import BytesIO from typing import Any from dropbox import Dropbox -from dropbox.exceptions import ApiError -from dropbox.exceptions import AuthError -from dropbox.files import FileMetadata -from dropbox.files import FolderMetadata +from dropbox.exceptions import ApiError, AuthError +from dropbox.files import FileMetadata, FolderMetadata from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialInvalidError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialInvalidError, + InsufficientPermissionsError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.file_processing.extract_file_text import extract_file_text from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/drupal_wiki/connector.py b/backend/onyx/connectors/drupal_wiki/connector.py index 8db8c008af..8816bf0ac3 100644 --- a/backend/onyx/connectors/drupal_wiki/connector.py +++ b/backend/onyx/connectors/drupal_wiki/connector.py @@ -5,42 +5,50 @@ from typing import Any import requests from typing_extensions import override -from onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE -from onyx.configs.app_configs import DRUPAL_WIKI_ATTACHMENT_SIZE_THRESHOLD -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.app_configs import ( + CONTINUE_ON_CONNECTOR_FAILURE, + DRUPAL_WIKI_ATTACHMENT_SIZE_THRESHOLD, + INDEX_BATCH_SIZE, +) +from onyx.configs.constants import DocumentSource, FileOrigin from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( datetime_from_utc_timestamp, ) -from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rl_requests -from onyx.connectors.cross_connector_utils.tabular_section_utils import is_tabular_file +from onyx.connectors.cross_connector_utils.rate_limit_wrapper import ( + rate_limit_builder, + rl_requests, +) from onyx.connectors.cross_connector_utils.tabular_section_utils import ( + is_tabular_file, tabular_file_to_sections, ) -from onyx.connectors.drupal_wiki.models import DrupalWikiCheckpoint -from onyx.connectors.drupal_wiki.models import DrupalWikiPage -from onyx.connectors.drupal_wiki.models import DrupalWikiPageResponse -from onyx.connectors.drupal_wiki.models import DrupalWikiSpaceResponse +from onyx.connectors.drupal_wiki.models import ( + DrupalWikiCheckpoint, + DrupalWikiPage, + DrupalWikiPageResponse, + DrupalWikiSpaceResponse, +) from onyx.connectors.drupal_wiki.utils import build_drupal_wiki_document_id from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import ConnectorFailure -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import extract_text_and_images -from onyx.file_processing.extract_file_text import get_file_ext +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + ConnectorFailure, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnector, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + DocumentFailure, + HierarchyNode, + ImageSection, + SlimDocument, + TabularSection, + TextSection, +) +from onyx.file_processing.extract_file_text import extract_text_and_images, get_file_ext from onyx.file_processing.file_types import OnyxFileExtensions from onyx.file_processing.html_utils import parse_html_page_basic from onyx.file_processing.image_utils import store_image_and_create_section diff --git a/backend/onyx/connectors/drupal_wiki/models.py b/backend/onyx/connectors/drupal_wiki/models.py index a89ea4c402..66b860c6f5 100644 --- a/backend/onyx/connectors/drupal_wiki/models.py +++ b/backend/onyx/connectors/drupal_wiki/models.py @@ -1,8 +1,5 @@ from enum import Enum -from typing import Generic -from typing import List -from typing import Optional -from typing import TypeVar +from typing import Generic, List, Optional, TypeVar from pydantic import BaseModel diff --git a/backend/onyx/connectors/egnyte/connector.py b/backend/onyx/connectors/egnyte/connector.py index 123345c519..dc1b1b492e 100644 --- a/backend/onyx/connectors/egnyte/connector.py +++ b/backend/onyx/connectors/egnyte/connector.py @@ -1,35 +1,41 @@ import io import os from collections.abc import Generator -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import IO +from datetime import datetime, timezone +from typing import Any, IO from urllib.parse import quote from pydantic import Field -from onyx.configs.app_configs import EGNYTE_CLIENT_ID -from onyx.configs.app_configs import EGNYTE_CLIENT_SECRET -from onyx.configs.app_configs import INDEX_BATCH_SIZE +from onyx.configs.app_configs import ( + EGNYTE_CLIENT_ID, + EGNYTE_CLIENT_SECRET, + INDEX_BATCH_SIZE, +) from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( get_oauth_callback_uri, ) -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import OAuthConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import detect_encoding -from onyx.file_processing.extract_file_text import extract_file_text -from onyx.file_processing.extract_file_text import get_file_ext -from onyx.file_processing.extract_file_text import read_text_file +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + OAuthConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) +from onyx.file_processing.extract_file_text import ( + detect_encoding, + extract_file_text, + get_file_ext, + read_text_file, +) from onyx.file_processing.file_types import OnyxFileExtensions from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import request_with_retries diff --git a/backend/onyx/connectors/factory.py b/backend/onyx/connectors/factory.py index 186d7c3771..477a65640a 100644 --- a/backend/onyx/connectors/factory.py +++ b/backend/onyx/connectors/factory.py @@ -1,6 +1,5 @@ import importlib -from typing import Any -from typing import Type +from typing import Any, Type from sqlalchemy.orm import Session @@ -8,19 +7,19 @@ from onyx.configs.app_configs import INTEGRATION_TESTS_MODE from onyx.configs.constants import DocumentSource from onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled from onyx.connectors.credentials_provider import OnyxDBCredentialsProvider -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import ValidationError -from onyx.connectors.interfaces import BaseConnector -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CredentialsConnector -from onyx.connectors.interfaces import EventConnector -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector +from onyx.connectors.exceptions import ConnectorValidationError, ValidationError +from onyx.connectors.interfaces import ( + BaseConnector, + CheckpointedConnector, + CredentialsConnector, + EventConnector, + LoadConnector, + PollConnector, +) from onyx.connectors.models import InputType from onyx.connectors.registry import CONNECTOR_CLASS_MAP from onyx.db.connector import fetch_connector_by_id -from onyx.db.credentials import backend_update_credential_json -from onyx.db.credentials import fetch_credential_by_id +from onyx.db.credentials import backend_update_credential_json, fetch_credential_by_id from onyx.db.enums import AccessType from onyx.db.models import Credential from onyx.file_store.staging import RawFileCallback diff --git a/backend/onyx/connectors/file/connector.py b/backend/onyx/connectors/file/connector.py index 36535fbce5..a41ff38e3e 100644 --- a/backend/onyx/connectors/file/connector.py +++ b/backend/onyx/connectors/file/connector.py @@ -1,31 +1,28 @@ import json import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from io import BytesIO from pathlib import Path -from typing import Any -from typing import IO +from typing import Any, IO from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.constants import DocumentSource, FileOrigin from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( process_onyx_metadata, ) -from onyx.connectors.cross_connector_utils.tabular_section_utils import is_tabular_file from onyx.connectors.cross_connector_utils.tabular_section_utils import ( + is_tabular_file, tabular_file_to_sections, ) -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import extract_text_and_images -from onyx.file_processing.extract_file_text import get_file_ext +from onyx.connectors.interfaces import GenerateDocumentsOutput, LoadConnector +from onyx.connectors.models import ( + Document, + HierarchyNode, + ImageSection, + TabularSection, + TextSection, +) +from onyx.file_processing.extract_file_text import extract_text_and_images, get_file_ext from onyx.file_processing.file_types import OnyxFileExtensions from onyx.file_processing.image_utils import store_image_and_create_section from onyx.file_store.file_store import get_default_file_store diff --git a/backend/onyx/connectors/fireflies/connector.py b/backend/onyx/connectors/fireflies/connector.py index df31716feb..cdc38907e0 100644 --- a/backend/onyx/connectors/fireflies/connector.py +++ b/backend/onyx/connectors/fireflies/connector.py @@ -1,25 +1,26 @@ import time from collections.abc import Iterator -from datetime import datetime -from datetime import timezone -from typing import cast -from typing import List +from datetime import datetime, timezone +from typing import cast, List import requests -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + ImageSection, + TextSection, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/freshdesk/connector.py b/backend/onyx/connectors/freshdesk/connector.py index afac6a861b..db1278da42 100644 --- a/backend/onyx/connectors/freshdesk/connector.py +++ b/backend/onyx/connectors/freshdesk/connector.py @@ -1,7 +1,6 @@ import json from collections.abc import Iterator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import List import requests @@ -9,14 +8,18 @@ import requests from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rl_requests -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.file_processing.html_utils import parse_html_page_basic from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/gitbook/connector.py b/backend/onyx/connectors/gitbook/connector.py index a425df260f..20e67bc2ab 100644 --- a/backend/onyx/connectors/gitbook/connector.py +++ b/backend/onyx/connectors/gitbook/connector.py @@ -1,21 +1,23 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from urllib.parse import urljoin import requests -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/github/connector.py b/backend/onyx/connectors/github/connector.py index 26d0a87707..972b27713e 100644 --- a/backend/onyx/connectors/github/connector.py +++ b/backend/onyx/connectors/github/connector.py @@ -1,20 +1,13 @@ import copy import os -from collections.abc import Callable -from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from collections.abc import Callable, Generator +from datetime import datetime, timedelta, timezone from enum import Enum from io import BytesIO -from typing import Any -from typing import cast +from typing import Any, cast -from github import Github -from github import RateLimitExceededException -from github import Repository -from github.GithubException import GithubException -from github.GithubException import UnknownObjectException +from github import Github, RateLimitExceededException, Repository +from github.GithubException import GithubException, UnknownObjectException from github.Issue import Issue from github.NamedUser import NamedUser from github.PaginatedList import PaginatedList @@ -25,35 +18,41 @@ from typing_extensions import override from onyx.access.models import ExternalAccess from onyx.configs.app_configs import GITHUB_CONNECTOR_BASE_URL from onyx.configs.constants import DocumentSource -from onyx.connectors.connector_runner import CheckpointOutputWrapper -from onyx.connectors.connector_runner import ConnectorRunner -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.exceptions import ValidationError +from onyx.connectors.connector_runner import CheckpointOutputWrapper, ConnectorRunner +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, + ValidationError, +) from onyx.connectors.github.models import SerializedRepository from onyx.connectors.github.rate_limit_utils import sleep_after_rate_limit_exception -from onyx.connectors.github.utils import deserialize_repository -from onyx.connectors.github.utils import get_external_access_permission -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import ConnectorCheckpoint -from onyx.connectors.interfaces import ConnectorFailure -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import IndexingHeartbeatInterface -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import file_io_to_text -from onyx.file_processing.extract_file_text import is_text_file +from onyx.connectors.github.utils import ( + deserialize_repository, + get_external_access_permission, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + ConnectorCheckpoint, + ConnectorFailure, + GenerateSlimDocumentOutput, + IndexingHeartbeatInterface, + SecondsSinceUnixEpoch, + SlimConnector, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + DocumentFailure, + EntityFailure, + HierarchyNode, + SlimDocument, + TextSection, +) +from onyx.file_processing.extract_file_text import file_io_to_text, is_text_file from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/github/rate_limit_utils.py b/backend/onyx/connectors/github/rate_limit_utils.py index c58f6c7c8e..e89cdf0d5d 100644 --- a/backend/onyx/connectors/github/rate_limit_utils.py +++ b/backend/onyx/connectors/github/rate_limit_utils.py @@ -1,7 +1,5 @@ import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from github import Github diff --git a/backend/onyx/connectors/github/utils.py b/backend/onyx/connectors/github/utils.py index 90e31b1500..b215df126f 100644 --- a/backend/onyx/connectors/github/utils.py +++ b/backend/onyx/connectors/github/utils.py @@ -7,8 +7,10 @@ from github.Repository import Repository from onyx.access.models import ExternalAccess from onyx.connectors.github.models import SerializedRepository from onyx.utils.logger import setup_logger -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) logger = setup_logger() diff --git a/backend/onyx/connectors/gitlab/connector.py b/backend/onyx/connectors/gitlab/connector.py index b05d09bc78..7e7be8ce9d 100644 --- a/backend/onyx/connectors/gitlab/connector.py +++ b/backend/onyx/connectors/gitlab/connector.py @@ -1,30 +1,33 @@ import fnmatch import itertools from collections import deque -from collections.abc import Iterable -from collections.abc import Iterator -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import TypeVar +from collections.abc import Iterable, Iterator +from datetime import datetime, timezone +from typing import Any, TypeVar import gitlab import pytz from gitlab.v4.objects import Project -from onyx.configs.app_configs import GITLAB_CONNECTOR_INCLUDE_CODE_FILES -from onyx.configs.app_configs import INDEX_BATCH_SIZE +from onyx.configs.app_configs import ( + GITLAB_CONNECTOR_INCLUDE_CODE_FILES, + INDEX_BATCH_SIZE, +) from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/gmail/connector.py b/backend/onyx/connectors/gmail/connector.py index e443137a0d..6fe544c110 100644 --- a/backend/onyx/connectors/gmail/connector.py +++ b/backend/onyx/connectors/gmail/connector.py @@ -1,9 +1,6 @@ from base64 import urlsafe_b64decode -from collections.abc import Callable -from collections.abc import Iterator -from typing import Any -from typing import cast -from typing import Dict +from collections.abc import Callable, Iterator +from typing import Any, cast, Dict from google.oauth2.credentials import Credentials as OAuthCredentials from google.oauth2.service_account import Credentials as ServiceAccountCredentials @@ -14,36 +11,42 @@ from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc from onyx.connectors.google_utils.google_auth import get_google_creds -from onyx.connectors.google_utils.google_utils import execute_paginated_retrieval from onyx.connectors.google_utils.google_utils import ( + execute_paginated_retrieval, execute_paginated_retrieval_with_max_pages, + execute_single_retrieval, + PAGE_TOKEN_KEY, +) +from onyx.connectors.google_utils.resources import ( + get_admin_service, + get_gmail_service, + GmailService, ) -from onyx.connectors.google_utils.google_utils import execute_single_retrieval -from onyx.connectors.google_utils.google_utils import PAGE_TOKEN_KEY -from onyx.connectors.google_utils.resources import get_admin_service -from onyx.connectors.google_utils.resources import get_gmail_service -from onyx.connectors.google_utils.resources import GmailService from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_PRIMARY_ADMIN_KEY, + MISSING_SCOPES_ERROR_STR, + ONYX_SCOPE_INSTRUCTIONS, + SLIM_BATCH_SIZE, + USER_FIELDS, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + ConnectorFailure, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorCheckpoint, + Document, + DocumentFailure, + HierarchyNode, + ImageSection, + SlimDocument, + TextSection, ) -from onyx.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR -from onyx.connectors.google_utils.shared_constants import ONYX_SCOPE_INSTRUCTIONS -from onyx.connectors.google_utils.shared_constants import SLIM_BATCH_SIZE -from onyx.connectors.google_utils.shared_constants import USER_FIELDS -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import ConnectorFailure -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/gong/connector.py b/backend/onyx/connectors/gong/connector.py index 152e4194ae..d411808ae8 100644 --- a/backend/onyx/connectors/gong/connector.py +++ b/backend/onyx/connectors/gong/connector.py @@ -2,11 +2,8 @@ import base64 import copy import time from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timedelta, timezone +from typing import Any, cast from urllib.parse import urlparse import requests @@ -16,15 +13,19 @@ from urllib3.util import Retry from onyx.configs.app_configs import GONG_CONNECTOR_START_TIME from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + TextSection, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/google_drive/connector.py b/backend/onyx/connectors/google_drive/connector.py index 4cf8cbdbc7..f97a12916f 100644 --- a/backend/onyx/connectors/google_drive/connector.py +++ b/backend/onyx/connectors/google_drive/connector.py @@ -3,17 +3,11 @@ import json import os import sys import threading -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Callable, Generator, Iterator from datetime import datetime from enum import Enum -from typing import Any -from typing import cast -from typing import Protocol -from urllib.parse import parse_qs -from urllib.parse import ParseResult -from urllib.parse import urlparse +from typing import Any, cast, Protocol +from urllib.parse import parse_qs, ParseResult, urlparse from google.auth.exceptions import RefreshError from google.oauth2.credentials import Credentials as OAuthCredentials @@ -22,84 +16,96 @@ from googleapiclient.errors import HttpError from typing_extensions import override from onyx.access.models import ExternalAccess -from onyx.configs.app_configs import GOOGLE_DRIVE_CONNECTOR_SIZE_THRESHOLD -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import MAX_DRIVE_WORKERS +from onyx.configs.app_configs import ( + GOOGLE_DRIVE_CONNECTOR_SIZE_THRESHOLD, + INDEX_BATCH_SIZE, + MAX_DRIVE_WORKERS, +) from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) from onyx.connectors.google_drive.doc_conversion import ( _FALLBACK_BINARY_WEB_VIEW_LINK_TEMPLATE, -) -from onyx.connectors.google_drive.doc_conversion import ( _FALLBACK_WEB_VIEW_LINK_TEMPLATES, + build_slim_document, + convert_drive_item_to_document, + onyx_document_id_from_drive_file, + PermissionSyncContext, + WEB_VIEW_LINK_KEY, ) -from onyx.connectors.google_drive.doc_conversion import build_slim_document -from onyx.connectors.google_drive.doc_conversion import convert_drive_item_to_document -from onyx.connectors.google_drive.doc_conversion import onyx_document_id_from_drive_file -from onyx.connectors.google_drive.doc_conversion import PermissionSyncContext -from onyx.connectors.google_drive.doc_conversion import WEB_VIEW_LINK_KEY -from onyx.connectors.google_drive.file_retrieval import crawl_folders_for_files -from onyx.connectors.google_drive.file_retrieval import DriveFileFieldType -from onyx.connectors.google_drive.file_retrieval import get_all_files_for_oauth from onyx.connectors.google_drive.file_retrieval import ( + crawl_folders_for_files, + DriveFileFieldType, + get_all_files_for_oauth, get_all_files_in_my_drive_and_shared, -) -from onyx.connectors.google_drive.file_retrieval import get_external_access_for_folder -from onyx.connectors.google_drive.file_retrieval import ( + get_external_access_for_folder, get_files_by_web_view_links_batch, + get_files_in_shared_drive, + get_folder_metadata, + get_root_folder_id, + get_shared_drive_name, + has_link_only_permission, +) +from onyx.connectors.google_drive.models import ( + DriveRetrievalStage, + GoogleDriveCheckpoint, + GoogleDriveFileType, + RetrievedDriveFile, + StageCompletion, ) -from onyx.connectors.google_drive.file_retrieval import get_files_in_shared_drive -from onyx.connectors.google_drive.file_retrieval import get_folder_metadata -from onyx.connectors.google_drive.file_retrieval import get_root_folder_id -from onyx.connectors.google_drive.file_retrieval import get_shared_drive_name -from onyx.connectors.google_drive.file_retrieval import has_link_only_permission -from onyx.connectors.google_drive.models import DriveRetrievalStage -from onyx.connectors.google_drive.models import GoogleDriveCheckpoint -from onyx.connectors.google_drive.models import GoogleDriveFileType -from onyx.connectors.google_drive.models import RetrievedDriveFile -from onyx.connectors.google_drive.models import StageCompletion from onyx.connectors.google_utils.google_auth import get_google_creds -from onyx.connectors.google_utils.google_utils import execute_paginated_retrieval -from onyx.connectors.google_utils.google_utils import get_file_owners -from onyx.connectors.google_utils.google_utils import GoogleFields -from onyx.connectors.google_utils.resources import get_admin_service -from onyx.connectors.google_utils.resources import get_drive_service -from onyx.connectors.google_utils.resources import GoogleDriveService -from onyx.connectors.google_utils.resources import ImpersonationError -from onyx.connectors.google_utils.resources import make_user_removal_checker +from onyx.connectors.google_utils.google_utils import ( + execute_paginated_retrieval, + get_file_owners, + GoogleFields, +) +from onyx.connectors.google_utils.resources import ( + get_admin_service, + get_drive_service, + GoogleDriveService, + ImpersonationError, + make_user_removal_checker, +) from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_PRIMARY_ADMIN_KEY, + MISSING_SCOPES_ERROR_STR, + ONYX_SCOPE_INSTRUCTIONS, + SLIM_BATCH_SIZE, + USER_FIELDS, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + GenerateSlimDocumentOutput, + NormalizationResult, + Resolver, + SecondsSinceUnixEpoch, + SlimConnector, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + EntityFailure, + HierarchyNode, + SlimDocument, ) -from onyx.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR -from onyx.connectors.google_utils.shared_constants import ONYX_SCOPE_INSTRUCTIONS -from onyx.connectors.google_utils.shared_constants import SLIM_BATCH_SIZE -from onyx.connectors.google_utils.shared_constants import USER_FIELDS -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import NormalizationResult -from onyx.connectors.interfaces import Resolver -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument from onyx.db.enums import HierarchyNodeType from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.batching import batch_generator from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder -from onyx.utils.threadpool_concurrency import parallel_yield -from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel -from onyx.utils.threadpool_concurrency import ThreadSafeDict -from onyx.utils.threadpool_concurrency import ThreadSafeSet +from onyx.utils.threadpool_concurrency import ( + parallel_yield, + run_functions_tuples_in_parallel, + ThreadSafeDict, + ThreadSafeSet, +) logger = setup_logger() # TODO: Improve this by using the batch utility: https://googleapis.github.io/google-api-python-client/docs/batch.html diff --git a/backend/onyx/connectors/google_drive/doc_conversion.py b/backend/onyx/connectors/google_drive/doc_conversion.py index 779da77a96..26950cfc31 100644 --- a/backend/onyx/connectors/google_drive/doc_conversion.py +++ b/backend/onyx/connectors/google_drive/doc_conversion.py @@ -1,10 +1,8 @@ import io from collections.abc import Callable from datetime import datetime -from typing import Any -from typing import cast -from urllib.parse import urlparse -from urllib.parse import urlunparse +from typing import Any, cast +from urllib.parse import urlparse, urlunparse from googleapiclient.errors import HttpError from googleapiclient.http import MediaIoBaseDownload @@ -12,48 +10,62 @@ from pydantic import BaseModel from onyx.access.models import ExternalAccess from onyx.configs.app_configs import GOOGLE_DRIVE_ADVANCED_PARSE_MAX_BYTES -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.constants import DocumentSource, FileOrigin from onyx.connectors.cross_connector_utils.section_utils import cap_sections_text from onyx.connectors.cross_connector_utils.tabular_section_utils import ( extract_and_stage_tabular_file, + is_tabular_file, +) +from onyx.connectors.google_drive.constants import ( + DRIVE_FOLDER_TYPE, + DRIVE_SHORTCUT_TYPE, +) +from onyx.connectors.google_drive.file_retrieval import ( + add_drive_resource_key_header, + DRIVE_RESOURCE_KEY_FIELD, +) +from onyx.connectors.google_drive.models import GDriveMimeType, GoogleDriveFileType +from onyx.connectors.google_drive.section_extraction import ( + get_document_sections, + HEADING_DELIMITER, +) +from onyx.connectors.google_utils.resources import ( + get_drive_service, + get_google_authorized_session, + GoogleDriveService, +) +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + ImageSection, + SlimDocument, + TabularSection, + TextSection, +) +from onyx.file_processing.extract_file_text import ( + extract_file_text, + get_file_ext, + read_docx_file, + read_pdf_file, + read_pptx_file, +) +from onyx.file_processing.file_types import ( + OnyxFileExtensions, + OnyxMimeTypes, + PRESENTATION_MIME_TYPE, + SPREADSHEET_MIME_TYPE, +) +from onyx.file_processing.image_utils import ( + make_image_callback, + store_image_and_create_section, ) -from onyx.connectors.cross_connector_utils.tabular_section_utils import is_tabular_file -from onyx.connectors.google_drive.constants import DRIVE_FOLDER_TYPE -from onyx.connectors.google_drive.constants import DRIVE_SHORTCUT_TYPE -from onyx.connectors.google_drive.file_retrieval import add_drive_resource_key_header -from onyx.connectors.google_drive.file_retrieval import DRIVE_RESOURCE_KEY_FIELD -from onyx.connectors.google_drive.models import GDriveMimeType -from onyx.connectors.google_drive.models import GoogleDriveFileType -from onyx.connectors.google_drive.section_extraction import get_document_sections -from onyx.connectors.google_drive.section_extraction import HEADING_DELIMITER -from onyx.connectors.google_utils.resources import get_drive_service -from onyx.connectors.google_utils.resources import get_google_authorized_session -from onyx.connectors.google_utils.resources import GoogleDriveService -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import ImageSection -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import extract_file_text -from onyx.file_processing.extract_file_text import get_file_ext -from onyx.file_processing.extract_file_text import read_docx_file -from onyx.file_processing.extract_file_text import read_pdf_file -from onyx.file_processing.extract_file_text import read_pptx_file -from onyx.file_processing.file_types import OnyxFileExtensions -from onyx.file_processing.file_types import OnyxMimeTypes -from onyx.file_processing.file_types import PRESENTATION_MIME_TYPE -from onyx.file_processing.file_types import SPREADSHEET_MIME_TYPE -from onyx.file_processing.image_utils import make_image_callback -from onyx.file_processing.image_utils import store_image_and_create_section from onyx.file_store.staging import RawFileCallback from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, + noop_fallback, ) -from onyx.utils.variable_functionality import noop_fallback logger = setup_logger() diff --git a/backend/onyx/connectors/google_drive/file_retrieval.py b/backend/onyx/connectors/google_drive/file_retrieval.py index 164eb6de4c..cf78261ac4 100644 --- a/backend/onyx/connectors/google_drive/file_retrieval.py +++ b/backend/onyx/connectors/google_drive/file_retrieval.py @@ -1,12 +1,8 @@ -from collections.abc import Callable -from collections.abc import Iterator -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Iterator +from datetime import datetime, timezone from enum import Enum -from typing import Any -from typing import cast -from urllib.parse import parse_qs -from urllib.parse import urlparse +from typing import Any, cast +from urllib.parse import parse_qs, urlparse from google.auth.exceptions import RefreshError from googleapiclient.discovery import Resource @@ -14,25 +10,29 @@ from googleapiclient.errors import HttpError from googleapiclient.http import BatchHttpRequest from onyx.access.models import ExternalAccess -from onyx.connectors.google_drive.constants import DRIVE_FOLDER_TYPE -from onyx.connectors.google_drive.constants import DRIVE_SHORTCUT_TYPE -from onyx.connectors.google_drive.models import DriveRetrievalStage -from onyx.connectors.google_drive.models import GoogleDriveFileType -from onyx.connectors.google_drive.models import RetrievedDriveFile -from onyx.connectors.google_utils.google_utils import execute_paginated_retrieval +from onyx.connectors.google_drive.constants import ( + DRIVE_FOLDER_TYPE, + DRIVE_SHORTCUT_TYPE, +) +from onyx.connectors.google_drive.models import ( + DriveRetrievalStage, + GoogleDriveFileType, + RetrievedDriveFile, +) from onyx.connectors.google_utils.google_utils import ( + execute_paginated_retrieval, execute_paginated_retrieval_with_max_pages, + GoogleFields, + ORDER_BY_KEY, + PAGE_TOKEN_KEY, ) -from onyx.connectors.google_utils.google_utils import GoogleFields -from onyx.connectors.google_utils.google_utils import ORDER_BY_KEY -from onyx.connectors.google_utils.google_utils import PAGE_TOKEN_KEY from onyx.connectors.google_utils.resources import GoogleDriveService from onyx.connectors.interfaces import SecondsSinceUnixEpoch from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, + noop_fallback, ) -from onyx.utils.variable_functionality import noop_fallback logger = setup_logger() diff --git a/backend/onyx/connectors/google_drive/models.py b/backend/onyx/connectors/google_drive/models.py index ce831c82c1..cd4131797e 100644 --- a/backend/onyx/connectors/google_drive/models.py +++ b/backend/onyx/connectors/google_drive/models.py @@ -1,16 +1,10 @@ from enum import Enum from typing import Any -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_serializer -from pydantic import field_validator +from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator -from onyx.connectors.interfaces import ConnectorCheckpoint -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.utils.threadpool_concurrency import ThreadSafeDict -from onyx.utils.threadpool_concurrency import ThreadSafeSet +from onyx.connectors.interfaces import ConnectorCheckpoint, SecondsSinceUnixEpoch +from onyx.utils.threadpool_concurrency import ThreadSafeDict, ThreadSafeSet class GDriveMimeType(str, Enum): diff --git a/backend/onyx/connectors/google_site/connector.py b/backend/onyx/connectors/google_site/connector.py index 260fc2f83e..82f394f305 100644 --- a/backend/onyx/connectors/google_site/connector.py +++ b/backend/onyx/connectors/google_site/connector.py @@ -1,20 +1,14 @@ import os import re -from typing import Any -from typing import cast +from typing import Any, cast -from bs4 import BeautifulSoup -from bs4 import Tag +from bs4 import BeautifulSoup, Tag from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection -from onyx.file_processing.extract_file_text import load_files_from_zip -from onyx.file_processing.extract_file_text import read_text_file +from onyx.connectors.interfaces import GenerateDocumentsOutput, LoadConnector +from onyx.connectors.models import Document, HierarchyNode, TextSection +from onyx.file_processing.extract_file_text import load_files_from_zip, read_text_file from onyx.file_processing.html_utils import web_html_cleanup from onyx.file_store.file_store import get_default_file_store from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/google_utils/google_auth.py b/backend/onyx/connectors/google_utils/google_auth.py index e97bdbf8fb..372dd8541a 100644 --- a/backend/onyx/connectors/google_utils/google_auth.py +++ b/backend/onyx/connectors/google_utils/google_auth.py @@ -5,21 +5,17 @@ from google.auth.transport.requests import Request from google.oauth2.credentials import Credentials as OAuthCredentials from google.oauth2.service_account import Credentials as ServiceAccountCredentials -from onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_ID -from onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_SECRET +from onyx.configs.app_configs import ( + OAUTH_GOOGLE_DRIVE_CLIENT_ID, + OAUTH_GOOGLE_DRIVE_CLIENT_SECRET, +) from onyx.configs.constants import DocumentSource from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, -) -from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY -from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_DICT_TOKEN_KEY, DB_CREDENTIALS_PRIMARY_ADMIN_KEY, -) -from onyx.connectors.google_utils.shared_constants import GOOGLE_SCOPES -from onyx.connectors.google_utils.shared_constants import ( + GOOGLE_SCOPES, GoogleOAuthAuthenticationMethod, ) from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/google_utils/google_kv.py b/backend/onyx/connectors/google_utils/google_kv.py index 363cc14f5e..03831f8c7d 100644 --- a/backend/onyx/connectors/google_utils/google_kv.py +++ b/backend/onyx/connectors/google_utils/google_kv.py @@ -1,52 +1,37 @@ import json -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import cast -from urllib.parse import parse_qs -from urllib.parse import ParseResult -from urllib.parse import urlparse +from datetime import datetime, timedelta, timezone +from typing import Any, cast +from urllib.parse import parse_qs, ParseResult, urlparse from google.oauth2.credentials import Credentials as OAuthCredentials from google_auth_oauthlib.flow import InstalledAppFlow from sqlalchemy.orm import Session from onyx.configs.app_configs import WEB_DOMAIN -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import KV_CRED_KEY -from onyx.connectors.google_utils.resources import get_drive_service -from onyx.connectors.google_utils.resources import get_gmail_service +from onyx.configs.constants import DocumentSource, KV_CRED_KEY +from onyx.connectors.google_utils.resources import get_drive_service, get_gmail_service from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_APP_CREDENTIAL_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, -) -from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY -from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_DICT_TOKEN_KEY, DB_CREDENTIALS_PRIMARY_ADMIN_KEY, -) -from onyx.connectors.google_utils.shared_constants import GOOGLE_SCOPES -from onyx.connectors.google_utils.shared_constants import ( + GOOGLE_SCOPES, GoogleOAuthAuthenticationMethod, + MISSING_SCOPES_ERROR_STR, + ONYX_SCOPE_INSTRUCTIONS, +) +from onyx.db.credentials import fetch_credential_by_id_for_user, update_credential_json +from onyx.db.encrypted_kv_store import ( + delete_encrypted_kv, + load_encrypted_kv, + upsert_encrypted_kv, ) -from onyx.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR -from onyx.connectors.google_utils.shared_constants import ONYX_SCOPE_INSTRUCTIONS -from onyx.db.credentials import fetch_credential_by_id_for_user -from onyx.db.credentials import update_credential_json -from onyx.db.encrypted_kv_store import delete_encrypted_kv -from onyx.db.encrypted_kv_store import load_encrypted_kv -from onyx.db.encrypted_kv_store import upsert_encrypted_kv from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.key_value_store.interface import KvKeyNotFoundError -from onyx.server.documents.models import CredentialBase -from onyx.server.documents.models import GoogleServiceAccountKey +from onyx.server.documents.models import CredentialBase, GoogleServiceAccountKey from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/google_utils/google_utils.py b/backend/onyx/connectors/google_utils/google_utils.py index 1584d07683..5b7af6234c 100644 --- a/backend/onyx/connectors/google_utils/google_utils.py +++ b/backend/onyx/connectors/google_utils/google_utils.py @@ -2,10 +2,8 @@ import re import socket import ssl import time -from collections.abc import Callable -from collections.abc import Iterator -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Iterator +from datetime import datetime, timezone from enum import Enum from typing import Any diff --git a/backend/onyx/connectors/google_utils/resources.py b/backend/onyx/connectors/google_utils/resources.py index d17856c27f..79a846eaa9 100644 --- a/backend/onyx/connectors/google_utils/resources.py +++ b/backend/onyx/connectors/google_utils/resources.py @@ -4,8 +4,7 @@ from google.auth.exceptions import RefreshError from google.auth.transport.requests import AuthorizedSession from google.oauth2.credentials import Credentials as OAuthCredentials from google.oauth2.service_account import Credentials as ServiceAccountCredentials -from googleapiclient.discovery import build -from googleapiclient.discovery import Resource +from googleapiclient.discovery import build, Resource from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/guru/connector.py b/backend/onyx/connectors/guru/connector.py index c170abe131..35aa9bf888 100644 --- a/backend/onyx/connectors/guru/connector.py +++ b/backend/onyx/connectors/guru/connector.py @@ -1,6 +1,5 @@ import json -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import requests @@ -8,15 +7,19 @@ import requests from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.file_processing.html_utils import parse_html_page_basic from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/highspot/client.py b/backend/onyx/connectors/highspot/client.py index 1bc40353cf..f605f4a9d3 100644 --- a/backend/onyx/connectors/highspot/client.py +++ b/backend/onyx/connectors/highspot/client.py @@ -1,15 +1,10 @@ import base64 -from typing import Any -from typing import Dict -from typing import List -from typing import Optional +from typing import Any, Dict, List, Optional from urllib.parse import urljoin import requests from requests.adapters import HTTPAdapter -from requests.exceptions import HTTPError -from requests.exceptions import RequestException -from requests.exceptions import Timeout +from requests.exceptions import HTTPError, RequestException, Timeout from urllib3.util.retry import Retry from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/highspot/connector.py b/backend/onyx/connectors/highspot/connector.py index 2681abd802..0eb67c23d4 100644 --- a/backend/onyx/connectors/highspot/connector.py +++ b/backend/onyx/connectors/highspot/connector.py @@ -1,29 +1,29 @@ import os from datetime import datetime from io import BytesIO -from typing import Any -from typing import Dict -from typing import List -from typing import Optional +from typing import Any, Dict, List, Optional from pydantic import BaseModel from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.highspot.client import HighspotClient -from onyx.connectors.highspot.client import HighspotClientError +from onyx.connectors.highspot.client import HighspotClient, HighspotClientError from onyx.connectors.highspot.utils import scrape_url_content -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + GenerateSlimDocumentOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + SlimDocument, + TextSection, +) from onyx.file_processing.extract_file_text import extract_file_text from onyx.file_processing.file_types import OnyxFileExtensions from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface diff --git a/backend/onyx/connectors/hubspot/connector.py b/backend/onyx/connectors/hubspot/connector.py index 2dbc8ce658..f382a177fc 100644 --- a/backend/onyx/connectors/hubspot/connector.py +++ b/backend/onyx/connectors/hubspot/connector.py @@ -1,12 +1,8 @@ import re import time -from collections.abc import Callable -from collections.abc import Generator -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast -from typing import TypeVar +from collections.abc import Callable, Generator +from datetime import datetime, timezone +from typing import Any, cast, TypeVar import requests from hubspot import HubSpot @@ -17,9 +13,7 @@ from hubspot.crm.companies.models import SimplePublicObjectId as CompanyObjectId from hubspot.crm.contacts.models import ( BatchReadInputSimplePublicObjectId as ContactsBatchReadInput, ) -from hubspot.crm.contacts.models import Filter -from hubspot.crm.contacts.models import FilterGroup -from hubspot.crm.contacts.models import PublicObjectSearchRequest +from hubspot.crm.contacts.models import Filter, FilterGroup, PublicObjectSearchRequest from hubspot.crm.contacts.models import SimplePublicObjectId as ContactObjectId from hubspot.crm.deals.models import ( BatchReadInputSimplePublicObjectId as DealsBatchReadInput, @@ -34,19 +28,22 @@ from hubspot.crm.tickets.models import ( ) from hubspot.crm.tickets.models import SimplePublicObjectId as TicketObjectId -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource from onyx.connectors.hubspot.rate_limit import HubSpotRateLimiter -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + ImageSection, + TextSection, +) from onyx.utils.logger import setup_logger HUBSPOT_BASE_URL = "https://app.hubspot.com" diff --git a/backend/onyx/connectors/hubspot/rate_limit.py b/backend/onyx/connectors/hubspot/rate_limit.py index af11c73c57..43527f258f 100644 --- a/backend/onyx/connectors/hubspot/rate_limit.py +++ b/backend/onyx/connectors/hubspot/rate_limit.py @@ -2,11 +2,10 @@ from __future__ import annotations import time from collections.abc import Callable -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar -from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder from onyx.connectors.cross_connector_utils.rate_limit_wrapper import ( + rate_limit_builder, RateLimitTriedTooManyTimesError, ) from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/imap/connector.py b/backend/onyx/connectors/imap/connector.py index 4e6741b607..34ef9e82e3 100644 --- a/backend/onyx/connectors/imap/connector.py +++ b/backend/onyx/connectors/imap/connector.py @@ -3,13 +3,11 @@ import email import imaplib import os import re -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from email.message import Message from email.utils import parseaddr from enum import Enum -from typing import Any -from typing import cast +from typing import Any, cast import bs4 from pydantic import BaseModel @@ -17,15 +15,19 @@ from pydantic import BaseModel from onyx.access.models import ExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.imap.models import EmailHeaders -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import CredentialsConnector -from onyx.connectors.interfaces import CredentialsProviderInterface -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + CredentialsConnector, + CredentialsProviderInterface, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorCheckpoint, + Document, + TextSection, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/interfaces.py b/backend/onyx/connectors/interfaces.py index edad8f6419..d1ae17d13f 100644 --- a/backend/onyx/connectors/interfaces.py +++ b/backend/onyx/connectors/interfaces.py @@ -1,21 +1,18 @@ import abc -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Generator, Iterator from types import TracebackType -from typing import Any -from typing import Generic -from typing import TypeAlias -from typing import TypeVar +from typing import Any, Generic, TypeAlias, TypeVar -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.configs.constants import DocumentSource -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + Document, + HierarchyNode, + SlimDocument, +) from onyx.file_store.staging import RawFileCallback from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop diff --git a/backend/onyx/connectors/jira/access.py b/backend/onyx/connectors/jira/access.py index 616626ce23..d542b1241c 100644 --- a/backend/onyx/connectors/jira/access.py +++ b/backend/onyx/connectors/jira/access.py @@ -8,8 +8,10 @@ from typing import cast from jira import JIRA from onyx.access.models import ExternalAccess -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) def get_project_permissions( diff --git a/backend/onyx/connectors/jira/connector.py b/backend/onyx/connectors/jira/connector.py index 8eb202875a..d09ee58968 100644 --- a/backend/onyx/connectors/jira/connector.py +++ b/backend/onyx/connectors/jira/connector.py @@ -1,13 +1,8 @@ import copy import json import os -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterable -from collections.abc import Iterator -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from collections.abc import Callable, Generator, Iterable, Iterator +from datetime import datetime, timedelta, timezone from typing import Any import requests @@ -17,41 +12,51 @@ from jira.resources import Issue from more_itertools import chunked from typing_extensions import override -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import JIRA_CONNECTOR_LABELS_TO_SKIP -from onyx.configs.app_configs import JIRA_CONNECTOR_MAX_TICKET_SIZE -from onyx.configs.app_configs import JIRA_SLIM_PAGE_SIZE +from onyx.configs.app_configs import ( + INDEX_BATCH_SIZE, + JIRA_CONNECTOR_LABELS_TO_SKIP, + JIRA_CONNECTOR_MAX_TICKET_SIZE, + JIRA_SLIM_PAGE_SIZE, +) from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( is_atlassian_date_error, + time_str_to_utc, +) +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnector, + SlimConnectorWithPermSync, ) -from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync from onyx.connectors.jira.access import get_project_permissions -from onyx.connectors.jira.utils import best_effort_basic_expert_info -from onyx.connectors.jira.utils import best_effort_get_field_from_issue -from onyx.connectors.jira.utils import build_jira_client -from onyx.connectors.jira.utils import build_jira_url -from onyx.connectors.jira.utils import extract_text_from_adf -from onyx.connectors.jira.utils import get_comment_strs -from onyx.connectors.jira.utils import JIRA_CLOUD_API_VERSION -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.jira.utils import ( + best_effort_basic_expert_info, + best_effort_get_field_from_issue, + build_jira_client, + build_jira_url, + extract_text_from_adf, + get_comment_strs, + JIRA_CLOUD_API_VERSION, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + HierarchyNode, + SlimDocument, + TextSection, +) from onyx.db.enums import HierarchyNodeType from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/jira/utils.py b/backend/onyx/connectors/jira/utils.py index 52bd91e3da..02eb36db83 100644 --- a/backend/onyx/connectors/jira/utils.py +++ b/backend/onyx/connectors/jira/utils.py @@ -1,14 +1,11 @@ """Module with custom fields processing functions""" import os -from typing import Any -from typing import List +from typing import Any, List from urllib.parse import urlparse from jira import JIRA -from jira.resources import CustomFieldOption -from jira.resources import Issue -from jira.resources import User +from jira.resources import CustomFieldOption, Issue, User from onyx.connectors.cross_connector_utils.miscellaneous_utils import scoped_url from onyx.connectors.models import BasicExpertInfo diff --git a/backend/onyx/connectors/linear/connector.py b/backend/onyx/connectors/linear/connector.py index 367a45bce0..b3d4b985ee 100644 --- a/backend/onyx/connectors/linear/connector.py +++ b/backend/onyx/connectors/linear/connector.py @@ -1,34 +1,38 @@ import os import re import time -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timezone +from typing import Any, cast from urllib.parse import urlparse import requests from typing_extensions import override -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import LINEAR_CLIENT_ID -from onyx.configs.app_configs import LINEAR_CLIENT_SECRET +from onyx.configs.app_configs import ( + INDEX_BATCH_SIZE, + LINEAR_CLIENT_ID, + LINEAR_CLIENT_SECRET, +) from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( get_oauth_callback_uri, + time_str_to_utc, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + NormalizationResult, + OAuthConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + ImageSection, + TextSection, ) -from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import NormalizationResult -from onyx.connectors.interfaces import OAuthConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import request_with_retries diff --git a/backend/onyx/connectors/loopio/connector.py b/backend/onyx/connectors/loopio/connector.py index a2ae7e5d38..5c5fa9ea43 100644 --- a/backend/onyx/connectors/loopio/connector.py +++ b/backend/onyx/connectors/loopio/connector.py @@ -1,7 +1,6 @@ import json from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from oauthlib.oauth2 import BackendApplicationClient @@ -10,17 +9,23 @@ from requests_oauthlib import OAuth2Session from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection -from onyx.file_processing.html_utils import parse_html_page_basic -from onyx.file_processing.html_utils import strip_excessive_newlines_and_spaces +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) +from onyx.file_processing.html_utils import ( + parse_html_page_basic, + strip_excessive_newlines_and_spaces, +) from onyx.utils.logger import setup_logger LOOPIO_API_BASE = "https://api.loopio.com/" diff --git a/backend/onyx/connectors/lumapps/connector.py b/backend/onyx/connectors/lumapps/connector.py index b8dc39b99b..a1b19e7df8 100644 --- a/backend/onyx/connectors/lumapps/connector.py +++ b/backend/onyx/connectors/lumapps/connector.py @@ -1,32 +1,38 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.lumapps.client import LumAppsClientError -from onyx.connectors.lumapps.client import OnyxLumApps +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnector, +) +from onyx.connectors.lumapps.client import LumAppsClientError, OnyxLumApps from onyx.connectors.lumapps.models import LumAppsCheckpoint -from onyx.connectors.lumapps.utils import extract_body_text -from onyx.connectors.lumapps.utils import pick_lang -from onyx.connectors.lumapps.utils import resolve_metadata_labels -from onyx.connectors.lumapps.utils import slugify_family_key -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.lumapps.utils import ( + extract_body_text, + pick_lang, + resolve_metadata_labels, + slugify_family_key, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + HierarchyNode, + SlimDocument, + TextSection, +) from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/mediawiki/family.py b/backend/onyx/connectors/mediawiki/family.py index 5313ae332b..5d8308f507 100644 --- a/backend/onyx/connectors/mediawiki/family.py +++ b/backend/onyx/connectors/mediawiki/family.py @@ -6,11 +6,9 @@ import itertools import tempfile from typing import Any from unittest import mock -from urllib.parse import urlparse -from urllib.parse import urlunparse +from urllib.parse import urlparse, urlunparse -from pywikibot import family -from pywikibot import pagegenerators +from pywikibot import family, pagegenerators from pywikibot.scripts import generate_family_file from pywikibot.scripts.generate_user_files import pywikibot diff --git a/backend/onyx/connectors/mediawiki/wiki.py b/backend/onyx/connectors/mediawiki/wiki.py index 68aac5db09..ff88cc8d52 100644 --- a/backend/onyx/connectors/mediawiki/wiki.py +++ b/backend/onyx/connectors/mediawiki/wiki.py @@ -4,26 +4,22 @@ import datetime import itertools import tempfile from collections.abc import Iterator -from typing import Any -from typing import cast -from typing import ClassVar +from typing import Any, cast, ClassVar import pywikibot.config import pywikibot.time -from pywikibot import pagegenerators -from pywikibot import textlib +from pywikibot import pagegenerators, textlib from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) from onyx.connectors.mediawiki.family import family_class_dispatch -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, HierarchyNode, ImageSection, TextSection from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/mock_connector/connector.py b/backend/onyx/connectors/mock_connector/connector.py index cd7b8bf0bf..ff70c97b4e 100644 --- a/backend/onyx/connectors/mock_connector/connector.py +++ b/backend/onyx/connectors/mock_connector/connector.py @@ -5,12 +5,12 @@ from pydantic import BaseModel from typing_extensions import override from onyx.access.models import ExternalAccess -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ConnectorCheckpoint, ConnectorFailure, Document from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/models.py b/backend/onyx/connectors/models.py index 9fc6834832..a93d46233c 100644 --- a/backend/onyx/connectors/models.py +++ b/backend/onyx/connectors/models.py @@ -4,21 +4,13 @@ import sys from collections.abc import Sequence from datetime import datetime from enum import Enum -from typing import Any -from typing import cast -from typing import Literal +from typing import Any, cast, Literal -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator -from pydantic import model_validator +from pydantic import BaseModel, Field, field_validator, model_validator from onyx.access.models import ExternalAccess -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import INDEX_SEPARATOR -from onyx.configs.constants import RETURN_SEPARATOR -from onyx.db.enums import HierarchyNodeType -from onyx.db.enums import IndexModelStatus +from onyx.configs.constants import DocumentSource, INDEX_SEPARATOR, RETURN_SEPARATOR +from onyx.db.enums import HierarchyNodeType, IndexModelStatus from onyx.utils.text_processing import make_url_compatible diff --git a/backend/onyx/connectors/notion/connector.py b/backend/onyx/connectors/notion/connector.py index af502dc70c..319b1792d1 100644 --- a/backend/onyx/connectors/notion/connector.py +++ b/backend/onyx/connectors/notion/connector.py @@ -1,34 +1,39 @@ import re from collections.abc import Generator from datetime import datetime -from typing import Any -from typing import cast -from typing import Optional -from urllib.parse import parse_qs -from urllib.parse import urlparse +from typing import Any, cast, Optional +from urllib.parse import parse_qs, urlparse import requests from pydantic import BaseModel from typing_extensions import override -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import NOTION_CONNECTOR_DISABLE_RECURSIVE_PAGE_LOOKUP +from onyx.configs.app_configs import ( + INDEX_BATCH_SIZE, + NOTION_CONNECTOR_DISABLE_RECURSIVE_PAGE_LOOKUP, +) from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rl_requests -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import NormalizationResult -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + NormalizationResult, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + ImageSection, + TextSection, +) from onyx.db.enums import HierarchyNodeType from onyx.utils.batching import batch_generator from onyx.utils.datetime import datetime_to_utc diff --git a/backend/onyx/connectors/outline/client.py b/backend/onyx/connectors/outline/client.py index 482ee4fff2..668fa35410 100644 --- a/backend/onyx/connectors/outline/client.py +++ b/backend/onyx/connectors/outline/client.py @@ -2,8 +2,7 @@ from typing import Any import requests from requests.exceptions import ConnectionError as RequestsConnectionError -from requests.exceptions import RequestException -from requests.exceptions import Timeout +from requests.exceptions import RequestException, Timeout from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS diff --git a/backend/onyx/connectors/outline/connector.py b/backend/onyx/connectors/outline/connector.py index bf702b6202..61b7a96d6c 100644 --- a/backend/onyx/connectors/outline/connector.py +++ b/backend/onyx/connectors/outline/connector.py @@ -6,19 +6,27 @@ from typing import Any from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection -from onyx.connectors.outline.client import OutlineApiClient -from onyx.connectors.outline.client import OutlineClientRequestFailedError +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) +from onyx.connectors.outline.client import ( + OutlineApiClient, + OutlineClientRequestFailedError, +) class OutlineConnector(LoadConnector, PollConnector): diff --git a/backend/onyx/connectors/productboard/connector.py b/backend/onyx/connectors/productboard/connector.py index 7f2a0c0ded..b743a78dff 100644 --- a/backend/onyx/connectors/productboard/connector.py +++ b/backend/onyx/connectors/productboard/connector.py @@ -1,22 +1,19 @@ from collections.abc import Generator from itertools import chain -from typing import Any -from typing import cast +from typing import Any, cast import requests from bs4 import BeautifulSoup -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import BasicExpertInfo, Document, HierarchyNode, TextSection from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/salesforce/connector.py b/backend/onyx/connectors/salesforce/connector.py index b24bd1f98a..e082dd2e33 100644 --- a/backend/onyx/connectors/salesforce/connector.py +++ b/backend/onyx/connectors/salesforce/connector.py @@ -12,32 +12,40 @@ from typing import Any from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection -from onyx.connectors.salesforce.doc_conversion import convert_sf_object_to_doc -from onyx.connectors.salesforce.doc_conversion import convert_sf_query_result_to_doc -from onyx.connectors.salesforce.doc_conversion import ID_PREFIX +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + GenerateSlimDocumentOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorCheckpoint, + ConnectorMissingCredentialError, + Document, + HierarchyNode, + SlimDocument, + TextSection, +) +from onyx.connectors.salesforce.doc_conversion import ( + convert_sf_object_to_doc, + convert_sf_query_result_to_doc, + ID_PREFIX, +) from onyx.connectors.salesforce.onyx_salesforce import OnyxSalesforce from onyx.connectors.salesforce.salesforce_calls import fetch_all_csvs_in_parallel from onyx.connectors.salesforce.sqlite_functions import OnyxSalesforceSQLite -from onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE -from onyx.connectors.salesforce.utils import CREATED_FIELD -from onyx.connectors.salesforce.utils import ID_FIELD -from onyx.connectors.salesforce.utils import MODIFIED_FIELD -from onyx.connectors.salesforce.utils import NAME_FIELD -from onyx.connectors.salesforce.utils import USER_OBJECT_TYPE -from onyx.connectors.salesforce.utils import validate_sf_identifier +from onyx.connectors.salesforce.utils import ( + ACCOUNT_OBJECT_TYPE, + CREATED_FIELD, + ID_FIELD, + MODIFIED_FIELD, + NAME_FIELD, + USER_OBJECT_TYPE, + validate_sf_identifier, +) from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/salesforce/doc_conversion.py b/backend/onyx/connectors/salesforce/doc_conversion.py index 734dc0a4b0..c6f10a7650 100644 --- a/backend/onyx/connectors/salesforce/doc_conversion.py +++ b/backend/onyx/connectors/salesforce/doc_conversion.py @@ -1,20 +1,18 @@ import re -from typing import Any -from typing import cast +from typing import Any, cast from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import Document -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import BasicExpertInfo, Document, ImageSection, TextSection from onyx.connectors.salesforce.onyx_salesforce import OnyxSalesforce from onyx.connectors.salesforce.sqlite_functions import OnyxSalesforceSQLite -from onyx.connectors.salesforce.utils import CREATED_FIELD -from onyx.connectors.salesforce.utils import ID_FIELD -from onyx.connectors.salesforce.utils import MODIFIED_FIELD -from onyx.connectors.salesforce.utils import NAME_FIELD -from onyx.connectors.salesforce.utils import SalesforceObject +from onyx.connectors.salesforce.utils import ( + CREATED_FIELD, + ID_FIELD, + MODIFIED_FIELD, + NAME_FIELD, + SalesforceObject, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/salesforce/onyx_salesforce.py b/backend/onyx/connectors/salesforce/onyx_salesforce.py index 7f0499f925..16978b1940 100644 --- a/backend/onyx/connectors/salesforce/onyx_salesforce.py +++ b/backend/onyx/connectors/salesforce/onyx_salesforce.py @@ -2,19 +2,19 @@ import time from typing import Any from requests.adapters import HTTPAdapter -from simple_salesforce import Salesforce -from simple_salesforce import SFType +from simple_salesforce import Salesforce, SFType from simple_salesforce.exceptions import SalesforceRefusedRequest from simple_salesforce.format import format_soql from urllib3.util.retry import Retry from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.salesforce.blacklist import SALESFORCE_BLACKLISTED_OBJECTS -from onyx.connectors.salesforce.blacklist import SALESFORCE_BLACKLISTED_PREFIXES -from onyx.connectors.salesforce.blacklist import SALESFORCE_BLACKLISTED_SUFFIXES +from onyx.connectors.salesforce.blacklist import ( + SALESFORCE_BLACKLISTED_OBJECTS, + SALESFORCE_BLACKLISTED_PREFIXES, + SALESFORCE_BLACKLISTED_SUFFIXES, +) from onyx.connectors.salesforce.salesforce_calls import get_object_by_id_query -from onyx.connectors.salesforce.utils import ID_FIELD -from onyx.connectors.salesforce.utils import validate_sf_identifier +from onyx.connectors.salesforce.utils import ID_FIELD, validate_sf_identifier from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/salesforce/salesforce_calls.py b/backend/onyx/connectors/salesforce/salesforce_calls.py index 087778f47d..738c220f5c 100644 --- a/backend/onyx/connectors/salesforce/salesforce_calls.py +++ b/backend/onyx/connectors/salesforce/salesforce_calls.py @@ -6,15 +6,13 @@ from datetime import datetime from pytz import UTC from simple_salesforce import Salesforce -from simple_salesforce.bulk2 import SFBulk2Handler -from simple_salesforce.bulk2 import SFBulk2Type +from simple_salesforce.bulk2 import SFBulk2Handler, SFBulk2Type from simple_salesforce.exceptions import SalesforceRefusedRequest from simple_salesforce.format import format_soql from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.salesforce.utils import MODIFIED_FIELD -from onyx.connectors.salesforce.utils import validate_sf_identifier +from onyx.connectors.salesforce.utils import MODIFIED_FIELD, validate_sf_identifier from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/connectors/salesforce/shelve_stuff/old_test_salesforce_shelves.py b/backend/onyx/connectors/salesforce/shelve_stuff/old_test_salesforce_shelves.py index 642e345a1d..84ccb578cd 100644 --- a/backend/onyx/connectors/salesforce/shelve_stuff/old_test_salesforce_shelves.py +++ b/backend/onyx/connectors/salesforce/shelve_stuff/old_test_salesforce_shelves.py @@ -2,17 +2,14 @@ import csv import os import shutil -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import find_ids_by_type from onyx.connectors.salesforce.shelve_stuff.shelve_functions import ( + find_ids_by_type, get_affected_parent_ids_by_type, -) -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_child_ids -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_record -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import ( + get_child_ids, + get_record, update_sf_db_with_csv, ) -from onyx.connectors.salesforce.utils import BASE_DATA_PATH -from onyx.connectors.salesforce.utils import get_object_type_path +from onyx.connectors.salesforce.utils import BASE_DATA_PATH, get_object_type_path _VALID_SALESFORCE_IDS = [ "001bm00000fd9Z3AAI", diff --git a/backend/onyx/connectors/salesforce/shelve_stuff/shelve_functions.py b/backend/onyx/connectors/salesforce/shelve_stuff/shelve_functions.py index 2eedff6ddd..2ed5383998 100644 --- a/backend/onyx/connectors/salesforce/shelve_stuff/shelve_functions.py +++ b/backend/onyx/connectors/salesforce/shelve_stuff/shelve_functions.py @@ -3,14 +3,11 @@ import shelve from onyx.connectors.salesforce.shelve_stuff.shelve_utils import ( get_child_to_parent_shelf_path, -) -from onyx.connectors.salesforce.shelve_stuff.shelve_utils import get_id_type_shelf_path -from onyx.connectors.salesforce.shelve_stuff.shelve_utils import get_object_shelf_path -from onyx.connectors.salesforce.shelve_stuff.shelve_utils import ( + get_id_type_shelf_path, + get_object_shelf_path, get_parent_to_child_shelf_path, ) -from onyx.connectors.salesforce.utils import SalesforceObject -from onyx.connectors.salesforce.utils import validate_salesforce_id +from onyx.connectors.salesforce.utils import SalesforceObject, validate_salesforce_id from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/connectors/salesforce/shelve_stuff/shelve_utils.py b/backend/onyx/connectors/salesforce/shelve_stuff/shelve_utils.py index 0ea937af12..ca84f846cd 100644 --- a/backend/onyx/connectors/salesforce/shelve_stuff/shelve_utils.py +++ b/backend/onyx/connectors/salesforce/shelve_stuff/shelve_utils.py @@ -1,7 +1,6 @@ import os -from onyx.connectors.salesforce.utils import BASE_DATA_PATH -from onyx.connectors.salesforce.utils import get_object_type_path +from onyx.connectors.salesforce.utils import BASE_DATA_PATH, get_object_type_path def get_object_shelf_path(object_type: str) -> str: diff --git a/backend/onyx/connectors/salesforce/shelve_stuff/test_salesforce_shelves.py b/backend/onyx/connectors/salesforce/shelve_stuff/test_salesforce_shelves.py index 642e345a1d..84ccb578cd 100644 --- a/backend/onyx/connectors/salesforce/shelve_stuff/test_salesforce_shelves.py +++ b/backend/onyx/connectors/salesforce/shelve_stuff/test_salesforce_shelves.py @@ -2,17 +2,14 @@ import csv import os import shutil -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import find_ids_by_type from onyx.connectors.salesforce.shelve_stuff.shelve_functions import ( + find_ids_by_type, get_affected_parent_ids_by_type, -) -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_child_ids -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_record -from onyx.connectors.salesforce.shelve_stuff.shelve_functions import ( + get_child_ids, + get_record, update_sf_db_with_csv, ) -from onyx.connectors.salesforce.utils import BASE_DATA_PATH -from onyx.connectors.salesforce.utils import get_object_type_path +from onyx.connectors.salesforce.utils import BASE_DATA_PATH, get_object_type_path _VALID_SALESFORCE_IDS = [ "001bm00000fd9Z3AAI", diff --git a/backend/onyx/connectors/salesforce/sqlite_functions.py b/backend/onyx/connectors/salesforce/sqlite_functions.py index abbf3f1454..8fd2a1e9e6 100644 --- a/backend/onyx/connectors/salesforce/sqlite_functions.py +++ b/backend/onyx/connectors/salesforce/sqlite_functions.py @@ -5,17 +5,18 @@ import sqlite3 import time from collections.abc import Iterator from pathlib import Path -from typing import Any -from typing import cast +from typing import Any, cast from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE -from onyx.connectors.salesforce.utils import ID_FIELD -from onyx.connectors.salesforce.utils import NAME_FIELD -from onyx.connectors.salesforce.utils import remove_sqlite_db_files -from onyx.connectors.salesforce.utils import SalesforceObject -from onyx.connectors.salesforce.utils import USER_OBJECT_TYPE -from onyx.connectors.salesforce.utils import validate_salesforce_id +from onyx.connectors.salesforce.utils import ( + ACCOUNT_OBJECT_TYPE, + ID_FIELD, + NAME_FIELD, + remove_sqlite_db_files, + SalesforceObject, + USER_OBJECT_TYPE, + validate_salesforce_id, +) from onyx.utils.logger import setup_logger from shared_configs.utils import batch_list diff --git a/backend/onyx/connectors/sharepoint/connector.py b/backend/onyx/connectors/sharepoint/connector.py index deb55168ed..2baf3ba144 100644 --- a/backend/onyx/connectors/sharepoint/connector.py +++ b/backend/onyx/connectors/sharepoint/connector.py @@ -8,22 +8,15 @@ import random import re import time from collections import deque -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterable -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Generator, Iterable +from datetime import datetime, timezone from enum import Enum -from typing import Any -from typing import cast -from urllib.parse import quote -from urllib.parse import unquote -from urllib.parse import urlsplit +from typing import Any, cast +from urllib.parse import quote, unquote, urlsplit import msal import requests -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.serialization import pkcs12 from office365.graph_client import GraphClient from office365.onedrive.driveitems.driveItem import DriveItem @@ -34,57 +27,60 @@ from office365.runtime.client_request import ClientRequestException from office365.runtime.paths.resource_path import ResourcePath from office365.runtime.queries.client_query import ClientQuery from office365.sharepoint.client_context import ClientContext -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from requests.exceptions import HTTPError from typing_extensions import override -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS -from onyx.configs.app_configs import SHAREPOINT_CONNECTOR_SIZE_THRESHOLD -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.app_configs import ( + INDEX_BATCH_SIZE, + REQUEST_TIMEOUT_SECONDS, + SHAREPOINT_CONNECTOR_SIZE_THRESHOLD, +) +from onyx.configs.constants import DocumentSource, FileOrigin from onyx.connectors.cross_connector_utils.tabular_section_utils import ( extract_and_stage_tabular_file, + is_tabular_file, ) -from onyx.connectors.cross_connector_utils.tabular_section_utils import is_tabular_file from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import IndexingHeartbeatInterface -from onyx.connectors.interfaces import Resolver -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + GenerateSlimDocumentOutput, + IndexingHeartbeatInterface, + Resolver, + SecondsSinceUnixEpoch, + SlimConnector, + SlimConnectorWithPermSync, +) from onyx.connectors.microsoft_graph_env import resolve_microsoft_environment -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import ExternalAccess -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + EntityFailure, + ExternalAccess, + HierarchyNode, + ImageSection, + SlimDocument, + TabularSection, + TextSection, +) from onyx.connectors.sharepoint.connector_utils import get_sharepoint_external_access from onyx.db.enums import HierarchyNodeType -from onyx.file_processing.extract_file_text import extract_text_and_images -from onyx.file_processing.extract_file_text import get_file_ext -from onyx.file_processing.file_types import OnyxFileExtensions -from onyx.file_processing.file_types import OnyxMimeTypes -from onyx.file_processing.image_utils import make_image_callback -from onyx.file_processing.image_utils import store_image_and_create_section +from onyx.file_processing.extract_file_text import extract_text_and_images, get_file_ext +from onyx.file_processing.file_types import OnyxFileExtensions, OnyxMimeTypes +from onyx.file_processing.image_utils import ( + make_image_callback, + store_image_and_create_section, +) from onyx.file_store.staging import RawFileCallback from onyx.utils.logger import setup_logger from onyx.utils.retry_after import parse_retry_after_seconds from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel -from onyx.utils.url import SSRFException -from onyx.utils.url import validate_outbound_http_url +from onyx.utils.url import SSRFException, validate_outbound_http_url logger = setup_logger() SLIM_BATCH_SIZE = 1000 diff --git a/backend/onyx/connectors/slab/connector.py b/backend/onyx/connectors/slab/connector.py index d0b99e6308..db01215add 100644 --- a/backend/onyx/connectors/slab/connector.py +++ b/backend/onyx/connectors/slab/connector.py @@ -1,8 +1,6 @@ import json -from collections.abc import Callable -from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Generator +from datetime import datetime, timezone from typing import Any from urllib.parse import urljoin @@ -13,17 +11,21 @@ from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + GenerateSlimDocumentOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + SlimDocument, + TextSection, +) from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/slack/access.py b/backend/onyx/connectors/slack/access.py index 92fc3f89cf..f9bb2fda30 100644 --- a/backend/onyx/connectors/slack/access.py +++ b/backend/onyx/connectors/slack/access.py @@ -6,8 +6,10 @@ from slack_sdk import WebClient from onyx.access.models import ExternalAccess from onyx.connectors.models import BasicExpertInfo from onyx.connectors.slack.models import ChannelType -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) def get_channel_access( diff --git a/backend/onyx/connectors/slack/connector.py b/backend/onyx/connectors/slack/connector.py index 999d8a07b0..642d2014e2 100644 --- a/backend/onyx/connectors/slack/connector.py +++ b/backend/onyx/connectors/slack/connector.py @@ -2,69 +2,70 @@ import contextvars import copy import itertools import re -from collections.abc import Callable -from collections.abc import Generator -from concurrent.futures import as_completed -from concurrent.futures import Future -from concurrent.futures import ThreadPoolExecutor -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Generator +from concurrent.futures import as_completed, Future, ThreadPoolExecutor +from datetime import datetime, timezone from enum import Enum -from http.client import IncompleteRead -from http.client import RemoteDisconnected -from typing import Any -from typing import cast +from http.client import IncompleteRead, RemoteDisconnected +from typing import Any, cast from urllib.error import URLError from urllib.parse import urlparse from pydantic import BaseModel from slack_sdk import WebClient from slack_sdk.errors import SlackApiError -from slack_sdk.http_retry import ConnectionErrorRetryHandler -from slack_sdk.http_retry import RetryHandler +from slack_sdk.http_retry import ConnectionErrorRetryHandler, RetryHandler from slack_sdk.http_retry.builtin_interval_calculators import ( FixedValueRetryIntervalCalculator, ) from typing_extensions import override from onyx.access.models import ExternalAccess -from onyx.configs.app_configs import ENABLE_EXPENSIVE_EXPERT_CALLS -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import SLACK_NUM_THREADS +from onyx.configs.app_configs import ( + ENABLE_EXPENSIVE_EXPERT_CALLS, + INDEX_BATCH_SIZE, + SLACK_NUM_THREADS, +) from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import CredentialsConnector -from onyx.connectors.interfaces import CredentialsProviderInterface -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import NormalizationResult -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + CredentialsConnector, + CredentialsProviderInterface, + GenerateSlimDocumentOutput, + NormalizationResult, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + DocumentFailure, + EntityFailure, + HierarchyNode, + SlimDocument, + TextSection, +) from onyx.connectors.slack.access import get_channel_access -from onyx.connectors.slack.models import ChannelType -from onyx.connectors.slack.models import MessageType -from onyx.connectors.slack.models import ThreadType +from onyx.connectors.slack.models import ChannelType, MessageType, ThreadType from onyx.connectors.slack.onyx_retry_handler import OnyxRedisSlackRetryHandler from onyx.connectors.slack.onyx_slack_web_client import OnyxSlackWebClient -from onyx.connectors.slack.utils import expert_info_from_slack_id -from onyx.connectors.slack.utils import fetch_team_user_emails -from onyx.connectors.slack.utils import get_message_link -from onyx.connectors.slack.utils import make_paginated_slack_api_call -from onyx.connectors.slack.utils import SlackTextCleaner +from onyx.connectors.slack.utils import ( + expert_info_from_slack_id, + fetch_team_user_emails, + get_message_link, + make_paginated_slack_api_call, + SlackTextCleaner, +) from onyx.db.enums import HierarchyNodeType from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.redis.redis_pool import get_redis_client diff --git a/backend/onyx/connectors/slack/onyx_slack_web_client.py b/backend/onyx/connectors/slack/onyx_slack_web_client.py index a1b976024b..a575472ec6 100644 --- a/backend/onyx/connectors/slack/onyx_slack_web_client.py +++ b/backend/onyx/connectors/slack/onyx_slack_web_client.py @@ -1,15 +1,16 @@ import threading import time -from typing import Any -from typing import Dict +from typing import Any, Dict from urllib.request import Request from redis.lock import Lock as RedisLock from slack_sdk import WebClient -from onyx.connectors.slack.utils import ONYX_SLACK_LOCK_BLOCKING_TIMEOUT -from onyx.connectors.slack.utils import ONYX_SLACK_LOCK_TOTAL_BLOCKING_TIMEOUT -from onyx.connectors.slack.utils import ONYX_SLACK_LOCK_TTL +from onyx.connectors.slack.utils import ( + ONYX_SLACK_LOCK_BLOCKING_TIMEOUT, + ONYX_SLACK_LOCK_TOTAL_BLOCKING_TIMEOUT, + ONYX_SLACK_LOCK_TTL, +) from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/slack/utils.py b/backend/onyx/connectors/slack/utils.py index a3fabd9c96..dee19b95a5 100644 --- a/backend/onyx/connectors/slack/utils.py +++ b/backend/onyx/connectors/slack/utils.py @@ -1,10 +1,7 @@ import re -from collections.abc import Callable -from collections.abc import Generator -from functools import lru_cache -from functools import wraps -from typing import Any -from typing import cast +from collections.abc import Callable, Generator +from functools import lru_cache, wraps +from typing import Any, cast from slack_sdk import WebClient from slack_sdk.errors import SlackApiError diff --git a/backend/onyx/connectors/teams/connector.py b/backend/onyx/connectors/teams/connector.py index 11621ba0fb..242ab9f9c5 100644 --- a/backend/onyx/connectors/teams/connector.py +++ b/backend/onyx/connectors/teams/connector.py @@ -1,8 +1,7 @@ import copy import os from collections.abc import Iterator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import msal @@ -13,30 +12,38 @@ from office365.teams.channels.channel import Channel from office365.teams.team import Team from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnectorWithPermSync, + CheckpointOutput, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) from onyx.connectors.microsoft_graph_env import resolve_microsoft_environment -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + EntityFailure, + HierarchyNode, + SlimDocument, + TextSection, +) from onyx.connectors.teams.models import Message -from onyx.connectors.teams.utils import execute_query_with_retry -from onyx.connectors.teams.utils import fetch_expert_infos -from onyx.connectors.teams.utils import fetch_external_access -from onyx.connectors.teams.utils import fetch_messages -from onyx.connectors.teams.utils import fetch_replies +from onyx.connectors.teams.utils import ( + execute_query_with_retry, + fetch_expert_infos, + fetch_external_access, + fetch_messages, + fetch_replies, +) from onyx.file_processing.html_utils import parse_html_page_basic from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/teams/models.py b/backend/onyx/connectors/teams/models.py index a03874a477..a2148f0b20 100644 --- a/backend/onyx/connectors/teams/models.py +++ b/backend/onyx/connectors/teams/models.py @@ -1,8 +1,6 @@ from datetime import datetime -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field from pydantic.alias_generators import to_camel diff --git a/backend/onyx/connectors/teams/utils.py b/backend/onyx/connectors/teams/utils.py index 972b452f94..5d27faafa7 100644 --- a/backend/onyx/connectors/teams/utils.py +++ b/backend/onyx/connectors/teams/utils.py @@ -1,15 +1,13 @@ import random import time from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from office365.graph_client import GraphClient from office365.runtime.client_request_exception import ClientRequestException from office365.runtime.queries.client_query import ClientQuery -from office365.teams.channels.channel import Channel -from office365.teams.channels.channel import ConversationMember +from office365.teams.channels.channel import Channel, ConversationMember from onyx.access.models import ExternalAccess from onyx.connectors.interfaces import SecondsSinceUnixEpoch diff --git a/backend/onyx/connectors/testrail/connector.py b/backend/onyx/connectors/testrail/connector.py index 5a9523e94d..f450df8430 100644 --- a/backend/onyx/connectors/testrail/connector.py +++ b/backend/onyx/connectors/testrail/connector.py @@ -1,29 +1,31 @@ from __future__ import annotations from collections.abc import Iterator -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import ClassVar -from typing import Optional +from datetime import datetime, timezone +from typing import Any, ClassVar, Optional import requests from bs4 import BeautifulSoup -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) from onyx.file_processing.html_utils import format_document_soup from onyx.utils.logger import setup_logger from onyx.utils.text_processing import remove_markdown_image_references @@ -549,9 +551,11 @@ class TestRailConnector(LoadConnector, PollConnector): if __name__ == "__main__": - from onyx.configs.app_configs import TESTRAIL_API_KEY - from onyx.configs.app_configs import TESTRAIL_BASE_URL - from onyx.configs.app_configs import TESTRAIL_USERNAME + from onyx.configs.app_configs import ( + TESTRAIL_API_KEY, + TESTRAIL_BASE_URL, + TESTRAIL_USERNAME, + ) connector = TestRailConnector() diff --git a/backend/onyx/connectors/web/connector.py b/backend/onyx/connectors/web/connector.py index f2631584f9..c4b9e2a8e7 100644 --- a/backend/onyx/connectors/web/connector.py +++ b/backend/onyx/connectors/web/connector.py @@ -3,35 +3,31 @@ import random import socket import time from enum import Enum -from typing import Any -from typing import cast -from urllib.parse import urljoin -from urllib.parse import urlparse +from typing import Any, cast +from urllib.parse import urljoin, urlparse import requests from bs4 import BeautifulSoup -from playwright.sync_api import BrowserContext -from playwright.sync_api import Playwright -from playwright.sync_api import TimeoutError +from playwright.sync_api import BrowserContext, Playwright, TimeoutError from typing_extensions import override from urllib3.exceptions import MaxRetryError -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS +from onyx.configs.app_configs import INDEX_BATCH_SIZE, REQUEST_TIMEOUT_SECONDS from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + GenerateSlimDocumentOutput, + LoadConnector, + SecondsSinceUnixEpoch, + SlimConnector, +) +from onyx.connectors.models import Document, HierarchyNode, SlimDocument, TextSection from onyx.file_processing.html_utils import web_html_cleanup from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.server.security.models import web_connector_ssrf_enforced @@ -40,12 +36,13 @@ from onyx.utils.logger import setup_logger # Re-exported for backwards compatibility with existing tests/callers that # patch these names on `onyx.connectors.web.connector`. -from onyx.utils.playwright_fetch import DEFAULT_HEADERS -from onyx.utils.playwright_fetch import DEFAULT_USER_AGENT # noqa: F401 -from onyx.utils.playwright_fetch import start_playwright +from onyx.utils.playwright_fetch import ( + DEFAULT_HEADERS, + DEFAULT_USER_AGENT, # noqa: F401 + start_playwright, +) from onyx.utils.sitemap import list_pages_for_site -from onyx.utils.web_content import extract_pdf_text -from onyx.utils.web_content import is_pdf_resource +from onyx.utils.web_content import extract_pdf_text, is_pdf_resource from shared_configs.configs import MULTI_TENANT logger = setup_logger() diff --git a/backend/onyx/connectors/xenforo/connector.py b/backend/onyx/connectors/xenforo/connector.py index 9d78f834d1..b8a0455eb5 100644 --- a/backend/onyx/connectors/xenforo/connector.py +++ b/backend/onyx/connectors/xenforo/connector.py @@ -12,23 +12,16 @@ can be used to specify a state from which to start loading documents. """ import re -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from urllib.parse import urlparse import requests -from bs4 import BeautifulSoup -from bs4 import Tag +from bs4 import BeautifulSoup, Tag from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import GenerateDocumentsOutput, LoadConnector +from onyx.connectors.models import BasicExpertInfo, Document, HierarchyNode, TextSection from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/connectors/zendesk/connector.py b/backend/onyx/connectors/zendesk/connector.py index 3f79771cd6..fcfb349034 100644 --- a/backend/onyx/connectors/zendesk/connector.py +++ b/backend/onyx/connectors/zendesk/connector.py @@ -1,36 +1,42 @@ import copy import time -from collections.abc import Callable -from collections.abc import Iterator -from typing import Any -from typing import cast +from collections.abc import Callable, Iterator +from typing import Any, cast import requests from pydantic import BaseModel from requests.exceptions import HTTPError from typing_extensions import override -from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS -from onyx.configs.app_configs import ZENDESK_CONNECTOR_SKIP_ARTICLE_LABELS +from onyx.configs.app_configs import ( + REQUEST_TIMEOUT_SECONDS, + ZENDESK_CONNECTOR_SKIP_ARTICLE_LABELS, +) from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import ConnectorFailure -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from onyx.connectors.models import TextSection +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + ConnectorFailure, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) +from onyx.connectors.models import ( + BasicExpertInfo, + ConnectorCheckpoint, + Document, + DocumentFailure, + HierarchyNode, + SlimDocument, + TextSection, +) from onyx.file_processing.html_utils import parse_html_page_basic from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.utils.retry_after import parse_retry_after_seconds diff --git a/backend/onyx/connectors/zulip/connector.py b/backend/onyx/connectors/zulip/connector.py index 73d1dcde31..38d60cae6c 100644 --- a/backend/onyx/connectors/zulip/connector.py +++ b/backend/onyx/connectors/zulip/connector.py @@ -2,31 +2,31 @@ import os import tempfile import urllib.parse from collections.abc import Generator -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import Dict -from typing import List -from typing import Tuple -from typing import Union +from datetime import datetime, timezone +from typing import Any, Dict, List, Tuple, Union from zulip import Client from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import GenerateDocumentsOutput -from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.interfaces import PollConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection -from onyx.connectors.zulip.schemas import GetMessagesResponse -from onyx.connectors.zulip.schemas import Message -from onyx.connectors.zulip.utils import build_search_narrow -from onyx.connectors.zulip.utils import call_api -from onyx.connectors.zulip.utils import encode_zulip_narrow_operand +from onyx.connectors.interfaces import ( + GenerateDocumentsOutput, + LoadConnector, + PollConnector, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, + TextSection, +) +from onyx.connectors.zulip.schemas import GetMessagesResponse, Message +from onyx.connectors.zulip.utils import ( + build_search_narrow, + call_api, + encode_zulip_narrow_operand, +) from onyx.utils.logger import setup_logger # Potential improvements diff --git a/backend/onyx/connectors/zulip/schemas.py b/backend/onyx/connectors/zulip/schemas.py index e999aace5b..2dcf376cbc 100644 --- a/backend/onyx/connectors/zulip/schemas.py +++ b/backend/onyx/connectors/zulip/schemas.py @@ -1,10 +1,6 @@ -from typing import Any -from typing import List -from typing import Optional -from typing import Union +from typing import Any, List, Optional, Union -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field class Message(BaseModel): diff --git a/backend/onyx/connectors/zulip/utils.py b/backend/onyx/connectors/zulip/utils.py index 76252d267c..ee2eda40af 100644 --- a/backend/onyx/connectors/zulip/utils.py +++ b/backend/onyx/connectors/zulip/utils.py @@ -1,8 +1,6 @@ import time from collections.abc import Callable -from typing import Any -from typing import Dict -from typing import Optional +from typing import Any, Dict, Optional from urllib.parse import quote from onyx.utils.logger import setup_logger diff --git a/backend/onyx/context/search/federated/slack_search.py b/backend/onyx/context/search/federated/slack_search.py index 3c7bf0c2e9..895f6c3e54 100644 --- a/backend/onyx/context/search/federated/slack_search.py +++ b/backend/onyx/context/search/federated/slack_search.py @@ -2,36 +2,35 @@ import json import re import time from datetime import datetime -from typing import Any -from typing import cast +from typing import Any, cast -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import ValidationError +from pydantic import BaseModel, ConfigDict, ValidationError from slack_sdk import WebClient from slack_sdk.errors import SlackApiError from sqlalchemy.orm import Session -from onyx.configs.app_configs import ENABLE_CONTEXTUAL_RAG -from onyx.configs.app_configs import MAX_SLACK_THREAD_CONTEXT_MESSAGES -from onyx.configs.app_configs import SLACK_THREAD_CONTEXT_BATCH_SIZE -from onyx.configs.chat_configs import DOC_TIME_DECAY -from onyx.connectors.models import IndexingDocument -from onyx.connectors.models import TextSection -from onyx.context.search.federated.models import ChannelMetadata -from onyx.context.search.federated.models import DirectThreadFetch -from onyx.context.search.federated.models import SlackMessage -from onyx.context.search.federated.slack_search_utils import ALL_CHANNEL_TYPES -from onyx.context.search.federated.slack_search_utils import build_channel_query_filter -from onyx.context.search.federated.slack_search_utils import build_slack_queries -from onyx.context.search.federated.slack_search_utils import get_channel_type -from onyx.context.search.federated.slack_search_utils import ( - get_channel_type_for_missing_scope, +from onyx.configs.app_configs import ( + ENABLE_CONTEXTUAL_RAG, + MAX_SLACK_THREAD_CONTEXT_MESSAGES, + SLACK_THREAD_CONTEXT_BATCH_SIZE, ) -from onyx.context.search.federated.slack_search_utils import is_recency_query -from onyx.context.search.federated.slack_search_utils import should_include_message -from onyx.context.search.models import ChunkIndexRequest -from onyx.context.search.models import InferenceChunk +from onyx.configs.chat_configs import DOC_TIME_DECAY +from onyx.connectors.models import IndexingDocument, TextSection +from onyx.context.search.federated.models import ( + ChannelMetadata, + DirectThreadFetch, + SlackMessage, +) +from onyx.context.search.federated.slack_search_utils import ( + ALL_CHANNEL_TYPES, + build_channel_query_filter, + build_slack_queries, + get_channel_type, + get_channel_type_for_missing_scope, + is_recency_query, + should_include_message, +) +from onyx.context.search.models import ChunkIndexRequest, InferenceChunk from onyx.db.document import DocumentSource from onyx.db.models import SearchSettings from onyx.db.search_settings import get_current_search_settings @@ -41,8 +40,7 @@ from onyx.indexing.chunker import Chunker from onyx.indexing.embedder import DefaultIndexingEmbedder from onyx.indexing.models import DocAwareChunk from onyx.llm.factory import get_default_llm -from onyx.onyxbot.slack.models import ChannelType -from onyx.onyxbot.slack.models import SlackContext +from onyx.onyxbot.slack.models import ChannelType, SlackContext from onyx.redis.redis_pool import get_redis_client from onyx.server.federated.models import FederatedConnectorDetail from onyx.utils.logger import setup_logger diff --git a/backend/onyx/context/search/federated/slack_search_utils.py b/backend/onyx/context/search/federated/slack_search_utils.py index a81bf51ac1..b09bf3f99d 100644 --- a/backend/onyx/context/search/federated/slack_search_utils.py +++ b/backend/onyx/context/search/federated/slack_search_utils.py @@ -1,16 +1,13 @@ import fnmatch import json import re -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from pydantic import ValidationError from onyx.configs.app_configs import MAX_SLACK_QUERY_EXPANSIONS -from onyx.context.search.federated.models import ChannelMetadata -from onyx.context.search.federated.models import DirectThreadFetch +from onyx.context.search.federated.models import ChannelMetadata, DirectThreadFetch from onyx.context.search.models import ChunkIndexRequest from onyx.federated_connectors.slack.models import SlackEntities from onyx.llm.interfaces import LLM @@ -18,11 +15,12 @@ from onyx.llm.models import UserMessage from onyx.llm.utils import llm_response_to_string from onyx.natural_language_processing.english_stopwords import ENGLISH_STOPWORDS_SET from onyx.onyxbot.slack.models import ChannelType -from onyx.prompts.federated_search import SLACK_DATE_EXTRACTION_PROMPT -from onyx.prompts.federated_search import SLACK_QUERY_EXPANSION_PROMPT +from onyx.prompts.federated_search import ( + SLACK_DATE_EXTRACTION_PROMPT, + SLACK_QUERY_EXPANSION_PROMPT, +) from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/context/search/models.py b/backend/onyx/context/search/models.py index adb14fe267..fec40a4723 100644 --- a/backend/onyx/context/search/models.py +++ b/backend/onyx/context/search/models.py @@ -1,18 +1,13 @@ from collections.abc import Sequence -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from enum import Enum from typing import Any -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator -from pydantic import model_validator +from pydantic import BaseModel, Field, field_validator, model_validator from onyx.configs.constants import DocumentSource from onyx.db.models import SearchSettings -from onyx.indexing.models import BaseChunk -from onyx.indexing.models import IndexingSetting +from onyx.indexing.models import BaseChunk, IndexingSetting from onyx.tools.tool_implementations.web_search.models import WEB_SEARCH_PREFIX diff --git a/backend/onyx/context/search/pipeline.py b/backend/onyx/context/search/pipeline.py index a024b9b187..ded669cbb2 100644 --- a/backend/onyx/context/search/pipeline.py +++ b/backend/onyx/context/search/pipeline.py @@ -3,14 +3,16 @@ from datetime import datetime from sqlalchemy.orm import Session -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import ChunkIndexRequest -from onyx.context.search.models import ChunkSearchRequest -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection -from onyx.context.search.models import PersonaSearchInfo -from onyx.context.search.models import TimeRange +from onyx.context.search.models import ( + BaseFilters, + ChunkIndexRequest, + ChunkSearchRequest, + IndexFilters, + InferenceChunk, + InferenceSection, + PersonaSearchInfo, + TimeRange, +) from onyx.context.search.preprocessing.access_filters import ( build_access_filters_for_user, ) diff --git a/backend/onyx/context/search/retrieval/search_runner.py b/backend/onyx/context/search/retrieval/search_runner.py index 8116cfbb5b..fcfbc9d44d 100644 --- a/backend/onyx/context/search/retrieval/search_runner.py +++ b/backend/onyx/context/search/retrieval/search_runner.py @@ -3,19 +3,18 @@ from uuid import UUID from sqlalchemy.orm import Session -from onyx.configs.chat_configs import HYBRID_ALPHA -from onyx.configs.chat_configs import NUM_RETURNED_HITS +from onyx.configs.chat_configs import HYBRID_ALPHA, NUM_RETURNED_HITS from onyx.context.search.enums import QueryType -from onyx.context.search.models import ChunkIndexRequest -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection -from onyx.context.search.utils import get_query_embedding -from onyx.context.search.utils import inference_section_from_chunks -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentSectionRequest -from onyx.federated_connectors.federated_retrieval import FederatedRetrievalInfo +from onyx.context.search.models import ( + ChunkIndexRequest, + IndexFilters, + InferenceChunk, + InferenceSection, +) +from onyx.context.search.utils import get_query_embedding, inference_section_from_chunks +from onyx.document_index.interfaces_new import DocumentIndex, DocumentSectionRequest from onyx.federated_connectors.federated_retrieval import ( + FederatedRetrievalInfo, get_federated_retrieval_functions, ) from onyx.natural_language_processing.search_nlp_models import EmbeddingModel diff --git a/backend/onyx/context/search/utils.py b/backend/onyx/context/search/utils.py index 2425a3944e..a3d54c2aab 100644 --- a/backend/onyx/context/search/utils.py +++ b/backend/onyx/context/search/utils.py @@ -4,28 +4,29 @@ from typing import TypeVar from sqlalchemy.orm import Session -from onyx.configs.app_configs import QUERY_EMBEDDING_CACHE_ENABLED -from onyx.configs.app_configs import QUERY_EMBEDDING_CACHE_TTL_S -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection -from onyx.context.search.models import SavedSearchDoc -from onyx.context.search.models import SavedSearchDocWithContent -from onyx.context.search.models import SearchDoc +from onyx.configs.app_configs import ( + QUERY_EMBEDDING_CACHE_ENABLED, + QUERY_EMBEDDING_CACHE_TTL_S, +) +from onyx.context.search.models import ( + InferenceChunk, + InferenceSection, + SavedSearchDoc, + SavedSearchDocWithContent, + SearchDoc, +) from onyx.db.document import get_document_id_to_file_id_map from onyx.db.models import SearchSettings from onyx.db.search_settings import get_current_search_settings from onyx.natural_language_processing.query_embedding_cache import ( cache_query_embeddings, -) -from onyx.natural_language_processing.query_embedding_cache import ( get_cached_query_embeddings, + record_cache_skipped, ) -from onyx.natural_language_processing.query_embedding_cache import record_cache_skipped from onyx.natural_language_processing.search_nlp_models import EmbeddingModel from onyx.utils.logger import setup_logger from onyx.utils.timing import log_function_time -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT +from shared_configs.configs import MODEL_SERVER_HOST, MODEL_SERVER_PORT from shared_configs.enums import EmbedTextType from shared_configs.model_server_models import Embedding diff --git a/backend/onyx/db/_deprecated/pg_file_store.py b/backend/onyx/db/_deprecated/pg_file_store.py index 4f17650b78..6395d19ed4 100644 --- a/backend/onyx/db/_deprecated/pg_file_store.py +++ b/backend/onyx/db/_deprecated/pg_file_store.py @@ -8,8 +8,7 @@ from psycopg2.extensions import connection from sqlalchemy import text # NEW: for SQL large-object helpers from sqlalchemy.orm import Session -from onyx.file_store.constants import MAX_IN_MEMORY_SIZE -from onyx.file_store.constants import STANDARD_CHUNK_SIZE +from onyx.file_store.constants import MAX_IN_MEMORY_SIZE, STANDARD_CHUNK_SIZE from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/admin_banner.py b/backend/onyx/db/admin_banner.py index 364db5175f..500ee77d5d 100644 --- a/backend/onyx/db/admin_banner.py +++ b/backend/onyx/db/admin_banner.py @@ -1,5 +1,4 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from pydantic import BaseModel diff --git a/backend/onyx/db/api_key.py b/backend/onyx/db/api_key.py index fbb8cbfdd9..2b4a2c9737 100644 --- a/backend/onyx/db/api_key.py +++ b/backend/onyx/db/api_key.py @@ -1,28 +1,26 @@ import uuid from fastapi_users.password import PasswordHelper -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import Session +from sqlalchemy.orm import joinedload, Session -from onyx.auth.api_key import ApiKeyDescriptor -from onyx.auth.api_key import build_displayable_api_key -from onyx.auth.api_key import generate_api_key -from onyx.auth.api_key import hash_api_key +from onyx.auth.api_key import ( + ApiKeyDescriptor, + build_displayable_api_key, + generate_api_key, + hash_api_key, +) from onyx.auth.schemas import UserRole -from onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN -from onyx.configs.constants import DANSWER_API_KEY_PREFIX -from onyx.configs.constants import UNNAMED_KEY_PLACEHOLDER +from onyx.configs.constants import ( + DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN, + DANSWER_API_KEY_PREFIX, + UNNAMED_KEY_PLACEHOLDER, +) from onyx.db.enums import AccountType -from onyx.db.models import ApiKey -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import ApiKey, User, User__UserGroup, UserGroup from onyx.db.permissions import recompute_user_permissions__no_commit -from onyx.db.users import assign_user_to_default_groups__no_commit -from onyx.db.users import delete_user_from_db +from onyx.db.users import assign_user_to_default_groups__no_commit, delete_user_from_db from onyx.server.api_key.models import APIKeyArgs from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/db/auth.py b/backend/onyx/db/auth.py index f1babd89a4..cbac6f059e 100644 --- a/backend/onyx/db/auth.py +++ b/backend/onyx/db/auth.py @@ -1,29 +1,23 @@ -from collections.abc import AsyncGenerator -from collections.abc import Callable -from typing import Any -from typing import Dict -from typing import TypeVar +from collections.abc import AsyncGenerator, Callable +from typing import Any, Dict, TypeVar from fastapi import Depends -from fastapi_users.models import ID -from fastapi_users.models import UP +from fastapi_users.models import ID, UP from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase from fastapi_users_db_sqlalchemy.access_token import SQLAlchemyAccessTokenDatabase -from sqlalchemy import func -from sqlalchemy import Select +from sqlalchemy import func, Select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.configs.constants import ANONYMOUS_USER_EMAIL -from onyx.configs.constants import NO_AUTH_PLACEHOLDER_USER_EMAIL +from onyx.configs.constants import ANONYMOUS_USER_EMAIL, NO_AUTH_PLACEHOLDER_USER_EMAIL from onyx.db.api_key import get_api_key_email_pattern -from onyx.db.engine.async_sql_engine import get_async_session -from onyx.db.engine.async_sql_engine import get_async_session_context_manager -from onyx.db.models import AccessToken -from onyx.db.models import OAuthAccount -from onyx.db.models import User +from onyx.db.engine.async_sql_engine import ( + get_async_session, + get_async_session_context_manager, +) +from onyx.db.models import AccessToken, OAuthAccount, User from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, ) diff --git a/backend/onyx/db/chat.py b/backend/onyx/db/chat.py index 50a902b74b..eee8706222 100644 --- a/backend/onyx/db/chat.py +++ b/backend/onyx/db/chat.py @@ -1,42 +1,31 @@ from collections.abc import Sequence -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Tuple from uuid import UUID from fastapi import HTTPException -from sqlalchemy import delete -from sqlalchemy import desc -from sqlalchemy import func -from sqlalchemy import nullsfirst -from sqlalchemy import or_ -from sqlalchemy import Row -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import delete, desc, func, nullsfirst, or_, Row, select, update from sqlalchemy.exc import MultipleResultsFound -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import joinedload, selectinload, Session from onyx.configs.chat_configs import HARD_DELETE_CHATS from onyx.configs.constants import MessageType -from onyx.context.search.models import InferenceSection -from onyx.context.search.models import SavedSearchDoc +from onyx.context.search.models import InferenceSection, SavedSearchDoc from onyx.context.search.models import SearchDoc as ServerSearchDoc -from onyx.db.models import ChatMessage -from onyx.db.models import ChatMessage__SearchDoc -from onyx.db.models import ChatSession -from onyx.db.models import ChatSessionSharedStatus -from onyx.db.models import Persona +from onyx.db.models import ( + ChatMessage, + ChatMessage__SearchDoc, + ChatSession, + ChatSessionSharedStatus, + Persona, + ToolCall, + User, +) from onyx.db.models import SearchDoc as DBSearchDoc -from onyx.db.models import ToolCall -from onyx.db.models import User from onyx.db.persona import get_best_persona_id_for_user from onyx.file_store.file_store import get_default_file_store from onyx.file_store.models import FileDescriptor -from onyx.llm.override_models import LLMOverride -from onyx.llm.override_models import PromptOverride +from onyx.llm.override_models import LLMOverride, PromptOverride from onyx.server.query_and_chat.models import ChatMessageDetail from onyx.utils.logger import setup_logger from onyx.utils.postgres_sanitization import sanitize_string diff --git a/backend/onyx/db/chat_search.py b/backend/onyx/db/chat_search.py index 8fb511680a..45c9f48597 100644 --- a/backend/onyx/db/chat_search.py +++ b/backend/onyx/db/chat_search.py @@ -1,18 +1,11 @@ -from typing import List -from typing import Optional -from typing import Tuple +from typing import List, Optional, Tuple from uuid import UUID -from sqlalchemy import column -from sqlalchemy import desc -from sqlalchemy import func -from sqlalchemy import select -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import Session +from sqlalchemy import column, desc, func, select +from sqlalchemy.orm import joinedload, Session from sqlalchemy.sql.expression import ColumnClause -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession +from onyx.db.models import ChatMessage, ChatSession def search_chat_sessions( diff --git a/backend/onyx/db/chunk.py b/backend/onyx/db/chunk.py index 3dd7272392..975f1f6408 100644 --- a/backend/onyx/db/chunk.py +++ b/backend/onyx/db/chunk.py @@ -1,5 +1,4 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from sqlalchemy import delete from sqlalchemy.orm import Session diff --git a/backend/onyx/db/connector.py b/backend/onyx/db/connector.py index e8ce0d7130..5cdd94e26f 100644 --- a/backend/onyx/db/connector.py +++ b/backend/onyx/db/connector.py @@ -1,24 +1,20 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from sqlalchemy import and_ -from sqlalchemy import exists -from sqlalchemy import func -from sqlalchemy import select -from sqlalchemy.orm import aliased -from sqlalchemy.orm import Session +from sqlalchemy import and_, exists, func, select +from sqlalchemy.orm import aliased, Session from onyx.configs.app_configs import DEFAULT_PRUNING_FREQ from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType from onyx.db.enums import IndexingMode -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import FederatedConnector -from onyx.db.models import IndexAttempt +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + FederatedConnector, + IndexAttempt, +) from onyx.kg.models import KGConnectorData -from onyx.server.documents.models import ConnectorBase -from onyx.server.documents.models import ObjectCreationIdResponse +from onyx.server.documents.models import ConnectorBase, ObjectCreationIdResponse from onyx.server.models import StatusResponse from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/connector_credential_pair.py b/backend/onyx/db/connector_credential_pair.py index 48bb7ef50e..51ec858bb6 100644 --- a/backend/onyx/db/connector_credential_pair.py +++ b/backend/onyx/db/connector_credential_pair.py @@ -4,38 +4,31 @@ from typing import TypeVarTuple from fastapi import HTTPException from pydantic import BaseModel -from sqlalchemy import delete -from sqlalchemy import desc -from sqlalchemy import exists -from sqlalchemy import func -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import update -from sqlalchemy.orm import aliased -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy import delete, desc, exists, func, Select, select, update +from sqlalchemy.orm import aliased, joinedload, selectinload, Session -from onyx.configs.constants import DEFAULT_CC_PAIR_ID -from onyx.configs.constants import DocumentSource +from onyx.configs.constants import DEFAULT_CC_PAIR_ID, DocumentSource from onyx.db.connector import fetch_connector_by_id -from onyx.db.credentials import fetch_credential_by_id -from onyx.db.credentials import fetch_credential_by_id_for_user +from onyx.db.credentials import fetch_credential_by_id, fetch_credential_by_id_for_user from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingMode -from onyx.db.enums import ProcessingMode -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexingStatus -from onyx.db.models import SearchSettings -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup__ConnectorCredentialPair -from onyx.db.models import UserRole +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + IndexingMode, + ProcessingMode, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + IndexAttempt, + IndexingStatus, + SearchSettings, + User, + User__UserGroup, + UserGroup__ConnectorCredentialPair, + UserRole, +) from onyx.server.models import StatusResponse from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop diff --git a/backend/onyx/db/credentials.py b/backend/onyx/db/credentials.py index 2e98bc9be5..2b656ead62 100644 --- a/backend/onyx/db/credentials.py +++ b/backend/onyx/db/credentials.py @@ -1,22 +1,20 @@ from typing import Any -from sqlalchemy import exists -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import exists, Select, select, update from sqlalchemy.orm import Session -from sqlalchemy.sql.expression import and_ -from sqlalchemy.sql.expression import or_ +from sqlalchemy.sql.expression import and_, or_ from onyx.auth.schemas import UserRole from onyx.configs.constants import DocumentSource from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import Credential__UserGroup -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import User -from onyx.db.models import User__UserGroup +from onyx.db.models import ( + ConnectorCredentialPair, + Credential, + Credential__UserGroup, + DocumentByConnectorCredentialPair, + User, + User__UserGroup, +) from onyx.server.documents.models import CredentialBase from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/deletion_attempt.py b/backend/onyx/db/deletion_attempt.py index 8cda15db34..ef0ed6e79d 100644 --- a/backend/onyx/db/deletion_attempt.py +++ b/backend/onyx/db/deletion_attempt.py @@ -1,8 +1,7 @@ from sqlalchemy.orm import Session from onyx.db.index_attempt import get_last_attempt -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexingStatus +from onyx.db.models import ConnectorCredentialPair, IndexingStatus from onyx.db.search_settings import get_current_search_settings diff --git a/backend/onyx/db/discord_bot.py b/backend/onyx/db/discord_bot.py index b88d0cb358..48f70e1299 100644 --- a/backend/onyx/db/discord_bot.py +++ b/backend/onyx/db/discord_bot.py @@ -1,25 +1,22 @@ """CRUD operations for Discord bot models.""" -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import Session +from sqlalchemy.orm import joinedload, Session -from onyx.auth.api_key import build_displayable_api_key -from onyx.auth.api_key import generate_api_key -from onyx.auth.api_key import hash_api_key +from onyx.auth.api_key import build_displayable_api_key, generate_api_key, hash_api_key from onyx.auth.schemas import UserRole from onyx.configs.constants import DISCORD_SERVICE_API_KEY_NAME from onyx.db.api_key import insert_api_key -from onyx.db.models import ApiKey -from onyx.db.models import DiscordBotConfig -from onyx.db.models import DiscordChannelConfig -from onyx.db.models import DiscordGuildConfig -from onyx.db.models import User +from onyx.db.models import ( + ApiKey, + DiscordBotConfig, + DiscordChannelConfig, + DiscordGuildConfig, + User, +) from onyx.db.utils import DiscordChannelView from onyx.server.api_key.models import APIKeyArgs from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/document.py b/backend/onyx/db/document.py index bbfeb911c4..70a6243d87 100644 --- a/backend/onyx/db/document.py +++ b/backend/onyx/db/document.py @@ -1,61 +1,59 @@ import contextlib import time -from collections.abc import Generator -from collections.abc import Iterable -from collections.abc import Sequence -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from collections.abc import Generator, Iterable, Sequence +from datetime import datetime, timedelta, timezone from typing import Any from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import CompoundSelect -from sqlalchemy import delete -from sqlalchemy import distinct -from sqlalchemy import exists -from sqlalchemy import func -from sqlalchemy import literal -from sqlalchemy import or_ -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import tuple_ -from sqlalchemy import update +from sqlalchemy import ( + and_, + CompoundSelect, + delete, + distinct, + exists, + func, + literal, + or_, + Select, + select, + tuple_, + update, +) from sqlalchemy.dialects.postgresql import insert from sqlalchemy.engine.util import TransactionalContext from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session from sqlalchemy.sql.expression import null -from onyx.configs.constants import DEFAULT_BOOST -from onyx.configs.constants import DocumentSource +from onyx.configs.constants import DEFAULT_BOOST, DocumentSource from onyx.configs.kg_configs import KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES from onyx.db.chunk import delete_chunk_stats_by_connector_credential_pair__no_commit from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.db.document_access import apply_document_access_filter -from onyx.db.entities import delete_from_kg_entities__no_commit -from onyx.db.entities import delete_from_kg_entities_extraction_staging__no_commit -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus +from onyx.db.entities import ( + delete_from_kg_entities__no_commit, + delete_from_kg_entities_extraction_staging__no_commit, +) +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus from onyx.db.feedback import delete_document_feedback_for_documents__no_commit from onyx.db.index_attempt_metrics import safe_record_single_event_if_set from onyx.db.index_attempt_metrics_models import IndexAttemptStage -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + DocumentByConnectorCredentialPair, + KGEntity, + KGRelationship, + User, +) from onyx.db.models import Document as DbDocument -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import KGEntity -from onyx.db.models import KGRelationship -from onyx.db.models import User -from onyx.db.relationships import delete_from_kg_relationships__no_commit from onyx.db.relationships import ( + delete_from_kg_relationships__no_commit, delete_from_kg_relationships_extraction_staging__no_commit, ) from onyx.db.tag import delete_document_tags_for_documents__no_commit -from onyx.db.utils import DocumentRow -from onyx.db.utils import model_to_dict -from onyx.db.utils import SortOrder +from onyx.db.utils import DocumentRow, model_to_dict, SortOrder from onyx.document_index.document_metadata import DocumentMetadata from onyx.file_store.staging import delete_files_best_effort from onyx.kg.models import KGStage diff --git a/backend/onyx/db/document_access.py b/backend/onyx/db/document_access.py index 776ea8945d..09533c4b85 100644 --- a/backend/onyx/db/document_access.py +++ b/backend/onyx/db/document_access.py @@ -10,22 +10,17 @@ This module provides reusable access filtering logic for documents based on: This is a standalone module to avoid circular imports between document.py and persona.py. """ -from sqlalchemy import and_ -from sqlalchemy import any_ -from sqlalchemy import cast -from sqlalchemy import or_ -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import String +from sqlalchemy import and_, any_, cast, or_, Select, select, String from sqlalchemy.dialects import postgresql from sqlalchemy.orm import Session from sqlalchemy.sql.elements import ColumnElement -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Document -from onyx.db.models import DocumentByConnectorCredentialPair +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus +from onyx.db.models import ( + ConnectorCredentialPair, + Document, + DocumentByConnectorCredentialPair, +) def apply_document_access_filter( diff --git a/backend/onyx/db/document_set.py b/backend/onyx/db/document_set.py index 626e733653..1eb098eb24 100644 --- a/backend/onyx/db/document_set.py +++ b/backend/onyx/db/document_set.py @@ -2,35 +2,32 @@ from collections.abc import Sequence from typing import cast from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import delete -from sqlalchemy import exists -from sqlalchemy import func -from sqlalchemy import or_ -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy.orm import aliased -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy import and_, delete, exists, func, or_, Select, select +from sqlalchemy.orm import aliased, selectinload, Session from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.db.connector_credential_pair import get_cc_pair_groups_for_ids -from onyx.db.connector_credential_pair import get_connector_credential_pairs -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus +from onyx.db.connector_credential_pair import ( + get_cc_pair_groups_for_ids, + get_connector_credential_pairs, +) +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus from onyx.db.federated import create_federated_connector_document_set_mapping -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Document -from onyx.db.models import DocumentByConnectorCredentialPair +from onyx.db.models import ( + ConnectorCredentialPair, + Document, + DocumentByConnectorCredentialPair, + DocumentSet__ConnectorCredentialPair, + DocumentSet__UserGroup, + FederatedConnector__DocumentSet, + User, + User__UserGroup, + UserRole, +) from onyx.db.models import DocumentSet as DocumentSetDBModel -from onyx.db.models import DocumentSet__ConnectorCredentialPair -from onyx.db.models import DocumentSet__UserGroup -from onyx.db.models import FederatedConnector__DocumentSet -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserRole -from onyx.server.features.document_set.models import DocumentSetCreationRequest -from onyx.server.features.document_set.models import DocumentSetUpdateRequest +from onyx.server.features.document_set.models import ( + DocumentSetCreationRequest, + DocumentSetUpdateRequest, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_versioned_implementation diff --git a/backend/onyx/db/encrypted_kv_store.py b/backend/onyx/db/encrypted_kv_store.py index 8ce8c7228c..1110eceb18 100644 --- a/backend/onyx/db/encrypted_kv_store.py +++ b/backend/onyx/db/encrypted_kv_store.py @@ -4,8 +4,7 @@ Backed by the ``encrypted_key_value_store`` table and never mirrored to the cach secrets stay encrypted at rest and out of Redis. Tenant context is picked up from the current-tenant contextvar, matching the plain KV store.""" -from typing import Any -from typing import cast +from typing import Any, cast from sqlalchemy.dialects.postgresql import insert as pg_insert diff --git a/backend/onyx/db/engine/async_sql_engine.py b/backend/onyx/db/engine/async_sql_engine.py index 0b1316fc3b..e0af1673d8 100644 --- a/backend/onyx/db/engine/async_sql_engine.py +++ b/backend/onyx/db/engine/async_sql_engine.py @@ -1,33 +1,32 @@ from collections.abc import AsyncGenerator from contextlib import asynccontextmanager -from typing import Any -from typing import AsyncContextManager +from typing import Any, AsyncContextManager from fastapi import HTTPException -from sqlalchemy import event -from sqlalchemy import pool -from sqlalchemy.ext.asyncio import AsyncEngine -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy import event, pool +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine -from onyx.configs.app_configs import AWS_REGION_NAME -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE -from onyx.configs.app_configs import POSTGRES_HOST -from onyx.configs.app_configs import POSTGRES_POOL_PRE_PING -from onyx.configs.app_configs import POSTGRES_POOL_RECYCLE -from onyx.configs.app_configs import POSTGRES_PORT -from onyx.configs.app_configs import POSTGRES_USE_NULL_POOL -from onyx.configs.app_configs import POSTGRES_USER +from onyx.configs.app_configs import ( + AWS_REGION_NAME, + POSTGRES_API_SERVER_POOL_OVERFLOW, + POSTGRES_API_SERVER_POOL_SIZE, + POSTGRES_HOST, + POSTGRES_POOL_PRE_PING, + POSTGRES_POOL_RECYCLE, + POSTGRES_PORT, + POSTGRES_USE_NULL_POOL, + POSTGRES_USER, +) from onyx.db.engine.iam_auth import get_iam_auth_token from onyx.db.engine.pg_ssl import create_pg_ssl_context -from onyx.db.engine.sql_engine import ASYNC_DB_API -from onyx.db.engine.sql_engine import build_connection_string -from onyx.db.engine.sql_engine import is_valid_schema_name -from onyx.db.engine.sql_engine import SqlEngine -from onyx.db.engine.sql_engine import USE_IAM_AUTH -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE +from onyx.db.engine.sql_engine import ( + ASYNC_DB_API, + build_connection_string, + is_valid_schema_name, + SqlEngine, + USE_IAM_AUTH, +) +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import get_current_tenant_id # Global so we don't create more than one engine per process diff --git a/backend/onyx/db/engine/iam_auth.py b/backend/onyx/db/engine/iam_auth.py index 6843839ba7..7f6668ead3 100644 --- a/backend/onyx/db/engine/iam_auth.py +++ b/backend/onyx/db/engine/iam_auth.py @@ -5,10 +5,12 @@ from typing import Any import boto3 -from onyx.configs.app_configs import POSTGRES_HOST -from onyx.configs.app_configs import POSTGRES_PORT -from onyx.configs.app_configs import POSTGRES_USER -from onyx.configs.app_configs import USE_IAM_AUTH +from onyx.configs.app_configs import ( + POSTGRES_HOST, + POSTGRES_PORT, + POSTGRES_USER, + USE_IAM_AUTH, +) from onyx.configs.constants import SSL_CERT_FILE diff --git a/backend/onyx/db/engine/pg_ssl.py b/backend/onyx/db/engine/pg_ssl.py index ec6c927931..44e8f18ff1 100644 --- a/backend/onyx/db/engine/pg_ssl.py +++ b/backend/onyx/db/engine/pg_ssl.py @@ -20,12 +20,14 @@ the explicit ``POSTGRES_SSL*`` settings only apply when IAM is off. import functools import ssl -from onyx.configs.app_configs import POSTGRES_SSLCERT -from onyx.configs.app_configs import POSTGRES_SSLKEY -from onyx.configs.app_configs import POSTGRES_SSLKEY_PASSWORD -from onyx.configs.app_configs import POSTGRES_SSLMODE -from onyx.configs.app_configs import POSTGRES_SSLROOTCERT -from onyx.configs.app_configs import USE_IAM_AUTH +from onyx.configs.app_configs import ( + POSTGRES_SSLCERT, + POSTGRES_SSLKEY, + POSTGRES_SSLKEY_PASSWORD, + POSTGRES_SSLMODE, + POSTGRES_SSLROOTCERT, + USE_IAM_AUTH, +) from onyx.db.engine.iam_auth import create_ssl_context_if_iam from onyx.utils.tls import build_ssl_context diff --git a/backend/onyx/db/engine/sql_engine.py b/backend/onyx/db/engine/sql_engine.py index 9079c0d1f8..6c9f9497d6 100644 --- a/backend/onyx/db/engine/sql_engine.py +++ b/backend/onyx/db/engine/sql_engine.py @@ -7,39 +7,43 @@ from contextlib import contextmanager from typing import Any from fastapi import HTTPException -from sqlalchemy import event -from sqlalchemy import pool -from sqlalchemy.engine import create_engine -from sqlalchemy.engine import Engine +from sqlalchemy import event, pool +from sqlalchemy.engine import create_engine, Engine from sqlalchemy.exc import DBAPIError from sqlalchemy.orm import Session -from onyx.configs.app_configs import DB_READONLY_PASSWORD -from onyx.configs.app_configs import DB_READONLY_USER -from onyx.configs.app_configs import LOG_POSTGRES_CONN_COUNTS -from onyx.configs.app_configs import LOG_POSTGRES_LATENCY -from onyx.configs.app_configs import POSTGRES_DB -from onyx.configs.app_configs import POSTGRES_HOST -from onyx.configs.app_configs import POSTGRES_PASSWORD -from onyx.configs.app_configs import POSTGRES_POOL_PRE_PING -from onyx.configs.app_configs import POSTGRES_POOL_RECYCLE -from onyx.configs.app_configs import POSTGRES_PORT -from onyx.configs.app_configs import POSTGRES_TCP_KEEPALIVES -from onyx.configs.app_configs import POSTGRES_TCP_KEEPALIVES_COUNT -from onyx.configs.app_configs import POSTGRES_TCP_KEEPALIVES_IDLE -from onyx.configs.app_configs import POSTGRES_TCP_KEEPALIVES_INTERVAL -from onyx.configs.app_configs import POSTGRES_USE_NULL_POOL -from onyx.configs.app_configs import POSTGRES_USER +from onyx.configs.app_configs import ( + DB_READONLY_PASSWORD, + DB_READONLY_USER, + LOG_POSTGRES_CONN_COUNTS, + LOG_POSTGRES_LATENCY, + POSTGRES_DB, + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_POOL_PRE_PING, + POSTGRES_POOL_RECYCLE, + POSTGRES_PORT, + POSTGRES_TCP_KEEPALIVES, + POSTGRES_TCP_KEEPALIVES_COUNT, + POSTGRES_TCP_KEEPALIVES_IDLE, + POSTGRES_TCP_KEEPALIVES_INTERVAL, + POSTGRES_USE_NULL_POOL, + POSTGRES_USER, +) from onyx.configs.constants import POSTGRES_UNKNOWN_APP_NAME from onyx.db.engine.iam_auth import provide_iam_token from onyx.db.engine.pg_ssl import pg_ssl_psycopg2_connect_args from onyx.server.utils import BasicAuthenticationError from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.configs import ( + MULTI_TENANT, + POSTGRES_DEFAULT_SCHEMA, + POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, +) +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) # Moved is_valid_schema_name here to avoid circular import diff --git a/backend/onyx/db/engine/tenant_utils.py b/backend/onyx/db/engine/tenant_utils.py index 4e71e3b5aa..8c3b49f249 100644 --- a/backend/onyx/db/engine/tenant_utils.py +++ b/backend/onyx/db/engine/tenant_utils.py @@ -2,11 +2,12 @@ import re from sqlalchemy import text -from onyx.db.engine.sql_engine import get_session_with_shared_schema -from onyx.db.engine.sql_engine import SqlEngine -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import TENANT_ID_PREFIX +from onyx.db.engine.sql_engine import get_session_with_shared_schema, SqlEngine +from shared_configs.configs import ( + MULTI_TENANT, + POSTGRES_DEFAULT_SCHEMA, + TENANT_ID_PREFIX, +) # Regex pattern for valid tenant IDs: # - UUID format: tenant_xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx diff --git a/backend/onyx/db/entities.py b/backend/onyx/db/entities.py index ade11bfd51..79ea582fa9 100644 --- a/backend/onyx/db/entities.py +++ b/backend/onyx/db/entities.py @@ -1,24 +1,16 @@ import uuid -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import List -from sqlalchemy import func -from sqlalchemy import literal -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import func, literal, select, update from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import Session import onyx.db.document as dbdocument from onyx.db.entity_type import UNGROUNDED_SOURCE_NAME -from onyx.db.models import Document -from onyx.db.models import KGEntity -from onyx.db.models import KGEntityExtractionStaging -from onyx.db.models import KGEntityType -from onyx.kg.models import KGGroundingType -from onyx.kg.models import KGStage +from onyx.db.models import Document, KGEntity, KGEntityExtractionStaging, KGEntityType +from onyx.kg.models import KGGroundingType, KGStage from onyx.kg.utils.formatting_utils import make_entity_id diff --git a/backend/onyx/db/entity_type.py b/backend/onyx/db/entity_type.py index 50e93d08d1..d902c8e7d3 100644 --- a/backend/onyx/db/entity_type.py +++ b/backend/onyx/db/entity_type.py @@ -5,8 +5,7 @@ from sqlalchemy.orm import Session from onyx.db.connector import fetch_unique_document_sources from onyx.db.document import DocumentSource -from onyx.db.models import Connector -from onyx.db.models import KGEntityType +from onyx.db.models import Connector, KGEntityType from onyx.kg.models import KGAttributeEntityOption from onyx.server.kg.models import EntityType diff --git a/backend/onyx/db/external_app.py b/backend/onyx/db/external_app.py index 73d6af802c..4e941dfa7e 100644 --- a/backend/onyx/db/external_app.py +++ b/backend/onyx/db/external_app.py @@ -1,28 +1,20 @@ import re -from typing import Any -from typing import cast -from uuid import UUID -from uuid import uuid4 +from typing import Any, cast +from uuid import UUID, uuid4 -from sqlalchemy import and_ -from sqlalchemy import delete -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import and_, delete, func, select from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.enums import SkillSharePermission -from onyx.db.models import ExternalApp -from onyx.db.models import ExternalAppPolicy -from onyx.db.models import ExternalAppUserCredential -from onyx.db.models import Skill -from onyx.db.models import User -from onyx.db.utils import is_set -from onyx.db.utils import UNSET -from onyx.db.utils import UnsetType +from onyx.db.enums import EndpointPolicy, ExternalAppType, SkillSharePermission +from onyx.db.models import ( + ExternalApp, + ExternalAppPolicy, + ExternalAppUserCredential, + Skill, + User, +) +from onyx.db.utils import is_set, UNSET, UnsetType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.skills.built_in import EXTERNAL_APP_BUILT_IN_SKILL_IDS @@ -295,8 +287,10 @@ def create_external_app( CUSTOM apps get a bundle-backed skill using ``slug``, or a generated ``custom-`` slug when omitted. """ - from onyx.db.skill import create_built_in_skill_row__no_commit - from onyx.db.skill import create_skill__no_commit + from onyx.db.skill import ( + create_built_in_skill_row__no_commit, + create_skill__no_commit, + ) # No existing app to restore from on create, so a masked value is rejected. organization_credentials = resolve_masked_credentials( diff --git a/backend/onyx/db/federated.py b/backend/onyx/db/federated.py index 6f12545e27..809f6142ce 100644 --- a/backend/onyx/db/federated.py +++ b/backend/onyx/db/federated.py @@ -3,16 +3,16 @@ from typing import Any from uuid import UUID from sqlalchemy import select -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import joinedload, selectinload, Session from onyx.configs.constants import FederatedConnectorSource from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import DocumentSet -from onyx.db.models import FederatedConnector -from onyx.db.models import FederatedConnector__DocumentSet -from onyx.db.models import FederatedConnectorOAuthToken +from onyx.db.models import ( + DocumentSet, + FederatedConnector, + FederatedConnector__DocumentSet, + FederatedConnectorOAuthToken, +) from onyx.federated_connectors.factory import get_federated_connector from onyx.utils.encryption import reject_masked_credentials from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/feedback.py b/backend/onyx/db/feedback.py index fae3eb5da7..95568be79d 100644 --- a/backend/onyx/db/feedback.py +++ b/backend/onyx/db/feedback.py @@ -1,31 +1,24 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import UUID from fastapi import HTTPException -from sqlalchemy import and_ -from sqlalchemy import asc -from sqlalchemy import delete -from sqlalchemy import desc -from sqlalchemy import exists -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy.orm import aliased -from sqlalchemy.orm import Session +from sqlalchemy import and_, asc, delete, desc, exists, Select, select +from sqlalchemy.orm import aliased, Session -from onyx.configs.constants import MessageType -from onyx.configs.constants import SearchFeedbackType +from onyx.configs.constants import MessageType, SearchFeedbackType from onyx.db.chat import get_chat_message from onyx.db.enums import AccessType -from onyx.db.models import ChatMessageFeedback -from onyx.db.models import ConnectorCredentialPair +from onyx.db.models import ( + ChatMessageFeedback, + ConnectorCredentialPair, + DocumentByConnectorCredentialPair, + DocumentRetrievalFeedback, + User, + User__UserGroup, + UserGroup__ConnectorCredentialPair, + UserRole, +) from onyx.db.models import Document as DbDocument -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import DocumentRetrievalFeedback -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup__ConnectorCredentialPair -from onyx.db.models import UserRole from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/file_record.py b/backend/onyx/db/file_record.py index aa50299539..8391e21673 100644 --- a/backend/onyx/db/file_record.py +++ b/backend/onyx/db/file_record.py @@ -1,16 +1,11 @@ -from sqlalchemy import and_ -from sqlalchemy import cast -from sqlalchemy import select -from sqlalchemy import String +from sqlalchemy import and_, cast, select, String from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session from onyx.background.task_utils import QUERY_REPORT_NAME_PREFIX -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import FileType +from onyx.configs.constants import FileOrigin, FileType from onyx.db.enums import IndexingStatus -from onyx.db.models import FileRecord -from onyx.db.models import IndexAttempt +from onyx.db.models import FileRecord, IndexAttempt def get_query_history_export_files( diff --git a/backend/onyx/db/hierarchy.py b/backend/onyx/db/hierarchy.py index c03adbde52..1586da1dbb 100644 --- a/backend/onyx/db/hierarchy.py +++ b/backend/onyx/db/hierarchy.py @@ -2,8 +2,7 @@ from collections import defaultdict -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.engine import CursorResult from sqlalchemy.exc import IntegrityError @@ -12,9 +11,11 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.models import HierarchyNode as PydanticHierarchyNode from onyx.db.enums import HierarchyNodeType -from onyx.db.models import Document -from onyx.db.models import HierarchyNode -from onyx.db.models import HierarchyNodeByConnectorCredentialPair +from onyx.db.models import ( + Document, + HierarchyNode, + HierarchyNodeByConnectorCredentialPair, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_versioned_implementation diff --git a/backend/onyx/db/hook.py b/backend/onyx/db/hook.py index f9d47c9771..156f8f5139 100644 --- a/backend/onyx/db/hook.py +++ b/backend/onyx/db/hook.py @@ -1,19 +1,14 @@ import datetime from uuid import UUID -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.engine import CursorResult from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session -from onyx.db.constants import UNSET -from onyx.db.constants import UnsetType -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint -from onyx.db.models import Hook -from onyx.db.models import HookExecutionLog +from onyx.db.constants import UNSET, UnsetType +from onyx.db.enums import HookFailStrategy, HookPoint +from onyx.db.models import Hook, HookExecutionLog from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/onyx/db/image_generation.py b/backend/onyx/db/image_generation.py index 63b81d31f8..8e76116c95 100644 --- a/backend/onyx/db/image_generation.py +++ b/backend/onyx/db/image_generation.py @@ -1,11 +1,7 @@ -from sqlalchemy import select -from sqlalchemy import update -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy import select, update +from sqlalchemy.orm import selectinload, Session -from onyx.db.models import ImageGenerationConfig -from onyx.db.models import LLMProvider -from onyx.db.models import ModelConfiguration +from onyx.db.models import ImageGenerationConfig, LLMProvider, ModelConfiguration from onyx.llm.model_capabilities import get_max_input_tokens from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/index_attempt.py b/backend/onyx/db/index_attempt.py index 7023e37c51..40b3efaa9f 100644 --- a/backend/onyx/db/index_attempt.py +++ b/backend/onyx/db/index_attempt.py @@ -1,34 +1,28 @@ from collections.abc import Sequence -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import TypeVarTuple -from sqlalchemy import and_ -from sqlalchemy import delete -from sqlalchemy import desc -from sqlalchemy import func -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import update -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import Session +from sqlalchemy import and_, delete, desc, func, Select, select, update +from sqlalchemy.orm import joinedload, Session from onyx.connectors.models import ConnectorFailure from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError -from onyx.db.models import SearchSettings +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + IndexingStatus, + IndexModelStatus, +) +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + IndexAttemptError, + SearchSettings, +) from onyx.redis.redis_docprocessing import RedisDocprocessing from onyx.redis.redis_pool import get_redis_client from onyx.server.documents.models import ConnectorCredentialPairIdentifier from onyx.utils.logger import setup_logger -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType +from onyx.utils.telemetry import optional_telemetry, RecordType logger = setup_logger() diff --git a/backend/onyx/db/index_attempt_metrics.py b/backend/onyx/db/index_attempt_metrics.py index 38a1c98984..a21dd1d051 100644 --- a/backend/onyx/db/index_attempt_metrics.py +++ b/backend/onyx/db/index_attempt_metrics.py @@ -20,11 +20,7 @@ import time from collections.abc import Generator from contextlib import contextmanager -from sqlalchemy import case -from sqlalchemy import cast -from sqlalchemy import Float -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import case, cast, Float, func, select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session diff --git a/backend/onyx/db/indexing_coordination.py b/backend/onyx/db/indexing_coordination.py index 1815dcbf2f..a70741c6e5 100644 --- a/backend/onyx/db/indexing_coordination.py +++ b/backend/onyx/db/indexing_coordination.py @@ -6,9 +6,11 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session from onyx.db.enums import IndexingStatus -from onyx.db.index_attempt import count_error_rows_for_index_attempt -from onyx.db.index_attempt import create_index_attempt -from onyx.db.index_attempt import get_index_attempt +from onyx.db.index_attempt import ( + count_error_rows_for_index_attempt, + create_index_attempt, + get_index_attempt, +) from onyx.db.models import IndexAttempt from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/input_prompt.py b/backend/onyx/db/input_prompt.py index f133619a19..678288a3e3 100644 --- a/backend/onyx/db/input_prompt.py +++ b/backend/onyx/db/input_prompt.py @@ -1,16 +1,12 @@ from uuid import UUID from fastapi import HTTPException -from sqlalchemy import or_ -from sqlalchemy import select +from sqlalchemy import or_, select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import aliased -from sqlalchemy.orm import Session +from sqlalchemy.orm import aliased, Session -from onyx.db.models import InputPrompt -from onyx.db.models import InputPrompt__User -from onyx.db.models import User +from onyx.db.models import InputPrompt, InputPrompt__User, User from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/llm.py b/backend/onyx/db/llm.py index 796f1882bd..162f39c408 100644 --- a/backend/onyx/db/llm.py +++ b/backend/onyx/db/llm.py @@ -1,35 +1,36 @@ -from sqlalchemy import delete -from sqlalchemy import or_ -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import delete, or_, select, update from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.orm import load_only -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import load_only, selectinload, Session from onyx.auth.schemas import UserRole from onyx.db.enums import LLMModelFlowType from onyx.db.models import CloudEmbeddingProvider as CloudEmbeddingProviderModel -from onyx.db.models import DocumentSet -from onyx.db.models import ImageGenerationConfig -from onyx.db.models import LLMModelFlow +from onyx.db.models import ( + DocumentSet, + ImageGenerationConfig, + LLMModelFlow, + LLMProvider__Persona, + LLMProvider__UserGroup, + ModelConfiguration, + Persona, + SearchSettings, + User, + User__UserGroup, + UserGroup, +) from onyx.db.models import LLMProvider as LLMProviderModel -from onyx.db.models import LLMProvider__Persona -from onyx.db.models import LLMProvider__UserGroup -from onyx.db.models import ModelConfiguration -from onyx.db.models import Persona -from onyx.db.models import SearchSettings from onyx.db.models import Tool as ToolModel -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup from onyx.llm.utils import model_supports_image_input from onyx.llm.well_known_providers.auto_update_models import LLMRecommendations -from onyx.server.manage.embedding.models import CloudEmbeddingProvider -from onyx.server.manage.embedding.models import CloudEmbeddingProviderCreationRequest -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import SyncModelEntry +from onyx.server.manage.embedding.models import ( + CloudEmbeddingProvider, + CloudEmbeddingProviderCreationRequest, +) +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + LLMProviderView, + SyncModelEntry, +) from onyx.utils.logger import setup_logger from shared_configs.enums import EmbeddingProvider diff --git a/backend/onyx/db/mcp.py b/backend/onyx/db/mcp.py index 1a713963c4..b65e4d9ef5 100644 --- a/backend/onyx/db/mcp.py +++ b/backend/onyx/db/mcp.py @@ -2,34 +2,31 @@ import datetime from typing import cast from uuid import UUID -from pydantic import BaseModel -from pydantic import ConfigDict -from sqlalchemy import and_ -from sqlalchemy import delete -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy.orm import aliased -from sqlalchemy.orm import Session +from pydantic import BaseModel, ConfigDict +from sqlalchemy import and_, delete, Select, select +from sqlalchemy.orm import aliased, Session from sqlalchemy.orm.attributes import flag_modified -from onyx.db.constants import UNSET -from onyx.db.constants import UnsetType -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPOAuthProviderMode -from onyx.db.enums import MCPServerStatus -from onyx.db.enums import MCPTransport -from onyx.db.models import MCPAuthenticationType -from onyx.db.models import MCPConnectionConfig -from onyx.db.models import MCPServer -from onyx.db.models import MCPServer__User -from onyx.db.models import MCPServer__UserGroup -from onyx.db.models import Persona -from onyx.db.models import Tool -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserRole -from onyx.server.features.mcp.models import DENYLISTED_MCP_HEADERS -from onyx.server.features.mcp.models import MCPConnectionData +from onyx.db.constants import UNSET, UnsetType +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPOAuthProviderMode, + MCPServerStatus, + MCPTransport, +) +from onyx.db.models import ( + MCPAuthenticationType, + MCPConnectionConfig, + MCPServer, + MCPServer__User, + MCPServer__UserGroup, + Persona, + Tool, + User, + User__UserGroup, + UserRole, +) +from onyx.server.features.mcp.models import DENYLISTED_MCP_HEADERS, MCPConnectionData from onyx.utils.logger import setup_logger from onyx.utils.sensitive import SensitiveValue diff --git a/backend/onyx/db/memory.py b/backend/onyx/db/memory.py index 316064e9dd..4037a23d9a 100644 --- a/backend/onyx/db/memory.py +++ b/backend/onyx/db/memory.py @@ -1,16 +1,15 @@ from uuid import UUID -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field from sqlalchemy import select from sqlalchemy.orm import Session -from onyx.auth.oauth_claims_capture import get_idp_profile_fields -from onyx.auth.oauth_claims_capture import get_idp_profile_placeholder_values +from onyx.auth.oauth_claims_capture import ( + get_idp_profile_fields, + get_idp_profile_placeholder_values, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none -from onyx.db.models import Memory -from onyx.db.models import User +from onyx.db.models import Memory, User MAX_MEMORIES_PER_USER = 10 diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index f56e98a2c4..7882222bec 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -1,124 +1,125 @@ import datetime import json -from typing import Any -from typing import Literal -from typing import NotRequired -from uuid import UUID -from uuid import uuid4 +from typing import Any, Literal, NotRequired +from uuid import UUID, uuid4 -from fastapi_users_db_sqlalchemy import SQLAlchemyBaseOAuthAccountTableUUID -from fastapi_users_db_sqlalchemy import SQLAlchemyBaseUserTableUUID +from fastapi_users_db_sqlalchemy import ( + SQLAlchemyBaseOAuthAccountTableUUID, + SQLAlchemyBaseUserTableUUID, +) from fastapi_users_db_sqlalchemy.access_token import SQLAlchemyBaseAccessTokenTableUUID from fastapi_users_db_sqlalchemy.generics import TIMESTAMPAware -from pydantic import BaseModel -from pydantic import ValidationError -from sqlalchemy import BigInteger -from sqlalchemy import Boolean -from sqlalchemy import CheckConstraint -from sqlalchemy import DateTime -from sqlalchemy import desc -from sqlalchemy import Enum -from sqlalchemy import event -from sqlalchemy import Float -from sqlalchemy import ForeignKey -from sqlalchemy import ForeignKeyConstraint -from sqlalchemy import func -from sqlalchemy import Index -from sqlalchemy import Integer -from sqlalchemy import Numeric -from sqlalchemy import PrimaryKeyConstraint -from sqlalchemy import Sequence -from sqlalchemy import String -from sqlalchemy import Text -from sqlalchemy import text -from sqlalchemy import true -from sqlalchemy import UniqueConstraint +from pydantic import BaseModel, ValidationError +from sqlalchemy import ( + BigInteger, + Boolean, + CheckConstraint, + DateTime, + desc, + Enum, + event, + Float, + ForeignKey, + ForeignKeyConstraint, + func, + Index, + Integer, + Numeric, + PrimaryKeyConstraint, + Sequence, + String, + Text, + text, + true, + UniqueConstraint, +) from sqlalchemy.dialects import postgresql from sqlalchemy.dialects.postgresql import JSONB as PGJSONB from sqlalchemy.dialects.postgresql import UUID as PGUUID from sqlalchemy.engine.interfaces import Dialect -from sqlalchemy.orm import DeclarativeBase -from sqlalchemy.orm import Mapped -from sqlalchemy.orm import mapped_column -from sqlalchemy.orm import Mapper -from sqlalchemy.orm import relationship -from sqlalchemy.orm import validates -from sqlalchemy.types import LargeBinary -from sqlalchemy.types import TypeDecorator +from sqlalchemy.orm import ( + DeclarativeBase, + Mapped, + mapped_column, + Mapper, + relationship, + validates, +) +from sqlalchemy.types import LargeBinary, TypeDecorator from typing_extensions import TypedDict # noreorder from onyx.auth.schemas import UserRole -from onyx.configs.constants import ANONYMOUS_USER_UUID -from onyx.configs.constants import DEFAULT_BOOST -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FederatedConnectorSource -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import MessageType -from onyx.configs.constants import MilestoneRecordType -from onyx.configs.constants import NotificationType -from onyx.configs.constants import SearchFeedbackType -from onyx.configs.constants import TokenRateLimitScope +from onyx.configs.constants import ( + ANONYMOUS_USER_UUID, + DEFAULT_BOOST, + DocumentSource, + FederatedConnectorSource, + FileOrigin, + MessageType, + MilestoneRecordType, + NotificationType, + SearchFeedbackType, + TokenRateLimitScope, +) from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import AccountType -from onyx.db.enums import ApprovalDecidedVia -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import ArtifactType -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import ChatSessionSharedStatus -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import DefaultAppMode -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.enums import GrantSource -from onyx.db.enums import HierarchyNodeType -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint -from onyx.db.enums import IndexingMode -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.enums import LLMModelFlowType -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPOAuthProviderMode -from onyx.db.enums import MCPServerStatus -from onyx.db.enums import MCPTransport -from onyx.db.enums import OpenSearchDocumentMigrationStatus -from onyx.db.enums import OpenSearchTenantMigrationStatus -from onyx.db.enums import PatType -from onyx.db.enums import Permission -from onyx.db.enums import PermissionSyncStatus -from onyx.db.enums import PersonaSharePermission -from onyx.db.enums import PortAttemptStatus -from onyx.db.enums import ProcessingMode -from onyx.db.enums import SandboxStatus -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import ScheduledTaskStatus -from onyx.db.enums import ScheduledTaskTriggerSource -from onyx.db.enums import SessionOrigin -from onyx.db.enums import SharingScope -from onyx.db.enums import SkillSharePermission -from onyx.db.enums import SSOProviderType -from onyx.db.enums import SwitchoverType -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType -from onyx.db.enums import TaskStatus -from onyx.db.enums import ThemePreference -from onyx.db.enums import UserFileStatus +from onyx.db.enums import ( + AccessType, + AccountType, + ApprovalDecidedVia, + ApprovalDecision, + ArtifactType, + BuildSessionStatus, + ChatSessionSharedStatus, + ConnectorCredentialPairStatus, + DefaultAppMode, + EmbeddingPrecision, + EndpointPolicy, + ExternalAppType, + GrantSource, + HierarchyNodeType, + HookFailStrategy, + HookPoint, + IndexingMode, + IndexingStatus, + IndexModelStatus, + LLMModelFlowType, + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPOAuthProviderMode, + MCPServerStatus, + MCPTransport, + OpenSearchDocumentMigrationStatus, + OpenSearchTenantMigrationStatus, + PatType, + Permission, + PermissionSyncStatus, + PersonaSharePermission, + PortAttemptStatus, + ProcessingMode, + SandboxStatus, + ScheduledTaskRunStatus, + ScheduledTaskStatus, + ScheduledTaskTriggerSource, + SessionOrigin, + SharingScope, + SkillSharePermission, + SSOProviderType, + SwitchoverType, + SyncStatus, + SyncType, + TaskStatus, + ThemePreference, + UserFileStatus, +) from onyx.db.index_attempt_metrics_models import IndexAttemptStage -from onyx.db.pydantic_type import PydanticListType -from onyx.db.pydantic_type import PydanticType +from onyx.db.pydantic_type import PydanticListType, PydanticType from onyx.external_apps.url_glob import UrlGlob from onyx.file_store.models import FileDescriptor -from onyx.kg.models import KGEntityTypeAttributes -from onyx.kg.models import KGStage -from onyx.llm.override_models import LLMOverride -from onyx.llm.override_models import PromptOverride +from onyx.kg.models import KGEntityTypeAttributes, KGStage +from onyx.llm.override_models import LLMOverride, PromptOverride from onyx.server.security.models import SSRFProtectionLevel from onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig -from onyx.utils.encryption import decrypt_bytes_to_string -from onyx.utils.encryption import encrypt_string_to_bytes +from onyx.utils.encryption import decrypt_bytes_to_string, encrypt_string_to_bytes from onyx.utils.headers import HeaderItemDict from onyx.utils.logger import setup_logger from onyx.utils.sensitive import SensitiveValue diff --git a/backend/onyx/db/notification.py b/backend/onyx/db/notification.py index c174223466..6fe6613b56 100644 --- a/backend/onyx/db/notification.py +++ b/backend/onyx/db/notification.py @@ -1,10 +1,7 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import UUID -from sqlalchemy import cast -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import cast, select, update from sqlalchemy.dialects import postgresql from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session @@ -13,8 +10,7 @@ from sqlalchemy.sql.elements import ColumnElement from onyx.auth.schemas import UserRole from onyx.configs.constants import NotificationType -from onyx.db.models import Notification -from onyx.db.models import User +from onyx.db.models import Notification, User def _notification_filters( diff --git a/backend/onyx/db/oauth_config.py b/backend/onyx/db/oauth_config.py index e85080e37f..ae67b7a85f 100644 --- a/backend/onyx/db/oauth_config.py +++ b/backend/onyx/db/oauth_config.py @@ -4,9 +4,7 @@ from uuid import UUID from sqlalchemy import select from sqlalchemy.orm import Session -from onyx.db.models import OAuthConfig -from onyx.db.models import OAuthUserToken -from onyx.db.models import Tool +from onyx.db.models import OAuthConfig, OAuthUserToken, Tool from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/opensearch_migration.py b/backend/onyx/db/opensearch_migration.py index 709b678d32..d322b5a591 100644 --- a/backend/onyx/db/opensearch_migration.py +++ b/backend/onyx/db/opensearch_migration.py @@ -5,26 +5,26 @@ from Vespa to OpenSearch. """ import json -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from sqlalchemy import select -from sqlalchemy import text +from sqlalchemy import select, text from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session from onyx.background.celery.tasks.opensearch_migration.constants import ( GET_VESPA_CHUNKS_SLICE_COUNT, -) -from onyx.background.celery.tasks.opensearch_migration.constants import ( TOTAL_ALLOWABLE_DOC_MIGRATION_ATTEMPTS_BEFORE_PERMANENT_FAILURE, ) -from onyx.configs.app_configs import ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX -from onyx.configs.app_configs import ONYX_DISABLE_VESPA +from onyx.configs.app_configs import ( + ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX, + ONYX_DISABLE_VESPA, +) from onyx.db.enums import OpenSearchDocumentMigrationStatus -from onyx.db.models import Document -from onyx.db.models import OpenSearchDocumentMigrationRecord -from onyx.db.models import OpenSearchTenantMigrationRecord +from onyx.db.models import ( + Document, + OpenSearchDocumentMigrationRecord, + OpenSearchTenantMigrationRecord, +) from onyx.document_index.vespa.shared_utils.utils import ( replace_invalid_doc_id_characters, ) diff --git a/backend/onyx/db/pat.py b/backend/onyx/db/pat.py index 7f037c9605..268aac985e 100644 --- a/backend/onyx/db/pat.py +++ b/backend/onyx/db/pat.py @@ -1,26 +1,23 @@ """Database operations for Personal Access Tokens.""" import asyncio -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import NamedTuple from uuid import UUID -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import contains_eager -from sqlalchemy.orm import Session +from sqlalchemy.orm import contains_eager, Session -from onyx.auth.pat import build_displayable_pat -from onyx.auth.pat import calculate_expiration -from onyx.auth.pat import generate_pat -from onyx.auth.pat import hash_pat +from onyx.auth.pat import ( + build_displayable_pat, + calculate_expiration, + generate_pat, + hash_pat, +) from onyx.db.engine.async_sql_engine import get_async_session_context_manager -from onyx.db.enums import PatType -from onyx.db.enums import Permission -from onyx.db.models import PersonalAccessToken -from onyx.db.models import User +from onyx.db.enums import PatType, Permission +from onyx.db.models import PersonalAccessToken, User from onyx.db.permissions import parse_permission_values from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/db/permission_sync_attempt.py b/backend/onyx/db/permission_sync_attempt.py index 454efb2766..4e0f3adb75 100644 --- a/backend/onyx/db/permission_sync_attempt.py +++ b/backend/onyx/db/permission_sync_attempt.py @@ -4,25 +4,22 @@ This module contains all CRUD operations for both DocPermissionSyncAttempt and ExternalGroupPermissionSyncAttempt models, along with shared utilities. """ -from typing import Any -from typing import cast +from typing import Any, cast -from sqlalchemy import delete -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import delete, func, select from sqlalchemy.engine.cursor import CursorResult -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import Session +from sqlalchemy.orm import joinedload, Session from onyx.configs.constants import DocumentSource from onyx.db.enums import PermissionSyncStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import DocPermissionSyncAttempt -from onyx.db.models import ExternalGroupPermissionSyncAttempt +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + DocPermissionSyncAttempt, + ExternalGroupPermissionSyncAttempt, +) from onyx.utils.logger import setup_logger -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType +from onyx.utils.telemetry import optional_telemetry, RecordType from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop logger = setup_logger() diff --git a/backend/onyx/db/permissions.py b/backend/onyx/db/permissions.py index 99e59c755a..773556a4bf 100644 --- a/backend/onyx/db/permissions.py +++ b/backend/onyx/db/permissions.py @@ -11,16 +11,12 @@ from collections import defaultdict from collections.abc import Iterable from uuid import UUID -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import select, update from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.enums import AccountType -from onyx.db.enums import Permission -from onyx.db.models import PermissionGrant -from onyx.db.models import User -from onyx.db.models import User__UserGroup +from onyx.db.enums import AccountType, Permission +from onyx.db.models import PermissionGrant, User, User__UserGroup from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/persona.py b/backend/onyx/db/persona.py index 26b52075c4..c5a1f66285 100644 --- a/backend/onyx/db/persona.py +++ b/backend/onyx/db/persona.py @@ -4,52 +4,48 @@ from enum import Enum from uuid import UUID from fastapi import HTTPException -from sqlalchemy import exists -from sqlalchemy import func -from sqlalchemy import not_ -from sqlalchemy import or_ -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import update -from sqlalchemy.orm import aliased -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy import exists, func, not_, or_, Select, select, update +from sqlalchemy.orm import aliased, selectinload, Session from onyx.access.hierarchy_access import get_user_external_group_ids from onyx.auth.schemas import UserRole from onyx.configs.app_configs import CURATORS_CANNOT_VIEW_OR_EDIT_NON_OWNED_ASSISTANTS -from onyx.configs.constants import DEFAULT_PERSONA_ID -from onyx.configs.constants import NotificationType +from onyx.configs.constants import DEFAULT_PERSONA_ID, NotificationType from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX from onyx.db.document_access import get_accessible_documents_by_ids from onyx.db.document_set import filter_document_set_ids_by_user_access -from onyx.db.enums import AccountType -from onyx.db.enums import PersonaSharePermission +from onyx.db.enums import AccountType, PersonaSharePermission from onyx.db.hierarchy import filter_accessible_hierarchy_node_ids -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Document -from onyx.db.models import DocumentSet -from onyx.db.models import FederatedConnector__DocumentSet -from onyx.db.models import HierarchyNode -from onyx.db.models import Persona -from onyx.db.models import Persona__User -from onyx.db.models import Persona__UserGroup -from onyx.db.models import PersonaLabel -from onyx.db.models import StarterMessage -from onyx.db.models import Tool -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserFile -from onyx.db.models import UserGroup +from onyx.db.models import ( + ConnectorCredentialPair, + Document, + DocumentSet, + FederatedConnector__DocumentSet, + HierarchyNode, + Persona, + Persona__User, + Persona__UserGroup, + PersonaLabel, + StarterMessage, + Tool, + User, + User__UserGroup, + UserFile, + UserGroup, +) from onyx.db.notification import create_notification -from onyx.db.persona_sharing import get_persona_access_level -from onyx.db.persona_sharing import get_user_group_ids_for_user -from onyx.db.persona_sharing import persona_ownership_is_vacant -from onyx.server.features.persona.models import FullPersonaSnapshot -from onyx.server.features.persona.models import MinimalPersonaSnapshot -from onyx.server.features.persona.models import PersonaSharedNotificationData -from onyx.server.features.persona.models import PersonaSnapshot -from onyx.server.features.persona.models import PersonaUpsertRequest +from onyx.db.persona_sharing import ( + get_persona_access_level, + get_user_group_ids_for_user, + persona_ownership_is_vacant, +) +from onyx.server.features.persona.models import ( + FullPersonaSnapshot, + MinimalPersonaSnapshot, + PersonaSharedNotificationData, + PersonaSnapshot, + PersonaUpsertRequest, +) from onyx.server.features.tool.tool_visibility import should_expose_tool_to_fe from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_versioned_implementation diff --git a/backend/onyx/db/persona_sharing.py b/backend/onyx/db/persona_sharing.py index 05a8247b5d..312dffd616 100644 --- a/backend/onyx/db/persona_sharing.py +++ b/backend/onyx/db/persona_sharing.py @@ -7,12 +7,12 @@ from sqlalchemy import select from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.enums import PersonaAccessLevel -from onyx.db.enums import PersonaSharePermission -from onyx.db.enums import PersonaSharingStatus -from onyx.db.models import Persona -from onyx.db.models import User -from onyx.db.models import User__UserGroup +from onyx.db.enums import ( + PersonaAccessLevel, + PersonaSharePermission, + PersonaSharingStatus, +) +from onyx.db.models import Persona, User, User__UserGroup def get_user_group_ids_for_user(db_session: Session, user_id: UUID) -> set[int]: diff --git a/backend/onyx/db/port_attempt.py b/backend/onyx/db/port_attempt.py index b96f009697..c8abbd0eff 100644 --- a/backend/onyx/db/port_attempt.py +++ b/backend/onyx/db/port_attempt.py @@ -8,22 +8,14 @@ history. Nothing here enqueues celery work — that is the caller's job. from datetime import datetime -from sqlalchemy import and_ -from sqlalchemy import exists -from sqlalchemy import func -from sqlalchemy import or_ -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import and_, exists, func, or_, select, update from sqlalchemy.orm import Session from onyx.db.connector_credential_pair import ( fetch_indexable_standard_connector_credential_pair_ids, ) -from onyx.db.enums import IndexModelStatus -from onyx.db.enums import PortAttemptStatus -from onyx.db.enums import SwitchoverType -from onyx.db.models import PortAttempt -from onyx.db.models import SearchSettings +from onyx.db.enums import IndexModelStatus, PortAttemptStatus, SwitchoverType +from onyx.db.models import PortAttempt, SearchSettings from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/projects.py b/backend/onyx/db/projects.py index ea10457019..557e65dd81 100644 --- a/backend/onyx/db/projects.py +++ b/backend/onyx/db/projects.py @@ -3,29 +3,27 @@ import uuid from typing import List from uuid import UUID -from fastapi import HTTPException -from fastapi import UploadFile -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field +from fastapi import HTTPException, UploadFile +from pydantic import BaseModel, ConfigDict, Field from sqlalchemy import func from sqlalchemy.orm import Session from starlette.background import BackgroundTasks from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.constants import CELERY_USER_FILE_PROCESSING_TASK_EXPIRES -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import ( + CELERY_USER_FILE_PROCESSING_TASK_EXPIRES, + FileOrigin, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, +) from onyx.db.enums import UserFileStatus -from onyx.db.models import Project__UserFile -from onyx.db.models import User -from onyx.db.models import UserFile -from onyx.db.models import UserProject +from onyx.db.models import Project__UserFile, User, UserFile, UserProject from onyx.server.documents.connector import upload_files -from onyx.server.features.projects.projects_file_utils import categorize_uploaded_files -from onyx.server.features.projects.projects_file_utils import RejectedFile +from onyx.server.features.projects.projects_file_utils import ( + categorize_uploaded_files, + RejectedFile, +) from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/db/pydantic_type.py b/backend/onyx/db/pydantic_type.py index cd0d021409..dbe674ad97 100644 --- a/backend/onyx/db/pydantic_type.py +++ b/backend/onyx/db/pydantic_type.py @@ -1,7 +1,5 @@ import json -from typing import Any -from typing import Optional -from typing import Type +from typing import Any, Optional, Type from pydantic import BaseModel from sqlalchemy.dialects.postgresql import JSONB diff --git a/backend/onyx/db/relationships.py b/backend/onyx/db/relationships.py index 201f1255cb..fbe61d8e4b 100644 --- a/backend/onyx/db/relationships.py +++ b/backend/onyx/db/relationships.py @@ -6,19 +6,23 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session import onyx.db.document as dbdocument -from onyx.db.models import KGEntity -from onyx.db.models import KGEntityExtractionStaging -from onyx.db.models import KGRelationship -from onyx.db.models import KGRelationshipExtractionStaging -from onyx.db.models import KGRelationshipType -from onyx.db.models import KGRelationshipTypeExtractionStaging -from onyx.db.models import KGStage -from onyx.kg.utils.formatting_utils import extract_relationship_type_id -from onyx.kg.utils.formatting_utils import format_relationship_id -from onyx.kg.utils.formatting_utils import get_entity_type -from onyx.kg.utils.formatting_utils import make_relationship_id -from onyx.kg.utils.formatting_utils import make_relationship_type_id -from onyx.kg.utils.formatting_utils import split_relationship_id +from onyx.db.models import ( + KGEntity, + KGEntityExtractionStaging, + KGRelationship, + KGRelationshipExtractionStaging, + KGRelationshipType, + KGRelationshipTypeExtractionStaging, + KGStage, +) +from onyx.kg.utils.formatting_utils import ( + extract_relationship_type_id, + format_relationship_id, + get_entity_type, + make_relationship_id, + make_relationship_type_id, + split_relationship_id, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/release_notes.py b/backend/onyx/db/release_notes.py index 359371b7ed..5d48f22a87 100644 --- a/backend/onyx/db/release_notes.py +++ b/backend/onyx/db/release_notes.py @@ -6,9 +6,11 @@ from sqlalchemy import select from sqlalchemy.orm import Session from onyx.configs.app_configs import INSTANCE_TYPE -from onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN -from onyx.configs.constants import NotificationType -from onyx.configs.constants import ONYX_UTM_SOURCE +from onyx.configs.constants import ( + DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN, + NotificationType, + ONYX_UTM_SOURCE, +) from onyx.db.enums import AccountType from onyx.db.models import User from onyx.db.notification import batch_create_notifications diff --git a/backend/onyx/db/rotate_encryption_key.py b/backend/onyx/db/rotate_encryption_key.py index 227941eb83..6701e9a133 100644 --- a/backend/onyx/db/rotate_encryption_key.py +++ b/backend/onyx/db/rotate_encryption_key.py @@ -12,15 +12,11 @@ safely resumed by re-running. import json from typing import Any -from sqlalchemy import LargeBinary -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import LargeBinary, select, update from sqlalchemy.orm import Session from onyx.configs.app_configs import ENCRYPTION_KEY_SECRET -from onyx.db.models import Base -from onyx.db.models import EncryptedJson -from onyx.db.models import EncryptedString +from onyx.db.models import Base, EncryptedJson, EncryptedString from onyx.utils.encryption import decrypt_bytes_to_string from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import global_version diff --git a/backend/onyx/db/saml.py b/backend/onyx/db/saml.py index 3b9afa1e22..4e14a6c964 100644 --- a/backend/onyx/db/saml.py +++ b/backend/onyx/db/saml.py @@ -2,12 +2,9 @@ import datetime from typing import cast from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import and_, func, select from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS from onyx.db.models import SamlAccount diff --git a/backend/onyx/db/scheduled_task.py b/backend/onyx/db/scheduled_task.py index 4c3839f8e1..4e88cbd66a 100644 --- a/backend/onyx/db/scheduled_task.py +++ b/backend/onyx/db/scheduled_task.py @@ -12,31 +12,27 @@ that claimed the task so concurrent beat ticks don't double-fire. from __future__ import annotations -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import desc -from sqlalchemy import literal -from sqlalchemy import select -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy import and_, desc, literal, select +from sqlalchemy.orm import selectinload, Session -from onyx.db.enums import ScheduledTaskErrorClass -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import ScheduledTaskSkipReason -from onyx.db.enums import ScheduledTaskStatus -from onyx.db.enums import ScheduledTaskTriggerSource -from onyx.db.models import ScheduledTask -from onyx.db.models import ScheduledTaskPreApprovedApp -from onyx.db.models import ScheduledTaskRun +from onyx.db.enums import ( + ScheduledTaskErrorClass, + ScheduledTaskRunStatus, + ScheduledTaskSkipReason, + ScheduledTaskStatus, + ScheduledTaskTriggerSource, +) +from onyx.db.models import ScheduledTask, ScheduledTaskPreApprovedApp, ScheduledTaskRun from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.build.scheduled_tasks.schedule import compute_next_run_at -from onyx.server.features.build.scheduled_tasks.schedule import EditorMode +from onyx.server.features.build.scheduled_tasks.schedule import ( + compute_next_run_at, + EditorMode, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/search_settings.py b/backend/onyx/db/search_settings.py index feccaf6667..6e0026e2b8 100644 --- a/backend/onyx/db/search_settings.py +++ b/backend/onyx/db/search_settings.py @@ -1,17 +1,19 @@ -from sqlalchemy import and_ -from sqlalchemy import delete +from sqlalchemy import and_, delete, select from sqlalchemy import inspect as sa_inspect -from sqlalchemy import select from sqlalchemy.orm import Session -from onyx.configs.model_configs import DEFAULT_DOCUMENT_ENCODER_MODEL -from onyx.configs.model_configs import DOCUMENT_ENCODER_MODEL +from onyx.configs.model_configs import ( + DEFAULT_DOCUMENT_ENCODER_MODEL, + DOCUMENT_ENCODER_MODEL, +) from onyx.context.search.models import SavedSearchSettings from onyx.db.llm import fetch_embedding_provider -from onyx.db.models import CloudEmbeddingProvider -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexModelStatus -from onyx.db.models import SearchSettings +from onyx.db.models import ( + CloudEmbeddingProvider, + IndexAttempt, + IndexModelStatus, + SearchSettings, +) from onyx.server.manage.embedding.models import ( CloudEmbeddingProvider as ServerCloudEmbeddingProvider, ) diff --git a/backend/onyx/db/seeding/chat_history_seeding.py b/backend/onyx/db/seeding/chat_history_seeding.py index 6f081897cc..e6e76d7206 100644 --- a/backend/onyx/db/seeding/chat_history_seeding.py +++ b/backend/onyx/db/seeding/chat_history_seeding.py @@ -1,14 +1,14 @@ import random -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from logging import getLogger from uuid import UUID from onyx.configs.constants import MessageType -from onyx.db.chat import create_chat_session -from onyx.db.chat import create_new_chat_message -from onyx.db.chat import get_or_create_root_message +from onyx.db.chat import ( + create_chat_session, + create_new_chat_message, + get_or_create_root_message, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import ChatSession diff --git a/backend/onyx/db/skill.py b/backend/onyx/db/skill.py index 1df25517f8..ac5bf10865 100644 --- a/backend/onyx/db/skill.py +++ b/backend/onyx/db/skill.py @@ -19,41 +19,41 @@ Callers normally control the transaction boundary so multi-step operations can roll back atomically. """ -from collections.abc import Iterable -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from dataclasses import dataclass from enum import Enum from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import ColumnElement -from sqlalchemy import delete -from sqlalchemy import exists -from sqlalchemy import func -from sqlalchemy import or_ -from sqlalchemy import Select -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import ( + and_, + ColumnElement, + delete, + exists, + func, + or_, + Select, + select, + update, +) from sqlalchemy.dialects.postgresql import insert from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session from onyx.auth.schemas import UserRole from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SkillSharePermission +from onyx.db.enums import SandboxStatus, SkillSharePermission from onyx.db.external_app import available_external_app_skill_ids_for_user -from onyx.db.models import ExternalApp -from onyx.db.models import Sandbox -from onyx.db.models import Skill -from onyx.db.models import Skill__User -from onyx.db.models import Skill__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserSkillPreference -from onyx.db.utils import is_fk_violation -from onyx.db.utils import is_unique_violation +from onyx.db.models import ( + ExternalApp, + Sandbox, + Skill, + Skill__User, + Skill__UserGroup, + User, + User__UserGroup, + UserSkillPreference, +) +from onyx.db.utils import is_fk_violation, is_unique_violation from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.skills.built_in import BUILT_IN_SKILLS diff --git a/backend/onyx/db/slack_channel_config.py b/backend/onyx/db/slack_channel_config.py index 9c6359d9bc..cb5c03624d 100644 --- a/backend/onyx/db/slack_channel_config.py +++ b/backend/onyx/db/slack_channel_config.py @@ -2,18 +2,20 @@ from collections.abc import Sequence from typing import Any from sqlalchemy import select -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import Session +from sqlalchemy.orm import joinedload, Session -from onyx.db.constants import DEFAULT_PERSONA_SLACK_CHANNEL_NAME -from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX -from onyx.db.models import ChannelConfig -from onyx.db.models import Persona -from onyx.db.models import Persona__DocumentSet -from onyx.db.models import SlackChannelConfig -from onyx.db.models import User -from onyx.db.persona import mark_persona_as_deleted -from onyx.db.persona import upsert_persona +from onyx.db.constants import ( + DEFAULT_PERSONA_SLACK_CHANNEL_NAME, + SLACK_BOT_PERSONA_PREFIX, +) +from onyx.db.models import ( + ChannelConfig, + Persona, + Persona__DocumentSet, + SlackChannelConfig, + User, +) +from onyx.db.persona import mark_persona_as_deleted, upsert_persona from onyx.db.tools import get_builtin_tool from onyx.tools.tool_implementations.search.search_tool import SearchTool from onyx.utils.errors import EERequiredError diff --git a/backend/onyx/db/sso_provider.py b/backend/onyx/db/sso_provider.py index 1ebd992083..0827f31a3b 100644 --- a/backend/onyx/db/sso_provider.py +++ b/backend/onyx/db/sso_provider.py @@ -4,15 +4,12 @@ import re from pathlib import Path from typing import Any -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import ValidationError +from pydantic import BaseModel, ConfigDict, ValidationError from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session -from onyx.configs.app_configs import SAML_CONF_DIR -from onyx.configs.app_configs import VALID_EMAIL_DOMAINS +from onyx.configs.app_configs import SAML_CONF_DIR, VALID_EMAIL_DOMAINS from onyx.db.enums import SSOProviderType from onyx.db.models import SSOProvider from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/swap_index.py b/backend/onyx/db/swap_index.py index 9fe463bdc7..d380e9808a 100644 --- a/backend/onyx/db/swap_index.py +++ b/backend/onyx/db/swap_index.py @@ -2,34 +2,42 @@ import time from sqlalchemy.orm import Session -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import VESPA_NUM_ATTEMPTS_ON_STARTUP +from onyx.configs.app_configs import DISABLE_VECTOR_DB, VESPA_NUM_ATTEMPTS_ON_STARTUP from onyx.configs.constants import KV_REINDEX_KEY from onyx.db.connector_credential_pair import ( fetch_indexable_standard_connector_credential_pair_ids, + get_connector_credential_pairs, + resync_cc_pair, +) +from onyx.db.document import ( + count_secondary_only_sync_pending_documents_for_cc_pairs, + delete_all_documents_for_connector_credential_pair, +) +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + IndexModelStatus, + SwitchoverType, ) -from onyx.db.connector_credential_pair import get_connector_credential_pairs -from onyx.db.connector_credential_pair import resync_cc_pair -from onyx.db.document import count_secondary_only_sync_pending_documents_for_cc_pairs -from onyx.db.document import delete_all_documents_for_connector_credential_pair -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.enums import SwitchoverType -from onyx.db.index_attempt import cancel_indexing_attempts_for_search_settings from onyx.db.index_attempt import ( + cancel_indexing_attempts_for_search_settings, count_unique_active_cc_pairs_with_successful_index_attempts, + count_unique_cc_pairs_with_successful_index_attempts, +) +from onyx.db.llm import ( + update_default_contextual_model, + update_no_default_contextual_rag_provider, +) +from onyx.db.models import ConnectorCredentialPair, SearchSettings +from onyx.db.port_attempt import ( + cancel_active_port_attempts, + get_active_port_attempt, + get_latest_port_attempt, +) +from onyx.db.search_settings import ( + get_current_search_settings, + get_secondary_search_settings, + update_search_settings_status, ) -from onyx.db.index_attempt import count_unique_cc_pairs_with_successful_index_attempts -from onyx.db.llm import update_default_contextual_model -from onyx.db.llm import update_no_default_contextual_rag_provider -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import SearchSettings -from onyx.db.port_attempt import cancel_active_port_attempts -from onyx.db.port_attempt import get_active_port_attempt -from onyx.db.port_attempt import get_latest_port_attempt -from onyx.db.search_settings import get_current_search_settings -from onyx.db.search_settings import get_secondary_search_settings -from onyx.db.search_settings import update_search_settings_status from onyx.document_index.factory import get_all_document_indices from onyx.key_value_store.factory import get_kv_store from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/sync_record.py b/backend/onyx/db/sync_record.py index f60ed7d19f..d02341e06c 100644 --- a/backend/onyx/db/sync_record.py +++ b/backend/onyx/db/sync_record.py @@ -1,12 +1,7 @@ -from sqlalchemy import and_ -from sqlalchemy import desc -from sqlalchemy import func -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import and_, desc, func, select, update from sqlalchemy.orm import Session -from onyx.db.enums import SyncStatus -from onyx.db.enums import SyncType +from onyx.db.enums import SyncStatus, SyncType from onyx.db.models import SyncRecord from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/tag.py b/backend/onyx/db/tag.py index f20c809f9f..bc386e6dcf 100644 --- a/backend/onyx/db/tag.py +++ b/backend/onyx/db/tag.py @@ -1,17 +1,11 @@ from typing import Any -from sqlalchemy import and_ -from sqlalchemy import delete -from sqlalchemy import or_ -from sqlalchemy import select +from sqlalchemy import and_, delete, or_, select from sqlalchemy.dialects.postgresql import insert as pg_insert -from sqlalchemy.orm import aliased -from sqlalchemy.orm import Session +from sqlalchemy.orm import aliased, Session from onyx.configs.constants import DocumentSource -from onyx.db.models import Document -from onyx.db.models import Document__Tag -from onyx.db.models import Tag +from onyx.db.models import Document, Document__Tag, Tag from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/targeted_reindex.py b/backend/onyx/db/targeted_reindex.py index da222d8d31..8fbb6610ca 100644 --- a/backend/onyx/db/targeted_reindex.py +++ b/backend/onyx/db/targeted_reindex.py @@ -18,22 +18,21 @@ return the job_id + per-request counts for the API response. from collections.abc import Sequence from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 from pydantic import BaseModel -from sqlalchemy import select -from sqlalchemy import tuple_ +from sqlalchemy import select, tuple_ from sqlalchemy.orm import Session -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import DocumentFailure +from onyx.connectors.models import ConnectorFailure, DocumentFailure from onyx.db.enums import IndexingStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError -from onyx.db.models import TargetedReindexJob -from onyx.db.models import TargetedReindexJobTarget +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + IndexAttemptError, + TargetedReindexJob, + TargetedReindexJobTarget, +) from onyx.db.search_settings import get_active_search_settings_list from onyx.utils.logger import setup_logger diff --git a/backend/onyx/db/tasks.py b/backend/onyx/db/tasks.py index 2f4ee3afae..43b0cdcae3 100644 --- a/backend/onyx/db/tasks.py +++ b/backend/onyx/db/tasks.py @@ -1,15 +1,12 @@ from datetime import datetime -from sqlalchemy import desc -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import desc, func, select from sqlalchemy.orm import Session from sqlalchemy.sql import delete from onyx.configs.app_configs import JOB_TIMEOUT from onyx.db.engine.time_utils import get_db_current_time -from onyx.db.models import TaskQueueState -from onyx.db.models import TaskStatus +from onyx.db.models import TaskQueueState, TaskStatus def get_latest_task( diff --git a/backend/onyx/db/tenant_invite_counter.py b/backend/onyx/db/tenant_invite_counter.py index cb778e0dc5..13a3b68f74 100644 --- a/backend/onyx/db/tenant_invite_counter.py +++ b/backend/onyx/db/tenant_invite_counter.py @@ -1,5 +1,4 @@ -from sqlalchemy import func -from sqlalchemy import update +from sqlalchemy import func, update from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session diff --git a/backend/onyx/db/token_limit.py b/backend/onyx/db/token_limit.py index f43dadb2ec..f968d0ec12 100644 --- a/backend/onyx/db/token_limit.py +++ b/backend/onyx/db/token_limit.py @@ -4,8 +4,7 @@ from sqlalchemy import select from sqlalchemy.orm import Session from onyx.configs.constants import TokenRateLimitScope -from onyx.db.models import TokenRateLimit -from onyx.db.models import TokenRateLimit__UserGroup +from onyx.db.models import TokenRateLimit, TokenRateLimit__UserGroup from onyx.server.token_rate_limits.models import TokenRateLimitArgs diff --git a/backend/onyx/db/tools.py b/backend/onyx/db/tools.py index e51409bb49..0557d94369 100644 --- a/backend/onyx/db/tools.py +++ b/backend/onyx/db/tools.py @@ -1,27 +1,17 @@ -from typing import Any -from typing import cast -from typing import Type -from typing import TYPE_CHECKING +from typing import Any, cast, Type, TYPE_CHECKING from uuid import UUID -from sqlalchemy import func -from sqlalchemy import or_ -from sqlalchemy import select +from sqlalchemy import func, or_, select from sqlalchemy.orm import Session -from onyx.db.constants import UNSET -from onyx.db.constants import UnsetType +from onyx.db.constants import UNSET, UnsetType from onyx.db.enums import MCPServerStatus -from onyx.db.models import MCPServer -from onyx.db.models import OAuthConfig -from onyx.db.models import Tool -from onyx.db.models import ToolCall +from onyx.db.models import MCPServer, OAuthConfig, Tool, ToolCall from onyx.server.features.tool.models import Header from onyx.tools.built_in_tools import BUILT_IN_TOOL_TYPES from onyx.utils.headers import HeaderItemDict from onyx.utils.logger import setup_logger -from onyx.utils.postgres_sanitization import sanitize_json_like -from onyx.utils.postgres_sanitization import sanitize_string +from onyx.utils.postgres_sanitization import sanitize_json_like, sanitize_string if TYPE_CHECKING: pass diff --git a/backend/onyx/db/usage.py b/backend/onyx/db/usage.py index ed247c486a..83315c58e1 100644 --- a/backend/onyx/db/usage.py +++ b/backend/onyx/db/usage.py @@ -1,7 +1,6 @@ """Database interactions for tenant usage tracking (cloud usage limits).""" -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from enum import Enum from pydantic import BaseModel diff --git a/backend/onyx/db/user_file.py b/backend/onyx/db/user_file.py index cf99b0757f..5bdeb48945 100644 --- a/backend/onyx/db/user_file.py +++ b/backend/onyx/db/user_file.py @@ -1,15 +1,10 @@ import datetime from uuid import UUID -from sqlalchemy import func -from sqlalchemy import select -from sqlalchemy.orm import joinedload -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy import func, select +from sqlalchemy.orm import joinedload, selectinload, Session -from onyx.db.models import Persona -from onyx.db.models import Project__UserFile -from onyx.db.models import UserFile +from onyx.db.models import Persona, Project__UserFile, UserFile def fetch_chunk_counts_for_user_files( diff --git a/backend/onyx/db/user_preferences.py b/backend/onyx/db/user_preferences.py index 39e2dc589c..6f77e80955 100644 --- a/backend/onyx/db/user_preferences.py +++ b/backend/onyx/db/user_preferences.py @@ -1,28 +1,22 @@ from collections.abc import Sequence from uuid import UUID -from sqlalchemy import Column -from sqlalchemy import delete -from sqlalchemy import desc -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import Column, delete, desc, select, update from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.enums import AccountType -from onyx.db.enums import DefaultAppMode -from onyx.db.enums import ThemePreference -from onyx.db.models import AccessToken -from onyx.db.models import Assistant__UserSpecificConfig -from onyx.db.models import Memory -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.enums import AccountType, DefaultAppMode, ThemePreference +from onyx.db.models import ( + AccessToken, + Assistant__UserSpecificConfig, + Memory, + User, + User__UserGroup, + UserGroup, +) from onyx.db.permissions import recompute_user_permissions__no_commit -from onyx.db.users import assign_user_to_default_groups__no_commit -from onyx.db.users import is_limited_user -from onyx.server.manage.models import MemoryItem -from onyx.server.manage.models import UserSpecificAssistantPreference +from onyx.db.users import assign_user_to_default_groups__no_commit, is_limited_user +from onyx.server.manage.models import MemoryItem, UserSpecificAssistantPreference from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/db/users.py b/backend/onyx/db/users.py index 012a6f84cf..9532612066 100644 --- a/backend/onyx/db/users.py +++ b/backend/onyx/db/users.py @@ -1,36 +1,34 @@ -from collections.abc import Callable -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import Any from uuid import UUID from fastapi import HTTPException from fastapi_users.password import PasswordHelper -from sqlalchemy import case -from sqlalchemy import func -from sqlalchemy import Select -from sqlalchemy import select +from sqlalchemy import case, func, Select, select from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session from sqlalchemy.sql import expression -from sqlalchemy.sql.elements import ColumnElement -from sqlalchemy.sql.elements import KeyedColumnElement +from sqlalchemy.sql.elements import ColumnElement, KeyedColumnElement from sqlalchemy.sql.expression import or_ from onyx.auth.invited_users import remove_user_from_invited_users from onyx.auth.schemas import UserRole -from onyx.configs.constants import ANONYMOUS_USER_EMAIL -from onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN -from onyx.configs.constants import NO_AUTH_PLACEHOLDER_USER_EMAIL +from onyx.configs.constants import ( + ANONYMOUS_USER_EMAIL, + DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN, + NO_AUTH_PLACEHOLDER_USER_EMAIL, +) from onyx.db.enums import AccountType -from onyx.db.models import DocumentSet -from onyx.db.models import DocumentSet__User -from onyx.db.models import Persona -from onyx.db.models import Persona__User -from onyx.db.models import SamlAccount -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import ( + DocumentSet, + DocumentSet__User, + Persona, + Persona__User, + SamlAccount, + User, + User__UserGroup, + UserGroup, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop diff --git a/backend/onyx/db/utils.py b/backend/onyx/db/utils.py index 517e01a6ea..974d2a9331 100644 --- a/backend/onyx/db/utils.py +++ b/backend/onyx/db/utils.py @@ -1,13 +1,8 @@ from enum import Enum -from typing import Any -from typing import Final -from typing import TypeGuard -from typing import TypeVar +from typing import Any, Final, TypeGuard, TypeVar -from psycopg2 import errorcodes -from psycopg2 import OperationalError -from psycopg2.errors import ForeignKeyViolation -from psycopg2.errors import UniqueViolation +from psycopg2 import errorcodes, OperationalError +from psycopg2.errors import ForeignKeyViolation, UniqueViolation from pydantic import BaseModel from sqlalchemy import inspect from sqlalchemy.exc import IntegrityError diff --git a/backend/onyx/db/voice.py b/backend/onyx/db/voice.py index 6f337966fb..b6fb87a010 100644 --- a/backend/onyx/db/voice.py +++ b/backend/onyx/db/voice.py @@ -1,12 +1,10 @@ from typing import Any from uuid import UUID -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import select, update from sqlalchemy.orm import Session -from onyx.db.models import User -from onyx.db.models import VoiceProvider +from onyx.db.models import User, VoiceProvider from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/onyx/db/web_search.py b/backend/onyx/db/web_search.py index 9c242ebe3e..86b51ac5ea 100644 --- a/backend/onyx/db/web_search.py +++ b/backend/onyx/db/web_search.py @@ -1,14 +1,11 @@ from __future__ import annotations -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import select, update from sqlalchemy.orm import Session -from onyx.db.models import InternetContentProvider -from onyx.db.models import InternetSearchProvider +from onyx.db.models import InternetContentProvider, InternetSearchProvider from onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig -from shared_configs.enums import WebContentProviderType -from shared_configs.enums import WebSearchProviderType +from shared_configs.enums import WebContentProviderType, WebSearchProviderType def fetch_web_search_providers(db_session: Session) -> list[InternetSearchProvider]: diff --git a/backend/onyx/deep_research/dr_loop.py b/backend/onyx/deep_research/dr_loop.py index 3fecf95832..a5bb38cd01 100644 --- a/backend/onyx/deep_research/dr_loop.py +++ b/backend/onyx/deep_research/dr_loop.py @@ -8,66 +8,69 @@ from collections.abc import Callable from typing import cast from onyx.chat.chat_state import ChatStateContainer -from onyx.chat.citation_processor import CitationMapping -from onyx.chat.citation_processor import DynamicCitationProcessor +from onyx.chat.citation_processor import CitationMapping, DynamicCitationProcessor from onyx.chat.emitter import Emitter from onyx.chat.llm_loop import construct_message_history -from onyx.chat.llm_step import run_llm_step -from onyx.chat.llm_step import run_llm_step_pkt_generator -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import FileToolMetadata -from onyx.chat.models import LlmStepResult -from onyx.chat.models import ToolCallSimple -from onyx.configs.chat_configs import DR_REPORT_LLM_TIMEOUT_S -from onyx.configs.chat_configs import SKIP_DEEP_RESEARCH_CLARIFICATION +from onyx.chat.llm_step import run_llm_step, run_llm_step_pkt_generator +from onyx.chat.models import ( + ChatMessageSimple, + FileToolMetadata, + LlmStepResult, + ToolCallSimple, +) +from onyx.configs.chat_configs import ( + DR_REPORT_LLM_TIMEOUT_S, + SKIP_DEEP_RESEARCH_CLARIFICATION, +) from onyx.configs.constants import MessageType from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.tools import get_tool_by_name -from onyx.deep_research.dr_mock_tools import get_clarification_tool_definitions -from onyx.deep_research.dr_mock_tools import get_orchestrator_tools -from onyx.deep_research.dr_mock_tools import RESEARCH_AGENT_TOOL_NAME -from onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_MESSAGE -from onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_TOKEN_COUNT -from onyx.deep_research.utils import check_special_tool_calls -from onyx.deep_research.utils import create_think_tool_token_processor -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMUserIdentity +from onyx.deep_research.dr_mock_tools import ( + get_clarification_tool_definitions, + get_orchestrator_tools, + RESEARCH_AGENT_TOOL_NAME, + THINK_TOOL_RESPONSE_MESSAGE, + THINK_TOOL_RESPONSE_TOKEN_COUNT, +) +from onyx.deep_research.utils import ( + check_special_tool_calls, + create_think_tool_token_processor, +) +from onyx.llm.interfaces import LLM, LLMUserIdentity from onyx.llm.model_capabilities import model_is_reasoning_model from onyx.llm.models import ToolChoiceOptions -from onyx.prompts.deep_research.orchestration_layer import CLARIFICATION_PROMPT -from onyx.prompts.deep_research.orchestration_layer import FINAL_REPORT_PROMPT -from onyx.prompts.deep_research.orchestration_layer import FIRST_CYCLE_REMINDER -from onyx.prompts.deep_research.orchestration_layer import FIRST_CYCLE_REMINDER_TOKENS from onyx.prompts.deep_research.orchestration_layer import ( + CLARIFICATION_PROMPT, + FINAL_REPORT_PROMPT, + FIRST_CYCLE_REMINDER, + FIRST_CYCLE_REMINDER_TOKENS, INTERNAL_SEARCH_CLARIFICATION_GUIDANCE, -) -from onyx.prompts.deep_research.orchestration_layer import ( INTERNAL_SEARCH_RESEARCH_TASK_GUIDANCE, + ORCHESTRATOR_PROMPT, + ORCHESTRATOR_PROMPT_REASONING, + RESEARCH_PLAN_PROMPT, + RESEARCH_PLAN_REMINDER, + USER_FINAL_REPORT_QUERY, ) -from onyx.prompts.deep_research.orchestration_layer import ORCHESTRATOR_PROMPT -from onyx.prompts.deep_research.orchestration_layer import ORCHESTRATOR_PROMPT_REASONING -from onyx.prompts.deep_research.orchestration_layer import RESEARCH_PLAN_PROMPT -from onyx.prompts.deep_research.orchestration_layer import RESEARCH_PLAN_REMINDER -from onyx.prompts.deep_research.orchestration_layer import USER_FINAL_REPORT_QUERY from onyx.prompts.prompt_utils import get_current_llm_day_time from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import DeepResearchPlanDelta -from onyx.server.query_and_chat.streaming_models import DeepResearchPlanStart -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import SectionEnd -from onyx.server.query_and_chat.streaming_models import TopLevelBranching +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, + DeepResearchPlanDelta, + DeepResearchPlanStart, + OverallStop, + Packet, + SectionEnd, + TopLevelBranching, +) from onyx.tools.fake_tools.research_agent import run_research_agent_calls from onyx.tools.interface import Tool -from onyx.tools.models import ToolCallInfo -from onyx.tools.models import ToolCallKickoff +from onyx.tools.models import ToolCallInfo, ToolCallKickoff from onyx.tools.tool_implementations.open_url.open_url_tool import OpenURLTool from onyx.tools.tool_implementations.search.search_tool import SearchTool from onyx.tools.tool_implementations.web_search.web_search_tool import WebSearchTool -from onyx.tracing.framework.create import function_span -from onyx.tracing.framework.create import trace +from onyx.tracing.framework.create import function_span, trace from onyx.utils.logger import setup_logger from onyx.utils.timing import log_function_time from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/deep_research/utils.py b/backend/onyx/deep_research/utils.py index 91a3b0cb85..dbba7f4abd 100644 --- a/backend/onyx/deep_research/utils.py +++ b/backend/onyx/deep_research/utils.py @@ -3,12 +3,9 @@ from typing import Any from pydantic import BaseModel -from onyx.deep_research.dr_mock_tools import GENERATE_REPORT_TOOL_NAME -from onyx.deep_research.dr_mock_tools import THINK_TOOL_NAME +from onyx.deep_research.dr_mock_tools import GENERATE_REPORT_TOOL_NAME, THINK_TOOL_NAME from onyx.deep_research.models import SpecialToolCalls -from onyx.llm.model_response import ChatCompletionDeltaToolCall -from onyx.llm.model_response import Delta -from onyx.llm.model_response import FunctionCall +from onyx.llm.model_response import ChatCompletionDeltaToolCall, Delta, FunctionCall from onyx.tools.models import ToolCallKickoff # JSON prefixes to detect in think_tool arguments diff --git a/backend/onyx/document_index/chunk_content_enrichment.py b/backend/onyx/document_index/chunk_content_enrichment.py index a56d5207f3..b4f19859a8 100644 --- a/backend/onyx/document_index/chunk_content_enrichment.py +++ b/backend/onyx/document_index/chunk_content_enrichment.py @@ -1,7 +1,6 @@ from onyx.configs.app_configs import BLURB_SIZE from onyx.configs.constants import RETURN_SEPARATOR -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceChunkUncleaned +from onyx.context.search.models import InferenceChunk, InferenceChunkUncleaned from onyx.indexing.models import DocAwareChunk diff --git a/backend/onyx/document_index/disabled.py b/backend/onyx/document_index/disabled.py index a0e9174e7d..290770375b 100644 --- a/backend/onyx/document_index/disabled.py +++ b/backend/onyx/document_index/disabled.py @@ -8,14 +8,15 @@ out against a nonexistent Vespa/OpenSearch instance. from collections.abc import Iterable from onyx.context.search.enums import QueryType -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import IndexFilters, InferenceChunk from onyx.db.enums import EmbeddingPrecision -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentInsertionRecord -from onyx.document_index.interfaces_new import DocumentSectionRequest -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import MetadataUpdateRequest +from onyx.document_index.interfaces_new import ( + DocumentIndex, + DocumentInsertionRecord, + DocumentSectionRequest, + IndexingMetadata, + MetadataUpdateRequest, +) from onyx.indexing.models import DocMetadataAwareIndexChunk from shared_configs.model_server_models import Embedding diff --git a/backend/onyx/document_index/document_index_utils.py b/backend/onyx/document_index/document_index_utils.py index d5660ab18e..f64e52999d 100644 --- a/backend/onyx/document_index/document_index_utils.py +++ b/backend/onyx/document_index/document_index_utils.py @@ -6,11 +6,12 @@ from sqlalchemy.orm import Session from onyx.configs.app_configs import ENABLE_MULTIPASS_INDEXING from onyx.db.models import SearchSettings -from onyx.db.search_settings import get_current_search_settings -from onyx.db.search_settings import get_secondary_search_settings +from onyx.db.search_settings import ( + get_current_search_settings, + get_secondary_search_settings, +) from onyx.document_index.vespa.internal_types import EnrichedDocumentIndexingInfo -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import MultipassConfig +from onyx.indexing.models import DocMetadataAwareIndexChunk, MultipassConfig from shared_configs.configs import MULTI_TENANT DEFAULT_BATCH_SIZE = 30 diff --git a/backend/onyx/document_index/factory.py b/backend/onyx/document_index/factory.py index 44bf51aba9..72e4b183c3 100644 --- a/backend/onyx/document_index/factory.py +++ b/backend/onyx/document_index/factory.py @@ -1,20 +1,23 @@ import httpx from sqlalchemy.orm import Session -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX -from onyx.configs.app_configs import ONYX_DISABLE_VESPA +from onyx.configs.app_configs import ( + DISABLE_VECTOR_DB, + ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + ONYX_DISABLE_VESPA, +) from onyx.db.models import SearchSettings from onyx.db.opensearch_migration import get_opensearch_retrieval_state from onyx.document_index.disabled import DisabledDocumentIndex -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces_new import DocumentIndex, TenantState from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, + OpenSearchIndexPair, +) +from onyx.document_index.vespa.vespa_document_index import ( + VespaDocumentIndex, + VespaIndexPair, ) -from onyx.document_index.opensearch.opensearch_document_index import OpenSearchIndexPair -from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex -from onyx.document_index.vespa.vespa_document_index import VespaIndexPair from onyx.indexing.models import IndexingSetting from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/document_index/interfaces_new.py b/backend/onyx/document_index/interfaces_new.py index 26c75a6ef4..a296dfd05d 100644 --- a/backend/onyx/document_index/interfaces_new.py +++ b/backend/onyx/document_index/interfaces_new.py @@ -3,14 +3,12 @@ from collections.abc import Iterable from datetime import datetime from typing import Self -from pydantic import BaseModel -from pydantic import model_validator +from pydantic import BaseModel, model_validator from onyx.access.models import DocumentAccess from onyx.configs.constants import PUBLIC_DOC_PAT from onyx.context.search.enums import QueryType -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import IndexFilters, InferenceChunk from onyx.db.enums import EmbeddingPrecision from onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE from onyx.indexing.models import DocMetadataAwareIndexChunk diff --git a/backend/onyx/document_index/opensearch/client.py b/backend/onyx/document_index/opensearch/client.py index 49e26dc7bb..6e986047b3 100644 --- a/backend/onyx/document_index/opensearch/client.py +++ b/backend/onyx/document_index/opensearch/client.py @@ -3,51 +3,58 @@ import logging import time from collections import Counter from collections.abc import Iterator -from contextlib import AbstractContextManager -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext from http import HTTPStatus -from typing import Any -from typing import Generic -from typing import TypeVar +from typing import Any, Generic, TypeVar import boto3 -from opensearchpy import NotFoundError -from opensearchpy import OpenSearch -from opensearchpy import TransportError -from opensearchpy import Urllib3AWSV4SignerAuth +from opensearchpy import ( + NotFoundError, + OpenSearch, + TransportError, + Urllib3AWSV4SignerAuth, +) from opensearchpy.helpers import bulk from pydantic import BaseModel -from onyx.configs.app_configs import DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S -from onyx.configs.app_configs import OPENSEARCH_ADMIN_PASSWORD -from onyx.configs.app_configs import OPENSEARCH_ADMIN_USERNAME -from onyx.configs.app_configs import OPENSEARCH_AUTH_METHOD -from onyx.configs.app_configs import OPENSEARCH_AWS_REGION -from onyx.configs.app_configs import OPENSEARCH_AWS_SERVICE -from onyx.configs.app_configs import OPENSEARCH_CA_CERTS -from onyx.configs.app_configs import OPENSEARCH_CLIENT_CERT -from onyx.configs.app_configs import OPENSEARCH_CLIENT_KEY -from onyx.configs.app_configs import OPENSEARCH_HOST -from onyx.configs.app_configs import OPENSEARCH_REST_API_PORT -from onyx.configs.app_configs import OPENSEARCH_USE_SSL -from onyx.configs.app_configs import OPENSEARCH_VERIFY_CERTS -from onyx.configs.app_configs import PIT_KEEP_ALIVE +from onyx.configs.app_configs import ( + DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S, + OPENSEARCH_ADMIN_PASSWORD, + OPENSEARCH_ADMIN_USERNAME, + OPENSEARCH_AUTH_METHOD, + OPENSEARCH_AWS_REGION, + OPENSEARCH_AWS_SERVICE, + OPENSEARCH_CA_CERTS, + OPENSEARCH_CLIENT_CERT, + OPENSEARCH_CLIENT_KEY, + OPENSEARCH_HOST, + OPENSEARCH_REST_API_PORT, + OPENSEARCH_USE_SSL, + OPENSEARCH_VERIFY_CERTS, + PIT_KEEP_ALIVE, +) from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE -from onyx.document_index.opensearch.constants import OpenSearchAuthMethod -from onyx.document_index.opensearch.constants import OpenSearchSearchType -from onyx.document_index.opensearch.schema import CHUNK_INDEX_FIELD_NAME -from onyx.document_index.opensearch.schema import CONTENT_VECTOR_FIELD_NAME -from onyx.document_index.opensearch.schema import DOCUMENT_ID_FIELD_NAME -from onyx.document_index.opensearch.schema import DocumentChunk -from onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors -from onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id -from onyx.document_index.opensearch.schema import MAX_CHUNK_SIZE_FIELD_NAME -from onyx.document_index.opensearch.schema import TITLE_VECTOR_FIELD_NAME +from onyx.document_index.opensearch.constants import ( + DEFAULT_MAX_CHUNK_SIZE, + OpenSearchAuthMethod, + OpenSearchSearchType, +) +from onyx.document_index.opensearch.schema import ( + CHUNK_INDEX_FIELD_NAME, + CONTENT_VECTOR_FIELD_NAME, + DOCUMENT_ID_FIELD_NAME, + DocumentChunk, + DocumentChunkWithoutVectors, + get_opensearch_doc_chunk_id, + MAX_CHUNK_SIZE_FIELD_NAME, + TITLE_VECTOR_FIELD_NAME, +) from onyx.document_index.opensearch.search import DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW -from onyx.server.metrics.opensearch_search import observe_opensearch_search -from onyx.server.metrics.opensearch_search import record_opensearch_search_error -from onyx.server.metrics.opensearch_search import track_opensearch_search +from onyx.server.metrics.opensearch_search import ( + observe_opensearch_search, + record_opensearch_search_error, + track_opensearch_search, +) from onyx.utils.logger import setup_logger from onyx.utils.timing import log_function_time diff --git a/backend/onyx/document_index/opensearch/opensearch_document_index.py b/backend/onyx/document_index/opensearch/opensearch_document_index.py index d68f086c37..1627c347b8 100644 --- a/backend/onyx/document_index/opensearch/opensearch_document_index.py +++ b/backend/onyx/document_index/opensearch/opensearch_document_index.py @@ -5,61 +5,65 @@ from typing import Any from opensearchpy.helpers.errors import BulkIndexError from onyx.access.models import DocumentAccess -from onyx.configs.app_configs import MAX_CHUNKS_PER_DOC_BATCH -from onyx.configs.app_configs import VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT -from onyx.configs.constants import OnyxRedisLocks -from onyx.configs.constants import PUBLIC_DOC_PAT +from onyx.configs.app_configs import ( + MAX_CHUNKS_PER_DOC_BATCH, + VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT, +) +from onyx.configs.constants import OnyxRedisLocks, PUBLIC_DOC_PAT from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( get_experts_stores_representations, ) from onyx.connectors.models import convert_metadata_list_of_strings_to_dict from onyx.context.search.enums import QueryType -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceChunkUncleaned +from onyx.context.search.models import ( + IndexFilters, + InferenceChunk, + InferenceChunkUncleaned, +) from onyx.db.enums import EmbeddingPrecision from onyx.db.models import DocumentSource -from onyx.document_index.chunk_content_enrichment import cleanup_content_for_chunks from onyx.document_index.chunk_content_enrichment import ( + cleanup_content_for_chunks, generate_enriched_content_for_chunk_text, ) -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentInsertionRecord -from onyx.document_index.interfaces_new import DocumentSectionRequest -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import MetadataUpdateRequest -from onyx.document_index.interfaces_new import SecondaryIndexDocumentMissingError -from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.opensearch.client import OpenSearchClient -from onyx.document_index.opensearch.client import OpenSearchDocumentMissingError -from onyx.document_index.opensearch.client import OpenSearchIndexClient -from onyx.document_index.opensearch.client import SearchHit +from onyx.document_index.interfaces_new import ( + DocumentIndex, + DocumentInsertionRecord, + DocumentSectionRequest, + IndexingMetadata, + MetadataUpdateRequest, + SecondaryIndexDocumentMissingError, + TenantState, +) +from onyx.document_index.opensearch.client import ( + OpenSearchClient, + OpenSearchDocumentMissingError, + OpenSearchIndexClient, + SearchHit, +) from onyx.document_index.opensearch.cluster_settings import OPENSEARCH_CLUSTER_SETTINGS from onyx.document_index.opensearch.constants import OpenSearchSearchType -from onyx.document_index.opensearch.schema import ACCESS_CONTROL_LIST_FIELD_NAME -from onyx.document_index.opensearch.schema import CONTENT_FIELD_NAME -from onyx.document_index.opensearch.schema import CREATED_AT_FIELD_NAME -from onyx.document_index.opensearch.schema import DOCUMENT_SETS_FIELD_NAME -from onyx.document_index.opensearch.schema import DocumentChunk -from onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors -from onyx.document_index.opensearch.schema import DocumentSchema -from onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id -from onyx.document_index.opensearch.schema import GLOBAL_BOOST_FIELD_NAME -from onyx.document_index.opensearch.schema import HIDDEN_FIELD_NAME -from onyx.document_index.opensearch.schema import PERSONAS_FIELD_NAME -from onyx.document_index.opensearch.schema import USER_PROJECTS_FIELD_NAME -from onyx.document_index.opensearch.search import DocumentQuery +from onyx.document_index.opensearch.schema import ( + ACCESS_CONTROL_LIST_FIELD_NAME, + CONTENT_FIELD_NAME, + CREATED_AT_FIELD_NAME, + DOCUMENT_SETS_FIELD_NAME, + DocumentChunk, + DocumentChunkWithoutVectors, + DocumentSchema, + get_opensearch_doc_chunk_id, + GLOBAL_BOOST_FIELD_NAME, + HIDDEN_FIELD_NAME, + PERSONAS_FIELD_NAME, + USER_PROJECTS_FIELD_NAME, +) from onyx.document_index.opensearch.search import ( + DocumentQuery, get_min_max_normalization_pipeline_name_and_config, -) -from onyx.document_index.opensearch.search import ( get_normalization_pipeline_name_and_config, -) -from onyx.document_index.opensearch.search import ( get_zscore_normalization_pipeline_name_and_config, ) -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import Document +from onyx.indexing.models import DocMetadataAwareIndexChunk, Document from onyx.redis.lock_context import redis_shared_lock from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/document_index/opensearch/port_copy.py b/backend/onyx/document_index/opensearch/port_copy.py index 30ec718ea1..0949bb9ed2 100644 --- a/backend/onyx/document_index/opensearch/port_copy.py +++ b/backend/onyx/document_index/opensearch/port_copy.py @@ -14,8 +14,7 @@ so they stream PIT pages. """ from collections import defaultdict -from collections.abc import Callable -from collections.abc import Iterable +from collections.abc import Callable, Iterable from onyx.db.models import SearchSettings from onyx.document_index.factory import build_opensearch_document_index @@ -25,15 +24,15 @@ from onyx.document_index.opensearch.opensearch_document_index import ( ) from onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors from onyx.indexing.chunker import DEFAULT_CONTEXTUAL_RAG_RESERVED_TOKENS -from onyx.indexing.embedder import DefaultIndexingEmbedder -from onyx.indexing.embedder import IndexingEmbedder -from onyx.indexing.port_reembed import AugmentationReembedContext -from onyx.indexing.port_reembed import re_embed_chunks -from onyx.indexing.port_reembed import ReembedStrategy -from onyx.indexing.port_reembed import select_reembed_strategy +from onyx.indexing.embedder import DefaultIndexingEmbedder, IndexingEmbedder +from onyx.indexing.port_reembed import ( + AugmentationReembedContext, + re_embed_chunks, + ReembedStrategy, + select_reembed_strategy, +) from onyx.llm.factory import get_contextual_rag_llm_for_search_settings -from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.natural_language_processing.utils import get_tokenizer +from onyx.natural_language_processing.utils import BaseTokenizer, get_tokenizer from shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE # Cap per bulk write so it can't run long unheartbeated and get a live port stall-failed. diff --git a/backend/onyx/document_index/opensearch/schema.py b/backend/onyx/document_index/opensearch/schema.py index 66ec3cfe68..1a7da773c3 100644 --- a/backend/onyx/document_index/opensearch/schema.py +++ b/backend/onyx/document_index/opensearch/schema.py @@ -1,32 +1,34 @@ import hashlib -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import Self +from datetime import datetime, timezone +from typing import Any, Self -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_serializer -from pydantic import field_validator -from pydantic import model_serializer -from pydantic import model_validator -from pydantic import SerializerFunctionWrapHandler -from pydantic import ValidationInfo +from pydantic import ( + BaseModel, + Field, + field_serializer, + field_validator, + model_serializer, + model_validator, + SerializerFunctionWrapHandler, + ValidationInfo, +) -from onyx.configs.app_configs import OPENSEARCH_INDEX_NUM_REPLICAS -from onyx.configs.app_configs import OPENSEARCH_INDEX_NUM_SHARDS -from onyx.configs.app_configs import OPENSEARCH_TEXT_ANALYZER -from onyx.configs.app_configs import USING_AWS_MANAGED_OPENSEARCH +from onyx.configs.app_configs import ( + OPENSEARCH_INDEX_NUM_REPLICAS, + OPENSEARCH_INDEX_NUM_SHARDS, + OPENSEARCH_TEXT_ANALYZER, + USING_AWS_MANAGED_OPENSEARCH, +) from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE -from onyx.document_index.opensearch.constants import EF_CONSTRUCTION -from onyx.document_index.opensearch.constants import EF_SEARCH -from onyx.document_index.opensearch.constants import M -from onyx.document_index.opensearch.string_filtering import DocumentIDTooLongError -from onyx.document_index.opensearch.string_filtering import ( - filter_and_validate_document_id, +from onyx.document_index.opensearch.constants import ( + DEFAULT_MAX_CHUNK_SIZE, + EF_CONSTRUCTION, + EF_SEARCH, + M, ) from onyx.document_index.opensearch.string_filtering import ( + DocumentIDTooLongError, + filter_and_validate_document_id, MAX_DOCUMENT_ID_ENCODED_LENGTH, ) from onyx.utils.datetime import datetime_to_utc diff --git a/backend/onyx/document_index/opensearch/search.py b/backend/onyx/document_index/opensearch/search.py index ecfce64e1e..03652625c2 100644 --- a/backend/onyx/document_index/opensearch/search.py +++ b/backend/onyx/document_index/opensearch/search.py @@ -1,55 +1,46 @@ import random -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import TypeAlias -from typing import TypeVar +from datetime import datetime, timedelta, timezone +from typing import Any, TypeAlias, TypeVar -from onyx.configs.app_configs import DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S -from onyx.configs.app_configs import OPENSEARCH_EXPLAIN_ENABLED -from onyx.configs.app_configs import OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED -from onyx.configs.app_configs import OPENSEARCH_PROFILING_DISABLED -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import INDEX_SEPARATOR -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import Tag -from onyx.context.search.models import TimeRange +from onyx.configs.app_configs import ( + DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S, + OPENSEARCH_EXPLAIN_ENABLED, + OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED, + OPENSEARCH_PROFILING_DISABLED, +) +from onyx.configs.constants import DocumentSource, INDEX_SEPARATOR +from onyx.context.search.models import IndexFilters, Tag, TimeRange from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.opensearch.constants import ASSUMED_DOCUMENT_AGE_DAYS from onyx.document_index.opensearch.constants import ( + ASSUMED_DOCUMENT_AGE_DAYS, DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES, -) -from onyx.document_index.opensearch.constants import ( DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW, -) -from onyx.document_index.opensearch.constants import ( HYBRID_SEARCH_NORMALIZATION_PIPELINE, -) -from onyx.document_index.opensearch.constants import ( HYBRID_SEARCH_SUBQUERY_CONFIGURATION, + HybridSearchNormalizationPipeline, + HybridSearchSubqueryConfiguration, +) +from onyx.document_index.opensearch.schema import ( + ACCESS_CONTROL_LIST_FIELD_NAME, + ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME, + CHUNK_INDEX_FIELD_NAME, + CONTENT_FIELD_NAME, + CONTENT_VECTOR_FIELD_NAME, + CREATED_AT_FIELD_NAME, + DOCUMENT_ID_FIELD_NAME, + DOCUMENT_SETS_FIELD_NAME, + HIDDEN_FIELD_NAME, + LAST_UPDATED_FIELD_NAME, + MAX_CHUNK_SIZE_FIELD_NAME, + METADATA_LIST_FIELD_NAME, + PERSONAS_FIELD_NAME, + PUBLIC_FIELD_NAME, + SOURCE_TYPE_FIELD_NAME, + TENANT_ID_FIELD_NAME, + TITLE_FIELD_NAME, + TITLE_VECTOR_FIELD_NAME, + USER_PROJECTS_FIELD_NAME, ) -from onyx.document_index.opensearch.constants import HybridSearchNormalizationPipeline -from onyx.document_index.opensearch.constants import HybridSearchSubqueryConfiguration -from onyx.document_index.opensearch.schema import ACCESS_CONTROL_LIST_FIELD_NAME -from onyx.document_index.opensearch.schema import ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME -from onyx.document_index.opensearch.schema import CHUNK_INDEX_FIELD_NAME -from onyx.document_index.opensearch.schema import CONTENT_FIELD_NAME -from onyx.document_index.opensearch.schema import CONTENT_VECTOR_FIELD_NAME -from onyx.document_index.opensearch.schema import CREATED_AT_FIELD_NAME -from onyx.document_index.opensearch.schema import DOCUMENT_ID_FIELD_NAME -from onyx.document_index.opensearch.schema import DOCUMENT_SETS_FIELD_NAME -from onyx.document_index.opensearch.schema import HIDDEN_FIELD_NAME -from onyx.document_index.opensearch.schema import LAST_UPDATED_FIELD_NAME -from onyx.document_index.opensearch.schema import MAX_CHUNK_SIZE_FIELD_NAME -from onyx.document_index.opensearch.schema import METADATA_LIST_FIELD_NAME -from onyx.document_index.opensearch.schema import PERSONAS_FIELD_NAME -from onyx.document_index.opensearch.schema import PUBLIC_FIELD_NAME -from onyx.document_index.opensearch.schema import SOURCE_TYPE_FIELD_NAME -from onyx.document_index.opensearch.schema import TENANT_ID_FIELD_NAME -from onyx.document_index.opensearch.schema import TITLE_FIELD_NAME -from onyx.document_index.opensearch.schema import TITLE_VECTOR_FIELD_NAME -from onyx.document_index.opensearch.schema import USER_PROJECTS_FIELD_NAME from onyx.utils.datetime import datetime_to_utc # See https://docs.opensearch.org/latest/query-dsl/term/terms/. diff --git a/backend/onyx/document_index/vespa/chunk_retrieval.py b/backend/onyx/document_index/vespa/chunk_retrieval.py index a318e6f3ba..0d1c1c0862 100644 --- a/backend/onyx/document_index/vespa/chunk_retrieval.py +++ b/backend/onyx/document_index/vespa/chunk_retrieval.py @@ -1,12 +1,9 @@ import json import string import time -from collections.abc import Callable -from collections.abc import Mapping -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast +from collections.abc import Callable, Mapping +from datetime import datetime, timezone +from typing import Any, cast import httpx @@ -16,49 +13,50 @@ from onyx.background.celery.tasks.opensearch_migration.constants import ( from onyx.background.celery.tasks.opensearch_migration.transformer import ( FIELDS_NEEDED_FOR_TRANSFORMATION, ) -from onyx.configs.app_configs import LOG_VESPA_TIMING_INFORMATION -from onyx.configs.app_configs import VESPA_LANGUAGE_OVERRIDE -from onyx.configs.app_configs import VESPA_MIGRATION_REQUEST_TIMEOUT_S -from onyx.configs.app_configs import VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunkUncleaned +from onyx.configs.app_configs import ( + LOG_VESPA_TIMING_INFORMATION, + VESPA_LANGUAGE_OVERRIDE, + VESPA_MIGRATION_REQUEST_TIMEOUT_S, + VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT, +) +from onyx.context.search.models import IndexFilters, InferenceChunkUncleaned from onyx.document_index.interfaces_new import TenantState from onyx.document_index.vespa.internal_types import VespaChunkRequest from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client from onyx.document_index.vespa.shared_utils.vespa_request_builders import ( build_vespa_filters, -) -from onyx.document_index.vespa.shared_utils.vespa_request_builders import ( build_vespa_id_based_retrieval_yql, ) -from onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST -from onyx.document_index.vespa_constants import BLURB -from onyx.document_index.vespa_constants import BOOST -from onyx.document_index.vespa_constants import CHUNK_CONTEXT -from onyx.document_index.vespa_constants import CHUNK_ID -from onyx.document_index.vespa_constants import CONTENT -from onyx.document_index.vespa_constants import CONTENT_SUMMARY -from onyx.document_index.vespa_constants import DOC_SUMMARY -from onyx.document_index.vespa_constants import DOC_UPDATED_AT -from onyx.document_index.vespa_constants import DOCUMENT_ID -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT -from onyx.document_index.vespa_constants import HIDDEN -from onyx.document_index.vespa_constants import IMAGE_FILE_NAME -from onyx.document_index.vespa_constants import LARGE_CHUNK_REFERENCE_IDS -from onyx.document_index.vespa_constants import MAX_ID_SEARCH_QUERY_SIZE -from onyx.document_index.vespa_constants import MAX_OR_CONDITIONS -from onyx.document_index.vespa_constants import METADATA -from onyx.document_index.vespa_constants import METADATA_SUFFIX -from onyx.document_index.vespa_constants import PRIMARY_OWNERS -from onyx.document_index.vespa_constants import SEARCH_ENDPOINT -from onyx.document_index.vespa_constants import SECONDARY_OWNERS -from onyx.document_index.vespa_constants import SECTION_CONTINUATION -from onyx.document_index.vespa_constants import SEMANTIC_IDENTIFIER -from onyx.document_index.vespa_constants import SOURCE_LINKS -from onyx.document_index.vespa_constants import SOURCE_TYPE -from onyx.document_index.vespa_constants import TENANT_ID -from onyx.document_index.vespa_constants import TITLE -from onyx.document_index.vespa_constants import YQL_BASE +from onyx.document_index.vespa_constants import ( + ACCESS_CONTROL_LIST, + BLURB, + BOOST, + CHUNK_CONTEXT, + CHUNK_ID, + CONTENT, + CONTENT_SUMMARY, + DOC_SUMMARY, + DOC_UPDATED_AT, + DOCUMENT_ID, + DOCUMENT_ID_ENDPOINT, + HIDDEN, + IMAGE_FILE_NAME, + LARGE_CHUNK_REFERENCE_IDS, + MAX_ID_SEARCH_QUERY_SIZE, + MAX_OR_CONDITIONS, + METADATA, + METADATA_SUFFIX, + PRIMARY_OWNERS, + SEARCH_ENDPOINT, + SECONDARY_OWNERS, + SECTION_CONTINUATION, + SEMANTIC_IDENTIFIER, + SOURCE_LINKS, + SOURCE_TYPE, + TENANT_ID, + TITLE, + YQL_BASE, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel diff --git a/backend/onyx/document_index/vespa/deletion.py b/backend/onyx/document_index/vespa/deletion.py index b31f86ca91..9215683950 100644 --- a/backend/onyx/document_index/vespa/deletion.py +++ b/backend/onyx/document_index/vespa/deletion.py @@ -3,8 +3,7 @@ from uuid import UUID import httpx -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT -from onyx.document_index.vespa_constants import NUM_THREADS +from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT, NUM_THREADS from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/document_index/vespa/indexing_utils.py b/backend/onyx/document_index/vespa/indexing_utils.py index a8659f45f1..05e683eab0 100644 --- a/backend/onyx/document_index/vespa/indexing_utils.py +++ b/backend/onyx/document_index/vespa/indexing_utils.py @@ -3,11 +3,9 @@ import json import random import time import uuid -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import Callable -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from http import HTTPStatus import httpx @@ -18,45 +16,49 @@ from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( from onyx.document_index.chunk_content_enrichment import ( generate_enriched_content_for_chunk_text, ) -from onyx.document_index.document_index_utils import get_uuid_from_chunk -from onyx.document_index.document_index_utils import get_uuid_from_chunk_info_old +from onyx.document_index.document_index_utils import ( + get_uuid_from_chunk, + get_uuid_from_chunk_info_old, +) from onyx.document_index.vespa.internal_types import MinimalDocumentIndexingInfo from onyx.document_index.vespa.shared_utils.utils import ( replace_invalid_doc_id_characters, ) -from onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST -from onyx.document_index.vespa_constants import AGGREGATED_CHUNK_BOOST_FACTOR -from onyx.document_index.vespa_constants import BLURB -from onyx.document_index.vespa_constants import BOOST -from onyx.document_index.vespa_constants import CHUNK_CONTEXT -from onyx.document_index.vespa_constants import CHUNK_ID -from onyx.document_index.vespa_constants import CONTENT -from onyx.document_index.vespa_constants import CONTENT_SUMMARY -from onyx.document_index.vespa_constants import DOC_SUMMARY -from onyx.document_index.vespa_constants import DOC_UPDATED_AT -from onyx.document_index.vespa_constants import DOCUMENT_ID -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT -from onyx.document_index.vespa_constants import DOCUMENT_SETS -from onyx.document_index.vespa_constants import EMBEDDINGS -from onyx.document_index.vespa_constants import FULL_CHUNK_EMBEDDING_KEY -from onyx.document_index.vespa_constants import IMAGE_FILE_NAME -from onyx.document_index.vespa_constants import LARGE_CHUNK_REFERENCE_IDS -from onyx.document_index.vespa_constants import METADATA -from onyx.document_index.vespa_constants import METADATA_LIST -from onyx.document_index.vespa_constants import METADATA_SUFFIX -from onyx.document_index.vespa_constants import NUM_THREADS -from onyx.document_index.vespa_constants import PERSONAS -from onyx.document_index.vespa_constants import PRIMARY_OWNERS -from onyx.document_index.vespa_constants import SECONDARY_OWNERS -from onyx.document_index.vespa_constants import SECTION_CONTINUATION -from onyx.document_index.vespa_constants import SEMANTIC_IDENTIFIER -from onyx.document_index.vespa_constants import SKIP_TITLE_EMBEDDING -from onyx.document_index.vespa_constants import SOURCE_LINKS -from onyx.document_index.vespa_constants import SOURCE_TYPE -from onyx.document_index.vespa_constants import TENANT_ID -from onyx.document_index.vespa_constants import TITLE -from onyx.document_index.vespa_constants import TITLE_EMBEDDING -from onyx.document_index.vespa_constants import USER_PROJECT +from onyx.document_index.vespa_constants import ( + ACCESS_CONTROL_LIST, + AGGREGATED_CHUNK_BOOST_FACTOR, + BLURB, + BOOST, + CHUNK_CONTEXT, + CHUNK_ID, + CONTENT, + CONTENT_SUMMARY, + DOC_SUMMARY, + DOC_UPDATED_AT, + DOCUMENT_ID, + DOCUMENT_ID_ENDPOINT, + DOCUMENT_SETS, + EMBEDDINGS, + FULL_CHUNK_EMBEDDING_KEY, + IMAGE_FILE_NAME, + LARGE_CHUNK_REFERENCE_IDS, + METADATA, + METADATA_LIST, + METADATA_SUFFIX, + NUM_THREADS, + PERSONAS, + PRIMARY_OWNERS, + SECONDARY_OWNERS, + SECTION_CONTINUATION, + SEMANTIC_IDENTIFIER, + SKIP_TITLE_EMBEDDING, + SOURCE_LINKS, + SOURCE_TYPE, + TENANT_ID, + TITLE, + TITLE_EMBEDDING, + USER_PROJECT, +) from onyx.indexing.models import DocMetadataAwareIndexChunk from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/document_index/vespa/kg_interactions.py b/backend/onyx/document_index/vespa/kg_interactions.py index 2bfae7a71c..f6c8f78dd2 100644 --- a/backend/onyx/document_index/vespa/kg_interactions.py +++ b/backend/onyx/document_index/vespa/kg_interactions.py @@ -1,9 +1,13 @@ -from onyx.db.document import get_document_kg_entities_and_relationships -from onyx.db.document import get_num_chunks_for_document +from onyx.db.document import ( + get_document_kg_entities_and_relationships, + get_num_chunks_for_document, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.vespa.vespa_document_index import KGUChunkUpdateRequest -from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex +from onyx.document_index.vespa.vespa_document_index import ( + KGUChunkUpdateRequest, + VespaDocumentIndex, +) from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/document_index/vespa/shared_utils/utils.py b/backend/onyx/document_index/vespa/shared_utils/utils.py index 38399eea62..a99c64d21e 100644 --- a/backend/onyx/document_index/vespa/shared_utils/utils.py +++ b/backend/onyx/document_index/vespa/shared_utils/utils.py @@ -3,10 +3,12 @@ from typing import cast import httpx -from onyx.configs.app_configs import MANAGED_VESPA -from onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH -from onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH -from onyx.configs.app_configs import VESPA_REQUEST_TIMEOUT +from onyx.configs.app_configs import ( + MANAGED_VESPA, + VESPA_CLOUD_CERT_PATH, + VESPA_CLOUD_KEY_PATH, + VESPA_REQUEST_TIMEOUT, +) from onyx.document_index.vespa_constants import VESPA_APP_CONTAINER_URL from onyx.utils.logger import setup_logger diff --git a/backend/onyx/document_index/vespa/shared_utils/vespa_request_builders.py b/backend/onyx/document_index/vespa/shared_utils/vespa_request_builders.py index 1f2100b7ad..d9f257a1c8 100644 --- a/backend/onyx/document_index/vespa/shared_utils/vespa_request_builders.py +++ b/backend/onyx/document_index/vespa/shared_utils/vespa_request_builders.py @@ -1,21 +1,21 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from onyx.configs.constants import INDEX_SEPARATOR from onyx.context.search.models import IndexFilters from onyx.document_index.vespa.internal_types import VespaChunkRequest -from onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST -from onyx.document_index.vespa_constants import CHUNK_ID -from onyx.document_index.vespa_constants import DOC_UPDATED_AT -from onyx.document_index.vespa_constants import DOCUMENT_ID -from onyx.document_index.vespa_constants import DOCUMENT_SETS -from onyx.document_index.vespa_constants import HIDDEN -from onyx.document_index.vespa_constants import METADATA_LIST -from onyx.document_index.vespa_constants import PERSONAS -from onyx.document_index.vespa_constants import SOURCE_TYPE -from onyx.document_index.vespa_constants import TENANT_ID -from onyx.document_index.vespa_constants import USER_PROJECT +from onyx.document_index.vespa_constants import ( + ACCESS_CONTROL_LIST, + CHUNK_ID, + DOC_UPDATED_AT, + DOCUMENT_ID, + DOCUMENT_SETS, + HIDDEN, + METADATA_LIST, + PERSONAS, + SOURCE_TYPE, + TENANT_ID, + USER_PROJECT, +) from onyx.kg.utils.formatting_utils import split_relationship_id from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/document_index/vespa/vespa_document_index.py b/backend/onyx/document_index/vespa/vespa_document_index.py index 1f325ef35a..b58ce35f02 100644 --- a/backend/onyx/document_index/vespa/vespa_document_index.py +++ b/backend/onyx/document_index/vespa/vespa_document_index.py @@ -7,13 +7,9 @@ import re import time import urllib.parse import zipfile -from collections.abc import Generator -from collections.abc import Iterable -from datetime import datetime -from datetime import timedelta -from typing import Any -from typing import BinaryIO -from typing import cast +from collections.abc import Generator, Iterable +from datetime import datetime, timedelta +from typing import Any, BinaryIO, cast from uuid import UUID import httpx @@ -21,59 +17,74 @@ import jinja2 import requests from pydantic import BaseModel -from onyx.configs.app_configs import MAX_CHUNKS_PER_DOC_BATCH -from onyx.configs.app_configs import RECENCY_BIAS_MULTIPLIER -from onyx.configs.app_configs import RERANK_COUNT -from onyx.configs.chat_configs import DOC_TIME_DECAY -from onyx.configs.chat_configs import HYBRID_ALPHA -from onyx.configs.chat_configs import TITLE_CONTENT_RATIO -from onyx.configs.chat_configs import VESPA_SEARCHER_THREADS +from onyx.configs.app_configs import ( + MAX_CHUNKS_PER_DOC_BATCH, + RECENCY_BIAS_MULTIPLIER, + RERANK_COUNT, +) +from onyx.configs.chat_configs import ( + DOC_TIME_DECAY, + HYBRID_ALPHA, + TITLE_CONTENT_RATIO, + VESPA_SEARCHER_THREADS, +) from onyx.configs.constants import KV_REINDEX_KEY from onyx.context.search.enums import QueryType -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import IndexFilters, InferenceChunk from onyx.db.enums import EmbeddingPrecision from onyx.document_index.chunk_content_enrichment import cleanup_content_for_chunks -from onyx.document_index.document_index_utils import get_document_chunk_ids -from onyx.document_index.document_index_utils import get_uuid_from_chunk_info -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentInsertionRecord -from onyx.document_index.interfaces_new import DocumentSectionRequest -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import MetadataUpdateRequest -from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.vespa.chunk_retrieval import batch_search_api_retrieval -from onyx.document_index.vespa.chunk_retrieval import get_all_chunks_paginated -from onyx.document_index.vespa.chunk_retrieval import get_chunks_via_visit_api -from onyx.document_index.vespa.chunk_retrieval import parallel_visit_api_retrieval -from onyx.document_index.vespa.chunk_retrieval import query_vespa +from onyx.document_index.document_index_utils import ( + get_document_chunk_ids, + get_uuid_from_chunk_info, +) +from onyx.document_index.interfaces_new import ( + DocumentIndex, + DocumentInsertionRecord, + DocumentSectionRequest, + IndexingMetadata, + MetadataUpdateRequest, + TenantState, +) +from onyx.document_index.vespa.chunk_retrieval import ( + batch_search_api_retrieval, + get_all_chunks_paginated, + get_chunks_via_visit_api, + parallel_visit_api_retrieval, + query_vespa, +) from onyx.document_index.vespa.deletion import delete_vespa_chunks -from onyx.document_index.vespa.indexing_utils import BaseHTTPXClientContext -from onyx.document_index.vespa.indexing_utils import batch_index_vespa_chunks -from onyx.document_index.vespa.indexing_utils import check_for_final_chunk_existence -from onyx.document_index.vespa.indexing_utils import clean_chunk_id_copy -from onyx.document_index.vespa.indexing_utils import GlobalHTTPXClientContext -from onyx.document_index.vespa.indexing_utils import TemporaryHTTPXClientContext -from onyx.document_index.vespa.internal_types import EnrichedDocumentIndexingInfo -from onyx.document_index.vespa.internal_types import MinimalDocumentIndexingInfo -from onyx.document_index.vespa.internal_types import VespaChunkRequest -from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client +from onyx.document_index.vespa.indexing_utils import ( + BaseHTTPXClientContext, + batch_index_vespa_chunks, + check_for_final_chunk_existence, + clean_chunk_id_copy, + GlobalHTTPXClientContext, + TemporaryHTTPXClientContext, +) +from onyx.document_index.vespa.internal_types import ( + EnrichedDocumentIndexingInfo, + MinimalDocumentIndexingInfo, + VespaChunkRequest, +) from onyx.document_index.vespa.shared_utils.utils import ( + get_vespa_http_client, replace_invalid_doc_id_characters, ) from onyx.document_index.vespa.shared_utils.vespa_request_builders import ( build_vespa_filters, ) -from onyx.document_index.vespa_constants import BATCH_SIZE -from onyx.document_index.vespa_constants import CHUNK_ID -from onyx.document_index.vespa_constants import CONTENT_SUMMARY -from onyx.document_index.vespa_constants import DOCUMENT_ID -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT -from onyx.document_index.vespa_constants import NUM_THREADS -from onyx.document_index.vespa_constants import SEARCH_ENDPOINT -from onyx.document_index.vespa_constants import VESPA_APPLICATION_ENDPOINT -from onyx.document_index.vespa_constants import VESPA_TIMEOUT -from onyx.document_index.vespa_constants import YQL_BASE +from onyx.document_index.vespa_constants import ( + BATCH_SIZE, + CHUNK_ID, + CONTENT_SUMMARY, + DOCUMENT_ID, + DOCUMENT_ID_ENDPOINT, + NUM_THREADS, + SEARCH_ENDPOINT, + VESPA_APPLICATION_ENDPOINT, + VESPA_TIMEOUT, + YQL_BASE, +) from onyx.indexing.models import DocMetadataAwareIndexChunk from onyx.key_value_store.factory import get_shared_kv_store from onyx.kg.utils.formatting_utils import split_relationship_id diff --git a/backend/onyx/document_index/vespa_constants.py b/backend/onyx/document_index/vespa_constants.py index 9b5e0512d0..3caf5945ed 100644 --- a/backend/onyx/document_index/vespa_constants.py +++ b/backend/onyx/document_index/vespa_constants.py @@ -1,8 +1,10 @@ -from onyx.configs.app_configs import VESPA_CLOUD_URL -from onyx.configs.app_configs import VESPA_CONFIG_SERVER_HOST -from onyx.configs.app_configs import VESPA_HOST -from onyx.configs.app_configs import VESPA_PORT -from onyx.configs.app_configs import VESPA_TENANT_PORT +from onyx.configs.app_configs import ( + VESPA_CLOUD_URL, + VESPA_CONFIG_SERVER_HOST, + VESPA_HOST, + VESPA_PORT, + VESPA_TENANT_PORT, +) from onyx.configs.constants import SOURCE_TYPE # config server diff --git a/backend/onyx/error_handling/exceptions.py b/backend/onyx/error_handling/exceptions.py index ad5cfac1e6..c6f5bd9510 100644 --- a/backend/onyx/error_handling/exceptions.py +++ b/backend/onyx/error_handling/exceptions.py @@ -22,8 +22,7 @@ use ``status_code_override``:: ) """ -from fastapi import FastAPI -from fastapi import Request +from fastapi import FastAPI, Request from fastapi.responses import JSONResponse from onyx.error_handling.error_codes import OnyxErrorCode diff --git a/backend/onyx/evals/eval.py b/backend/onyx/evals/eval.py index eee611c3b2..8c3fc59ce3 100644 --- a/backend/onyx/evals/eval.py +++ b/backend/onyx/evals/eval.py @@ -1,37 +1,37 @@ import time -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager from typing import Any -from sqlalchemy import Engine -from sqlalchemy import event -from sqlalchemy.orm import Session -from sqlalchemy.orm import sessionmaker +from sqlalchemy import Engine, event +from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.orm.session import SessionTransaction from onyx.chat.chat_state import ChatStateContainer from onyx.chat.models import ChatFullResponse -from onyx.chat.process_message import gather_stream_full -from onyx.chat.process_message import handle_stream_message_objects +from onyx.chat.process_message import gather_stream_full, handle_stream_message_objects from onyx.configs.constants import DEFAULT_PERSONA_ID from onyx.db.chat import create_chat_session from onyx.db.engine.sql_engine import get_sqlalchemy_engine from onyx.db.users import get_user_by_email -from onyx.evals.models import ChatFullEvalResult -from onyx.evals.models import EvalationAck -from onyx.evals.models import EvalConfigurationOptions -from onyx.evals.models import EvalMessage -from onyx.evals.models import EvalProvider -from onyx.evals.models import EvalTimings -from onyx.evals.models import EvalToolResult -from onyx.evals.models import MultiTurnEvalResult -from onyx.evals.models import ToolAssertion +from onyx.evals.models import ( + ChatFullEvalResult, + EvalationAck, + EvalConfigurationOptions, + EvalMessage, + EvalProvider, + EvalTimings, + EvalToolResult, + MultiTurnEvalResult, + ToolAssertion, +) from onyx.evals.provider import get_provider from onyx.llm.override_models import LLMOverride -from onyx.server.query_and_chat.models import AUTO_PLACE_AFTER_LATEST_MESSAGE -from onyx.server.query_and_chat.models import ChatSessionCreationRequest -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.server.query_and_chat.models import ( + AUTO_PLACE_AFTER_LATEST_MESSAGE, + ChatSessionCreationRequest, + SendMessageRequest, +) from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/evals/eval_cli.py b/backend/onyx/evals/eval_cli.py index 7f15cbe161..d28d611c83 100644 --- a/backend/onyx/evals/eval_cli.py +++ b/backend/onyx/evals/eval_cli.py @@ -12,13 +12,14 @@ from typing import Any import braintrust import requests -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE +from onyx.configs.app_configs import ( + POSTGRES_API_SERVER_POOL_OVERFLOW, + POSTGRES_API_SERVER_POOL_SIZE, +) from onyx.configs.constants import POSTGRES_WEB_APP_NAME from onyx.db.engine.sql_engine import SqlEngine from onyx.evals.eval import run_eval -from onyx.evals.models import EvalationAck -from onyx.evals.models import EvalConfigurationOptions +from onyx.evals.models import EvalationAck, EvalConfigurationOptions from onyx.evals.provider import get_provider from onyx.tracing.setup import setup_tracing diff --git a/backend/onyx/evals/models.py b/backend/onyx/evals/models.py index 59c72169e9..b0f2984db0 100644 --- a/backend/onyx/evals/models.py +++ b/backend/onyx/evals/models.py @@ -1,10 +1,8 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import Callable from typing import Any -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from sqlalchemy.orm import Session from onyx.db.tools import get_builtin_tool diff --git a/backend/onyx/evals/one_off/create_braintrust_dataset.py b/backend/onyx/evals/one_off/create_braintrust_dataset.py index 9739ee67c2..97bf9e360b 100644 --- a/backend/onyx/evals/one_off/create_braintrust_dataset.py +++ b/backend/onyx/evals/one_off/create_braintrust_dataset.py @@ -16,9 +16,7 @@ import argparse import csv import os import sys -from typing import Any -from typing import Dict -from typing import List +from typing import Any, Dict, List from onyx.configs.app_configs import BRAINTRUST_API_KEY diff --git a/backend/onyx/evals/providers/braintrust.py b/backend/onyx/evals/providers/braintrust.py index 91ecb26ae7..46149d0422 100644 --- a/backend/onyx/evals/providers/braintrust.py +++ b/backend/onyx/evals/providers/braintrust.py @@ -1,19 +1,16 @@ from collections.abc import Callable -from typing import Any -from typing import Union +from typing import Any, Union -from braintrust import Eval -from braintrust import EvalCase -from braintrust import init_dataset -from braintrust import Score +from braintrust import Eval, EvalCase, init_dataset, Score -from onyx.configs.app_configs import BRAINTRUST_MAX_CONCURRENCY -from onyx.configs.app_configs import BRAINTRUST_PROJECT -from onyx.evals.models import EvalationAck -from onyx.evals.models import EvalConfigurationOptions -from onyx.evals.models import EvalProvider -from onyx.evals.models import EvalToolResult -from onyx.evals.models import MultiTurnEvalResult +from onyx.configs.app_configs import BRAINTRUST_MAX_CONCURRENCY, BRAINTRUST_PROJECT +from onyx.evals.models import ( + EvalationAck, + EvalConfigurationOptions, + EvalProvider, + EvalToolResult, + MultiTurnEvalResult, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/evals/providers/local.py b/backend/onyx/evals/providers/local.py index 0b9a5170fb..8442e81241 100644 --- a/backend/onyx/evals/providers/local.py +++ b/backend/onyx/evals/providers/local.py @@ -6,11 +6,13 @@ No external dependencies like Braintrust required. from collections.abc import Callable from typing import Any -from onyx.evals.models import EvalationAck -from onyx.evals.models import EvalConfigurationOptions -from onyx.evals.models import EvalProvider -from onyx.evals.models import EvalToolResult -from onyx.evals.models import MultiTurnEvalResult +from onyx.evals.models import ( + EvalationAck, + EvalConfigurationOptions, + EvalProvider, + EvalToolResult, + MultiTurnEvalResult, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/external_apps/credentials.py b/backend/onyx/external_apps/credentials.py index a81db794a1..6f3fac2128 100644 --- a/backend/onyx/external_apps/credentials.py +++ b/backend/onyx/external_apps/credentials.py @@ -3,8 +3,10 @@ from uuid import UUID from sqlalchemy.orm import Session -from onyx.db.external_app import get_external_app_by_id -from onyx.db.external_app import get_external_app_user_credential +from onyx.db.external_app import ( + get_external_app_by_id, + get_external_app_user_credential, +) from onyx.db.models import ExternalApp diff --git a/backend/onyx/external_apps/matching/engine.py b/backend/onyx/external_apps/matching/engine.py index 08e04fb487..952e2691aa 100644 --- a/backend/onyx/external_apps/matching/engine.py +++ b/backend/onyx/external_apps/matching/engine.py @@ -1,27 +1,22 @@ """Composes recognition + policy resolution into a single verdict for an outbound request.""" -from collections.abc import Iterable -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from typing import Any -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from sqlalchemy.orm import Session -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.enums import POLICY_SEVERITY +from onyx.db.enums import EndpointPolicy, ExternalAppType, POLICY_SEVERITY from onyx.db.external_app import get_policies from onyx.db.models import ExternalApp -from onyx.external_apps.matching.request import MatchContext -from onyx.external_apps.matching.request import ProxiedRequest +from onyx.external_apps.matching.request import MatchContext, ProxiedRequest from onyx.external_apps.matching.rules import rule_matches -from onyx.external_apps.providers.registry import effective_policy -from onyx.external_apps.providers.registry import get_endpoint_catalog -from onyx.external_apps.providers.registry import get_provider_for_app +from onyx.external_apps.providers.registry import ( + effective_policy, + get_endpoint_catalog, + get_provider_for_app, +) # action_type for a domain-matched request that hit no catalog action. WHOLE_DOMAIN_ACTION_TYPE = "unspecified" diff --git a/backend/onyx/external_apps/matching/graphql_parsing.py b/backend/onyx/external_apps/matching/graphql_parsing.py index 1ade3f50fb..25460a0f20 100644 --- a/backend/onyx/external_apps/matching/graphql_parsing.py +++ b/backend/onyx/external_apps/matching/graphql_parsing.py @@ -8,12 +8,14 @@ the rule-matching layer. import json -from graphql import FieldNode -from graphql import FragmentDefinitionNode -from graphql import FragmentSpreadNode -from graphql import GraphQLSyntaxError -from graphql import InlineFragmentNode -from graphql import OperationDefinitionNode +from graphql import ( + FieldNode, + FragmentDefinitionNode, + FragmentSpreadNode, + GraphQLSyntaxError, + InlineFragmentNode, + OperationDefinitionNode, +) from graphql import parse as parse_graphql from graphql.language import SelectionSetNode diff --git a/backend/onyx/external_apps/matching/rules.py b/backend/onyx/external_apps/matching/rules.py index 3bf59e0973..c7041a5e02 100644 --- a/backend/onyx/external_apps/matching/rules.py +++ b/backend/onyx/external_apps/matching/rules.py @@ -4,19 +4,16 @@ registry. Adding a new rule kind is "write a ``RuleMatcher`` + register it" — the dispatch in ``rule_matches`` never changes (open/closed). """ -from abc import ABC -from abc import abstractmethod -from typing import Any -from typing import ClassVar -from typing import Generic -from typing import get_args -from typing import TypeVar +from abc import ABC, abstractmethod +from typing import Any, ClassVar, Generic, get_args, TypeVar from onyx.external_apps.matching.request import MatchContext -from onyx.external_apps.providers.actions import GraphQLOp -from onyx.external_apps.providers.actions import MatchRule -from onyx.external_apps.providers.actions import path_matches -from onyx.external_apps.providers.actions import RestRoute +from onyx.external_apps.providers.actions import ( + GraphQLOp, + MatchRule, + path_matches, + RestRoute, +) RuleT = TypeVar("RuleT") diff --git a/backend/onyx/external_apps/models.py b/backend/onyx/external_apps/models.py index 16676840eb..09c8129865 100644 --- a/backend/onyx/external_apps/models.py +++ b/backend/onyx/external_apps/models.py @@ -1,7 +1,6 @@ from pydantic import BaseModel -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType class ActionPolicyView(BaseModel): diff --git a/backend/onyx/external_apps/presentation/payload_decoders.py b/backend/onyx/external_apps/presentation/payload_decoders.py index 28540d994a..0194b70bcb 100644 --- a/backend/onyx/external_apps/presentation/payload_decoders.py +++ b/backend/onyx/external_apps/presentation/payload_decoders.py @@ -2,13 +2,10 @@ from __future__ import annotations import base64 import binascii -from email import message_from_bytes -from email import policy +from email import message_from_bytes, policy from email.message import EmailMessage -from email.utils import formataddr -from email.utils import getaddresses -from typing import Any -from typing import Protocol +from email.utils import formataddr, getaddresses +from typing import Any, Protocol class PayloadDecoder(Protocol): diff --git a/backend/onyx/external_apps/providers/actions.py b/backend/onyx/external_apps/providers/actions.py index c9534716c7..ac0ee92078 100644 --- a/backend/onyx/external_apps/providers/actions.py +++ b/backend/onyx/external_apps/providers/actions.py @@ -1,12 +1,7 @@ from enum import Enum -from typing import Annotated -from typing import Literal +from typing import Annotated, Literal -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_validator -from pydantic import InstanceOf +from pydantic import BaseModel, ConfigDict, Field, field_validator, InstanceOf from onyx.db.enums import EndpointPolicy diff --git a/backend/onyx/external_apps/providers/base.py b/backend/onyx/external_apps/providers/base.py index 33530d0966..78b9dfc2c4 100644 --- a/backend/onyx/external_apps/providers/base.py +++ b/backend/onyx/external_apps/providers/base.py @@ -1,12 +1,9 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import Any -from typing import ClassVar +from typing import Any, ClassVar import requests -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from onyx.db.enums import ExternalAppType from onyx.external_apps.presentation.payload_decoders import PayloadDecoder diff --git a/backend/onyx/external_apps/providers/github.py b/backend/onyx/external_apps/providers/github.py index 40342bd427..6897b864ad 100644 --- a/backend/onyx/external_apps/providers/github.py +++ b/backend/onyx/external_apps/providers/github.py @@ -2,23 +2,28 @@ from typing import Any import requests -from onyx.configs.app_configs import EXT_APP_GITHUB_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_GITHUB_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.configs.app_configs import ( + EXT_APP_GITHUB_CLIENT_ID, + EXT_APP_GITHUB_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import GraphQLOp -from onyx.external_apps.providers.actions import RestRoute -from onyx.external_apps.providers.base import AdminDescriptorSpec -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import OAuthFlowSpec -from onyx.external_apps.providers.base import OAuthProviderSpec -from onyx.external_apps.providers.base import OnyxManagedExtApp -from onyx.external_apps.providers.base import OrgCredentialField -from onyx.external_apps.providers.base import token_response_error +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + GraphQLOp, + RestRoute, +) +from onyx.external_apps.providers.base import ( + AdminDescriptorSpec, + OAuthExternalAppProvider, + OAuthFlowSpec, + OAuthProviderSpec, + OnyxManagedExtApp, + OrgCredentialField, + token_response_error, +) class GitHubAction(ExternalAppAction): diff --git a/backend/onyx/external_apps/providers/gmail.py b/backend/onyx/external_apps/providers/gmail.py index 01aad23fa0..ec3d2eb863 100644 --- a/backend/onyx/external_apps/providers/gmail.py +++ b/backend/onyx/external_apps/providers/gmail.py @@ -1,14 +1,19 @@ from collections.abc import Mapping -from onyx.configs.app_configs import EXT_APP_GMAIL_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_GMAIL_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.presentation.payload_decoders import GmailRawMimeDecoder -from onyx.external_apps.presentation.payload_decoders import PayloadDecoder -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import RestRoute +from onyx.configs.app_configs import ( + EXT_APP_GMAIL_CLIENT_ID, + EXT_APP_GMAIL_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.presentation.payload_decoders import ( + GmailRawMimeDecoder, + PayloadDecoder, +) +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + RestRoute, +) from onyx.external_apps.providers.base import OnyxManagedExtApp from onyx.external_apps.providers.google_base import GoogleOAuthProvider diff --git a/backend/onyx/external_apps/providers/google_base.py b/backend/onyx/external_apps/providers/google_base.py index 8c858fa616..ce028feb16 100644 --- a/backend/onyx/external_apps/providers/google_base.py +++ b/backend/onyx/external_apps/providers/google_base.py @@ -4,11 +4,13 @@ from onyx.db.enums import ExternalAppType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.base import AdminDescriptorSpec -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import OAuthFlowSpec -from onyx.external_apps.providers.base import OAuthProviderSpec -from onyx.external_apps.providers.base import OrgCredentialField +from onyx.external_apps.providers.base import ( + AdminDescriptorSpec, + OAuthExternalAppProvider, + OAuthFlowSpec, + OAuthProviderSpec, + OrgCredentialField, +) # Google's OAuth 2.0 endpoints are shared across all Google APIs. _AUTHORIZE_URL = "https://accounts.google.com/o/oauth2/v2/auth" diff --git a/backend/onyx/external_apps/providers/google_calendar.py b/backend/onyx/external_apps/providers/google_calendar.py index 6e1b9faa60..cda89c50a6 100644 --- a/backend/onyx/external_apps/providers/google_calendar.py +++ b/backend/onyx/external_apps/providers/google_calendar.py @@ -1,10 +1,13 @@ -from onyx.configs.app_configs import EXT_APP_GOOGLE_CALENDAR_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_GOOGLE_CALENDAR_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import RestRoute +from onyx.configs.app_configs import ( + EXT_APP_GOOGLE_CALENDAR_CLIENT_ID, + EXT_APP_GOOGLE_CALENDAR_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + RestRoute, +) from onyx.external_apps.providers.base import OnyxManagedExtApp from onyx.external_apps.providers.google_base import GoogleOAuthProvider diff --git a/backend/onyx/external_apps/providers/google_drive.py b/backend/onyx/external_apps/providers/google_drive.py index a956fa313c..4bdd2e7976 100644 --- a/backend/onyx/external_apps/providers/google_drive.py +++ b/backend/onyx/external_apps/providers/google_drive.py @@ -1,10 +1,13 @@ -from onyx.configs.app_configs import EXT_APP_GOOGLE_DRIVE_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_GOOGLE_DRIVE_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import RestRoute +from onyx.configs.app_configs import ( + EXT_APP_GOOGLE_DRIVE_CLIENT_ID, + EXT_APP_GOOGLE_DRIVE_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + RestRoute, +) from onyx.external_apps.providers.base import OnyxManagedExtApp from onyx.external_apps.providers.google_base import GoogleOAuthProvider diff --git a/backend/onyx/external_apps/providers/hubspot.py b/backend/onyx/external_apps/providers/hubspot.py index ecff4cf72e..e1a728189e 100644 --- a/backend/onyx/external_apps/providers/hubspot.py +++ b/backend/onyx/external_apps/providers/hubspot.py @@ -3,23 +3,28 @@ from urllib.parse import quote import requests -from onyx.configs.app_configs import EXT_APP_HUBSPOT_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_HUBSPOT_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.configs.app_configs import ( + EXT_APP_HUBSPOT_CLIENT_ID, + EXT_APP_HUBSPOT_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import RestRoute -from onyx.external_apps.providers.base import AdminDescriptorSpec -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import OAuthFlowSpec -from onyx.external_apps.providers.base import OAuthProviderSpec -from onyx.external_apps.providers.base import OnyxManagedExtApp -from onyx.external_apps.providers.base import OrgCredentialField -from onyx.external_apps.providers.base import parse_granted_scopes -from onyx.external_apps.providers.base import token_response_error +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + RestRoute, +) +from onyx.external_apps.providers.base import ( + AdminDescriptorSpec, + OAuthExternalAppProvider, + OAuthFlowSpec, + OAuthProviderSpec, + OnyxManagedExtApp, + OrgCredentialField, + parse_granted_scopes, + token_response_error, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/external_apps/providers/linear.py b/backend/onyx/external_apps/providers/linear.py index a093b7da58..f94e6db1cd 100644 --- a/backend/onyx/external_apps/providers/linear.py +++ b/backend/onyx/external_apps/providers/linear.py @@ -1,20 +1,25 @@ from typing import Any -from onyx.configs.app_configs import EXT_APP_LINEAR_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_LINEAR_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.configs.app_configs import ( + EXT_APP_LINEAR_CLIENT_ID, + EXT_APP_LINEAR_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import GraphQLOp -from onyx.external_apps.providers.base import AdminDescriptorSpec -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import OAuthFlowSpec -from onyx.external_apps.providers.base import OAuthProviderSpec -from onyx.external_apps.providers.base import OnyxManagedExtApp -from onyx.external_apps.providers.base import OrgCredentialField +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + GraphQLOp, +) +from onyx.external_apps.providers.base import ( + AdminDescriptorSpec, + OAuthExternalAppProvider, + OAuthFlowSpec, + OAuthProviderSpec, + OnyxManagedExtApp, + OrgCredentialField, +) class LinearAction(ExternalAppAction): diff --git a/backend/onyx/external_apps/providers/notion.py b/backend/onyx/external_apps/providers/notion.py index 09609f1d89..eb6ac8c589 100644 --- a/backend/onyx/external_apps/providers/notion.py +++ b/backend/onyx/external_apps/providers/notion.py @@ -1,22 +1,27 @@ import base64 from typing import Any -from onyx.configs.app_configs import EXT_APP_NOTION_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_NOTION_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.configs.app_configs import ( + EXT_APP_NOTION_CLIENT_ID, + EXT_APP_NOTION_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import RestRoute -from onyx.external_apps.providers.base import AdminDescriptorSpec -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import OAuthFlowSpec -from onyx.external_apps.providers.base import OAuthProviderSpec -from onyx.external_apps.providers.base import OnyxManagedExtApp -from onyx.external_apps.providers.base import OrgCredentialField -from onyx.external_apps.providers.base import TokenExchangeRequest +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + RestRoute, +) +from onyx.external_apps.providers.base import ( + AdminDescriptorSpec, + OAuthExternalAppProvider, + OAuthFlowSpec, + OAuthProviderSpec, + OnyxManagedExtApp, + OrgCredentialField, + TokenExchangeRequest, +) # Pinned across the provider and the sandbox skill so request-shaping stays # consistent with what the OAuth exchange negotiated. 2025-09-03 is the current diff --git a/backend/onyx/external_apps/providers/registry.py b/backend/onyx/external_apps/providers/registry.py index 5351a1d5d2..46d2beb943 100644 --- a/backend/onyx/external_apps/providers/registry.py +++ b/backend/onyx/external_apps/providers/registry.py @@ -1,15 +1,15 @@ -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.db.models import ExternalApp from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.models import ActionPolicyView -from onyx.external_apps.models import BuiltInExternalAppDescriptor -from onyx.external_apps.models import EndpointDescriptor -from onyx.external_apps.models import OrgCredentialFieldDescriptor +from onyx.external_apps.models import ( + ActionPolicyView, + BuiltInExternalAppDescriptor, + EndpointDescriptor, + OrgCredentialFieldDescriptor, +) from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.base import ExternalAppProvider -from onyx.external_apps.providers.base import OnyxManagedExtApp +from onyx.external_apps.providers.base import ExternalAppProvider, OnyxManagedExtApp from onyx.external_apps.providers.github import GitHubProvider from onyx.external_apps.providers.gmail import GmailProvider from onyx.external_apps.providers.google_calendar import GoogleCalendarProvider diff --git a/backend/onyx/external_apps/providers/slack.py b/backend/onyx/external_apps/providers/slack.py index 84847557bb..cdf1c1e775 100644 --- a/backend/onyx/external_apps/providers/slack.py +++ b/backend/onyx/external_apps/providers/slack.py @@ -1,21 +1,26 @@ from typing import Any -from onyx.configs.app_configs import EXT_APP_SLACK_CLIENT_ID -from onyx.configs.app_configs import EXT_APP_SLACK_CLIENT_SECRET -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.configs.app_configs import ( + EXT_APP_SLACK_CLIENT_ID, + EXT_APP_SLACK_CLIENT_SECRET, +) +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import RestRoute -from onyx.external_apps.providers.base import AdminDescriptorSpec -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import OAuthFlowSpec -from onyx.external_apps.providers.base import OAuthProviderSpec -from onyx.external_apps.providers.base import OnyxManagedExtApp -from onyx.external_apps.providers.base import OrgCredentialField -from onyx.external_apps.providers.base import parse_granted_scopes +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + RestRoute, +) +from onyx.external_apps.providers.base import ( + AdminDescriptorSpec, + OAuthExternalAppProvider, + OAuthFlowSpec, + OAuthProviderSpec, + OnyxManagedExtApp, + OrgCredentialField, + parse_granted_scopes, +) class SlackAction(ExternalAppAction): diff --git a/backend/onyx/external_apps/token_refresh.py b/backend/onyx/external_apps/token_refresh.py index 20773f119f..cc0c0d8239 100644 --- a/backend/onyx/external_apps/token_refresh.py +++ b/backend/onyx/external_apps/token_refresh.py @@ -6,8 +6,7 @@ behind :func:`ensure_fresh_credentials`. Each step takes its own short session, no connection is held across the lock wait or the POST. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from uuid import UUID @@ -16,19 +15,21 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.external_app import delete_external_app_user_credential -from onyx.db.external_app import get_external_app_by_id -from onyx.db.external_app import get_external_app_user_credential -from onyx.db.external_app import upsert_external_app_user_credential +from onyx.db.external_app import ( + delete_external_app_user_credential, + get_external_app_by_id, + get_external_app_user_credential, + upsert_external_app_user_credential, +) from onyx.db.models import ExternalApp -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import TokenRefreshTerminalError -from onyx.external_apps.providers.base import TokenRefreshTransientError +from onyx.external_apps.providers.base import ( + OAuthExternalAppProvider, + TokenRefreshTerminalError, + TokenRefreshTransientError, +) from onyx.external_apps.providers.registry import get_provider_for_app -from onyx.external_apps.token_utils import needs_refresh -from onyx.external_apps.token_utils import stamp_expires_at -from onyx.redis.lock_context import redis_shared_lock -from onyx.redis.lock_context import RedisSharedLockAcquisitionError +from onyx.external_apps.token_utils import needs_refresh, stamp_expires_at +from onyx.redis.lock_context import redis_shared_lock, RedisSharedLockAcquisitionError from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/external_apps/token_utils.py b/backend/onyx/external_apps/token_utils.py index 3615f54776..1f116ee9f5 100644 --- a/backend/onyx/external_apps/token_utils.py +++ b/backend/onyx/external_apps/token_utils.py @@ -2,8 +2,7 @@ from providers/orchestrator) so callers can use it without an import cycle. """ -from datetime import datetime -from datetime import timedelta +from datetime import datetime, timedelta from typing import Any from onyx.utils.datetime import datetime_to_utc diff --git a/backend/onyx/external_apps/url_glob.py b/backend/onyx/external_apps/url_glob.py index af51e3ec7d..b6f3abd742 100644 --- a/backend/onyx/external_apps/url_glob.py +++ b/backend/onyx/external_apps/url_glob.py @@ -1,8 +1,7 @@ import re from functools import lru_cache -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/onyx/feature_flags/factory.py b/backend/onyx/feature_flags/factory.py index 6251503e24..965487446f 100644 --- a/backend/onyx/feature_flags/factory.py +++ b/backend/onyx/feature_flags/factory.py @@ -1,6 +1,5 @@ from onyx.configs.app_configs import DEV_MODE -from onyx.feature_flags.interface import FeatureFlagProvider -from onyx.feature_flags.interface import NoOpFeatureFlagProvider +from onyx.feature_flags.interface import FeatureFlagProvider, NoOpFeatureFlagProvider from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, ) diff --git a/backend/onyx/federated_connectors/factory.py b/backend/onyx/federated_connectors/factory.py index d63222d404..e4acc76e0c 100644 --- a/backend/onyx/federated_connectors/factory.py +++ b/backend/onyx/federated_connectors/factory.py @@ -1,8 +1,7 @@ """Factory for creating federated connector instances.""" import importlib -from typing import Any -from typing import Type +from typing import Any, Type from onyx.configs.constants import FederatedConnectorSource from onyx.federated_connectors.interfaces import FederatedConnector diff --git a/backend/onyx/federated_connectors/federated_retrieval.py b/backend/onyx/federated_connectors/federated_retrieval.py index 30dec55dd8..6a86d36626 100644 --- a/backend/onyx/federated_connectors/federated_retrieval.py +++ b/backend/onyx/federated_connectors/federated_retrieval.py @@ -3,18 +3,15 @@ from collections.abc import Callable from typing import Any from uuid import UUID -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from sqlalchemy.orm import Session -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FederatedConnectorSource -from onyx.context.search.models import ChunkIndexRequest -from onyx.context.search.models import InferenceChunk +from onyx.configs.constants import DocumentSource, FederatedConnectorSource +from onyx.context.search.models import ChunkIndexRequest, InferenceChunk from onyx.db.federated import ( get_federated_connector_document_set_mappings_by_document_set_names, + list_federated_connector_oauth_tokens, ) -from onyx.db.federated import list_federated_connector_oauth_tokens from onyx.db.models import FederatedConnector__DocumentSet from onyx.db.slack_bot import fetch_slack_bots from onyx.federated_connectors.factory import get_federated_connector diff --git a/backend/onyx/federated_connectors/interfaces.py b/backend/onyx/federated_connectors/interfaces.py index b48cdca3b1..2c3fe8671e 100644 --- a/backend/onyx/federated_connectors/interfaces.py +++ b/backend/onyx/federated_connectors/interfaces.py @@ -1,13 +1,8 @@ -from abc import ABC -from abc import abstractmethod -from typing import Any -from typing import Dict +from abc import ABC, abstractmethod +from typing import Any, Dict -from onyx.context.search.models import ChunkIndexRequest -from onyx.context.search.models import InferenceChunk -from onyx.federated_connectors.models import CredentialField -from onyx.federated_connectors.models import EntityField -from onyx.federated_connectors.models import OAuthResult +from onyx.context.search.models import ChunkIndexRequest, InferenceChunk +from onyx.federated_connectors.models import CredentialField, EntityField, OAuthResult from onyx.onyxbot.slack.models import SlackContext diff --git a/backend/onyx/federated_connectors/models.py b/backend/onyx/federated_connectors/models.py index 8f22220a2c..02945f7634 100644 --- a/backend/onyx/federated_connectors/models.py +++ b/backend/onyx/federated_connectors/models.py @@ -1,10 +1,7 @@ from datetime import datetime -from typing import Any -from typing import Dict -from typing import Optional +from typing import Any, Dict, Optional -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field class FieldSpec(BaseModel): diff --git a/backend/onyx/federated_connectors/slack/federated_connector.py b/backend/onyx/federated_connectors/slack/federated_connector.py index 6d18020577..74bad708f8 100644 --- a/backend/onyx/federated_connectors/slack/federated_connector.py +++ b/backend/onyx/federated_connectors/slack/federated_connector.py @@ -1,6 +1,4 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from urllib.parse import urlencode @@ -10,15 +8,11 @@ from slack_sdk import WebClient from typing_extensions import override from onyx.context.search.federated.slack_search import slack_retrieval -from onyx.context.search.models import ChunkIndexRequest -from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import ChunkIndexRequest, InferenceChunk from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.federated_connectors.interfaces import FederatedConnector -from onyx.federated_connectors.models import CredentialField -from onyx.federated_connectors.models import EntityField -from onyx.federated_connectors.models import OAuthResult -from onyx.federated_connectors.slack.models import SlackCredentials -from onyx.federated_connectors.slack.models import SlackEntities +from onyx.federated_connectors.models import CredentialField, EntityField, OAuthResult +from onyx.federated_connectors.slack.models import SlackCredentials, SlackEntities from onyx.onyxbot.slack.models import SlackContext from onyx.utils.logger import setup_logger diff --git a/backend/onyx/federated_connectors/slack/models.py b/backend/onyx/federated_connectors/slack/models.py index 7aef330a69..9b9eba1474 100644 --- a/backend/onyx/federated_connectors/slack/models.py +++ b/backend/onyx/federated_connectors/slack/models.py @@ -1,9 +1,6 @@ from typing import Optional -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator -from pydantic import model_validator +from pydantic import BaseModel, Field, field_validator, model_validator class SlackEntities(BaseModel): diff --git a/backend/onyx/file_processing/extract_file_text.py b/backend/onyx/file_processing/extract_file_text.py index 918442ec64..54700a260f 100644 --- a/backend/onyx/file_processing/extract_file_text.py +++ b/backend/onyx/file_processing/extract_file_text.py @@ -6,18 +6,11 @@ import os import re import tempfile import zipfile -from collections.abc import Callable -from collections.abc import Iterator -from collections.abc import Sequence +from collections.abc import Callable, Iterator, Sequence from email.parser import Parser as EmailParser from io import BytesIO from pathlib import Path -from typing import Any -from typing import cast -from typing import IO -from typing import NamedTuple -from typing import Optional -from typing import TYPE_CHECKING +from typing import Any, cast, IO, NamedTuple, Optional, TYPE_CHECKING from zipfile import BadZipFile import chardet @@ -25,21 +18,26 @@ import openpyxl from openpyxl.worksheet._read_only import ReadOnlyWorksheet from PIL import Image -from onyx.configs.app_configs import MAX_EMBEDDED_IMAGES_PER_FILE -from onyx.configs.app_configs import MAX_XLSX_CELLS_PER_SHEET -from onyx.configs.app_configs import PDF_TEXT_EXTRACTION_TIMEOUT_SECONDS +from onyx.configs.app_configs import ( + MAX_EMBEDDED_IMAGES_PER_FILE, + MAX_XLSX_CELLS_PER_SHEET, + PDF_TEXT_EXTRACTION_TIMEOUT_SECONDS, +) from onyx.configs.constants import ONYX_METADATA_FILENAME from onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled -from onyx.file_processing.file_types import OnyxFileExtensions -from onyx.file_processing.file_types import OnyxMimeTypes -from onyx.file_processing.file_types import PRESENTATION_MIME_TYPE -from onyx.file_processing.file_types import WORD_PROCESSING_MIME_TYPE +from onyx.file_processing.file_types import ( + OnyxFileExtensions, + OnyxMimeTypes, + PRESENTATION_MIME_TYPE, + WORD_PROCESSING_MIME_TYPE, +) from onyx.file_processing.html_utils import parse_html_page_basic -from onyx.file_processing.unstructured import get_unstructured_api_key -from onyx.file_processing.unstructured import unstructured_to_text +from onyx.file_processing.unstructured import ( + get_unstructured_api_key, + unstructured_to_text, +) from onyx.utils.logger import setup_logger -from onyx.utils.process_isolation import IsolatedProcessError -from onyx.utils.process_isolation import run_in_isolated_process +from onyx.utils.process_isolation import IsolatedProcessError, run_in_isolated_process if TYPE_CHECKING: from markitdown import MarkItDown @@ -488,9 +486,11 @@ def read_docx_file( The images list returned is empty in this case. """ md = get_markitdown_converter() - from markitdown import FileConversionException - from markitdown import StreamInfo - from markitdown import UnsupportedFormatException + from markitdown import ( + FileConversionException, + StreamInfo, + UnsupportedFormatException, + ) try: doc = md.convert( @@ -546,9 +546,11 @@ def extract_pptx_images(pptx_bytes: IO[Any]) -> Iterator[tuple[bytes, str]]: def pptx_to_text(file: IO[Any], file_name: str = "") -> str: md = get_markitdown_converter() - from markitdown import FileConversionException - from markitdown import StreamInfo - from markitdown import UnsupportedFormatException + from markitdown import ( + FileConversionException, + StreamInfo, + UnsupportedFormatException, + ) stream_info = StreamInfo( mimetype=PRESENTATION_MIME_TYPE, filename=file_name or None, extension=".pptx" diff --git a/backend/onyx/file_processing/html_utils.py b/backend/onyx/file_processing/html_utils.py index 0371df4c24..57c7d9d106 100644 --- a/backend/onyx/file_processing/html_utils.py +++ b/backend/onyx/file_processing/html_utils.py @@ -6,10 +6,12 @@ from typing import IO import bs4 -from onyx.configs.app_configs import HTML_BASED_CONNECTOR_TRANSFORM_LINKS_STRATEGY -from onyx.configs.app_configs import PARSE_WITH_TRAFILATURA -from onyx.configs.app_configs import WEB_CONNECTOR_IGNORED_CLASSES -from onyx.configs.app_configs import WEB_CONNECTOR_IGNORED_ELEMENTS +from onyx.configs.app_configs import ( + HTML_BASED_CONNECTOR_TRANSFORM_LINKS_STRATEGY, + PARSE_WITH_TRAFILATURA, + WEB_CONNECTOR_IGNORED_CLASSES, + WEB_CONNECTOR_IGNORED_ELEMENTS, +) from onyx.file_processing.enums import HtmlBasedConnectorTransformLinksStrategy from onyx.utils.logger import setup_logger diff --git a/backend/onyx/file_processing/image_summarization.py b/backend/onyx/file_processing/image_summarization.py index 129aecd49e..162e5737a9 100644 --- a/backend/onyx/file_processing/image_summarization.py +++ b/backend/onyx/file_processing/image_summarization.py @@ -3,21 +3,24 @@ from io import BytesIO from PIL import Image -from onyx.configs.app_configs import IMAGE_SUMMARIZATION_SYSTEM_PROMPT -from onyx.configs.app_configs import IMAGE_SUMMARIZATION_USER_PROMPT +from onyx.configs.app_configs import ( + IMAGE_SUMMARIZATION_SYSTEM_PROMPT, + IMAGE_SUMMARIZATION_USER_PROMPT, +) from onyx.llm.interfaces import LLM -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import ContentPart -from onyx.llm.models import ImageContentPart -from onyx.llm.models import ImageUrlDetail -from onyx.llm.models import SystemMessage -from onyx.llm.models import TextContentPart -from onyx.llm.models import UserMessage +from onyx.llm.models import ( + ChatCompletionMessage, + ContentPart, + ImageContentPart, + ImageUrlDetail, + SystemMessage, + TextContentPart, + UserMessage, +) from onyx.llm.utils import llm_response_to_string from onyx.server.metrics.image_processing import track_image_summarization from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.b64 import get_image_type_from_bytes from onyx.utils.logger import setup_logger diff --git a/backend/onyx/file_processing/image_utils.py b/backend/onyx/file_processing/image_utils.py index c31331b737..478c2b95a1 100644 --- a/backend/onyx/file_processing/image_utils.py +++ b/backend/onyx/file_processing/image_utils.py @@ -3,9 +3,7 @@ from io import BytesIO from typing import Tuple from onyx.configs.constants import FileOrigin -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import ImageSection, TabularSection, TextSection from onyx.file_processing.file_types import OnyxMimeTypes from onyx.file_store.file_store import get_default_file_store from onyx.utils.b64 import get_image_type_from_bytes diff --git a/backend/onyx/file_processing/password_validation.py b/backend/onyx/file_processing/password_validation.py index cea56afff3..94e8b2995a 100644 --- a/backend/onyx/file_processing/password_validation.py +++ b/backend/onyx/file_processing/password_validation.py @@ -1,8 +1,6 @@ -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager -from typing import Any -from typing import IO +from typing import Any, IO from onyx.file_processing.extract_file_text import get_file_ext from onyx.utils.logger import setup_logger diff --git a/backend/onyx/file_processing/unstructured.py b/backend/onyx/file_processing/unstructured.py index 84f7db3e16..bb91ec5e21 100644 --- a/backend/onyx/file_processing/unstructured.py +++ b/backend/onyx/file_processing/unstructured.py @@ -1,7 +1,4 @@ -from typing import Any -from typing import cast -from typing import IO -from typing import TYPE_CHECKING +from typing import Any, cast, IO, TYPE_CHECKING from onyx.configs.constants import KV_UNSTRUCTURED_API_KEY from onyx.key_value_store.factory import get_kv_store @@ -36,8 +33,7 @@ def delete_unstructured_api_key() -> None: def _sdk_partition_request( file: IO[Any], file_name: str, **kwargs: Any ) -> "operations.PartitionRequest": - from unstructured_client.models import operations - from unstructured_client.models import shared + from unstructured_client.models import operations, shared file.seek(0, 0) try: diff --git a/backend/onyx/file_store/document_batch_storage.py b/backend/onyx/file_store/document_batch_storage.py index 38d3d1898c..1417dfa1de 100644 --- a/backend/onyx/file_store/document_batch_storage.py +++ b/backend/onyx/file_store/document_batch_storage.py @@ -1,20 +1,14 @@ import json -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from enum import Enum from io import StringIO -from typing import List -from typing import Optional -from typing import TypeAlias +from typing import List, Optional, TypeAlias from pydantic import BaseModel from onyx.configs.constants import FileOrigin -from onyx.connectors.models import DocExtractionContext -from onyx.connectors.models import DocIndexingContext -from onyx.connectors.models import Document -from onyx.file_store.file_store import FileStore -from onyx.file_store.file_store import get_default_file_store +from onyx.connectors.models import DocExtractionContext, DocIndexingContext, Document +from onyx.file_store.file_store import FileStore, get_default_file_store from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/file_store/file_store.py b/backend/onyx/file_store/file_store.py index e5f8954a4d..03270c3dc9 100644 --- a/backend/onyx/file_store/file_store.py +++ b/backend/onyx/file_store/file_store.py @@ -1,15 +1,9 @@ import hashlib import tempfile import uuid -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from io import BytesIO -from typing import Any -from typing import cast -from typing import IO -from typing import NotRequired -from typing import TYPE_CHECKING -from typing import TypedDict +from typing import Any, cast, IO, NotRequired, TYPE_CHECKING, TypedDict import boto3 import puremagic @@ -17,22 +11,28 @@ from botocore.config import Config from botocore.exceptions import ClientError from sqlalchemy.orm import Session -from onyx.configs.app_configs import AWS_REGION_NAME -from onyx.configs.app_configs import S3_AWS_ACCESS_KEY_ID -from onyx.configs.app_configs import S3_AWS_SECRET_ACCESS_KEY -from onyx.configs.app_configs import S3_ENDPOINT_URL -from onyx.configs.app_configs import S3_FILE_STORE_BUCKET_NAME -from onyx.configs.app_configs import S3_FILE_STORE_PREFIX -from onyx.configs.app_configs import S3_GENERATE_LOCAL_CHECKSUM -from onyx.configs.app_configs import S3_VERIFY_SSL +from onyx.configs.app_configs import ( + AWS_REGION_NAME, + S3_AWS_ACCESS_KEY_ID, + S3_AWS_SECRET_ACCESS_KEY, + S3_ENDPOINT_URL, + S3_FILE_STORE_BUCKET_NAME, + S3_FILE_STORE_PREFIX, + S3_GENERATE_LOCAL_CHECKSUM, + S3_VERIFY_SSL, +) from onyx.configs.constants import FileOrigin -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none -from onyx.db.file_record import delete_filerecord_by_file_id -from onyx.db.file_record import get_filerecord_by_file_id -from onyx.db.file_record import get_filerecord_by_file_id_optional -from onyx.db.file_record import get_filerecord_by_prefix -from onyx.db.file_record import upsert_filerecord +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_current_tenant_if_none, +) +from onyx.db.file_record import ( + delete_filerecord_by_file_id, + get_filerecord_by_file_id, + get_filerecord_by_file_id_optional, + get_filerecord_by_prefix, + upsert_filerecord, +) from onyx.db.models import FileRecord from onyx.db.models import FileRecord as FileStoreModel from onyx.file_store.s3_key_utils import generate_s3_key @@ -608,11 +608,13 @@ def get_s3_file_store() -> S3BackedFileStore: def get_gcs_file_store() -> "GCSBackedFileStore": """Returns the GCS file store implementation.""" - from onyx.configs.app_configs import GCS_FILE_STORE_BUCKET_NAME - from onyx.configs.app_configs import GCS_FILE_STORE_PREFIX - from onyx.configs.app_configs import GCS_PROJECT_ID - from onyx.configs.app_configs import GCS_SERVICE_ACCOUNT_KEY_JSON - from onyx.configs.app_configs import GCS_SERVICE_ACCOUNT_KEY_PATH + from onyx.configs.app_configs import ( + GCS_FILE_STORE_BUCKET_NAME, + GCS_FILE_STORE_PREFIX, + GCS_PROJECT_ID, + GCS_SERVICE_ACCOUNT_KEY_JSON, + GCS_SERVICE_ACCOUNT_KEY_PATH, + ) from onyx.file_store.gcs_file_store import GCSBackedFileStore bucket_name = GCS_FILE_STORE_BUCKET_NAME diff --git a/backend/onyx/file_store/gcs_file_store.py b/backend/onyx/file_store/gcs_file_store.py index 87b58ea627..6e57a167ef 100644 --- a/backend/onyx/file_store/gcs_file_store.py +++ b/backend/onyx/file_store/gcs_file_store.py @@ -4,10 +4,7 @@ import json import tempfile import uuid from io import BytesIO -from typing import Any -from typing import cast -from typing import IO -from typing import TYPE_CHECKING +from typing import Any, cast, IO, TYPE_CHECKING import puremagic from sqlalchemy.orm import Session @@ -16,13 +13,17 @@ if TYPE_CHECKING: from google.cloud.storage import Client as GCSClient from onyx.configs.constants import FileOrigin -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none -from onyx.db.file_record import delete_filerecord_by_file_id -from onyx.db.file_record import get_filerecord_by_file_id -from onyx.db.file_record import get_filerecord_by_file_id_optional -from onyx.db.file_record import get_filerecord_by_prefix -from onyx.db.file_record import upsert_filerecord +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_current_tenant_if_none, +) +from onyx.db.file_record import ( + delete_filerecord_by_file_id, + get_filerecord_by_file_id, + get_filerecord_by_file_id_optional, + get_filerecord_by_prefix, + upsert_filerecord, +) from onyx.db.models import FileRecord from onyx.file_store.file_store import FileStore from onyx.file_store.s3_key_utils import generate_s3_key @@ -124,8 +125,7 @@ class GCSBackedFileStore(FileStore): def initialize(self) -> None: """Initialize the GCS file store by ensuring the bucket exists.""" - from google.api_core.exceptions import Forbidden - from google.api_core.exceptions import NotFound + from google.api_core.exceptions import Forbidden, NotFound client = self._get_gcs_client() try: diff --git a/backend/onyx/file_store/models.py b/backend/onyx/file_store/models.py index 23c5685269..0b3f59bba6 100644 --- a/backend/onyx/file_store/models.py +++ b/backend/onyx/file_store/models.py @@ -1,8 +1,7 @@ import base64 import threading from enum import Enum -from typing import Callable -from typing import NotRequired +from typing import Callable, NotRequired from pydantic import BaseModel from typing_extensions import TypedDict # noreorder diff --git a/backend/onyx/file_store/postgres_file_store.py b/backend/onyx/file_store/postgres_file_store.py index d816494123..e29c28c98c 100644 --- a/backend/onyx/file_store/postgres_file_store.py +++ b/backend/onyx/file_store/postgres_file_store.py @@ -7,27 +7,31 @@ eliminating the need for an external S3/MinIO service. import tempfile import uuid from io import BytesIO -from typing import Any -from typing import cast -from typing import IO +from typing import Any, cast, IO import puremagic from psycopg2.extensions import connection as Psycopg2Connection from sqlalchemy.orm import Session from onyx.configs.constants import FileOrigin -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none -from onyx.db.file_content import delete_file_content_by_file_id -from onyx.db.file_content import get_file_content_by_file_id -from onyx.db.file_content import get_file_content_by_file_id_optional -from onyx.db.file_content import transfer_file_content_file_id -from onyx.db.file_content import upsert_file_content -from onyx.db.file_record import delete_filerecord_by_file_id -from onyx.db.file_record import get_filerecord_by_file_id -from onyx.db.file_record import get_filerecord_by_file_id_optional -from onyx.db.file_record import get_filerecord_by_prefix -from onyx.db.file_record import upsert_filerecord +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_current_tenant_if_none, +) +from onyx.db.file_content import ( + delete_file_content_by_file_id, + get_file_content_by_file_id, + get_file_content_by_file_id_optional, + transfer_file_content_file_id, + upsert_file_content, +) +from onyx.db.file_record import ( + delete_filerecord_by_file_id, + get_filerecord_by_file_id, + get_filerecord_by_file_id_optional, + get_filerecord_by_prefix, + upsert_filerecord, +) from onyx.db.models import FileRecord from onyx.db.models import FileRecord as FileStoreModel from onyx.file_store.file_store import FileStore diff --git a/backend/onyx/file_store/staging.py b/backend/onyx/file_store/staging.py index a154baebe3..6df8d9e872 100644 --- a/backend/onyx/file_store/staging.py +++ b/backend/onyx/file_store/staging.py @@ -1,13 +1,14 @@ from collections.abc import Callable -from typing import Any -from typing import IO +from typing import Any, IO from sqlalchemy.orm import Session from onyx.configs.constants import FileOrigin -from onyx.db.file_record import get_staged_file_ids_by_index_attempt_id -from onyx.db.file_record import get_staged_file_ids_for_cc_pair_excluding_attempt -from onyx.db.file_record import update_filerecord_origin +from onyx.db.file_record import ( + get_staged_file_ids_by_index_attempt_id, + get_staged_file_ids_for_cc_pair_excluding_attempt, + update_filerecord_origin, +) from onyx.file_store.file_store import get_default_file_store from onyx.utils.logger import setup_logger diff --git a/backend/onyx/file_store/utils.py b/backend/onyx/file_store/utils.py index 77c846ebe1..282e4453ed 100644 --- a/backend/onyx/file_store/utils.py +++ b/backend/onyx/file_store/utils.py @@ -11,9 +11,7 @@ from onyx.configs.app_configs import WEB_DOMAIN from onyx.configs.constants import FileOrigin from onyx.db.models import UserFile from onyx.file_store.file_store import get_default_file_store -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import FileDescriptor -from onyx.file_store.models import InMemoryChatFile +from onyx.file_store.models import ChatFileType, FileDescriptor, InMemoryChatFile from onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type from onyx.utils.b64 import get_image_type from onyx.utils.logger import setup_logger diff --git a/backend/onyx/hooks/executor.py b/backend/onyx/hooks/executor.py index 014ec185ff..2aa6a7b642 100644 --- a/backend/onyx/hooks/executor.py +++ b/backend/onyx/hooks/executor.py @@ -9,8 +9,7 @@ via fetch_versioned_implementation so that: - EE: ee.onyx.hooks.executor._execute_hook_impl → real HTTP call """ -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar from pydantic import BaseModel from sqlalchemy.orm import Session diff --git a/backend/onyx/hooks/models.py b/backend/onyx/hooks/models.py index ceb3d0f3e8..f54e0859f2 100644 --- a/backend/onyx/hooks/models.py +++ b/backend/onyx/hooks/models.py @@ -1,16 +1,10 @@ from datetime import datetime from enum import Enum -from typing import Annotated -from typing import Any +from typing import Annotated, Any -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator -from pydantic import model_validator -from pydantic import SecretStr +from pydantic import BaseModel, Field, field_validator, model_validator, SecretStr -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint +from onyx.db.enums import HookFailStrategy, HookPoint NonEmptySecretStr = Annotated[SecretStr, Field(min_length=1)] diff --git a/backend/onyx/hooks/points/base.py b/backend/onyx/hooks/points/base.py index 9be41a73e2..d89af71c23 100644 --- a/backend/onyx/hooks/points/base.py +++ b/backend/onyx/hooks/points/base.py @@ -1,10 +1,8 @@ -from typing import Any -from typing import ClassVar +from typing import Any, ClassVar from pydantic import BaseModel -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint +from onyx.db.enums import HookFailStrategy, HookPoint _REQUIRED_ATTRS = ( "hook_point", diff --git a/backend/onyx/hooks/points/document_ingestion.py b/backend/onyx/hooks/points/document_ingestion.py index 654b5dbb62..9d7e61dd90 100644 --- a/backend/onyx/hooks/points/document_ingestion.py +++ b/backend/onyx/hooks/points/document_ingestion.py @@ -1,8 +1,6 @@ -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint +from onyx.db.enums import HookFailStrategy, HookPoint from onyx.hooks.points.base import HookPointSpec diff --git a/backend/onyx/hooks/points/document_push.py b/backend/onyx/hooks/points/document_push.py index f385cd1ced..12bff7c5c1 100644 --- a/backend/onyx/hooks/points/document_push.py +++ b/backend/onyx/hooks/points/document_push.py @@ -1,8 +1,6 @@ -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint +from onyx.db.enums import HookFailStrategy, HookPoint from onyx.hooks.points.base import HookPointSpec -from onyx.indexing.document_push import DocumentPushPayload -from onyx.indexing.document_push import DocumentPushResponse +from onyx.indexing.document_push import DocumentPushPayload, DocumentPushResponse __all__ = ["DocumentPushPayload", "DocumentPushResponse", "DocumentPushSpec"] diff --git a/backend/onyx/hooks/points/query_processing.py b/backend/onyx/hooks/points/query_processing.py index 090e7f4d96..69b9215e49 100644 --- a/backend/onyx/hooks/points/query_processing.py +++ b/backend/onyx/hooks/points/query_processing.py @@ -1,9 +1,6 @@ -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint +from onyx.db.enums import HookFailStrategy, HookPoint from onyx.hooks.points.base import HookPointSpec diff --git a/backend/onyx/image_gen/factory.py b/backend/onyx/image_gen/factory.py index e4a8926d34..7668669c4c 100644 --- a/backend/onyx/image_gen/factory.py +++ b/backend/onyx/image_gen/factory.py @@ -1,7 +1,9 @@ from enum import Enum -from onyx.image_gen.interfaces import ImageGenerationProvider -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials +from onyx.image_gen.interfaces import ( + ImageGenerationProvider, + ImageGenerationProviderCredentials, +) from onyx.image_gen.providers.azure_img_gen import AzureImageGenerationProvider from onyx.image_gen.providers.openai_img_gen import OpenAIImageGenerationProvider from onyx.image_gen.providers.vertex_img_gen import VertexImageGenerationProvider diff --git a/backend/onyx/image_gen/generation.py b/backend/onyx/image_gen/generation.py index a5729b37f9..7ed94fdeee 100644 --- a/backend/onyx/image_gen/generation.py +++ b/backend/onyx/image_gen/generation.py @@ -4,12 +4,13 @@ from sqlalchemy.orm import Session from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.image_generation import get_default_image_generation_config from onyx.image_gen.exceptions import ImageGenerationNotConfiguredError -from onyx.image_gen.factory import get_image_generation_provider -from onyx.image_gen.factory import validate_credentials -from onyx.image_gen.interfaces import ImageGenerationProvider -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.image_gen.interfaces import ImageShape -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.factory import get_image_generation_provider, validate_credentials +from onyx.image_gen.interfaces import ( + ImageGenerationProvider, + ImageGenerationProviderCredentials, + ImageShape, + ReferenceImage, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/image_gen/interfaces.py b/backend/onyx/image_gen/interfaces.py index 8897d8e262..1a0e483dca 100644 --- a/backend/onyx/image_gen/interfaces.py +++ b/backend/onyx/image_gen/interfaces.py @@ -2,8 +2,7 @@ from __future__ import annotations import abc from enum import Enum -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING from pydantic import BaseModel diff --git a/backend/onyx/image_gen/providers/azure_img_gen.py b/backend/onyx/image_gen/providers/azure_img_gen.py index ef752e755c..91aeddf030 100644 --- a/backend/onyx/image_gen/providers/azure_img_gen.py +++ b/backend/onyx/image_gen/providers/azure_img_gen.py @@ -1,11 +1,12 @@ from __future__ import annotations -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING -from onyx.image_gen.interfaces import ImageGenerationProvider -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.interfaces import ( + ImageGenerationProvider, + ImageGenerationProviderCredentials, + ReferenceImage, +) from onyx.tracing.flows import LLMFlow from onyx.tracing.llm_utils import traced_llm_call diff --git a/backend/onyx/image_gen/providers/openai_img_gen.py b/backend/onyx/image_gen/providers/openai_img_gen.py index 97487675f8..89a6b4ba11 100644 --- a/backend/onyx/image_gen/providers/openai_img_gen.py +++ b/backend/onyx/image_gen/providers/openai_img_gen.py @@ -1,11 +1,12 @@ from __future__ import annotations -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING -from onyx.image_gen.interfaces import ImageGenerationProvider -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.interfaces import ( + ImageGenerationProvider, + ImageGenerationProviderCredentials, + ReferenceImage, +) from onyx.tracing.flows import LLMFlow from onyx.tracing.llm_utils import traced_llm_call diff --git a/backend/onyx/image_gen/providers/vertex_img_gen.py b/backend/onyx/image_gen/providers/vertex_img_gen.py index e58739bb66..04a1b1672c 100644 --- a/backend/onyx/image_gen/providers/vertex_img_gen.py +++ b/backend/onyx/image_gen/providers/vertex_img_gen.py @@ -3,21 +3,24 @@ from __future__ import annotations import base64 import json from datetime import datetime -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING from pydantic import BaseModel from onyx.image_gen.exceptions import ImageProviderCredentialsError -from onyx.image_gen.interfaces import ImageGenerationProvider -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.image_gen.interfaces import ReferenceImage -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_SERVICE_ACCOUNT -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_WORKLOAD_IDENTITY -from onyx.llm.well_known_providers.constants import VERTEX_CREDENTIALS_FILE_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_LOCATION_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_PROJECT_KWARG +from onyx.image_gen.interfaces import ( + ImageGenerationProvider, + ImageGenerationProviderCredentials, + ReferenceImage, +) +from onyx.llm.well_known_providers.constants import ( + VERTEX_AUTH_METHOD_KWARG, + VERTEX_AUTH_METHOD_SERVICE_ACCOUNT, + VERTEX_AUTH_METHOD_WORKLOAD_IDENTITY, + VERTEX_CREDENTIALS_FILE_KWARG, + VERTEX_LOCATION_KWARG, + VERTEX_PROJECT_KWARG, +) from onyx.tracing.flows import LLMFlow from onyx.tracing.llm_utils import traced_llm_call @@ -129,8 +132,7 @@ class VertexImageGenerationProvider(ImageGenerationProvider): ) -> ImageGenerationResponse: from google import genai from google.genai import types as genai_types - from litellm.types.utils import ImageObject - from litellm.types.utils import ImageResponse + from litellm.types.utils import ImageObject, ImageResponse credentials: Any if self._use_workload_identity: diff --git a/backend/onyx/indexing/adapters/document_indexing_adapter.py b/backend/onyx/indexing/adapters/document_indexing_adapter.py index 5e6f17083c..fa6fbbac7f 100644 --- a/backend/onyx/indexing/adapters/document_indexing_adapter.py +++ b/backend/onyx/indexing/adapters/document_indexing_adapter.py @@ -7,28 +7,33 @@ from sqlalchemy.orm import Session from onyx.access.access import get_access_for_documents from onyx.access.models import DocumentAccess from onyx.configs.constants import DEFAULT_BOOST -from onyx.connectors.models import Document -from onyx.connectors.models import IndexAttemptMetadata +from onyx.connectors.models import Document, IndexAttemptMetadata from onyx.db.chunk import update_chunk_boost_components__no_commit -from onyx.db.document import fetch_chunk_counts_for_documents -from onyx.db.document import mark_document_as_indexed_for_cc_pair__no_commit -from onyx.db.document import prepare_to_modify_documents -from onyx.db.document import update_docs_chunk_count__no_commit -from onyx.db.document import update_docs_created_at__no_commit -from onyx.db.document import update_docs_last_modified__no_commit -from onyx.db.document import update_docs_updated_at__no_commit +from onyx.db.document import ( + fetch_chunk_counts_for_documents, + mark_document_as_indexed_for_cc_pair__no_commit, + prepare_to_modify_documents, + update_docs_chunk_count__no_commit, + update_docs_created_at__no_commit, + update_docs_last_modified__no_commit, + update_docs_updated_at__no_commit, +) from onyx.db.document_set import fetch_document_sets_for_documents from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.index_attempt_metrics import safe_record_single_event_if_set from onyx.db.index_attempt_metrics_models import IndexAttemptStage -from onyx.indexing.indexing_pipeline import DocumentBatchPrepareContext -from onyx.indexing.indexing_pipeline import index_doc_batch_prepare -from onyx.indexing.models import ChunkEnrichmentContext -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import IndexChunk -from onyx.indexing.models import IndexingBatchAdapter -from onyx.indexing.models import UpdatableChunkData +from onyx.indexing.indexing_pipeline import ( + DocumentBatchPrepareContext, + index_doc_batch_prepare, +) +from onyx.indexing.models import ( + ChunkEnrichmentContext, + DocAwareChunk, + DocMetadataAwareIndexChunk, + IndexChunk, + IndexingBatchAdapter, + UpdatableChunkData, +) from onyx.redis.redis_hierarchy import get_ancestors_from_raw_id from onyx.redis.redis_pool import get_redis_client from onyx.utils.logger import setup_logger diff --git a/backend/onyx/indexing/adapters/user_file_indexing_adapter.py b/backend/onyx/indexing/adapters/user_file_indexing_adapter.py index 912fa84b0c..23469a26f6 100644 --- a/backend/onyx/indexing/adapters/user_file_indexing_adapter.py +++ b/backend/onyx/indexing/adapters/user_file_indexing_adapter.py @@ -9,31 +9,31 @@ from uuid import UUID from sqlalchemy import select from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session from onyx.access.access import get_access_for_user_files from onyx.access.models import DocumentAccess -from onyx.configs.constants import DEFAULT_BOOST -from onyx.configs.constants import NotificationType +from onyx.configs.constants import DEFAULT_BOOST, NotificationType from onyx.connectors.models import Document from onyx.db.enums import UserFileStatus -from onyx.db.models import Persona -from onyx.db.models import UserFile +from onyx.db.models import Persona, UserFile from onyx.db.notification import create_notification -from onyx.db.user_file import fetch_chunk_counts_for_user_files -from onyx.db.user_file import fetch_persona_ids_for_user_files -from onyx.db.user_file import fetch_user_project_ids_for_user_files +from onyx.db.user_file import ( + fetch_chunk_counts_for_user_files, + fetch_persona_ids_for_user_files, + fetch_user_project_ids_for_user_files, +) from onyx.file_store.utils import store_user_file_plaintext from onyx.indexing.indexing_pipeline import DocumentBatchPrepareContext -from onyx.indexing.models import ChunkEnrichmentContext -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import IndexChunk -from onyx.indexing.models import UpdatableChunkData +from onyx.indexing.models import ( + ChunkEnrichmentContext, + DocAwareChunk, + DocMetadataAwareIndexChunk, + IndexChunk, + UpdatableChunkData, +) from onyx.llm.factory import get_default_llm -from onyx.natural_language_processing.utils import count_tokens -from onyx.natural_language_processing.utils import get_tokenizer +from onyx.natural_language_processing.utils import count_tokens, get_tokenizer from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/indexing/chunker.py b/backend/onyx/indexing/chunker.py index bd56a1e6d7..2876c8e022 100644 --- a/backend/onyx/indexing/chunker.py +++ b/backend/onyx/indexing/chunker.py @@ -1,21 +1,20 @@ from chonkie import SentenceChunker -from onyx.configs.app_configs import AVERAGE_SUMMARY_EMBEDDINGS -from onyx.configs.app_configs import BLURB_SIZE -from onyx.configs.app_configs import LARGE_CHUNK_RATIO -from onyx.configs.app_configs import MINI_CHUNK_SIZE -from onyx.configs.app_configs import SKIP_METADATA_IN_CHUNK -from onyx.configs.app_configs import USE_CHUNK_SUMMARY -from onyx.configs.app_configs import USE_DOCUMENT_SUMMARY -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import RETURN_SEPARATOR -from onyx.configs.constants import SECTION_SEPARATOR +from onyx.configs.app_configs import ( + AVERAGE_SUMMARY_EMBEDDINGS, + BLURB_SIZE, + LARGE_CHUNK_RATIO, + MINI_CHUNK_SIZE, + SKIP_METADATA_IN_CHUNK, + USE_CHUNK_SUMMARY, + USE_DOCUMENT_SUMMARY, +) +from onyx.configs.constants import DocumentSource, RETURN_SEPARATOR, SECTION_SEPARATOR from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( get_metadata_keys_to_ignore, ) from onyx.connectors.models import IndexingDocument -from onyx.indexing.chunking import DocumentChunker -from onyx.indexing.chunking import extract_blurb +from onyx.indexing.chunking import DocumentChunker, extract_blurb from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.indexing.models import DocAwareChunk from onyx.llm.utils import MAX_CONTEXT_TOKENS diff --git a/backend/onyx/indexing/chunking/document_chunker.py b/backend/onyx/indexing/chunking/document_chunker.py index e6d8fc3c29..324575857d 100644 --- a/backend/onyx/indexing/chunking/document_chunker.py +++ b/backend/onyx/indexing/chunking/document_chunker.py @@ -1,13 +1,17 @@ from chonkie import SentenceChunker -from onyx.connectors.models import IndexingDocument -from onyx.connectors.models import Section -from onyx.connectors.models import SectionType -from onyx.connectors.models import TabularSection +from onyx.connectors.models import ( + IndexingDocument, + Section, + SectionType, + TabularSection, +) from onyx.indexing.chunking.image_section_chunker import ImageChunker -from onyx.indexing.chunking.section_chunker import AccumulatorState -from onyx.indexing.chunking.section_chunker import ChunkPayload -from onyx.indexing.chunking.section_chunker import SectionChunker +from onyx.indexing.chunking.section_chunker import ( + AccumulatorState, + ChunkPayload, + SectionChunker, +) from onyx.indexing.chunking.tabular_section_chunker import TabularChunker from onyx.indexing.chunking.text_section_chunker import TextChunker from onyx.indexing.models import DocAwareChunk diff --git a/backend/onyx/indexing/chunking/image_section_chunker.py b/backend/onyx/indexing/chunking/image_section_chunker.py index e3f32459d6..34d6991fac 100644 --- a/backend/onyx/indexing/chunking/image_section_chunker.py +++ b/backend/onyx/indexing/chunking/image_section_chunker.py @@ -1,8 +1,10 @@ from onyx.connectors.models import Section -from onyx.indexing.chunking.section_chunker import AccumulatorState -from onyx.indexing.chunking.section_chunker import ChunkPayload -from onyx.indexing.chunking.section_chunker import SectionChunker -from onyx.indexing.chunking.section_chunker import SectionChunkerOutput +from onyx.indexing.chunking.section_chunker import ( + AccumulatorState, + ChunkPayload, + SectionChunker, + SectionChunkerOutput, +) from onyx.utils.text_processing import clean_text diff --git a/backend/onyx/indexing/chunking/section_chunker.py b/backend/onyx/indexing/chunking/section_chunker.py index 310c73f077..b9ef9a5245 100644 --- a/backend/onyx/indexing/chunking/section_chunker.py +++ b/backend/onyx/indexing/chunking/section_chunker.py @@ -1,14 +1,11 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import Sequence from typing import cast from chonkie import SentenceChunker -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field -from onyx.connectors.models import IndexingDocument -from onyx.connectors.models import Section +from onyx.connectors.models import IndexingDocument, Section from onyx.indexing.models import DocAwareChunk diff --git a/backend/onyx/indexing/chunking/tabular_section_chunker/analysis.py b/backend/onyx/indexing/chunking/tabular_section_chunker/analysis.py index 106495118e..e7467d4ae0 100644 --- a/backend/onyx/indexing/chunking/tabular_section_chunker/analysis.py +++ b/backend/onyx/indexing/chunking/tabular_section_chunker/analysis.py @@ -3,9 +3,7 @@ from collections.abc import Iterable from datetime import date from dateutil.parser import parse as parse_dt -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field from onyx.utils.csv_utils import ParsedRow diff --git a/backend/onyx/indexing/chunking/tabular_section_chunker/sheet_descriptor.py b/backend/onyx/indexing/chunking/tabular_section_chunker/sheet_descriptor.py index 45741e80cc..191a586f95 100644 --- a/backend/onyx/indexing/chunking/tabular_section_chunker/sheet_descriptor.py +++ b/backend/onyx/indexing/chunking/tabular_section_chunker/sheet_descriptor.py @@ -1,6 +1,5 @@ from onyx.indexing.chunking.tabular_section_chunker.analysis import SheetAnalysis -from onyx.indexing.chunking.tabular_section_chunker.util import label -from onyx.indexing.chunking.tabular_section_chunker.util import pack_lines +from onyx.indexing.chunking.tabular_section_chunker.util import label, pack_lines from onyx.natural_language_processing.utils import BaseTokenizer MAX_NUMERIC_COLS = 12 diff --git a/backend/onyx/indexing/chunking/tabular_section_chunker/tabular_section_chunker.py b/backend/onyx/indexing/chunking/tabular_section_chunker/tabular_section_chunker.py index 4d2375c966..85d432772e 100644 --- a/backend/onyx/indexing/chunking/tabular_section_chunker/tabular_section_chunker.py +++ b/backend/onyx/indexing/chunking/tabular_section_chunker/tabular_section_chunker.py @@ -4,27 +4,30 @@ from itertools import chain from pydantic import BaseModel -from onyx.connectors.models import Section -from onyx.connectors.models import TabularSection +from onyx.connectors.models import Section, TabularSection from onyx.file_store.file_store import get_default_file_store -from onyx.indexing.chunking.section_chunker import AccumulatorState -from onyx.indexing.chunking.section_chunker import ChunkPayload -from onyx.indexing.chunking.section_chunker import SectionChunker -from onyx.indexing.chunking.section_chunker import SectionChunkerOutput -from onyx.indexing.chunking.tabular_section_chunker.analysis import analyze_sheet -from onyx.indexing.chunking.tabular_section_chunker.analysis import SheetAnalysis +from onyx.indexing.chunking.section_chunker import ( + AccumulatorState, + ChunkPayload, + SectionChunker, + SectionChunkerOutput, +) +from onyx.indexing.chunking.tabular_section_chunker.analysis import ( + analyze_sheet, + SheetAnalysis, +) from onyx.indexing.chunking.tabular_section_chunker.sheet_descriptor import ( build_sheet_descriptor_chunks, ) from onyx.indexing.chunking.tabular_section_chunker.total_descriptor import ( build_total_descriptor_chunks, ) -from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.natural_language_processing.utils import count_tokens -from onyx.natural_language_processing.utils import split_text_by_tokens -from onyx.utils.csv_utils import parse_csv_stream -from onyx.utils.csv_utils import ParsedRow -from onyx.utils.csv_utils import read_csv_header +from onyx.natural_language_processing.utils import ( + BaseTokenizer, + count_tokens, + split_text_by_tokens, +) +from onyx.utils.csv_utils import parse_csv_stream, ParsedRow, read_csv_header from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/indexing/chunking/tabular_section_chunker/total_descriptor.py b/backend/onyx/indexing/chunking/tabular_section_chunker/total_descriptor.py index 57105b9d40..373f1f8620 100644 --- a/backend/onyx/indexing/chunking/tabular_section_chunker/total_descriptor.py +++ b/backend/onyx/indexing/chunking/tabular_section_chunker/total_descriptor.py @@ -1,9 +1,10 @@ from collections import Counter -from onyx.indexing.chunking.tabular_section_chunker.analysis import NumericAggregate -from onyx.indexing.chunking.tabular_section_chunker.analysis import SheetAnalysis -from onyx.indexing.chunking.tabular_section_chunker.util import label -from onyx.indexing.chunking.tabular_section_chunker.util import pack_lines +from onyx.indexing.chunking.tabular_section_chunker.analysis import ( + NumericAggregate, + SheetAnalysis, +) +from onyx.indexing.chunking.tabular_section_chunker.util import label, pack_lines from onyx.natural_language_processing.utils import BaseTokenizer TOTALS_HEADER = ( diff --git a/backend/onyx/indexing/chunking/tabular_section_chunker/util.py b/backend/onyx/indexing/chunking/tabular_section_chunker/util.py index a64690c092..780816230e 100644 --- a/backend/onyx/indexing/chunking/tabular_section_chunker/util.py +++ b/backend/onyx/indexing/chunking/tabular_section_chunker/util.py @@ -1,5 +1,4 @@ -from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.natural_language_processing.utils import count_tokens +from onyx.natural_language_processing.utils import BaseTokenizer, count_tokens def label(name: str) -> str: diff --git a/backend/onyx/indexing/chunking/text_section_chunker.py b/backend/onyx/indexing/chunking/text_section_chunker.py index 8abbc0f958..9aa37d6810 100644 --- a/backend/onyx/indexing/chunking/text_section_chunker.py +++ b/backend/onyx/indexing/chunking/text_section_chunker.py @@ -4,15 +4,18 @@ from chonkie import SentenceChunker from onyx.configs.constants import SECTION_SEPARATOR from onyx.connectors.models import Section -from onyx.indexing.chunking.section_chunker import AccumulatorState -from onyx.indexing.chunking.section_chunker import ChunkPayload -from onyx.indexing.chunking.section_chunker import SectionChunker -from onyx.indexing.chunking.section_chunker import SectionChunkerOutput -from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.natural_language_processing.utils import count_tokens -from onyx.natural_language_processing.utils import split_text_by_tokens -from onyx.utils.text_processing import clean_text -from onyx.utils.text_processing import shared_precompare_cleanup +from onyx.indexing.chunking.section_chunker import ( + AccumulatorState, + ChunkPayload, + SectionChunker, + SectionChunkerOutput, +) +from onyx.natural_language_processing.utils import ( + BaseTokenizer, + count_tokens, + split_text_by_tokens, +) +from onyx.utils.text_processing import clean_text, shared_precompare_cleanup from shared_configs.configs import STRICT_CHUNK_TOKEN_LIMIT diff --git a/backend/onyx/indexing/document_push.py b/backend/onyx/indexing/document_push.py index 5163572543..a2e4d26c28 100644 --- a/backend/onyx/indexing/document_push.py +++ b/backend/onyx/indexing/document_push.py @@ -25,17 +25,16 @@ Push failures never fail the indexing batch: failures are logged and swallowed. from functools import lru_cache -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field -from onyx.configs.app_configs import DOCUMENT_PUSH_API_KEY -from onyx.configs.app_configs import DOCUMENT_PUSH_ENDPOINT_URL -from onyx.configs.app_configs import DOCUMENT_PUSH_TIMEOUT_SECONDS -from onyx.utils.external_endpoint import ExternalEndpointConfig -from onyx.utils.external_endpoint import post_json_to_endpoint +from onyx.configs.app_configs import ( + DOCUMENT_PUSH_API_KEY, + DOCUMENT_PUSH_ENDPOINT_URL, + DOCUMENT_PUSH_TIMEOUT_SECONDS, +) +from onyx.utils.external_endpoint import ExternalEndpointConfig, post_json_to_endpoint from onyx.utils.logger import setup_logger -from onyx.utils.url import SSRFException -from onyx.utils.url import validate_outbound_http_url +from onyx.utils.url import SSRFException, validate_outbound_http_url logger = setup_logger() diff --git a/backend/onyx/indexing/embedder.py b/backend/onyx/indexing/embedder.py index 4a8c967aab..d65c74d426 100644 --- a/backend/onyx/indexing/embedder.py +++ b/backend/onyx/indexing/embedder.py @@ -1,29 +1,29 @@ import time -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections import defaultdict import sentry_sdk -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorStopSignal -from onyx.connectors.models import DocumentFailure +from onyx.connectors.models import ( + ConnectorFailure, + ConnectorStopSignal, + DocumentFailure, +) from onyx.db.models import SearchSettings from onyx.document_index.chunk_content_enrichment import ( generate_enriched_content_for_chunk_embedding, ) from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import IndexChunk +from onyx.indexing.models import ChunkEmbedding, DocAwareChunk, IndexChunk from onyx.natural_language_processing.search_nlp_models import EmbeddingModel from onyx.utils.logger import setup_logger from onyx.utils.pydantic_util import shallow_model_dump from onyx.utils.timing import log_function_time -from shared_configs.configs import INDEXING_MODEL_SERVER_HOST -from shared_configs.configs import INDEXING_MODEL_SERVER_PORT -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType +from shared_configs.configs import ( + INDEXING_MODEL_SERVER_HOST, + INDEXING_MODEL_SERVER_PORT, +) +from shared_configs.enums import EmbeddingProvider, EmbedTextType from shared_configs.model_server_models import Embedding logger = setup_logger() diff --git a/backend/onyx/indexing/indexing_heartbeat.py b/backend/onyx/indexing/indexing_heartbeat.py index 7adc76e771..581ce53c42 100644 --- a/backend/onyx/indexing/indexing_heartbeat.py +++ b/backend/onyx/indexing/indexing_heartbeat.py @@ -1,5 +1,4 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod class IndexingHeartbeatInterface(ABC): diff --git a/backend/onyx/indexing/indexing_pipeline.py b/backend/onyx/indexing/indexing_pipeline.py index cc5c4a3f91..6fd577e051 100644 --- a/backend/onyx/indexing/indexing_pipeline.py +++ b/backend/onyx/indexing/indexing_pipeline.py @@ -1,98 +1,109 @@ import time from collections import defaultdict -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Callable, Generator, Iterator from contextlib import contextmanager -from typing import NamedTuple -from typing import Protocol +from typing import NamedTuple, Protocol import sentry_sdk -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from sqlalchemy.orm import Session -from onyx.configs.app_configs import ENABLE_CONTEXTUAL_RAG -from onyx.configs.app_configs import MAX_CHUNKS_PER_DOC_BATCH -from onyx.configs.app_configs import MAX_DOCUMENT_CHARS -from onyx.configs.app_configs import MAX_TOKENS_FOR_FULL_INCLUSION -from onyx.configs.app_configs import USE_CHUNK_SUMMARY -from onyx.configs.app_configs import USE_DOCUMENT_SUMMARY +from onyx.configs.app_configs import ( + ENABLE_CONTEXTUAL_RAG, + MAX_CHUNKS_PER_DOC_BATCH, + MAX_DOCUMENT_CHARS, + MAX_TOKENS_FOR_FULL_INCLUSION, + USE_CHUNK_SUMMARY, + USE_DOCUMENT_SUMMARY, +) from onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( get_experts_stores_representations, ) -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorStopSignal -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import ImageSection -from onyx.connectors.models import IndexAttemptMetadata -from onyx.connectors.models import IndexingDocument -from onyx.connectors.models import Section -from onyx.connectors.models import SectionType -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + ConnectorFailure, + ConnectorStopSignal, + Document, + DocumentFailure, + ImageSection, + IndexAttemptMetadata, + IndexingDocument, + Section, + SectionType, + TextSection, +) from onyx.db.connector_credential_pair import get_connector_credential_pair -from onyx.db.document import get_documents_by_ids -from onyx.db.document import update_docs_content_hash__no_commit -from onyx.db.document import upsert_document_by_connector_credential_pair -from onyx.db.document import upsert_documents +from onyx.db.document import ( + get_documents_by_ids, + update_docs_content_hash__no_commit, + upsert_document_by_connector_credential_pair, + upsert_documents, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import HookPoint +from onyx.db.enums import AccessType, HookPoint from onyx.db.hierarchy import link_hierarchy_nodes_to_documents -from onyx.db.index_attempt_metrics import IndexAttemptStage -from onyx.db.index_attempt_metrics import safe_record_single_event_if_set -from onyx.db.index_attempt_metrics import time_stage_if_set +from onyx.db.index_attempt_metrics import ( + IndexAttemptStage, + safe_record_single_event_if_set, + time_stage_if_set, +) from onyx.db.models import Document as DBDocument from onyx.db.models import IndexModelStatus from onyx.db.search_settings import get_active_search_settings from onyx.db.tag import upsert_document_tags from onyx.document_index.document_index_utils import get_multipass_config from onyx.document_index.document_metadata import DocumentMetadata -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentInsertionRecord -from onyx.document_index.interfaces_new import IndexingMetadata +from onyx.document_index.interfaces_new import ( + DocumentIndex, + DocumentInsertionRecord, + IndexingMetadata, +) from onyx.file_processing.image_summarization import summarize_image_with_error_handling from onyx.file_store.file_store import get_default_file_store from onyx.file_store.staging import promote_staged_file -from onyx.hooks.executor import execute_hook -from onyx.hooks.executor import HookSkipped -from onyx.hooks.executor import HookSoftFailed -from onyx.hooks.points.document_ingestion import DocumentIngestionOwner -from onyx.hooks.points.document_ingestion import DocumentIngestionPayload -from onyx.hooks.points.document_ingestion import DocumentIngestionResponse -from onyx.hooks.points.document_ingestion import DocumentIngestionSection +from onyx.hooks.executor import execute_hook, HookSkipped, HookSoftFailed +from onyx.hooks.points.document_ingestion import ( + DocumentIngestionOwner, + DocumentIngestionPayload, + DocumentIngestionResponse, + DocumentIngestionSection, +) from onyx.indexing.chunk_batch_store import ChunkBatchStore from onyx.indexing.chunker import Chunker -from onyx.indexing.document_push import DocumentPushPayload -from onyx.indexing.document_push import DocumentPushResponse -from onyx.indexing.document_push import get_document_push_config -from onyx.indexing.document_push import push_document_via_config -from onyx.indexing.embedder import embed_chunks_with_failure_handling -from onyx.indexing.embedder import IndexingEmbedder -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import IndexingBatchAdapter -from onyx.indexing.models import UpdatableChunkData +from onyx.indexing.document_push import ( + DocumentPushPayload, + DocumentPushResponse, + get_document_push_config, + push_document_via_config, +) +from onyx.indexing.embedder import embed_chunks_with_failure_handling, IndexingEmbedder +from onyx.indexing.models import ( + DocAwareChunk, + DocMetadataAwareIndexChunk, + IndexingBatchAdapter, + UpdatableChunkData, +) from onyx.indexing.vector_db_insertion import write_chunks_to_vector_db_with_backoff -from onyx.llm.factory import get_contextual_rag_llm_for_search_settings -from onyx.llm.factory import get_default_llm_with_vision +from onyx.llm.factory import ( + get_contextual_rag_llm_for_search_settings, + get_default_llm_with_vision, +) from onyx.llm.interfaces import LLM -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import UserMessage +from onyx.llm.models import ReasoningEffort, UserMessage from onyx.llm.multi_llm import LLMRateLimitError -from onyx.llm.utils import llm_response_to_string -from onyx.llm.utils import MAX_CONTEXT_TOKENS -from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.natural_language_processing.utils import get_tokenizer -from onyx.natural_language_processing.utils import tokenizer_trim_middle -from onyx.prompts.contextual_retrieval import CONTEXTUAL_RAG_PROMPT1 -from onyx.prompts.contextual_retrieval import CONTEXTUAL_RAG_PROMPT2 -from onyx.prompts.contextual_retrieval import DOCUMENT_SUMMARY_PROMPT +from onyx.llm.utils import llm_response_to_string, MAX_CONTEXT_TOKENS +from onyx.natural_language_processing.utils import ( + BaseTokenizer, + get_tokenizer, + tokenizer_trim_middle, +) +from onyx.prompts.contextual_retrieval import ( + CONTEXTUAL_RAG_PROMPT1, + CONTEXTUAL_RAG_PROMPT2, + DOCUMENT_SUMMARY_PROMPT, +) from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.batching import batch_generator from onyx.utils.logger import setup_logger from onyx.utils.postgres_sanitization import sanitize_documents_for_postgres diff --git a/backend/onyx/indexing/models.py b/backend/onyx/indexing/models.py index 19e897f413..c7cd7175b5 100644 --- a/backend/onyx/indexing/models.py +++ b/backend/onyx/indexing/models.py @@ -1,16 +1,12 @@ import contextlib from collections.abc import Generator -from typing import Optional -from typing import Protocol -from typing import TYPE_CHECKING +from typing import Optional, Protocol, TYPE_CHECKING -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.access.models import DocumentAccess from onyx.connectors.models import Document -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import SwitchoverType +from onyx.db.enums import EmbeddingPrecision, SwitchoverType from onyx.utils.logger import setup_logger from onyx.utils.pydantic_util import shallow_model_dump from shared_configs.enums import EmbeddingProvider diff --git a/backend/onyx/indexing/persistent_indexing.py b/backend/onyx/indexing/persistent_indexing.py index e7fca800f3..4b1434fd86 100644 --- a/backend/onyx/indexing/persistent_indexing.py +++ b/backend/onyx/indexing/persistent_indexing.py @@ -10,10 +10,12 @@ ConnectorFailure objects with proper context for known error conditions. import sentry_sdk from onyx.configs.constants import DocumentSource -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + EntityFailure, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.index_attempt import create_index_attempt_error from onyx.utils.logger import setup_logger diff --git a/backend/onyx/indexing/port_reembed.py b/backend/onyx/indexing/port_reembed.py index 73f7d98447..5722e0c1d9 100644 --- a/backend/onyx/indexing/port_reembed.py +++ b/backend/onyx/indexing/port_reembed.py @@ -38,23 +38,27 @@ from collections import defaultdict from dataclasses import dataclass from typing import TYPE_CHECKING -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import RETURN_SEPARATOR -from onyx.connectors.models import convert_metadata_list_of_strings_to_dict -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.configs.constants import DocumentSource, RETURN_SEPARATOR +from onyx.connectors.models import ( + convert_metadata_list_of_strings_to_dict, + Document, + TextSection, +) from onyx.db.models import SearchSettings -from onyx.document_index.chunk_content_enrichment import cleanup_content_for_chunks from onyx.document_index.chunk_content_enrichment import ( + cleanup_content_for_chunks, generate_enriched_content_for_chunk_text, ) -from onyx.document_index.opensearch.schema import DocumentChunk -from onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors -from onyx.indexing.chunker import get_metadata_suffix_for_document_index -from onyx.indexing.chunker import MAX_METADATA_PERCENTAGE +from onyx.document_index.opensearch.schema import ( + DocumentChunk, + DocumentChunkWithoutVectors, +) +from onyx.indexing.chunker import ( + get_metadata_suffix_for_document_index, + MAX_METADATA_PERCENTAGE, +) from onyx.indexing.embedder import IndexingEmbedder -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import IndexChunk +from onyx.indexing.models import DocAwareChunk, IndexChunk if TYPE_CHECKING: from onyx.llm.interfaces import LLM diff --git a/backend/onyx/indexing/vector_db_insertion.py b/backend/onyx/indexing/vector_db_insertion.py index 96cfad1b90..af532cac6a 100644 --- a/backend/onyx/indexing/vector_db_insertion.py +++ b/backend/onyx/indexing/vector_db_insertion.py @@ -1,18 +1,17 @@ import time -from collections.abc import Callable -from collections.abc import Iterable +from collections.abc import Callable, Iterable from http import HTTPStatus -from itertools import chain -from itertools import groupby +from itertools import chain, groupby import httpx import sentry_sdk -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import DocumentFailure -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentInsertionRecord -from onyx.document_index.interfaces_new import IndexingMetadata +from onyx.connectors.models import ConnectorFailure, DocumentFailure +from onyx.document_index.interfaces_new import ( + DocumentIndex, + DocumentInsertionRecord, + IndexingMetadata, +) from onyx.indexing.models import DocMetadataAwareIndexChunk from onyx.utils.logger import setup_logger diff --git a/backend/onyx/key_value_store/store.py b/backend/onyx/key_value_store/store.py index 51ce0f761e..10c24db90e 100644 --- a/backend/onyx/key_value_store/store.py +++ b/backend/onyx/key_value_store/store.py @@ -4,8 +4,7 @@ from typing import cast from onyx.cache.interface import CacheBackend from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import KVStore -from onyx.key_value_store.interface import KeyValueStore -from onyx.key_value_store.interface import KvKeyNotFoundError +from onyx.key_value_store.interface import KeyValueStore, KvKeyNotFoundError from onyx.utils.logger import setup_logger from onyx.utils.special_types import JSON_ro diff --git a/backend/onyx/kg/clustering/clustering.py b/backend/onyx/kg/clustering/clustering.py index 87efa0fd1c..0b9041a1aa 100644 --- a/backend/onyx/kg/clustering/clustering.py +++ b/backend/onyx/kg/clustering/clustering.py @@ -4,31 +4,37 @@ from typing import cast from rapidfuzz.fuzz import ratio from redis.lock import Lock as RedisLock -from sqlalchemy import func -from sqlalchemy import text +from sqlalchemy import func, text from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.kg_configs import KG_CLUSTERING_RETRIEVE_THRESHOLD -from onyx.configs.kg_configs import KG_CLUSTERING_THRESHOLD +from onyx.configs.kg_configs import ( + KG_CLUSTERING_RETRIEVE_THRESHOLD, + KG_CLUSTERING_THRESHOLD, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.entities import KGEntity -from onyx.db.entities import KGEntityExtractionStaging -from onyx.db.entities import merge_entities -from onyx.db.entities import transfer_entity -from onyx.db.kg_config import get_kg_config_settings -from onyx.db.kg_config import validate_kg_settings -from onyx.db.models import Document -from onyx.db.models import KGEntityType -from onyx.db.models import KGRelationshipExtractionStaging -from onyx.db.models import KGRelationshipTypeExtractionStaging -from onyx.db.relationships import transfer_relationship -from onyx.db.relationships import transfer_relationship_type -from onyx.db.relationships import upsert_relationship -from onyx.db.relationships import upsert_relationship_type +from onyx.db.entities import ( + KGEntity, + KGEntityExtractionStaging, + merge_entities, + transfer_entity, +) +from onyx.db.kg_config import get_kg_config_settings, validate_kg_settings +from onyx.db.models import ( + Document, + KGEntityType, + KGRelationshipExtractionStaging, + KGRelationshipTypeExtractionStaging, +) +from onyx.db.relationships import ( + transfer_relationship, + transfer_relationship_type, + upsert_relationship, + upsert_relationship_type, +) from onyx.document_index.vespa.kg_interactions import ( get_kg_vespa_info_update_requests_for_document, + update_kg_chunks_vespa_info, ) -from onyx.document_index.vespa.kg_interactions import update_kg_chunks_vespa_info from onyx.kg.models import KGGroundingType from onyx.kg.utils.formatting_utils import make_relationship_id from onyx.kg.utils.lock_utils import extend_lock diff --git a/backend/onyx/kg/clustering/normalizations.py b/backend/onyx/kg/clustering/normalizations.py index a429bd207d..d145f3a156 100644 --- a/backend/onyx/kg/clustering/normalizations.py +++ b/backend/onyx/kg/clustering/normalizations.py @@ -4,32 +4,29 @@ from typing import cast import numpy as np from rapidfuzz.distance.DamerauLevenshtein import normalized_similarity -from sqlalchemy import desc -from sqlalchemy import Float -from sqlalchemy import func -from sqlalchemy import MetaData -from sqlalchemy import select -from sqlalchemy import String -from sqlalchemy import Table +from sqlalchemy import desc, Float, func, MetaData, select, String, Table from sqlalchemy.dialects.postgresql import ARRAY -from onyx.configs.kg_configs import KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT -from onyx.configs.kg_configs import KG_NORMALIZATION_RERANK_NGRAM_WEIGHTS -from onyx.configs.kg_configs import KG_NORMALIZATION_RERANK_THRESHOLD -from onyx.configs.kg_configs import KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT +from onyx.configs.kg_configs import ( + KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT, + KG_NORMALIZATION_RERANK_NGRAM_WEIGHTS, + KG_NORMALIZATION_RERANK_THRESHOLD, + KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import KGEntity from onyx.db.relationships import get_relationships_for_entity_type_pairs -from onyx.kg.models import NormalizedEntities -from onyx.kg.models import NormalizedRelationships +from onyx.kg.models import NormalizedEntities, NormalizedRelationships from onyx.kg.utils.embeddings import encode_string_batch -from onyx.kg.utils.formatting_utils import format_entity_id_for_models -from onyx.kg.utils.formatting_utils import get_attributes -from onyx.kg.utils.formatting_utils import get_entity_type -from onyx.kg.utils.formatting_utils import make_entity_w_attributes -from onyx.kg.utils.formatting_utils import make_relationship_id -from onyx.kg.utils.formatting_utils import split_entity_id -from onyx.kg.utils.formatting_utils import split_relationship_id +from onyx.kg.utils.formatting_utils import ( + format_entity_id_for_models, + get_attributes, + get_entity_type, + make_entity_w_attributes, + make_relationship_id, + split_entity_id, + split_relationship_id, +) from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA diff --git a/backend/onyx/kg/extractions/extraction_processing.py b/backend/onyx/kg/extractions/extraction_processing.py index 9ca9dc775d..4c9849ed15 100644 --- a/backend/onyx/kg/extractions/extraction_processing.py +++ b/backend/onyx/kg/extractions/extraction_processing.py @@ -5,36 +5,43 @@ from redis.lock import Lock as RedisLock from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT from onyx.db.connector import get_kg_enabled_connectors -from onyx.db.document import get_document_updated_at -from onyx.db.document import get_skipped_kg_documents -from onyx.db.document import get_unprocessed_kg_document_batch_for_connector -from onyx.db.document import update_document_kg_info -from onyx.db.document import update_document_kg_stage +from onyx.db.document import ( + get_document_updated_at, + get_skipped_kg_documents, + get_unprocessed_kg_document_batch_for_connector, + update_document_kg_info, + update_document_kg_stage, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.entities import delete_from_kg_entities__no_commit -from onyx.db.entities import upsert_staging_entity +from onyx.db.entities import delete_from_kg_entities__no_commit, upsert_staging_entity from onyx.db.entity_type import get_entity_types -from onyx.db.kg_config import get_kg_config_settings -from onyx.db.kg_config import validate_kg_settings -from onyx.db.models import Document -from onyx.db.models import KGStage -from onyx.db.relationships import delete_from_kg_relationships__no_commit -from onyx.db.relationships import upsert_staging_relationship -from onyx.db.relationships import upsert_staging_relationship_type -from onyx.kg.models import KGClassificationInstructions -from onyx.kg.models import KGDocumentDeepExtractionResults -from onyx.kg.models import KGEnhancedDocumentMetadata -from onyx.kg.models import KGEntityTypeInstructions -from onyx.kg.models import KGExtractionInstructions -from onyx.kg.models import KGImpliedExtractionResults -from onyx.kg.utils.extraction_utils import EntityTypeMetadataTracker -from onyx.kg.utils.extraction_utils import get_batch_documents_metadata -from onyx.kg.utils.extraction_utils import kg_deep_extraction -from onyx.kg.utils.extraction_utils import kg_implied_extraction -from onyx.kg.utils.formatting_utils import extract_relationship_type_id -from onyx.kg.utils.formatting_utils import get_entity_type -from onyx.kg.utils.formatting_utils import split_entity_id -from onyx.kg.utils.formatting_utils import split_relationship_id +from onyx.db.kg_config import get_kg_config_settings, validate_kg_settings +from onyx.db.models import Document, KGStage +from onyx.db.relationships import ( + delete_from_kg_relationships__no_commit, + upsert_staging_relationship, + upsert_staging_relationship_type, +) +from onyx.kg.models import ( + KGClassificationInstructions, + KGDocumentDeepExtractionResults, + KGEnhancedDocumentMetadata, + KGEntityTypeInstructions, + KGExtractionInstructions, + KGImpliedExtractionResults, +) +from onyx.kg.utils.extraction_utils import ( + EntityTypeMetadataTracker, + get_batch_documents_metadata, + kg_deep_extraction, + kg_implied_extraction, +) +from onyx.kg.utils.formatting_utils import ( + extract_relationship_type_id, + get_entity_type, + split_entity_id, + split_relationship_id, +) from onyx.kg.utils.lock_utils import extend_lock from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel diff --git a/backend/onyx/kg/resets/reset_index.py b/backend/onyx/kg/resets/reset_index.py index e4eef0f15d..5d797699ff 100644 --- a/backend/onyx/kg/resets/reset_index.py +++ b/backend/onyx/kg/resets/reset_index.py @@ -1,14 +1,16 @@ from sqlalchemy.orm import Session from onyx.db.document import reset_all_document_kg_stages -from onyx.db.models import Connector -from onyx.db.models import KGEntity -from onyx.db.models import KGEntityExtractionStaging -from onyx.db.models import KGEntityType -from onyx.db.models import KGRelationship -from onyx.db.models import KGRelationshipExtractionStaging -from onyx.db.models import KGRelationshipType -from onyx.db.models import KGRelationshipTypeExtractionStaging +from onyx.db.models import ( + Connector, + KGEntity, + KGEntityExtractionStaging, + KGEntityType, + KGRelationship, + KGRelationshipExtractionStaging, + KGRelationshipType, + KGRelationshipTypeExtractionStaging, +) def reset_full_kg_index__commit(db_session: Session) -> None: diff --git a/backend/onyx/kg/resets/reset_source.py b/backend/onyx/kg/resets/reset_source.py index 64bfa31ef0..9eae3c97c3 100644 --- a/backend/onyx/kg/resets/reset_source.py +++ b/backend/onyx/kg/resets/reset_source.py @@ -3,17 +3,19 @@ from sqlalchemy import or_ from onyx.configs.constants import DocumentSource from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import Connector -from onyx.db.models import Document -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import KGEntity -from onyx.db.models import KGEntityExtractionStaging -from onyx.db.models import KGEntityType -from onyx.db.models import KGRelationship -from onyx.db.models import KGRelationshipExtractionStaging -from onyx.db.models import KGRelationshipType -from onyx.db.models import KGRelationshipTypeExtractionStaging -from onyx.db.models import KGStage +from onyx.db.models import ( + Connector, + Document, + DocumentByConnectorCredentialPair, + KGEntity, + KGEntityExtractionStaging, + KGEntityType, + KGRelationship, + KGRelationshipExtractionStaging, + KGRelationshipType, + KGRelationshipTypeExtractionStaging, + KGStage, +) from onyx.kg.resets.reset_index import reset_full_kg_index__commit from onyx.kg.resets.reset_vespa import reset_vespa_kg_index diff --git a/backend/onyx/kg/resets/reset_vespa.py b/backend/onyx/kg/resets/reset_vespa.py index ac7bb5ae40..1e15f097dd 100644 --- a/backend/onyx/kg/resets/reset_vespa.py +++ b/backend/onyx/kg/resets/reset_vespa.py @@ -3,17 +3,16 @@ from typing import Any from redis.lock import Lock as RedisLock -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import DocumentSource +from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT, DocumentSource from onyx.db.document import get_num_chunks_for_document from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import Connector -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import KGEntityType +from onyx.db.models import Connector, DocumentByConnectorCredentialPair, KGEntityType from onyx.document_index.document_index_utils import get_uuid_from_chunk_info from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.vespa.vespa_document_index import KGVespaChunkUpdateRequest -from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex +from onyx.document_index.vespa.vespa_document_index import ( + KGVespaChunkUpdateRequest, + VespaDocumentIndex, +) from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT from onyx.kg.utils.lock_utils import extend_lock from onyx.utils.logger import setup_logger diff --git a/backend/onyx/kg/setup/kg_default_entity_definitions.py b/backend/onyx/kg/setup/kg_default_entity_definitions.py index 377f41fdbc..5cd77af876 100644 --- a/backend/onyx/kg/setup/kg_default_entity_definitions.py +++ b/backend/onyx/kg/setup/kg_default_entity_definitions.py @@ -4,15 +4,16 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.db.entity_type import KGEntityType -from onyx.db.kg_config import get_kg_config_settings -from onyx.db.kg_config import validate_kg_settings -from onyx.kg.models import KGAttributeEntityOption -from onyx.kg.models import KGAttributeImplicationProperty -from onyx.kg.models import KGAttributeProperty -from onyx.kg.models import KGEntityTypeAttributes -from onyx.kg.models import KGEntityTypeClassificationInfo -from onyx.kg.models import KGEntityTypeDefinition -from onyx.kg.models import KGGroundingType +from onyx.db.kg_config import get_kg_config_settings, validate_kg_settings +from onyx.kg.models import ( + KGAttributeEntityOption, + KGAttributeImplicationProperty, + KGAttributeProperty, + KGEntityTypeAttributes, + KGEntityTypeClassificationInfo, + KGEntityTypeDefinition, + KGGroundingType, +) def get_default_entity_types(vendor_name: str) -> dict[str, KGEntityTypeDefinition]: diff --git a/backend/onyx/kg/utils/embeddings.py b/backend/onyx/kg/utils/embeddings.py index 79de0cb3c0..40256fc5e4 100644 --- a/backend/onyx/kg/utils/embeddings.py +++ b/backend/onyx/kg/utils/embeddings.py @@ -4,10 +4,11 @@ import numpy as np from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.search_settings import get_current_search_settings -from onyx.natural_language_processing.search_nlp_models import EmbeddingModel -from onyx.natural_language_processing.search_nlp_models import EmbedTextType -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT +from onyx.natural_language_processing.search_nlp_models import ( + EmbeddingModel, + EmbedTextType, +) +from shared_configs.configs import MODEL_SERVER_HOST, MODEL_SERVER_PORT def encode_string_batch(strings: List[str]) -> np.ndarray: diff --git a/backend/onyx/kg/utils/extraction_utils.py b/backend/onyx/kg/utils/extraction_utils.py index a33cd5cd81..1476dd917d 100644 --- a/backend/onyx/kg/utils/extraction_utils.py +++ b/backend/onyx/kg/utils/extraction_utils.py @@ -1,43 +1,45 @@ import json -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCallTypes +from onyx.configs.constants import DocumentSource, OnyxCallTypes from onyx.configs.kg_configs import KG_METADATA_TRACKING_THRESHOLD from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.entities import get_kg_entity_by_document from onyx.db.entity_type import get_entity_types from onyx.db.kg_config import KGConfigSettings -from onyx.db.models import Document -from onyx.db.models import KGEntityType -from onyx.db.models import KGRelationshipType +from onyx.db.models import Document, KGEntityType, KGRelationshipType from onyx.db.tag import get_structured_tags_for_document -from onyx.kg.models import KGAttributeEntityOption -from onyx.kg.models import KGAttributeTrackInfo -from onyx.kg.models import KGAttributeTrackType -from onyx.kg.models import KGChunkFormat -from onyx.kg.models import KGClassificationInstructions -from onyx.kg.models import KGClassificationResult -from onyx.kg.models import KGDocumentDeepExtractionResults -from onyx.kg.models import KGEnhancedDocumentMetadata -from onyx.kg.models import KGImpliedExtractionResults -from onyx.kg.models import KGMetadataContent -from onyx.kg.utils.formatting_utils import extract_email -from onyx.kg.utils.formatting_utils import get_entity_type -from onyx.kg.utils.formatting_utils import kg_email_processing -from onyx.kg.utils.formatting_utils import make_entity_id -from onyx.kg.utils.formatting_utils import make_relationship_id -from onyx.kg.utils.formatting_utils import make_relationship_type_id +from onyx.kg.models import ( + KGAttributeEntityOption, + KGAttributeTrackInfo, + KGAttributeTrackType, + KGChunkFormat, + KGClassificationInstructions, + KGClassificationResult, + KGDocumentDeepExtractionResults, + KGEnhancedDocumentMetadata, + KGImpliedExtractionResults, + KGMetadataContent, +) +from onyx.kg.utils.formatting_utils import ( + extract_email, + get_entity_type, + kg_email_processing, + make_entity_id, + make_relationship_id, + make_relationship_type_id, +) from onyx.kg.vespa.vespa_interactions import get_document_vespa_contents from onyx.llm.factory import get_default_llm from onyx.llm.models import UserMessage from onyx.llm.utils import llm_response_to_string -from onyx.prompts.kg_prompts import CALL_CHUNK_PREPROCESSING_PROMPT -from onyx.prompts.kg_prompts import CALL_DOCUMENT_CLASSIFICATION_PROMPT -from onyx.prompts.kg_prompts import GENERAL_CHUNK_PREPROCESSING_PROMPT -from onyx.prompts.kg_prompts import MASTER_EXTRACTION_PROMPT +from onyx.prompts.kg_prompts import ( + CALL_CHUNK_PREPROCESSING_PROMPT, + CALL_DOCUMENT_CLASSIFICATION_PROMPT, + GENERAL_CHUNK_PREPROCESSING_PROMPT, + MASTER_EXTRACTION_PROMPT, +) from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/llm/factory.py b/backend/onyx/llm/factory.py index 7362f196f9..8e315c591a 100644 --- a/backend/onyx/llm/factory.py +++ b/backend/onyx/llm/factory.py @@ -7,24 +7,26 @@ from onyx.auth.schemas import UserRole from onyx.configs.model_configs import GEN_AI_TEMPERATURE from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import LLMModelFlowType -from onyx.db.llm import can_user_access_llm_provider -from onyx.db.llm import fetch_default_contextual_rag_model -from onyx.db.llm import fetch_default_llm_model -from onyx.db.llm import fetch_default_vision_model -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import fetch_existing_models -from onyx.db.llm import fetch_model_configuration_by_id -from onyx.db.llm import fetch_user_group_ids +from onyx.db.llm import ( + can_user_access_llm_provider, + fetch_default_contextual_rag_model, + fetch_default_llm_model, + fetch_default_vision_model, + fetch_existing_llm_provider, + fetch_existing_models, + fetch_model_configuration_by_id, + fetch_user_group_ids, +) from onyx.db.models import LLMProvider as LLMProviderModel -from onyx.db.models import Persona -from onyx.db.models import SearchSettings -from onyx.db.models import User +from onyx.db.models import Persona, SearchSettings, User from onyx.llm.constants import LlmProviderNames from onyx.llm.interfaces import LLM from onyx.llm.multi_llm import LitellmLLM from onyx.llm.override_models import LLMOverride -from onyx.llm.utils import get_max_input_tokens_from_llm_provider -from onyx.llm.utils import model_supports_image_input +from onyx.llm.utils import ( + get_max_input_tokens_from_llm_provider, + model_supports_image_input, +) from onyx.llm.well_known_providers.constants import ( PROVIDERS_WITH_SPECIAL_API_KEY_HANDLING, ) diff --git a/backend/onyx/llm/interfaces.py b/backend/onyx/llm/interfaces.py index d9579fa632..7f2f9505b6 100644 --- a/backend/onyx/llm/interfaces.py +++ b/backend/onyx/llm/interfaces.py @@ -1,17 +1,12 @@ import abc -from collections.abc import Callable -from collections.abc import Iterator +from collections.abc import Callable, Iterator from typing import Any from pydantic import BaseModel -from onyx.llm.model_response import ModelResponse -from onyx.llm.model_response import ModelResponseStream -from onyx.llm.models import LanguageModelInput -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import ToolChoiceOptions -from onyx.llm.tracing_wrap import wrap_invoke -from onyx.llm.tracing_wrap import wrap_stream +from onyx.llm.model_response import ModelResponse, ModelResponseStream +from onyx.llm.models import LanguageModelInput, ReasoningEffort, ToolChoiceOptions +from onyx.llm.tracing_wrap import wrap_invoke, wrap_stream from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/llm/litellm_singleton/monkey_patches.py b/backend/onyx/llm/litellm_singleton/monkey_patches.py index 65582471e7..2d5797ad28 100644 --- a/backend/onyx/llm/litellm_singleton/monkey_patches.py +++ b/backend/onyx/llm/litellm_singleton/monkey_patches.py @@ -61,21 +61,15 @@ Status checked against LiteLLM v1.85.1 (2026-05-26): import time import uuid -from typing import Any -from typing import cast -from typing import List -from typing import Optional +from typing import Any, cast, List, Optional from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, -) -from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, ) from litellm.llms.ollama.chat.transformation import OllamaChatCompletionResponseIterator from litellm.llms.ollama.common_utils import OllamaError -from litellm.types.utils import ChatCompletionUsageBlock -from litellm.types.utils import ModelResponseStream +from litellm.types.utils import ChatCompletionUsageBlock, ModelResponseStream # Original upstream chunk_parser, saved before any patching for fallback use _original_responses_chunk_parser = ( @@ -122,8 +116,7 @@ def _patch_ollama_chunk_parser() -> None: - return finish_reason when done is true - return usage when done is true """ - from litellm.types.utils import Delta - from litellm.types.utils import StreamingChoices + from litellm.types.utils import Delta, StreamingChoices # process tool calls - if complete function arg - add id to tool call tool_calls = chunk["message"].get("tool_calls") @@ -260,9 +253,7 @@ def _patch_responses_reasoning_summary_newlines() -> None: self: Any, chunk: dict ) -> "ModelResponseStream": from litellm.types.llms.openai import ResponsesAPIStreamEvents - from litellm.types.utils import Delta - from litellm.types.utils import ModelResponseStream - from litellm.types.utils import StreamingChoices + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices from pydantic import BaseModel parsed_chunk = chunk @@ -446,8 +437,7 @@ def _patch_responses_api_usage_format() -> None: - litellm/llms/volcengine/responses/transformation.py (lines 225, 277) - litellm/completion_extras/litellm_responses_transformation/handler.py (line 51) """ - from litellm.types.llms.openai import ResponseAPIUsage - from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse original_model_construct = ResponsesAPIResponse.model_construct @@ -515,13 +505,14 @@ def _patch_logging_assembled_streaming_response() -> None: """ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.responses.utils import ResponseAPILoggingUtils - from litellm.types.llms.openai import ResponseAPIUsage - from litellm.types.llms.openai import ResponseCompletedEvent - from litellm.types.llms.openai import ResponseFailedEvent - from litellm.types.llms.openai import ResponseIncompleteEvent - from litellm.types.llms.openai import ResponsesAPIResponse - from litellm.types.utils import ModelResponse - from litellm.types.utils import TextCompletionResponse + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ) + from litellm.types.utils import ModelResponse, TextCompletionResponse original_method = LiteLLMLoggingObj._get_assembled_streaming_response diff --git a/backend/onyx/llm/model_capabilities.py b/backend/onyx/llm/model_capabilities.py index 50814d5c64..3f9d37171b 100644 --- a/backend/onyx/llm/model_capabilities.py +++ b/backend/onyx/llm/model_capabilities.py @@ -14,14 +14,14 @@ import re import threading import time from functools import lru_cache -from typing import Any -from typing import cast +from typing import Any, cast -from onyx.configs.model_configs import GEN_AI_MAX_TOKENS -from onyx.configs.model_configs import GEN_AI_MODEL_FALLBACK_MAX_TOKENS -from onyx.configs.model_configs import GEN_AI_NUM_RESERVED_OUTPUT_TOKENS -from onyx.llm.constants import BEDROCK_MODEL_TOKEN_LIMITS -from onyx.llm.constants import LlmProviderNames +from onyx.configs.model_configs import ( + GEN_AI_MAX_TOKENS, + GEN_AI_MODEL_FALLBACK_MAX_TOKENS, + GEN_AI_NUM_RESERVED_OUTPUT_TOKENS, +) +from onyx.llm.constants import BEDROCK_MODEL_TOKEN_LIMITS, LlmProviderNames from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/llm/model_name_parser.py b/backend/onyx/llm/model_name_parser.py index 4f1c25d67c..4381d0748e 100644 --- a/backend/onyx/llm/model_name_parser.py +++ b/backend/onyx/llm/model_name_parser.py @@ -20,12 +20,14 @@ from functools import lru_cache from pydantic import BaseModel -from onyx.llm.constants import AGGREGATOR_PROVIDERS -from onyx.llm.constants import HYPHENATED_MODEL_NAMES -from onyx.llm.constants import LlmProviderNames -from onyx.llm.constants import MODEL_PREFIX_TO_VENDOR -from onyx.llm.constants import PROVIDER_DISPLAY_NAMES -from onyx.llm.constants import VENDOR_BRAND_NAMES +from onyx.llm.constants import ( + AGGREGATOR_PROVIDERS, + HYPHENATED_MODEL_NAMES, + LlmProviderNames, + MODEL_PREFIX_TO_VENDOR, + PROVIDER_DISPLAY_NAMES, + VENDOR_BRAND_NAMES, +) class ParsedModelName(BaseModel): diff --git a/backend/onyx/llm/model_response.py b/backend/onyx/llm/model_response.py index 1c2c4ac623..41a70b1afa 100644 --- a/backend/onyx/llm/model_response.py +++ b/backend/onyx/llm/model_response.py @@ -1,11 +1,8 @@ from __future__ import annotations -from typing import Any -from typing import List -from typing import TYPE_CHECKING +from typing import Any, List, TYPE_CHECKING -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field class FunctionCall(BaseModel): diff --git a/backend/onyx/llm/multi_llm.py b/backend/onyx/llm/multi_llm.py index 88d97f9195..0321128f2c 100644 --- a/backend/onyx/llm/multi_llm.py +++ b/backend/onyx/llm/multi_llm.py @@ -4,62 +4,56 @@ import re import time from collections.abc import Iterator from contextlib import contextmanager -from typing import Any -from typing import cast -from typing import TYPE_CHECKING -from typing import Union +from typing import Any, cast, TYPE_CHECKING, Union from readerwriterlock import rwlock -from onyx.configs.app_configs import MOCK_LLM_RESPONSE -from onyx.configs.app_configs import SEND_USER_METADATA_TO_LLM_PROVIDER -from onyx.configs.chat_configs import LLM_FIRST_CHUNK_MAX_RETRIES -from onyx.configs.chat_configs import LLM_SOCKET_READ_TIMEOUT -from onyx.configs.model_configs import GEN_AI_TEMPERATURE -from onyx.configs.model_configs import LITELLM_EXTRA_BODY +from onyx.configs.app_configs import ( + MOCK_LLM_RESPONSE, + SEND_USER_METADATA_TO_LLM_PROVIDER, +) +from onyx.configs.chat_configs import ( + LLM_FIRST_CHUNK_MAX_RETRIES, + LLM_SOCKET_READ_TIMEOUT, +) +from onyx.configs.model_configs import GEN_AI_TEMPERATURE, LITELLM_EXTRA_BODY from onyx.llm.constants import LlmProviderNames from onyx.llm.cost import calculate_llm_cost_cents -from onyx.llm.interfaces import LanguageModelInput -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMConfig -from onyx.llm.interfaces import LLMUserIdentity -from onyx.llm.interfaces import ReasoningEffort -from onyx.llm.interfaces import ToolChoiceOptions -from onyx.llm.model_capabilities import is_true_openai_model -from onyx.llm.model_capabilities import model_is_reasoning_model -from onyx.llm.model_response import ModelResponse -from onyx.llm.model_response import ModelResponseStream -from onyx.llm.model_response import Usage -from onyx.llm.models import ANTHROPIC_ADAPTIVE_REASONING_EFFORT -from onyx.llm.models import ANTHROPIC_REASONING_EFFORT_BUDGET -from onyx.llm.models import OPENAI_REASONING_EFFORT +from onyx.llm.interfaces import ( + LanguageModelInput, + LLM, + LLMConfig, + LLMUserIdentity, + ReasoningEffort, + ToolChoiceOptions, +) +from onyx.llm.model_capabilities import is_true_openai_model, model_is_reasoning_model +from onyx.llm.model_response import ModelResponse, ModelResponseStream, Usage +from onyx.llm.models import ( + ANTHROPIC_ADAPTIVE_REASONING_EFFORT, + ANTHROPIC_REASONING_EFFORT_BUDGET, + OPENAI_REASONING_EFFORT, +) from onyx.llm.request_context import get_llm_mock_response from onyx.llm.utils import build_litellm_passthrough_kwargs -from onyx.llm.well_known_providers.constants import AWS_ACCESS_KEY_ID_KWARG from onyx.llm.well_known_providers.constants import ( + AWS_ACCESS_KEY_ID_KWARG, AWS_ACCESS_KEY_ID_KWARG_ENV_VAR_FORMAT, -) -from onyx.llm.well_known_providers.constants import ( AWS_BEARER_TOKEN_BEDROCK_KWARG_ENV_VAR_FORMAT, -) -from onyx.llm.well_known_providers.constants import AWS_REGION_NAME_KWARG -from onyx.llm.well_known_providers.constants import AWS_REGION_NAME_KWARG_ENV_VAR_FORMAT -from onyx.llm.well_known_providers.constants import AWS_SECRET_ACCESS_KEY_KWARG -from onyx.llm.well_known_providers.constants import ( + AWS_REGION_NAME_KWARG, + AWS_REGION_NAME_KWARG_ENV_VAR_FORMAT, + AWS_SECRET_ACCESS_KEY_KWARG, AWS_SECRET_ACCESS_KEY_KWARG_ENV_VAR_FORMAT, -) -from onyx.llm.well_known_providers.constants import LM_STUDIO_API_KEY_CONFIG_KEY -from onyx.llm.well_known_providers.constants import OLLAMA_API_KEY_CONFIG_KEY -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_WORKLOAD_IDENTITY -from onyx.llm.well_known_providers.constants import VERTEX_CREDENTIALS_FILE_KWARG -from onyx.llm.well_known_providers.constants import ( + LM_STUDIO_API_KEY_CONFIG_KEY, + OLLAMA_API_KEY_CONFIG_KEY, + VERTEX_AUTH_METHOD_KWARG, + VERTEX_AUTH_METHOD_WORKLOAD_IDENTITY, + VERTEX_CREDENTIALS_FILE_KWARG, VERTEX_CREDENTIALS_FILE_KWARG_ENV_VAR_FORMAT, + VERTEX_LOCATION_KWARG, + VERTEX_PROJECT_KWARG, ) -from onyx.llm.well_known_providers.constants import VERTEX_LOCATION_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_PROJECT_KWARG -from onyx.utils.encryption import mask_env_value_for_logging -from onyx.utils.encryption import mask_string +from onyx.utils.encryption import mask_env_value_for_logging, mask_string from onyx.utils.logger import setup_logger logger = setup_logger() @@ -71,8 +65,7 @@ logger = setup_logger() _env_rwlock = rwlock.RWLockWrite() if TYPE_CHECKING: - from litellm import CustomStreamWrapper - from litellm import HTTPHandler + from litellm import CustomStreamWrapper, HTTPHandler _LLM_PROMPT_LONG_TERM_LOG_CATEGORY = "llm_prompt" @@ -486,8 +479,7 @@ class LitellmLLM(LLM): return # Import here to avoid circular imports from onyx.db.engine.sql_engine import get_session_with_current_tenant - from onyx.db.usage import increment_usage - from onyx.db.usage import UsageType + from onyx.db.usage import increment_usage, UsageType # Calculate cost in cents cost_cents = calculate_llm_cost_cents( @@ -522,8 +514,7 @@ class LitellmLLM(LLM): client: "HTTPHandler | None" = None, ) -> Union["ModelResponse", "CustomStreamWrapper"]: # Lazy loading to avoid memory bloat for non-inference flows - from litellm.exceptions import RateLimitError - from litellm.exceptions import Timeout + from litellm.exceptions import RateLimitError, Timeout from onyx.llm.litellm_singleton import litellm diff --git a/backend/onyx/llm/prompt_cache/__init__.py b/backend/onyx/llm/prompt_cache/__init__.py index fe91ed6cfe..9867b1eed2 100644 --- a/backend/onyx/llm/prompt_cache/__init__.py +++ b/backend/onyx/llm/prompt_cache/__init__.py @@ -5,8 +5,7 @@ LLM providers. It supports both implicit caching (automatic provider-side cachin and explicit caching (with cache metadata management). """ -from onyx.llm.prompt_cache.cache_manager import CacheManager -from onyx.llm.prompt_cache.cache_manager import generate_cache_key_hash +from onyx.llm.prompt_cache.cache_manager import CacheManager, generate_cache_key_hash from onyx.llm.prompt_cache.models import CacheMetadata from onyx.llm.prompt_cache.processor import process_with_prompt_cache from onyx.llm.prompt_cache.providers.anthropic import AnthropicPromptCacheProvider @@ -15,8 +14,10 @@ from onyx.llm.prompt_cache.providers.factory import get_provider_adapter from onyx.llm.prompt_cache.providers.noop import NoOpPromptCacheProvider from onyx.llm.prompt_cache.providers.openai import OpenAIPromptCacheProvider from onyx.llm.prompt_cache.providers.vertex import VertexAIPromptCacheProvider -from onyx.llm.prompt_cache.utils import combine_messages_with_continuation -from onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform +from onyx.llm.prompt_cache.utils import ( + combine_messages_with_continuation, + prepare_messages_with_cacheable_transform, +) __all__ = [ "AnthropicPromptCacheProvider", diff --git a/backend/onyx/llm/prompt_cache/cache_manager.py b/backend/onyx/llm/prompt_cache/cache_manager.py index 51e191b463..4b56f0d786 100644 --- a/backend/onyx/llm/prompt_cache/cache_manager.py +++ b/backend/onyx/llm/prompt_cache/cache_manager.py @@ -2,8 +2,7 @@ import hashlib import json -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from onyx.configs.model_configs import PROMPT_CACHE_REDIS_TTL_MULTIPLIER from onyx.key_value_store.store import PgRedisKVStore diff --git a/backend/onyx/llm/prompt_cache/processor.py b/backend/onyx/llm/prompt_cache/processor.py index b848e6789c..050d54cd86 100644 --- a/backend/onyx/llm/prompt_cache/processor.py +++ b/backend/onyx/llm/prompt_cache/processor.py @@ -1,7 +1,6 @@ """Main processor for prompt caching.""" -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from onyx.configs.model_configs import ENABLE_PROMPT_CACHING from onyx.llm.interfaces import LLMConfig diff --git a/backend/onyx/llm/prompt_cache/providers/anthropic.py b/backend/onyx/llm/prompt_cache/providers/anthropic.py index cc393423bd..545a5e3545 100644 --- a/backend/onyx/llm/prompt_cache/providers/anthropic.py +++ b/backend/onyx/llm/prompt_cache/providers/anthropic.py @@ -6,8 +6,10 @@ from onyx.llm.interfaces import LanguageModelInput from onyx.llm.models import ChatCompletionMessage from onyx.llm.prompt_cache.models import CacheMetadata from onyx.llm.prompt_cache.providers.base import PromptCacheProvider -from onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform -from onyx.llm.prompt_cache.utils import revalidate_message_from_original +from onyx.llm.prompt_cache.utils import ( + prepare_messages_with_cacheable_transform, + revalidate_message_from_original, +) def _add_anthropic_cache_control( diff --git a/backend/onyx/llm/prompt_cache/providers/base.py b/backend/onyx/llm/prompt_cache/providers/base.py index 7e495eb5e8..0c26baa28d 100644 --- a/backend/onyx/llm/prompt_cache/providers/base.py +++ b/backend/onyx/llm/prompt_cache/providers/base.py @@ -1,7 +1,6 @@ """Base interface for provider-specific prompt caching adapters.""" -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from onyx.llm.interfaces import LanguageModelInput from onyx.llm.prompt_cache.models import CacheMetadata diff --git a/backend/onyx/llm/prompt_cache/providers/vertex.py b/backend/onyx/llm/prompt_cache/providers/vertex.py index 1d9a1abb18..7b88788bf6 100644 --- a/backend/onyx/llm/prompt_cache/providers/vertex.py +++ b/backend/onyx/llm/prompt_cache/providers/vertex.py @@ -6,8 +6,10 @@ from onyx.llm.interfaces import LanguageModelInput from onyx.llm.models import ChatCompletionMessage from onyx.llm.prompt_cache.models import CacheMetadata from onyx.llm.prompt_cache.providers.base import PromptCacheProvider -from onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform -from onyx.llm.prompt_cache.utils import revalidate_message_from_original +from onyx.llm.prompt_cache.utils import ( + prepare_messages_with_cacheable_transform, + revalidate_message_from_original, +) class VertexAIPromptCacheProvider(PromptCacheProvider): diff --git a/backend/onyx/llm/prompt_cache/utils.py b/backend/onyx/llm/prompt_cache/utils.py index 0da3b2f90d..39e27a6bb1 100644 --- a/backend/onyx/llm/prompt_cache/utils.py +++ b/backend/onyx/llm/prompt_cache/utils.py @@ -2,12 +2,10 @@ """Utility functions for prompt caching.""" import json -from collections.abc import Callable -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import Any -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import LanguageModelInput +from onyx.llm.models import ChatCompletionMessage, LanguageModelInput from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/llm/tracing_wrap.py b/backend/onyx/llm/tracing_wrap.py index 3b4765fe84..03fbe89344 100644 --- a/backend/onyx/llm/tracing_wrap.py +++ b/backend/onyx/llm/tracing_wrap.py @@ -17,21 +17,17 @@ from __future__ import annotations import functools import inspect -from collections.abc import Callable -from collections.abc import Iterator -from typing import Any -from typing import TYPE_CHECKING +from collections.abc import Callable, Iterator +from typing import Any, TYPE_CHECKING -from onyx.llm.model_response import ChatCompletionDeltaToolCall +from onyx.llm.model_response import ChatCompletionDeltaToolCall, Usage from onyx.llm.model_response import FunctionCall as DeltaFunctionCall -from onyx.llm.model_response import Usage from onyx.tracing.framework.create import get_current_span from onyx.tracing.framework.span_data import GenerationSpanData if TYPE_CHECKING: from onyx.llm.interfaces import LLM - from onyx.llm.model_response import ModelResponse - from onyx.llm.model_response import ModelResponseStream + from onyx.llm.model_response import ModelResponse, ModelResponseStream from onyx.llm.models import ToolCall @@ -141,8 +137,7 @@ def wrap_invoke( return invoke_fn(self, *args, **kwargs) from onyx.tracing.flows import LLMFlow - from onyx.tracing.llm_utils import llm_generation_span - from onyx.tracing.llm_utils import record_llm_response + from onyx.tracing.llm_utils import llm_generation_span, record_llm_response prompt, tools = _extract_prompt_and_tools(sig, self, args, kwargs) with llm_generation_span( @@ -198,8 +193,7 @@ def wrap_stream( return from onyx.tracing.flows import LLMFlow - from onyx.tracing.llm_utils import llm_generation_span - from onyx.tracing.llm_utils import record_llm_span_output + from onyx.tracing.llm_utils import llm_generation_span, record_llm_span_output prompt, tools = _extract_prompt_and_tools(sig, self, args, kwargs) with llm_generation_span( diff --git a/backend/onyx/llm/utils.py b/backend/onyx/llm/utils.py index a2d4382a2e..dc95f14759 100644 --- a/backend/onyx/llm/utils.py +++ b/backend/onyx/llm/utils.py @@ -1,29 +1,31 @@ import copy import re -from collections.abc import Callable -from collections.abc import Iterable -from typing import Any -from typing import TYPE_CHECKING +from collections.abc import Callable, Iterable +from typing import Any, TYPE_CHECKING from sqlalchemy import select -from onyx.configs.app_configs import LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS -from onyx.configs.app_configs import MAX_TOKENS_FOR_FULL_INCLUSION -from onyx.configs.app_configs import SEND_USER_METADATA_TO_LLM_PROVIDER -from onyx.configs.app_configs import USE_CHUNK_SUMMARY -from onyx.configs.app_configs import USE_DOCUMENT_SUMMARY +from onyx.configs.app_configs import ( + LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS, + MAX_TOKENS_FOR_FULL_INCLUSION, + SEND_USER_METADATA_TO_LLM_PROVIDER, + USE_CHUNK_SUMMARY, + USE_DOCUMENT_SUMMARY, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import LLMModelFlowType -from onyx.db.models import LLMProvider -from onyx.db.models import ModelConfiguration -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMUserIdentity -from onyx.llm.model_capabilities import get_max_input_tokens -from onyx.llm.model_capabilities import litellm_thinks_model_supports_image_input +from onyx.db.models import LLMProvider, ModelConfiguration +from onyx.llm.interfaces import LLM, LLMUserIdentity +from onyx.llm.model_capabilities import ( + get_max_input_tokens, + litellm_thinks_model_supports_image_input, +) from onyx.llm.model_response import ModelResponse from onyx.llm.models import UserMessage -from onyx.prompts.contextual_retrieval import CONTEXTUAL_RAG_TOKEN_ESTIMATE -from onyx.prompts.contextual_retrieval import DOCUMENT_SUMMARY_TOKEN_ESTIMATE +from onyx.prompts.contextual_retrieval import ( + CONTEXTUAL_RAG_TOKEN_ESTIMATE, + DOCUMENT_SUMMARY_TOKEN_ESTIMATE, +) from onyx.utils.logger import setup_logger from shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE @@ -126,19 +128,21 @@ def litellm_exception_to_error_msg( - error_code: Categorized error code for frontend display - is_retryable: Whether the user should try again """ - from litellm.exceptions import APIConnectionError - from litellm.exceptions import APIError - from litellm.exceptions import AuthenticationError - from litellm.exceptions import BadRequestError - from litellm.exceptions import BudgetExceededError - from litellm.exceptions import ContentPolicyViolationError - from litellm.exceptions import ContextWindowExceededError - from litellm.exceptions import NotFoundError - from litellm.exceptions import PermissionDeniedError - from litellm.exceptions import RateLimitError - from litellm.exceptions import ServiceUnavailableError - from litellm.exceptions import Timeout - from litellm.exceptions import UnprocessableEntityError + from litellm.exceptions import ( + APIConnectionError, + APIError, + AuthenticationError, + BadRequestError, + BudgetExceededError, + ContentPolicyViolationError, + ContextWindowExceededError, + NotFoundError, + PermissionDeniedError, + RateLimitError, + ServiceUnavailableError, + Timeout, + UnprocessableEntityError, + ) core_exception = _unwrap_nested_exception(e) error_msg = str(core_exception) diff --git a/backend/onyx/llm/well_known_providers/auto_update_models.py b/backend/onyx/llm/well_known_providers/auto_update_models.py index ecb509bb7b..e609d2f32e 100644 --- a/backend/onyx/llm/well_known_providers/auto_update_models.py +++ b/backend/onyx/llm/well_known_providers/auto_update_models.py @@ -3,8 +3,7 @@ from datetime import datetime from typing import Any -from pydantic import BaseModel -from pydantic import field_validator +from pydantic import BaseModel, field_validator from onyx.llm.well_known_providers.models import SimpleKnownModel diff --git a/backend/onyx/llm/well_known_providers/auto_update_service.py b/backend/onyx/llm/well_known_providers/auto_update_service.py index 24cd5666f1..7b33dd3962 100644 --- a/backend/onyx/llm/well_known_providers/auto_update_service.py +++ b/backend/onyx/llm/well_known_providers/auto_update_service.py @@ -15,8 +15,7 @@ from sqlalchemy.orm import Session from onyx.cache.factory import get_cache_backend from onyx.configs.app_configs import AUTO_LLM_CONFIG_URL -from onyx.db.llm import fetch_auto_mode_providers -from onyx.db.llm import sync_auto_mode_models +from onyx.db.llm import fetch_auto_mode_providers, sync_auto_mode_models from onyx.llm.well_known_providers.auto_update_models import LLMRecommendations from onyx.utils.logger import setup_logger diff --git a/backend/onyx/llm/well_known_providers/llm_provider_options.py b/backend/onyx/llm/well_known_providers/llm_provider_options.py index e27e035313..06e8753786 100644 --- a/backend/onyx/llm/well_known_providers/llm_provider_options.py +++ b/backend/onyx/llm/well_known_providers/llm_provider_options.py @@ -3,27 +3,31 @@ import pathlib import threading import time -from onyx.llm.constants import LlmProviderNames -from onyx.llm.constants import PROVIDER_DISPLAY_NAMES -from onyx.llm.constants import WELL_KNOWN_PROVIDER_NAMES +from onyx.llm.constants import ( + LlmProviderNames, + PROVIDER_DISPLAY_NAMES, + WELL_KNOWN_PROVIDER_NAMES, +) from onyx.llm.model_capabilities import get_max_input_tokens from onyx.llm.utils import model_supports_image_input from onyx.llm.well_known_providers.auto_update_models import LLMRecommendations from onyx.llm.well_known_providers.auto_update_service import ( fetch_llm_recommendations_from_github, ) -from onyx.llm.well_known_providers.constants import ANTHROPIC_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import AZURE_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import BEDROCK_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import BIFROST_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import LITELLM_PROXY_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import LM_STUDIO_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import NEBIUS_TOKENFACTORY_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import OLLAMA_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import OPENAI_COMPATIBLE_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import OPENAI_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import OPENROUTER_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import VERTEXAI_PROVIDER_NAME +from onyx.llm.well_known_providers.constants import ( + ANTHROPIC_PROVIDER_NAME, + AZURE_PROVIDER_NAME, + BEDROCK_PROVIDER_NAME, + BIFROST_PROVIDER_NAME, + LITELLM_PROXY_PROVIDER_NAME, + LM_STUDIO_PROVIDER_NAME, + NEBIUS_TOKENFACTORY_PROVIDER_NAME, + OLLAMA_PROVIDER_NAME, + OPENAI_COMPATIBLE_PROVIDER_NAME, + OPENAI_PROVIDER_NAME, + OPENROUTER_PROVIDER_NAME, + VERTEXAI_PROVIDER_NAME, +) from onyx.llm.well_known_providers.models import WellKnownLLMProviderDescriptor from onyx.server.manage.llm.models import ModelConfigurationView from onyx.utils.logger import setup_logger diff --git a/backend/onyx/llm/well_known_providers/models.py b/backend/onyx/llm/well_known_providers/models.py index 63121abc2c..de7dfdf047 100644 --- a/backend/onyx/llm/well_known_providers/models.py +++ b/backend/onyx/llm/well_known_providers/models.py @@ -1,7 +1,6 @@ from enum import Enum -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.server.manage.llm.models import ModelConfigurationView diff --git a/backend/onyx/main.py b/backend/onyx/main.py index 35c6e9c1b8..b0cb7525a4 100644 --- a/backend/onyx/main.py +++ b/backend/onyx/main.py @@ -4,16 +4,11 @@ import traceback import warnings from collections.abc import AsyncGenerator from contextlib import asynccontextmanager -from typing import Any -from typing import cast +from typing import Any, cast import uvicorn from anyio import to_thread -from fastapi import APIRouter -from fastapi import FastAPI -from fastapi import HTTPException -from fastapi import Request -from fastapi import status +from fastapi import APIRouter, FastAPI, HTTPException, Request, status from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse @@ -24,49 +19,51 @@ from sentry_sdk.integrations.starlette import StarletteIntegration from starlette.types import Lifespan from onyx import __version__ -from onyx.auth.schemas import UserCreate -from onyx.auth.schemas import UserRead -from onyx.auth.schemas import UserUpdate -from onyx.auth.users import auth_backend -from onyx.auth.users import create_onyx_oauth_router -from onyx.auth.users import fastapi_users -from onyx.auth.users import mobile_auth_backend -from onyx.auth.users import verify_auth_setting -from onyx.auth.users import verify_user_auth_secret +from onyx.auth.schemas import UserCreate, UserRead, UserUpdate +from onyx.auth.users import ( + auth_backend, + create_onyx_oauth_router, + fastapi_users, + mobile_auth_backend, + verify_auth_setting, + verify_user_auth_secret, +) from onyx.cache.interface import CacheBackendType -from onyx.configs.app_configs import API_SERVER_THREADPOOL_SIZE -from onyx.configs.app_configs import APP_API_PREFIX -from onyx.configs.app_configs import APP_HOST -from onyx.configs.app_configs import APP_PORT -from onyx.configs.app_configs import CACHE_BACKEND -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import ENABLE_PUBLIC_DOCS -from onyx.configs.app_configs import GOOGLE_LOGIN_BASE_SCOPES -from onyx.configs.app_configs import GOOGLE_OAUTH_SCOPE_OVERRIDE -from onyx.configs.app_configs import LOG_ENDPOINT_LATENCY -from onyx.configs.app_configs import OAUTH_CLIENT_ID -from onyx.configs.app_configs import OAUTH_CLIENT_SECRET -from onyx.configs.app_configs import OAUTH_ENABLED -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE -from onyx.configs.app_configs import POSTGRES_API_SERVER_READ_ONLY_POOL_OVERFLOW -from onyx.configs.app_configs import POSTGRES_API_SERVER_READ_ONLY_POOL_SIZE -from onyx.configs.app_configs import SYSTEM_RECURSION_LIMIT -from onyx.configs.app_configs import USER_AUTH_SECRET -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.configs.app_configs import ( + API_SERVER_THREADPOOL_SIZE, + APP_API_PREFIX, + APP_HOST, + APP_PORT, + CACHE_BACKEND, + DISABLE_VECTOR_DB, + ENABLE_PUBLIC_DOCS, + GOOGLE_LOGIN_BASE_SCOPES, + GOOGLE_OAUTH_SCOPE_OVERRIDE, + LOG_ENDPOINT_LATENCY, + OAUTH_CLIENT_ID, + OAUTH_CLIENT_SECRET, + OAUTH_ENABLED, + POSTGRES_API_SERVER_POOL_OVERFLOW, + POSTGRES_API_SERVER_POOL_SIZE, + POSTGRES_API_SERVER_READ_ONLY_POOL_OVERFLOW, + POSTGRES_API_SERVER_READ_ONLY_POOL_SIZE, + SYSTEM_RECURSION_LIMIT, + USER_AUTH_SECRET, + WEB_DOMAIN, +) from onyx.configs.constants import POSTGRES_WEB_APP_NAME -from onyx.db.engine.async_sql_engine import get_sqlalchemy_async_engine -from onyx.db.engine.async_sql_engine import reset_sqlalchemy_async_engine +from onyx.db.engine.async_sql_engine import ( + get_sqlalchemy_async_engine, + reset_sqlalchemy_async_engine, +) from onyx.db.engine.connection_warmup import warm_up_connections -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from onyx.db.sso_provider import seed_saml_provider_from_conf_dir from onyx.error_handling.exceptions import register_onyx_exception_handlers from onyx.file_store.file_store import get_default_file_store from onyx.hooks.registry import validate_registry from onyx.server.api_key.api import router as api_key_router -from onyx.server.auth.captcha_api import CaptchaCookieMiddleware -from onyx.server.auth.captcha_api import LoginCaptchaMiddleware +from onyx.server.auth.captcha_api import CaptchaCookieMiddleware, LoginCaptchaMiddleware from onyx.server.auth.captcha_api import router as captcha_router from onyx.server.auth.mobile import router as mobile_auth_router from onyx.server.auth_check import check_router_auth @@ -98,9 +95,8 @@ from onyx.server.features.oauth_config.api import ( ) from onyx.server.features.oauth_config.api import router as oauth_config_router from onyx.server.features.password.api import router as password_router -from onyx.server.features.persona.api import admin_agents_router +from onyx.server.features.persona.api import admin_agents_router, agents_router from onyx.server.features.persona.api import admin_router as admin_persona_router -from onyx.server.features.persona.api import agents_router from onyx.server.features.persona.api import basic_router as persona_router from onyx.server.features.projects.api import router as projects_router from onyx.server.features.search.api import router as search_api_router @@ -143,10 +139,12 @@ from onyx.server.metrics.postgres_connection_pool import ( ) from onyx.server.metrics.prometheus_setup import setup_prometheus_metrics from onyx.server.middleware.latency_logging import add_latency_logging_middleware -from onyx.server.middleware.rate_limiting import close_auth_limiter -from onyx.server.middleware.rate_limiting import get_auth_rate_limiters -from onyx.server.middleware.rate_limiting import RATE_LIMITING_ENABLED -from onyx.server.middleware.rate_limiting import setup_auth_limiter +from onyx.server.middleware.rate_limiting import ( + close_auth_limiter, + get_auth_rate_limiters, + RATE_LIMITING_ENABLED, + setup_auth_limiter, +) from onyx.server.oidc_multi import router as oidc_multi_router from onyx.server.onyx_api.ingestion import router as onyx_api_router from onyx.server.pat.api import router as pat_router @@ -159,27 +157,29 @@ from onyx.server.settings.api import admin_router as settings_admin_router from onyx.server.settings.api import basic_router as settings_router from onyx.server.token_rate_limits.api import router as token_rate_limit_settings_router from onyx.server.utils import BasicAuthenticationError -from onyx.setup import setup_multitenant_onyx -from onyx.setup import setup_onyx +from onyx.setup import setup_multitenant_onyx, setup_onyx from onyx.tracing.setup import setup_tracing from onyx.utils.client_ip import ClientIPMiddleware -from onyx.utils.logger import setup_logger -from onyx.utils.logger import setup_uvicorn_logger -from onyx.utils.middleware import add_endpoint_context_middleware -from onyx.utils.middleware import add_onyx_request_id_middleware -from onyx.utils.telemetry import get_or_generate_uuid -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType -from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version -from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable -from shared_configs.configs import CORS_ALLOW_CREDENTIALS -from shared_configs.configs import CORS_ALLOWED_ORIGIN -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import SENTRY_DSN -from shared_configs.configs import SENTRY_TRACES_SAMPLE_RATE +from onyx.utils.logger import setup_logger, setup_uvicorn_logger +from onyx.utils.middleware import ( + add_endpoint_context_middleware, + add_onyx_request_id_middleware, +) +from onyx.utils.telemetry import get_or_generate_uuid, optional_telemetry, RecordType +from onyx.utils.variable_functionality import ( + fetch_ee_implementation_or_noop, + fetch_versioned_implementation, + global_version, + set_is_ee_based_on_env_variable, +) +from shared_configs.configs import ( + CORS_ALLOW_CREDENTIALS, + CORS_ALLOWED_ORIGIN, + MULTI_TENANT, + POSTGRES_DEFAULT_SCHEMA, + SENTRY_DSN, + SENTRY_TRACES_SAMPLE_RATE, +) from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR warnings.filterwarnings( @@ -413,8 +413,10 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # noqa: ARG001 await setup_auth_limiter() if DISABLE_VECTOR_DB: - from onyx.background.periodic_poller import recover_stuck_user_files - from onyx.background.periodic_poller import start_periodic_poller + from onyx.background.periodic_poller import ( + recover_stuck_user_files, + start_periodic_poller, + ) recover_stuck_user_files(POSTGRES_DEFAULT_SCHEMA) start_periodic_poller(POSTGRES_DEFAULT_SCHEMA) diff --git a/backend/onyx/mcp_server/api.py b/backend/onyx/mcp_server/api.py index 0749ed7dde..8df017a75d 100644 --- a/backend/onyx/mcp_server/api.py +++ b/backend/onyx/mcp_server/api.py @@ -3,24 +3,22 @@ from collections.abc import AsyncGenerator from contextlib import asynccontextmanager -from fastapi import FastAPI -from fastapi import Request +from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import JSONResponse -from fastapi.responses import Response +from fastapi.responses import JSONResponse, Response from fastmcp import FastMCP from starlette.datastructures import MutableHeaders from starlette.middleware.base import RequestResponseEndpoint -from starlette.types import Receive -from starlette.types import Scope -from starlette.types import Send +from starlette.types import Receive, Scope, Send from onyx.configs.app_configs import MCP_SERVER_CORS_ORIGINS from onyx.error_handling.exceptions import register_onyx_exception_handlers from onyx.mcp_server.auth import OnyxTokenVerifier from onyx.mcp_server.utils import shutdown_http_client -from onyx.server.metrics.prometheus_setup import create_prometheus_instrumentator -from onyx.server.metrics.prometheus_setup import expose_prometheus_metrics +from onyx.server.metrics.prometheus_setup import ( + create_prometheus_instrumentator, + expose_prometheus_metrics, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable from shared_configs.configs import cors_allow_credentials diff --git a/backend/onyx/mcp_server/auth.py b/backend/onyx/mcp_server/auth.py index b65e215af5..351c3a2a0d 100644 --- a/backend/onyx/mcp_server/auth.py +++ b/backend/onyx/mcp_server/auth.py @@ -2,12 +2,10 @@ from typing import Optional -from fastmcp.server.auth.auth import AccessToken -from fastmcp.server.auth.auth import TokenVerifier +from fastmcp.server.auth.auth import AccessToken, TokenVerifier from onyx.mcp_server.utils import get_http_client -from onyx.server.metrics.mcp_server import MCPAuthResult -from onyx.server.metrics.mcp_server import record_mcp_auth_result +from onyx.server.metrics.mcp_server import MCPAuthResult, record_mcp_auth_result from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import build_api_server_url_for_http_requests diff --git a/backend/onyx/mcp_server/resources/__init__.py b/backend/onyx/mcp_server/resources/__init__.py index dd480454fe..aa4eef3ba8 100644 --- a/backend/onyx/mcp_server/resources/__init__.py +++ b/backend/onyx/mcp_server/resources/__init__.py @@ -1,5 +1,7 @@ """Resource registrations for the Onyx MCP server.""" # Import resource modules so decorators execute when the package loads. -from onyx.mcp_server.resources import document_sets # noqa: F401 -from onyx.mcp_server.resources import indexed_sources # noqa: F401 +from onyx.mcp_server.resources import ( + document_sets, # noqa: F401 + indexed_sources, # noqa: F401 +) diff --git a/backend/onyx/mcp_server/resources/document_sets.py b/backend/onyx/mcp_server/resources/document_sets.py index 093fd3c1a5..1d668ca5cf 100644 --- a/backend/onyx/mcp_server/resources/document_sets.py +++ b/backend/onyx/mcp_server/resources/document_sets.py @@ -5,8 +5,7 @@ from __future__ import annotations import json from onyx.mcp_server.api import mcp_server -from onyx.mcp_server.utils import get_accessible_document_sets -from onyx.mcp_server.utils import require_access_token +from onyx.mcp_server.utils import get_accessible_document_sets, require_access_token from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/mcp_server/resources/indexed_sources.py b/backend/onyx/mcp_server/resources/indexed_sources.py index 36e40f8029..f04f31b771 100644 --- a/backend/onyx/mcp_server/resources/indexed_sources.py +++ b/backend/onyx/mcp_server/resources/indexed_sources.py @@ -5,8 +5,7 @@ from __future__ import annotations import json from onyx.mcp_server.api import mcp_server -from onyx.mcp_server.utils import get_indexed_sources -from onyx.mcp_server.utils import require_access_token +from onyx.mcp_server.utils import get_indexed_sources, require_access_token from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/mcp_server/tools/search.py b/backend/onyx/mcp_server/tools/search.py index 49d5d588f8..cf50bd2fb2 100644 --- a/backend/onyx/mcp_server/tools/search.py +++ b/backend/onyx/mcp_server/tools/search.py @@ -6,28 +6,34 @@ from typing import Any import httpx from fastmcp.server.auth.auth import AccessToken -from pydantic import BaseModel -from pydantic import TypeAdapter -from pydantic import ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from onyx.configs.constants import DocumentSource from onyx.mcp_server.api import mcp_server -from onyx.mcp_server.utils import get_http_client -from onyx.mcp_server.utils import get_indexed_sources -from onyx.mcp_server.utils import require_access_token -from onyx.server.features.search.models import SearchRequest -from onyx.server.features.search.models import SearchResponse -from onyx.server.features.search.models import SearchResult -from onyx.server.features.web_search.models import OpenUrlsToolRequest -from onyx.server.features.web_search.models import OpenUrlsToolResponse -from onyx.server.features.web_search.models import WebSearchToolRequest -from onyx.server.features.web_search.models import WebSearchToolResponse +from onyx.mcp_server.utils import ( + get_http_client, + get_indexed_sources, + require_access_token, +) +from onyx.server.features.search.models import ( + SearchRequest, + SearchResponse, + SearchResult, +) +from onyx.server.features.web_search.models import ( + OpenUrlsToolRequest, + OpenUrlsToolResponse, + WebSearchToolRequest, + WebSearchToolResponse, +) from onyx.server.metrics.mcp_common import MCPToolCallStatus -from onyx.server.metrics.mcp_server import MCPServerToolName -from onyx.server.metrics.mcp_server import record_mcp_search_results -from onyx.server.metrics.mcp_server import record_mcp_search_source -from onyx.server.metrics.mcp_server import record_mcp_server_tool_outcome -from onyx.server.metrics.mcp_server import UNKNOWN_SOURCE_LABEL +from onyx.server.metrics.mcp_server import ( + MCPServerToolName, + record_mcp_search_results, + record_mcp_search_source, + record_mcp_server_tool_outcome, + UNKNOWN_SOURCE_LABEL, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import build_api_server_url_for_http_requests diff --git a/backend/onyx/mcp_server/utils.py b/backend/onyx/mcp_server/utils.py index e3061390b6..dad51bd9d4 100644 --- a/backend/onyx/mcp_server/utils.py +++ b/backend/onyx/mcp_server/utils.py @@ -5,8 +5,7 @@ from __future__ import annotations import httpx from fastmcp.server.auth.auth import AccessToken from fastmcp.server.dependencies import get_access_token -from pydantic import BaseModel -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import build_api_server_url_for_http_requests diff --git a/backend/onyx/mcp_server_main.py b/backend/onyx/mcp_server_main.py index fc7b726b97..b093e9c71d 100644 --- a/backend/onyx/mcp_server_main.py +++ b/backend/onyx/mcp_server_main.py @@ -2,9 +2,11 @@ import uvicorn -from onyx.configs.app_configs import MCP_SERVER_ENABLED -from onyx.configs.app_configs import MCP_SERVER_HOST -from onyx.configs.app_configs import MCP_SERVER_PORT +from onyx.configs.app_configs import ( + MCP_SERVER_ENABLED, + MCP_SERVER_HOST, + MCP_SERVER_PORT, +) from onyx.tracing.setup import setup_tracing from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable diff --git a/backend/onyx/natural_language_processing/constants.py b/backend/onyx/natural_language_processing/constants.py index 0424475b80..a62f090f11 100644 --- a/backend/onyx/natural_language_processing/constants.py +++ b/backend/onyx/natural_language_processing/constants.py @@ -5,8 +5,7 @@ This file contains constants moved from model_server to support the gradual migr of API-based calls to bypass the model server. """ -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType +from shared_configs.enums import EmbeddingProvider, EmbedTextType # Default model names for different providers DEFAULT_OPENAI_MODEL = "text-embedding-3-small" diff --git a/backend/onyx/natural_language_processing/query_embedding_cache.py b/backend/onyx/natural_language_processing/query_embedding_cache.py index 945b82c57d..3338daea1b 100644 --- a/backend/onyx/natural_language_processing/query_embedding_cache.py +++ b/backend/onyx/natural_language_processing/query_embedding_cache.py @@ -20,10 +20,12 @@ import struct from onyx.cache.factory import get_cache_backend from onyx.cache.interface import CACHE_TRANSIENT_ERRORS -from onyx.server.metrics.embedding import observe_query_embedding_cache_lookup -from onyx.server.metrics.embedding import observe_query_embedding_cache_write -from onyx.server.metrics.embedding import QueryEmbeddingCacheLookupOutcome -from onyx.server.metrics.embedding import QueryEmbeddingCacheWriteOutcome +from onyx.server.metrics.embedding import ( + observe_query_embedding_cache_lookup, + observe_query_embedding_cache_write, + QueryEmbeddingCacheLookupOutcome, + QueryEmbeddingCacheWriteOutcome, +) from onyx.utils.logger import setup_logger from shared_configs.enums import EmbeddingProvider from shared_configs.model_server_models import Embedding diff --git a/backend/onyx/natural_language_processing/search_nlp_models.py b/backend/onyx/natural_language_processing/search_nlp_models.py index fc53b56d2c..ee43c92c35 100644 --- a/backend/onyx/natural_language_processing/search_nlp_models.py +++ b/backend/onyx/natural_language_processing/search_nlp_models.py @@ -4,12 +4,10 @@ import os import threading import time from collections.abc import Callable -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, ThreadPoolExecutor from functools import wraps from types import TracebackType -from typing import Any -from typing import cast +from typing import Any, cast import aioboto3 import httpx @@ -19,62 +17,71 @@ from cohere import AsyncClient as CohereAsyncClient from cohere.core.api_error import ApiError from google.oauth2 import service_account from httpx import HTTPError -from requests import JSONDecodeError -from requests import RequestException -from requests import Response -from tenacity import retry -from tenacity import retry_if_exception_type -from tenacity import stop_after_attempt -from tenacity import wait_exponential -from tenacity import wait_fixed -from tenacity import wait_random +from requests import JSONDecodeError, RequestException, Response +from tenacity import ( + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, + wait_fixed, + wait_random, +) -from onyx.configs.app_configs import INDEXING_EMBEDDING_MODEL_NUM_THREADS -from onyx.configs.app_configs import LARGE_CHUNK_RATIO -from onyx.configs.model_configs import BATCH_SIZE_ENCODE_CHUNKS +from onyx.configs.app_configs import ( + INDEXING_EMBEDDING_MODEL_NUM_THREADS, + LARGE_CHUNK_RATIO, +) from onyx.configs.model_configs import ( + BATCH_SIZE_ENCODE_CHUNKS, BATCH_SIZE_ENCODE_CHUNKS_FOR_API_EMBEDDING_SERVICES, ) from onyx.connectors.models import ConnectorStopSignal from onyx.db.models import SearchSettings from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface -from onyx.natural_language_processing.constants import DEFAULT_COHERE_MODEL -from onyx.natural_language_processing.constants import DEFAULT_OPENAI_MODEL -from onyx.natural_language_processing.constants import DEFAULT_VERTEX_MODEL -from onyx.natural_language_processing.constants import DEFAULT_VOYAGE_MODEL -from onyx.natural_language_processing.constants import EmbeddingModelTextType -from onyx.natural_language_processing.exceptions import CohereBillingLimitError -from onyx.natural_language_processing.exceptions import ModelServerRateLimitError -from onyx.natural_language_processing.utils import get_tokenizer -from onyx.natural_language_processing.utils import tokenizer_trim_content -from onyx.server.metrics.embedding import observe_embedding_client -from onyx.server.metrics.embedding import track_embedding_in_progress +from onyx.natural_language_processing.constants import ( + DEFAULT_COHERE_MODEL, + DEFAULT_OPENAI_MODEL, + DEFAULT_VERTEX_MODEL, + DEFAULT_VOYAGE_MODEL, + EmbeddingModelTextType, +) +from onyx.natural_language_processing.exceptions import ( + CohereBillingLimitError, + ModelServerRateLimitError, +) +from onyx.natural_language_processing.utils import get_tokenizer, tokenizer_trim_content +from onyx.server.metrics.embedding import ( + observe_embedding_client, + track_embedding_in_progress, +) from onyx.tracing.flows import LLMFlow from onyx.tracing.llm_utils import traced_llm_call from onyx.utils.logger import setup_logger from onyx.utils.search_nlp_models_utils import pass_aws_key from onyx.utils.text_processing import remove_invalid_unicode_chars from onyx.utils.timing import log_function_time -from shared_configs.configs import API_BASED_EMBEDDING_TIMEOUT -from shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE -from shared_configs.configs import INDEXING_ONLY -from shared_configs.configs import MODEL_SERVER_CONNECT_TIMEOUT -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT -from shared_configs.configs import MODEL_SERVER_READ_TIMEOUT -from shared_configs.configs import OPENAI_EMBEDDING_TIMEOUT -from shared_configs.configs import SKIP_WARM_UP -from shared_configs.configs import VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType -from shared_configs.enums import RerankerProvider -from shared_configs.model_server_models import Embedding -from shared_configs.model_server_models import EmbedRequest -from shared_configs.model_server_models import EmbedResponse -from shared_configs.model_server_models import IntentRequest -from shared_configs.model_server_models import IntentResponse -from shared_configs.model_server_models import RerankRequest -from shared_configs.model_server_models import RerankResponse +from shared_configs.configs import ( + API_BASED_EMBEDDING_TIMEOUT, + DOC_EMBEDDING_CONTEXT_SIZE, + INDEXING_ONLY, + MODEL_SERVER_CONNECT_TIMEOUT, + MODEL_SERVER_HOST, + MODEL_SERVER_PORT, + MODEL_SERVER_READ_TIMEOUT, + OPENAI_EMBEDDING_TIMEOUT, + SKIP_WARM_UP, + VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE, +) +from shared_configs.enums import EmbeddingProvider, EmbedTextType, RerankerProvider +from shared_configs.model_server_models import ( + Embedding, + EmbedRequest, + EmbedResponse, + IntentRequest, + IntentResponse, + RerankRequest, + RerankResponse, +) from shared_configs.utils import batch_list logger = setup_logger() diff --git a/backend/onyx/natural_language_processing/utils.py b/backend/onyx/natural_language_processing/utils.py index 5847b79e76..74631b4e5e 100644 --- a/backend/onyx/natural_language_processing/utils.py +++ b/backend/onyx/natural_language_processing/utils.py @@ -1,10 +1,8 @@ import os -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from copy import copy -from tokenizers import Encoding -from tokenizers import Tokenizer +from tokenizers import Encoding, Tokenizer from onyx.configs.model_configs import DOCUMENT_ENCODER_MODEL from onyx.context.search.models import InferenceChunk diff --git a/backend/onyx/onyxbot/discord/api_client.py b/backend/onyx/onyxbot/discord/api_client.py index a9b4087b03..30b695c32c 100644 --- a/backend/onyx/onyxbot/discord/api_client.py +++ b/backend/onyx/onyxbot/discord/api_client.py @@ -4,12 +4,16 @@ import aiohttp from onyx.chat.models import ChatFullResponse from onyx.onyxbot.discord.constants import API_REQUEST_TIMEOUT -from onyx.onyxbot.discord.exceptions import APIConnectionError -from onyx.onyxbot.discord.exceptions import APIResponseError -from onyx.onyxbot.discord.exceptions import APITimeoutError -from onyx.server.query_and_chat.models import ChatSessionCreationRequest -from onyx.server.query_and_chat.models import MessageOrigin -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.onyxbot.discord.exceptions import ( + APIConnectionError, + APIResponseError, + APITimeoutError, +) +from onyx.server.query_and_chat.models import ( + ChatSessionCreationRequest, + MessageOrigin, + SendMessageRequest, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import build_api_server_url_for_http_requests diff --git a/backend/onyx/onyxbot/discord/cache.py b/backend/onyx/onyxbot/discord/cache.py index b68ee7b476..726083121a 100644 --- a/backend/onyx/onyxbot/discord/cache.py +++ b/backend/onyx/onyxbot/discord/cache.py @@ -3,8 +3,7 @@ import asyncio from typing import NamedTuple -from onyx.db.discord_bot import get_guild_configs -from onyx.db.discord_bot import get_or_create_discord_service_api_key +from onyx.db.discord_bot import get_guild_configs, get_or_create_discord_service_api_key from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.engine.tenant_utils import get_all_tenant_ids from onyx.onyxbot.discord.exceptions import CacheError diff --git a/backend/onyx/onyxbot/discord/client.py b/backend/onyx/onyxbot/discord/client.py index 67fbe00759..9a229f76f4 100644 --- a/backend/onyx/onyxbot/discord/client.py +++ b/backend/onyx/onyxbot/discord/client.py @@ -10,11 +10,12 @@ from onyx.configs.app_configs import DISCORD_BOT_INVOKE_CHAR from onyx.onyxbot.discord.api_client import OnyxAPIClient from onyx.onyxbot.discord.cache import DiscordCacheManager from onyx.onyxbot.discord.constants import CACHE_REFRESH_INTERVAL -from onyx.onyxbot.discord.handle_commands import handle_dm -from onyx.onyxbot.discord.handle_commands import handle_registration_command -from onyx.onyxbot.discord.handle_commands import handle_sync_channels_command -from onyx.onyxbot.discord.handle_message import process_chat_message -from onyx.onyxbot.discord.handle_message import should_respond +from onyx.onyxbot.discord.handle_commands import ( + handle_dm, + handle_registration_command, + handle_sync_channels_command, +) +from onyx.onyxbot.discord.handle_message import process_chat_message, should_respond from onyx.onyxbot.discord.utils import get_bot_token from onyx.utils.logger import setup_logger diff --git a/backend/onyx/onyxbot/discord/handle_commands.py b/backend/onyx/onyxbot/discord/handle_commands.py index aacd2bab8e..45095914ea 100644 --- a/backend/onyx/onyxbot/discord/handle_commands.py +++ b/backend/onyx/onyxbot/discord/handle_commands.py @@ -1,25 +1,24 @@ """Discord bot command handlers for registration and channel sync.""" import asyncio -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import discord from onyx.configs.app_configs import DISCORD_BOT_INVOKE_CHAR from onyx.configs.constants import ONYX_DISCORD_URL -from onyx.db.discord_bot import bulk_create_channel_configs -from onyx.db.discord_bot import get_guild_config_by_discord_id -from onyx.db.discord_bot import get_guild_config_by_internal_id -from onyx.db.discord_bot import get_guild_config_by_registration_key -from onyx.db.discord_bot import sync_channel_configs +from onyx.db.discord_bot import ( + bulk_create_channel_configs, + get_guild_config_by_discord_id, + get_guild_config_by_internal_id, + get_guild_config_by_registration_key, + sync_channel_configs, +) from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.utils import DiscordChannelView from onyx.onyxbot.discord.cache import DiscordCacheManager -from onyx.onyxbot.discord.constants import REGISTER_COMMAND -from onyx.onyxbot.discord.constants import SYNC_CHANNELS_COMMAND -from onyx.onyxbot.discord.exceptions import RegistrationError -from onyx.onyxbot.discord.exceptions import SyncChannelsError +from onyx.onyxbot.discord.constants import REGISTER_COMMAND, SYNC_CHANNELS_COMMAND +from onyx.onyxbot.discord.exceptions import RegistrationError, SyncChannelsError from onyx.server.manage.discord_bot.utils import parse_discord_registration_key from onyx.utils.logger import setup_logger from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/onyx/onyxbot/discord/handle_message.py b/backend/onyx/onyxbot/discord/handle_message.py index 280bfb9033..6e1caf3b08 100644 --- a/backend/onyx/onyxbot/discord/handle_message.py +++ b/backend/onyx/onyxbot/discord/handle_message.py @@ -6,15 +6,18 @@ import discord from pydantic import BaseModel from onyx.chat.models import ChatFullResponse -from onyx.db.discord_bot import get_channel_config_by_discord_ids -from onyx.db.discord_bot import get_guild_config_by_discord_id +from onyx.db.discord_bot import ( + get_channel_config_by_discord_ids, + get_guild_config_by_discord_id, +) from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.models import DiscordChannelConfig -from onyx.db.models import DiscordGuildConfig +from onyx.db.models import DiscordChannelConfig, DiscordGuildConfig from onyx.onyxbot.discord.api_client import OnyxAPIClient -from onyx.onyxbot.discord.constants import MAX_CONTEXT_MESSAGES -from onyx.onyxbot.discord.constants import MAX_MESSAGE_LENGTH -from onyx.onyxbot.discord.constants import THINKING_EMOJI +from onyx.onyxbot.discord.constants import ( + MAX_CONTEXT_MESSAGES, + MAX_MESSAGE_LENGTH, + THINKING_EMOJI, +) from onyx.onyxbot.discord.exceptions import APIError from onyx.utils.logger import setup_logger diff --git a/backend/onyx/onyxbot/discord/utils.py b/backend/onyx/onyxbot/discord/utils.py index 33e3a9cad9..80301fdbad 100644 --- a/backend/onyx/onyxbot/discord/utils.py +++ b/backend/onyx/onyxbot/discord/utils.py @@ -3,8 +3,7 @@ from onyx.db.discord_bot import get_discord_bot_config from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.utils.logger import setup_logger from onyx.utils.sensitive import SensitiveValue -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA logger = setup_logger() diff --git a/backend/onyx/onyxbot/slack/blocks.py b/backend/onyx/onyxbot/slack/blocks.py index 51cc228974..2bcaf35fb2 100644 --- a/backend/onyx/onyxbot/slack/blocks.py +++ b/backend/onyx/onyxbot/slack/blocks.py @@ -3,45 +3,52 @@ from typing import cast import pytz import timeago -from slack_sdk.models.blocks import ActionsBlock -from slack_sdk.models.blocks import Block -from slack_sdk.models.blocks import ButtonElement -from slack_sdk.models.blocks import ContextBlock -from slack_sdk.models.blocks import DividerBlock -from slack_sdk.models.blocks import HeaderBlock -from slack_sdk.models.blocks import Option -from slack_sdk.models.blocks import RadioButtonsElement -from slack_sdk.models.blocks import SectionBlock +from slack_sdk.models.blocks import ( + ActionsBlock, + Block, + ButtonElement, + ContextBlock, + DividerBlock, + HeaderBlock, + Option, + RadioButtonsElement, + SectionBlock, +) from slack_sdk.models.blocks.basic_components import MarkdownTextObject from slack_sdk.models.blocks.block_elements import ImageElement from onyx.chat.models import ChatBasicResponse from onyx.configs.app_configs import WEB_DOMAIN -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import SearchFeedbackType +from onyx.configs.constants import DocumentSource, SearchFeedbackType from onyx.configs.onyxbot_configs import ONYX_BOT_NUM_DOCS_TO_DISPLAY from onyx.context.search.models import SearchDoc from onyx.db.chat import get_chat_session_by_message_id from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import ChannelConfig -from onyx.onyxbot.slack.constants import CONTINUE_IN_WEB_UI_ACTION_ID -from onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID -from onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_ACTION_ID -from onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_RESOLVED_ACTION_ID -from onyx.onyxbot.slack.constants import IMMEDIATE_RESOLVED_BUTTON_ACTION_ID -from onyx.onyxbot.slack.constants import KEEP_TO_YOURSELF_ACTION_ID -from onyx.onyxbot.slack.constants import LIKE_BLOCK_ACTION_ID -from onyx.onyxbot.slack.constants import SHOW_EVERYONE_ACTION_ID +from onyx.onyxbot.slack.constants import ( + CONTINUE_IN_WEB_UI_ACTION_ID, + DISLIKE_BLOCK_ACTION_ID, + FOLLOWUP_BUTTON_ACTION_ID, + FOLLOWUP_BUTTON_RESOLVED_ACTION_ID, + IMMEDIATE_RESOLVED_BUTTON_ACTION_ID, + KEEP_TO_YOURSELF_ACTION_ID, + LIKE_BLOCK_ACTION_ID, + SHOW_EVERYONE_ACTION_ID, +) from onyx.onyxbot.slack.formatting import format_slack_message from onyx.onyxbot.slack.icons import source_to_github_img_link -from onyx.onyxbot.slack.models import ActionValuesEphemeralMessage -from onyx.onyxbot.slack.models import ActionValuesEphemeralMessageChannelConfig -from onyx.onyxbot.slack.models import ActionValuesEphemeralMessageMessageInfo -from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.onyxbot.slack.utils import build_continue_in_web_ui_id -from onyx.onyxbot.slack.utils import build_feedback_id -from onyx.onyxbot.slack.utils import build_publish_ephemeral_message_id -from onyx.onyxbot.slack.utils import remove_slack_text_interactions +from onyx.onyxbot.slack.models import ( + ActionValuesEphemeralMessage, + ActionValuesEphemeralMessageChannelConfig, + ActionValuesEphemeralMessageMessageInfo, + SlackMessageInfo, +) +from onyx.onyxbot.slack.utils import ( + build_continue_in_web_ui_id, + build_feedback_id, + build_publish_ephemeral_message_id, + remove_slack_text_interactions, +) from onyx.utils.datetime import datetime_to_utc from onyx.utils.text_processing import decode_escapes diff --git a/backend/onyx/onyxbot/slack/config.py b/backend/onyx/onyxbot/slack/config.py index 80107ae1f0..1f9bad4067 100644 --- a/backend/onyx/onyxbot/slack/config.py +++ b/backend/onyx/onyxbot/slack/config.py @@ -5,8 +5,8 @@ from sqlalchemy.orm import Session from onyx.db.models import SlackChannelConfig from onyx.db.slack_channel_config import ( fetch_slack_channel_config_for_channel_or_default, + fetch_slack_channel_configs, ) -from onyx.db.slack_channel_config import fetch_slack_channel_configs VALID_SLACK_FILTERS = [ "answerable_prefilter", diff --git a/backend/onyx/onyxbot/slack/formatting.py b/backend/onyx/onyxbot/slack/formatting.py index 6107106c2a..f5023e38eb 100644 --- a/backend/onyx/onyxbot/slack/formatting.py +++ b/backend/onyx/onyxbot/slack/formatting.py @@ -2,8 +2,7 @@ import re from collections.abc import Callable from typing import Any -from mistune import create_markdown -from mistune import HTMLRenderer +from mistune import create_markdown, HTMLRenderer # Tags that should be replaced with a newline (line-break and block-level elements) _HTML_NEWLINE_TAG_PATTERN = re.compile( diff --git a/backend/onyx/onyxbot/slack/handlers/handle_buttons.py b/backend/onyx/onyxbot/slack/handlers/handle_buttons.py index 745f6ded41..c4fe801865 100644 --- a/backend/onyx/onyxbot/slack/handlers/handle_buttons.py +++ b/backend/onyx/onyxbot/slack/handlers/handle_buttons.py @@ -1,6 +1,5 @@ import json -from typing import Any -from typing import cast +from typing import Any, cast from slack_sdk import WebClient from slack_sdk.models.blocks import SectionBlock @@ -10,43 +9,45 @@ from slack_sdk.webhook import WebhookClient from onyx.chat.models import ChatBasicResponse from onyx.chat.process_message import remove_answer_citations -from onyx.configs.constants import MessageType -from onyx.configs.constants import SearchFeedbackType +from onyx.configs.constants import MessageType, SearchFeedbackType from onyx.configs.onyxbot_configs import ONYX_BOT_FOLLOWUP_EMOJI from onyx.connectors.slack.utils import expert_info_from_slack_id -from onyx.context.search.models import SavedSearchDoc -from onyx.context.search.models import SearchDoc -from onyx.db.chat import get_chat_message -from onyx.db.chat import translate_db_message_to_chat_message_detail +from onyx.context.search.models import SavedSearchDoc, SearchDoc +from onyx.db.chat import get_chat_message, translate_db_message_to_chat_message_detail from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.feedback import create_chat_message_feedback -from onyx.db.feedback import create_doc_retrieval_feedback +from onyx.db.feedback import create_chat_message_feedback, create_doc_retrieval_feedback from onyx.db.users import get_user_by_email -from onyx.onyxbot.slack.blocks import build_follow_up_resolved_blocks -from onyx.onyxbot.slack.blocks import build_slack_response_blocks -from onyx.onyxbot.slack.blocks import get_document_feedback_blocks +from onyx.onyxbot.slack.blocks import ( + build_follow_up_resolved_blocks, + build_slack_response_blocks, + get_document_feedback_blocks, +) from onyx.onyxbot.slack.config import get_slack_channel_config_for_bot_and_channel -from onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID -from onyx.onyxbot.slack.constants import FeedbackVisibility -from onyx.onyxbot.slack.constants import KEEP_TO_YOURSELF_ACTION_ID -from onyx.onyxbot.slack.constants import LIKE_BLOCK_ACTION_ID -from onyx.onyxbot.slack.constants import SHOW_EVERYONE_ACTION_ID -from onyx.onyxbot.slack.constants import VIEW_DOC_FEEDBACK_ID +from onyx.onyxbot.slack.constants import ( + DISLIKE_BLOCK_ACTION_ID, + FeedbackVisibility, + KEEP_TO_YOURSELF_ACTION_ID, + LIKE_BLOCK_ACTION_ID, + SHOW_EVERYONE_ACTION_ID, + VIEW_DOC_FEEDBACK_ID, +) from onyx.onyxbot.slack.handlers.handle_message import ( remove_scheduled_feedback_reminder, ) from onyx.onyxbot.slack.handlers.handle_regular_answer import handle_regular_answer from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.onyxbot.slack.utils import build_feedback_id -from onyx.onyxbot.slack.utils import decompose_action_id -from onyx.onyxbot.slack.utils import fetch_group_ids_from_names -from onyx.onyxbot.slack.utils import fetch_slack_user_ids_from_emails -from onyx.onyxbot.slack.utils import get_channel_name_from_id -from onyx.onyxbot.slack.utils import get_feedback_visibility -from onyx.onyxbot.slack.utils import read_slack_thread -from onyx.onyxbot.slack.utils import respond_in_thread_or_channel -from onyx.onyxbot.slack.utils import TenantSocketModeClient -from onyx.onyxbot.slack.utils import update_emote_react +from onyx.onyxbot.slack.utils import ( + build_feedback_id, + decompose_action_id, + fetch_group_ids_from_names, + fetch_slack_user_ids_from_emails, + get_channel_name_from_id, + get_feedback_visibility, + read_slack_thread, + respond_in_thread_or_channel, + TenantSocketModeClient, + update_emote_react, +) from onyx.server.query_and_chat.models import ChatMessageDetail from onyx.server.query_and_chat.streaming_models import CitationInfo from onyx.utils.logger import setup_logger diff --git a/backend/onyx/onyxbot/slack/handlers/handle_message.py b/backend/onyx/onyxbot/slack/handlers/handle_message.py index 874f8e21d0..b9e42789c5 100644 --- a/backend/onyx/onyxbot/slack/handlers/handle_message.py +++ b/backend/onyx/onyxbot/slack/handlers/handle_message.py @@ -4,26 +4,28 @@ from slack_sdk import WebClient from slack_sdk.errors import SlackApiError from sqlalchemy.orm import Session -from onyx.configs.onyxbot_configs import ONYX_BOT_FEEDBACK_REMINDER -from onyx.configs.onyxbot_configs import ONYX_BOT_REACT_EMOJI +from onyx.configs.onyxbot_configs import ( + ONYX_BOT_FEEDBACK_REMINDER, + ONYX_BOT_REACT_EMOJI, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import AccountType -from onyx.db.models import ChannelConfig -from onyx.db.models import SlackChannelConfig +from onyx.db.models import ChannelConfig, SlackChannelConfig from onyx.db.user_preferences import activate_user -from onyx.db.users import add_slack_user_if_not_exists -from onyx.db.users import get_user_by_email +from onyx.db.users import add_slack_user_if_not_exists, get_user_by_email from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.onyxbot.slack.blocks import get_feedback_reminder_blocks from onyx.onyxbot.slack.handlers.handle_regular_answer import handle_regular_answer from onyx.onyxbot.slack.handlers.handle_standard_answers import handle_standard_answers from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.onyxbot.slack.utils import fetch_slack_user_ids_from_emails -from onyx.onyxbot.slack.utils import fetch_user_ids_from_groups -from onyx.onyxbot.slack.utils import respond_in_thread_or_channel -from onyx.onyxbot.slack.utils import slack_usage_report -from onyx.onyxbot.slack.utils import update_emote_react +from onyx.onyxbot.slack.utils import ( + fetch_slack_user_ids_from_emails, + fetch_user_ids_from_groups, + respond_in_thread_or_channel, + slack_usage_report, + update_emote_react, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop from shared_configs.configs import SLACK_CHANNEL_ID diff --git a/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py b/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py index 6bfebb5a2f..3647a436a2 100644 --- a/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py +++ b/backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py @@ -1,40 +1,39 @@ import functools from collections.abc import Callable -from typing import Any -from typing import Optional -from typing import TypeVar +from typing import Any, Optional, TypeVar from slack_sdk import WebClient from sqlalchemy.orm import Session from onyx.auth.users import get_anonymous_user from onyx.chat.models import ChatBasicResponse -from onyx.chat.process_message import gather_stream -from onyx.chat.process_message import handle_stream_message_objects -from onyx.configs.constants import DEFAULT_PERSONA_ID -from onyx.configs.constants import MessageType -from onyx.configs.onyxbot_configs import ONYX_BOT_DISABLE_DOCS_ONLY_ANSWER -from onyx.configs.onyxbot_configs import ONYX_BOT_DISPLAY_ERROR_MSGS -from onyx.configs.onyxbot_configs import ONYX_BOT_NUM_RETRIES -from onyx.configs.onyxbot_configs import ONYX_BOT_REACT_EMOJI -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import Tag -from onyx.db.models import SlackChannelConfig -from onyx.db.models import User +from onyx.chat.process_message import gather_stream, handle_stream_message_objects +from onyx.configs.constants import DEFAULT_PERSONA_ID, MessageType +from onyx.configs.onyxbot_configs import ( + ONYX_BOT_DISABLE_DOCS_ONLY_ANSWER, + ONYX_BOT_DISPLAY_ERROR_MSGS, + ONYX_BOT_NUM_RETRIES, + ONYX_BOT_REACT_EMOJI, +) +from onyx.context.search.models import BaseFilters, Tag +from onyx.db.models import SlackChannelConfig, User from onyx.db.persona import get_persona_by_id from onyx.db.users import get_user_by_email from onyx.onyxbot.slack.blocks import build_slack_response_blocks from onyx.onyxbot.slack.constants import SLACK_CHANNEL_REF_PATTERN -from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.onyxbot.slack.models import ThreadMessage -from onyx.onyxbot.slack.utils import get_channel_from_id -from onyx.onyxbot.slack.utils import get_channel_name_from_id -from onyx.onyxbot.slack.utils import respond_in_thread_or_channel -from onyx.onyxbot.slack.utils import SlackRateLimiter -from onyx.onyxbot.slack.utils import update_emote_react -from onyx.server.query_and_chat.models import ChatSessionCreationRequest -from onyx.server.query_and_chat.models import MessageOrigin -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.onyxbot.slack.models import SlackMessageInfo, ThreadMessage +from onyx.onyxbot.slack.utils import ( + get_channel_from_id, + get_channel_name_from_id, + respond_in_thread_or_channel, + SlackRateLimiter, + update_emote_react, +) +from onyx.server.query_and_chat.models import ( + ChatSessionCreationRequest, + MessageOrigin, + SendMessageRequest, +) from onyx.utils.logger import OnyxLoggingAdapter from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/onyxbot/slack/handlers/handle_standard_answers.py b/backend/onyx/onyxbot/slack/handlers/handle_standard_answers.py index 1f72f59fe5..5f07222787 100644 --- a/backend/onyx/onyxbot/slack/handlers/handle_standard_answers.py +++ b/backend/onyx/onyxbot/slack/handlers/handle_standard_answers.py @@ -3,8 +3,7 @@ from sqlalchemy.orm import Session from onyx.db.models import SlackChannelConfig from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.utils.logger import OnyxLoggingAdapter -from onyx.utils.logger import setup_logger +from onyx.utils.logger import OnyxLoggingAdapter, setup_logger from onyx.utils.variable_functionality import fetch_versioned_implementation logger = setup_logger() diff --git a/backend/onyx/onyxbot/slack/listener.py b/backend/onyx/onyxbot/slack/listener.py index ff7056c5ef..f804c4ab1b 100644 --- a/backend/onyx/onyxbot/slack/listener.py +++ b/backend/onyx/onyxbot/slack/listener.py @@ -7,98 +7,107 @@ from collections.abc import Callable from contextvars import Token from threading import Event from types import FrameType -from typing import Any -from typing import cast -from typing import Dict +from typing import Any, cast, Dict import psycopg2.errors -from prometheus_client import Gauge -from prometheus_client import start_http_server +from prometheus_client import Gauge, start_http_server from redis.exceptions import LockNotOwnedError from redis.lock import Lock from redis.lock import Lock as RedisLock from slack_sdk import WebClient from slack_sdk.errors import SlackApiError -from slack_sdk.http_retry import ConnectionErrorRetryHandler -from slack_sdk.http_retry import RateLimitErrorRetryHandler -from slack_sdk.http_retry import RetryHandler +from slack_sdk.http_retry import ( + ConnectionErrorRetryHandler, + RateLimitErrorRetryHandler, + RetryHandler, +) from slack_sdk.socket_mode.request import SocketModeRequest from slack_sdk.socket_mode.response import SocketModeResponse from sqlalchemy.exc import ProgrammingError from sqlalchemy.orm import Session -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import POD_NAME -from onyx.configs.app_configs import POD_NAMESPACE -from onyx.configs.constants import MessageType -from onyx.configs.constants import OnyxRedisLocks +from onyx.configs.app_configs import DEV_MODE, POD_NAME, POD_NAMESPACE +from onyx.configs.constants import MessageType, OnyxRedisLocks from onyx.configs.onyxbot_configs import NOTIFY_SLACKBOT_NO_ANSWER from onyx.connectors.slack.utils import expert_info_from_slack_id -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import ( + get_session_with_current_tenant, + get_session_with_tenant, + SqlEngine, +) from onyx.db.engine.tenant_utils import get_all_tenant_ids from onyx.db.models import SlackBot from onyx.db.search_settings import get_current_search_settings -from onyx.db.slack_bot import fetch_slack_bot -from onyx.db.slack_bot import fetch_slack_bots +from onyx.db.slack_bot import fetch_slack_bot, fetch_slack_bots from onyx.key_value_store.interface import KvKeyNotFoundError -from onyx.natural_language_processing.search_nlp_models import EmbeddingModel -from onyx.natural_language_processing.search_nlp_models import warm_up_bi_encoder -from onyx.onyxbot.slack.config import get_slack_channel_config_for_bot_and_channel -from onyx.onyxbot.slack.config import MAX_TENANTS_PER_POD -from onyx.onyxbot.slack.config import TENANT_ACQUISITION_INTERVAL -from onyx.onyxbot.slack.config import TENANT_HEARTBEAT_EXPIRATION -from onyx.onyxbot.slack.config import TENANT_HEARTBEAT_INTERVAL -from onyx.onyxbot.slack.config import TENANT_LOCK_EXPIRATION -from onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID -from onyx.onyxbot.slack.constants import FEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID -from onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_ACTION_ID -from onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_RESOLVED_ACTION_ID -from onyx.onyxbot.slack.constants import GENERATE_ANSWER_BUTTON_ACTION_ID -from onyx.onyxbot.slack.constants import IMMEDIATE_RESOLVED_BUTTON_ACTION_ID -from onyx.onyxbot.slack.constants import KEEP_TO_YOURSELF_ACTION_ID -from onyx.onyxbot.slack.constants import LIKE_BLOCK_ACTION_ID -from onyx.onyxbot.slack.constants import SHOW_EVERYONE_ACTION_ID -from onyx.onyxbot.slack.constants import VIEW_DOC_FEEDBACK_ID -from onyx.onyxbot.slack.handlers.handle_buttons import handle_doc_feedback_button -from onyx.onyxbot.slack.handlers.handle_buttons import handle_followup_button -from onyx.onyxbot.slack.handlers.handle_buttons import handle_followup_resolved_button -from onyx.onyxbot.slack.handlers.handle_buttons import handle_generate_answer_button +from onyx.natural_language_processing.search_nlp_models import ( + EmbeddingModel, + warm_up_bi_encoder, +) +from onyx.onyxbot.slack.config import ( + get_slack_channel_config_for_bot_and_channel, + MAX_TENANTS_PER_POD, + TENANT_ACQUISITION_INTERVAL, + TENANT_HEARTBEAT_EXPIRATION, + TENANT_HEARTBEAT_INTERVAL, + TENANT_LOCK_EXPIRATION, +) +from onyx.onyxbot.slack.constants import ( + DISLIKE_BLOCK_ACTION_ID, + FEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID, + FOLLOWUP_BUTTON_ACTION_ID, + FOLLOWUP_BUTTON_RESOLVED_ACTION_ID, + GENERATE_ANSWER_BUTTON_ACTION_ID, + IMMEDIATE_RESOLVED_BUTTON_ACTION_ID, + KEEP_TO_YOURSELF_ACTION_ID, + LIKE_BLOCK_ACTION_ID, + SHOW_EVERYONE_ACTION_ID, + VIEW_DOC_FEEDBACK_ID, +) from onyx.onyxbot.slack.handlers.handle_buttons import ( + handle_doc_feedback_button, + handle_followup_button, + handle_followup_resolved_button, + handle_generate_answer_button, handle_publish_ephemeral_message_button, + handle_slack_feedback, ) -from onyx.onyxbot.slack.handlers.handle_buttons import handle_slack_feedback -from onyx.onyxbot.slack.handlers.handle_message import handle_message from onyx.onyxbot.slack.handlers.handle_message import ( + handle_message, remove_scheduled_feedback_reminder, + schedule_feedback_reminder, +) +from onyx.onyxbot.slack.models import SlackContext, SlackMessageInfo, ThreadMessage +from onyx.onyxbot.slack.utils import ( + check_message_limit, + decompose_action_id, + get_channel_name_from_id, + get_channel_type_from_id, + get_onyx_bot_auth_ids, + read_slack_thread, + remove_onyx_bot_tag, + respond_in_thread_or_channel, + TenantSocketModeClient, ) -from onyx.onyxbot.slack.handlers.handle_message import schedule_feedback_reminder -from onyx.onyxbot.slack.models import SlackContext -from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.onyxbot.slack.models import ThreadMessage -from onyx.onyxbot.slack.utils import check_message_limit -from onyx.onyxbot.slack.utils import decompose_action_id -from onyx.onyxbot.slack.utils import get_channel_name_from_id -from onyx.onyxbot.slack.utils import get_channel_type_from_id -from onyx.onyxbot.slack.utils import get_onyx_bot_auth_ids -from onyx.onyxbot.slack.utils import read_slack_thread -from onyx.onyxbot.slack.utils import remove_onyx_bot_tag -from onyx.onyxbot.slack.utils import respond_in_thread_or_channel -from onyx.onyxbot.slack.utils import TenantSocketModeClient from onyx.redis.redis_pool import get_redis_client from onyx.server.manage.models import SlackBotTokens from onyx.tracing.setup import setup_tracing from onyx.utils.logger import setup_logger -from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop -from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable -from shared_configs.configs import DISALLOWED_SLACK_BOT_TENANT_LIST -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import SLACK_CHANNEL_ID -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from onyx.utils.variable_functionality import ( + fetch_ee_implementation_or_noop, + set_is_ee_based_on_env_variable, +) +from shared_configs.configs import ( + DISALLOWED_SLACK_BOT_TENANT_LIST, + MODEL_SERVER_HOST, + MODEL_SERVER_PORT, + POSTGRES_DEFAULT_SCHEMA, + SLACK_CHANNEL_ID, +) +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) logger = setup_logger() diff --git a/backend/onyx/onyxbot/slack/utils.py b/backend/onyx/onyxbot/slack/utils.py index 7e61ec89ed..26a451f537 100644 --- a/backend/onyx/onyxbot/slack/utils.py +++ b/backend/onyx/onyxbot/slack/utils.py @@ -6,35 +6,32 @@ import time import uuid from collections.abc import Generator from contextlib import contextmanager -from typing import Any -from typing import cast +from typing import Any, cast from slack_sdk import WebClient from slack_sdk.errors import SlackApiError -from slack_sdk.models.blocks import Block -from slack_sdk.models.blocks import SectionBlock +from slack_sdk.models.blocks import Block, SectionBlock from slack_sdk.models.metadata import Metadata from slack_sdk.socket_mode import SocketModeClient from onyx.configs.app_configs import DISABLE_TELEMETRY -from onyx.configs.constants import ID_SEPARATOR -from onyx.configs.constants import MessageType -from onyx.configs.onyxbot_configs import ONYX_BOT_FEEDBACK_VISIBILITY -from onyx.configs.onyxbot_configs import ONYX_BOT_MAX_QPM -from onyx.configs.onyxbot_configs import ONYX_BOT_MAX_WAIT_TIME -from onyx.configs.onyxbot_configs import ONYX_BOT_NUM_RETRIES -from onyx.configs.onyxbot_configs import ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD -from onyx.configs.onyxbot_configs import ONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS +from onyx.configs.constants import ID_SEPARATOR, MessageType +from onyx.configs.onyxbot_configs import ( + ONYX_BOT_FEEDBACK_VISIBILITY, + ONYX_BOT_MAX_QPM, + ONYX_BOT_MAX_WAIT_TIME, + ONYX_BOT_NUM_RETRIES, + ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD, + ONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS, +) from onyx.connectors.slack.utils import SlackTextCleaner from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.users import get_user_by_email from onyx.onyxbot.slack.constants import FeedbackVisibility -from onyx.onyxbot.slack.models import ChannelType -from onyx.onyxbot.slack.models import ThreadMessage +from onyx.onyxbot.slack.models import ChannelType, ThreadMessage from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType +from onyx.utils.telemetry import optional_telemetry, RecordType from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR logger = setup_logger() diff --git a/backend/onyx/prompts/coding_agent/coding_agent.py b/backend/onyx/prompts/coding_agent/coding_agent.py index a369f0673a..80417123b8 100644 --- a/backend/onyx/prompts/coding_agent/coding_agent.py +++ b/backend/onyx/prompts/coding_agent/coding_agent.py @@ -1,5 +1,4 @@ -from onyx.coding_agent.mock_tools import BASH_TOOL_NAME -from onyx.coding_agent.mock_tools import GENERATE_ANSWER_TOOL_NAME +from onyx.coding_agent.mock_tools import BASH_TOOL_NAME, GENERATE_ANSWER_TOOL_NAME from onyx.deep_research.dr_mock_tools import THINK_TOOL_NAME MAX_CODING_AGENT_CYCLES = 20 diff --git a/backend/onyx/prompts/deep_research/orchestration_layer.py b/backend/onyx/prompts/deep_research/orchestration_layer.py index 183d907c94..58f27f18f9 100644 --- a/backend/onyx/prompts/deep_research/orchestration_layer.py +++ b/backend/onyx/prompts/deep_research/orchestration_layer.py @@ -1,7 +1,9 @@ -from onyx.prompts.deep_research.dr_tool_prompts import GENERATE_PLAN_TOOL_NAME -from onyx.prompts.deep_research.dr_tool_prompts import GENERATE_REPORT_TOOL_NAME -from onyx.prompts.deep_research.dr_tool_prompts import RESEARCH_AGENT_TOOL_NAME -from onyx.prompts.deep_research.dr_tool_prompts import THINK_TOOL_NAME +from onyx.prompts.deep_research.dr_tool_prompts import ( + GENERATE_PLAN_TOOL_NAME, + GENERATE_REPORT_TOOL_NAME, + RESEARCH_AGENT_TOOL_NAME, + THINK_TOOL_NAME, +) # ruff: noqa: E501, W605 start CLARIFICATION_PROMPT = f""" diff --git a/backend/onyx/prompts/deep_research/research_agent.py b/backend/onyx/prompts/deep_research/research_agent.py index 78c2324991..7cb3fedcd2 100644 --- a/backend/onyx/prompts/deep_research/research_agent.py +++ b/backend/onyx/prompts/deep_research/research_agent.py @@ -1,5 +1,7 @@ -from onyx.prompts.deep_research.dr_tool_prompts import GENERATE_REPORT_TOOL_NAME -from onyx.prompts.deep_research.dr_tool_prompts import THINK_TOOL_NAME +from onyx.prompts.deep_research.dr_tool_prompts import ( + GENERATE_REPORT_TOOL_NAME, + THINK_TOOL_NAME, +) MAX_RESEARCH_CYCLES = 8 diff --git a/backend/onyx/prompts/prompt_utils.py b/backend/onyx/prompts/prompt_utils.py index 8a8a41f8b3..c67d09497b 100644 --- a/backend/onyx/prompts/prompt_utils.py +++ b/backend/onyx/prompts/prompt_utils.py @@ -6,15 +6,16 @@ from langchain_core.messages import BaseMessage from onyx.auth.oauth_claims_capture import IDP_PLACEHOLDER_KEYS from onyx.configs.constants import DocumentSource -from onyx.prompts.chat_prompts import ADDITIONAL_INFO -from onyx.prompts.chat_prompts import CITATION_GUIDANCE_REPLACEMENT_PAT -from onyx.prompts.chat_prompts import COMPANY_DESCRIPTION_BLOCK -from onyx.prompts.chat_prompts import COMPANY_NAME_BLOCK -from onyx.prompts.chat_prompts import DATETIME_REPLACEMENT_PAT -from onyx.prompts.chat_prompts import REMINDER_TAG_REPLACEMENT_PAT -from onyx.prompts.chat_prompts import REQUIRE_CITATION_GUIDANCE -from onyx.prompts.constants import CODE_BLOCK_PAT -from onyx.prompts.constants import REMINDER_TAG_DESCRIPTION +from onyx.prompts.chat_prompts import ( + ADDITIONAL_INFO, + CITATION_GUIDANCE_REPLACEMENT_PAT, + COMPANY_DESCRIPTION_BLOCK, + COMPANY_NAME_BLOCK, + DATETIME_REPLACEMENT_PAT, + REMINDER_TAG_REPLACEMENT_PAT, + REQUIRE_CITATION_GUIDANCE, +) +from onyx.prompts.constants import CODE_BLOCK_PAT, REMINDER_TAG_DESCRIPTION from onyx.server.settings.store import load_settings from onyx.utils.logger import setup_logger diff --git a/backend/onyx/redis/lock_context.py b/backend/onyx/redis/lock_context.py index 6b0b8c7058..6201b3baf8 100644 --- a/backend/onyx/redis/lock_context.py +++ b/backend/onyx/redis/lock_context.py @@ -1,8 +1,7 @@ import time from collections.abc import Generator from contextlib import contextmanager -from logging import Logger -from logging import LoggerAdapter +from logging import Logger, LoggerAdapter from redis.lock import Lock diff --git a/backend/onyx/redis/redis_connector_delete.py b/backend/onyx/redis/redis_connector_delete.py index 1b9d2a3fa1..ce6eb1865e 100644 --- a/backend/onyx/redis/redis_connector_delete.py +++ b/backend/onyx/redis/redis_connector_delete.py @@ -9,11 +9,13 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.configs.app_configs import DB_YIELD_PER_DEFAULT -from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants +from onyx.configs.constants import ( + CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, +) from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.db.document import construct_document_id_select_for_connector_credential_pair from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/onyx/redis/redis_connector_doc_perm_sync.py b/backend/onyx/redis/redis_connector_doc_perm_sync.py index 70baf3dbd3..918f1416ab 100644 --- a/backend/onyx/redis/redis_connector_doc_perm_sync.py +++ b/backend/onyx/redis/redis_connector_doc_perm_sync.py @@ -1,18 +1,17 @@ import time from datetime import datetime from logging import Logger -from typing import Any -from typing import cast -from typing import NamedTuple +from typing import Any, cast, NamedTuple from pydantic import BaseModel from redis.lock import Lock as RedisLock -from onyx.access.models import DocExternalAccess -from onyx.access.models import ElementExternalAccess -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT -from onyx.configs.constants import OnyxRedisConstants +from onyx.access.models import DocExternalAccess, ElementExternalAccess +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT, + OnyxRedisConstants, +) from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.server.metrics.perm_sync_metrics import ( diff --git a/backend/onyx/redis/redis_connector_prune.py b/backend/onyx/redis/redis_connector_prune.py index 9fe46e9460..190833002c 100644 --- a/backend/onyx/redis/redis_connector_prune.py +++ b/backend/onyx/redis/redis_connector_prune.py @@ -9,12 +9,14 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.configs.app_configs import PRUNE_FAILURE_BACKOFF_SECONDS -from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import CELERY_PRUNING_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants +from onyx.configs.constants import ( + CELERY_GENERIC_BEAT_LOCK_TIMEOUT, + CELERY_PRUNING_LOCK_TIMEOUT, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, +) from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/onyx/redis/redis_connector_utils.py b/backend/onyx/redis/redis_connector_utils.py index 6e8e6aa15b..d407edc077 100644 --- a/backend/onyx/redis/redis_connector_utils.py +++ b/backend/onyx/redis/redis_connector_utils.py @@ -1,8 +1,7 @@ from sqlalchemy.orm import Session from onyx.db.connector_credential_pair import get_connector_credential_pair -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import TaskStatus +from onyx.db.enums import ConnectorCredentialPairStatus, TaskStatus from onyx.db.models import TaskQueueState from onyx.redis.redis_connector import RedisConnector from onyx.server.documents.models import DeletionAttemptSnapshot diff --git a/backend/onyx/redis/redis_document_set.py b/backend/onyx/redis/redis_document_set.py index dbc8bccf57..6f718de941 100644 --- a/backend/onyx/redis/redis_document_set.py +++ b/backend/onyx/redis/redis_document_set.py @@ -7,11 +7,13 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.configs.app_configs import DB_YIELD_PER_DEFAULT -from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants +from onyx.configs.constants import ( + CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, +) from onyx.db.document_set import construct_document_id_select_by_docset from onyx.redis.redis_object_helper import RedisObjectHelper from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/onyx/redis/redis_hierarchy.py b/backend/onyx/redis/redis_hierarchy.py index 2e6c5b76c5..b3838d7361 100644 --- a/backend/onyx/redis/redis_hierarchy.py +++ b/backend/onyx/redis/redis_hierarchy.py @@ -16,8 +16,7 @@ Cache Strategy: using only the SOURCE-type node as the ancestor """ -from typing import cast -from typing import TYPE_CHECKING +from typing import cast, TYPE_CHECKING from pydantic import BaseModel from redis.lock import Lock as RedisLock diff --git a/backend/onyx/redis/redis_object_helper.py b/backend/onyx/redis/redis_object_helper.py index b1ef0ee9bf..686af95cbd 100644 --- a/backend/onyx/redis/redis_object_helper.py +++ b/backend/onyx/redis/redis_object_helper.py @@ -1,5 +1,4 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from celery import Celery from redis.lock import Lock as RedisLock diff --git a/backend/onyx/redis/redis_pool.py b/backend/onyx/redis/redis_pool.py index 3db9b248c4..9a683d22f3 100644 --- a/backend/onyx/redis/redis_pool.py +++ b/backend/onyx/redis/redis_pool.py @@ -1,9 +1,7 @@ import asyncio import json import threading -from typing import Any -from typing import cast -from typing import Optional +from typing import Any, cast, Optional import redis from fastapi import Request @@ -18,30 +16,38 @@ from redis.lock import Lock as RedisLock from redis.retry import Retry from redis.sentinel import Sentinel -from onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME -from onyx.auth.constants import API_KEY_HEADER_NAME -from onyx.auth.constants import BEARER_PREFIX -from onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX -from onyx.configs.app_configs import REDIS_DB_NUMBER -from onyx.configs.app_configs import REDIS_HEALTH_CHECK_INTERVAL -from onyx.configs.app_configs import REDIS_HOST -from onyx.configs.app_configs import REDIS_PASSWORD -from onyx.configs.app_configs import REDIS_POOL_MAX_CONNECTIONS -from onyx.configs.app_configs import REDIS_PORT -from onyx.configs.app_configs import REDIS_REPLICA_HOST -from onyx.configs.app_configs import REDIS_SENTINEL_HOSTS -from onyx.configs.app_configs import REDIS_SENTINEL_MASTER_NAME -from onyx.configs.app_configs import REDIS_SENTINEL_PASSWORD -from onyx.configs.app_configs import REDIS_SSL -from onyx.configs.app_configs import REDIS_SSL_CA_CERTS -from onyx.configs.app_configs import REDIS_SSL_CERT_REQS -from onyx.configs.app_configs import REDIS_SSL_CERTFILE -from onyx.configs.app_configs import REDIS_SSL_KEYFILE -from onyx.configs.app_configs import USE_REDIS_IAM_AUTH -from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME -from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS -from onyx.redis.iam_auth import configure_redis_iam_auth -from onyx.redis.iam_auth import create_redis_ssl_context_if_iam +from onyx.auth.constants import ( + API_KEY_HEADER_ALTERNATIVE_NAME, + API_KEY_HEADER_NAME, + BEARER_PREFIX, +) +from onyx.configs.app_configs import ( + REDIS_AUTH_KEY_PREFIX, + REDIS_DB_NUMBER, + REDIS_HEALTH_CHECK_INTERVAL, + REDIS_HOST, + REDIS_PASSWORD, + REDIS_POOL_MAX_CONNECTIONS, + REDIS_PORT, + REDIS_REPLICA_HOST, + REDIS_SENTINEL_HOSTS, + REDIS_SENTINEL_MASTER_NAME, + REDIS_SENTINEL_PASSWORD, + REDIS_SSL, + REDIS_SSL_CA_CERTS, + REDIS_SSL_CERT_REQS, + REDIS_SSL_CERTFILE, + REDIS_SSL_KEYFILE, + USE_REDIS_IAM_AUTH, +) +from onyx.configs.constants import ( + FASTAPI_USERS_AUTH_COOKIE_NAME, + REDIS_SOCKET_KEEPALIVE_OPTIONS, +) +from onyx.redis.iam_auth import ( + configure_redis_iam_auth, + create_redis_ssl_context_if_iam, +) from onyx.redis.tenant_redis_client import TenantRedisClient from onyx.utils.logger import setup_logger from shared_configs.configs import DEFAULT_REDIS_PREFIX diff --git a/backend/onyx/redis/redis_tenant_work_gating.py b/backend/onyx/redis/redis_tenant_work_gating.py index 07cdf7b65b..5666ab662d 100644 --- a/backend/onyx/redis/redis_tenant_work_gating.py +++ b/backend/onyx/redis/redis_tenant_work_gating.py @@ -11,8 +11,7 @@ All public functions no-op in single-tenant mode (`MULTI_TENANT=False`). import time from typing import cast -from prometheus_client import Counter -from prometheus_client import Gauge +from prometheus_client import Counter, Gauge from onyx.configs.constants import ONYX_CLOUD_TENANT_ID from onyx.redis.redis_pool import get_redis_client diff --git a/backend/onyx/redis/redis_usergroup.py b/backend/onyx/redis/redis_usergroup.py index add01edb7b..672c16254e 100644 --- a/backend/onyx/redis/redis_usergroup.py +++ b/backend/onyx/redis/redis_usergroup.py @@ -7,15 +7,19 @@ from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session from onyx.configs.app_configs import DB_YIELD_PER_DEFAULT -from onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import OnyxRedisConstants +from onyx.configs.constants import ( + CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + OnyxRedisConstants, +) from onyx.redis.redis_object_helper import RedisObjectHelper from onyx.redis.tenant_redis_client import TenantRedisClient -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) class RedisUserGroup(RedisObjectHelper): diff --git a/backend/onyx/redis/tenant_redis_client.py b/backend/onyx/redis/tenant_redis_client.py index 43a9b8a12d..e99047104a 100644 --- a/backend/onyx/redis/tenant_redis_client.py +++ b/backend/onyx/redis/tenant_redis_client.py @@ -12,10 +12,8 @@ typing error, not a silent cross-tenant write. # annotations like ``-> set[bytes]``. from __future__ import annotations -from collections.abc import Generator -from collections.abc import Mapping -from typing import Any -from typing import cast +from collections.abc import Generator, Mapping +from typing import Any, cast import redis from redis.client import Pipeline diff --git a/backend/onyx/sandbox_proxy/addons/gate.py b/backend/onyx/sandbox_proxy/addons/gate.py index 2154426972..2a0b14c478 100644 --- a/backend/onyx/sandbox_proxy/addons/gate.py +++ b/backend/onyx/sandbox_proxy/addons/gate.py @@ -17,49 +17,50 @@ from typing import Protocol from urllib.parse import urlparse from uuid import UUID -from cachetools import cachedmethod -from cachetools import TTLCache +from cachetools import cachedmethod, TTLCache from mitmproxy import http from mitmproxy.proxy import server_hooks from sqlalchemy.orm import Session -from onyx.cache.interface import CACHE_TRANSIENT_ERRORS -from onyx.cache.interface import CacheBackend +from onyx.cache.interface import CACHE_TRANSIENT_ERRORS, CacheBackend from onyx.configs.constants import NotificationType from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.enums import ApprovalDecidedVia -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import EndpointPolicy +from onyx.db.enums import ApprovalDecidedVia, ApprovalDecision, EndpointPolicy from onyx.db.notification import create_notification -from onyx.db.scheduled_task import get_live_scheduled_run_grants -from onyx.db.scheduled_task import ScheduledRunGrants -from onyx.external_apps.matching.engine import actions_requiring_approval -from onyx.external_apps.matching.engine import AllMatchedActions +from onyx.db.scheduled_task import get_live_scheduled_run_grants, ScheduledRunGrants +from onyx.external_apps.matching.engine import ( + actions_requiring_approval, + AllMatchedActions, +) from onyx.sandbox_proxy import approval_cache -from onyx.sandbox_proxy.credential_injection import CredentialInjectionDispatcher -from onyx.sandbox_proxy.credential_injection import InjectionContext -from onyx.sandbox_proxy.credential_injection import InjectionOutcome -from onyx.sandbox_proxy.errors import http_403 -from onyx.sandbox_proxy.errors import SandboxProxyError -from onyx.sandbox_proxy.identity import ResolvedSandbox -from onyx.sandbox_proxy.identity import SessionContext -from onyx.sandbox_proxy.logging_utils import approval_decided_args -from onyx.sandbox_proxy.logging_utils import APPROVAL_DECIDED_FIELDS -from onyx.sandbox_proxy.logging_utils import credential_outcome_label -from onyx.sandbox_proxy.logging_utils import egress_approval_matched_args -from onyx.sandbox_proxy.logging_utils import EGRESS_APPROVAL_MATCHED_FIELDS -from onyx.sandbox_proxy.logging_utils import egress_matched_args -from onyx.sandbox_proxy.logging_utils import EGRESS_MATCHED_FIELDS -from onyx.sandbox_proxy.logging_utils import egress_session_matched_args -from onyx.sandbox_proxy.logging_utils import EGRESS_SESSION_MATCHED_FIELDS -from onyx.sandbox_proxy.logging_utils import egress_target_args -from onyx.sandbox_proxy.logging_utils import EGRESS_TARGET_FIELDS -from onyx.sandbox_proxy.logging_utils import full_log_id -from onyx.sandbox_proxy.logging_utils import sandbox_log_label -from onyx.sandbox_proxy.logging_utils import short_log_id +from onyx.sandbox_proxy.credential_injection import ( + CredentialInjectionDispatcher, + InjectionContext, + InjectionOutcome, +) +from onyx.sandbox_proxy.errors import http_403, SandboxProxyError +from onyx.sandbox_proxy.identity import ResolvedSandbox, SessionContext +from onyx.sandbox_proxy.logging_utils import ( + approval_decided_args, + APPROVAL_DECIDED_FIELDS, + credential_outcome_label, + egress_approval_matched_args, + EGRESS_APPROVAL_MATCHED_FIELDS, + egress_matched_args, + EGRESS_MATCHED_FIELDS, + egress_session_matched_args, + EGRESS_SESSION_MATCHED_FIELDS, + egress_target_args, + EGRESS_TARGET_FIELDS, + full_log_id, + sandbox_log_label, + short_log_id, +) from onyx.sandbox_proxy.request_evaluator import RequestEvaluator -from onyx.server.features.build.configs import SANDBOX_API_SERVER_URL -from onyx.server.features.build.configs import SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS +from onyx.server.features.build.configs import ( + SANDBOX_API_SERVER_URL, + SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS, +) from onyx.server.features.build.db import action_approval from onyx.utils.logger import setup_logger diff --git a/backend/onyx/sandbox_proxy/backend.py b/backend/onyx/sandbox_proxy/backend.py index 12947a3a6f..148b7449d8 100644 --- a/backend/onyx/sandbox_proxy/backend.py +++ b/backend/onyx/sandbox_proxy/backend.py @@ -13,8 +13,7 @@ best and crash-on-missing-config at worst. from onyx.sandbox_proxy.ca import CAStore from onyx.sandbox_proxy.identity import SandboxIPLookup -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend def build_ca_store() -> CAStore: diff --git a/backend/onyx/sandbox_proxy/ca.py b/backend/onyx/sandbox_proxy/ca.py index f49f892a2d..ba366ffc30 100644 --- a/backend/onyx/sandbox_proxy/ca.py +++ b/backend/onyx/sandbox_proxy/ca.py @@ -7,8 +7,7 @@ from pathlib import Path from typing import Protocol from cryptography import x509 -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID diff --git a/backend/onyx/sandbox_proxy/ca_docker.py b/backend/onyx/sandbox_proxy/ca_docker.py index 821cc0e597..ca3bfe7e53 100644 --- a/backend/onyx/sandbox_proxy/ca_docker.py +++ b/backend/onyx/sandbox_proxy/ca_docker.py @@ -32,8 +32,7 @@ orphaned cert). Operator recovery is: delete ``ca.crt``, restart proxy. import os from pathlib import Path -from onyx.sandbox_proxy.ca import CAStore -from onyx.sandbox_proxy.ca import CAStoreConflictError +from onyx.sandbox_proxy.ca import CAStore, CAStoreConflictError from onyx.server.features.build.configs import SANDBOX_PROXY_CA_VOLUME_PATH from onyx.utils.logger import setup_logger diff --git a/backend/onyx/sandbox_proxy/ca_k8s.py b/backend/onyx/sandbox_proxy/ca_k8s.py index aa7f90dd11..47ec057864 100644 --- a/backend/onyx/sandbox_proxy/ca_k8s.py +++ b/backend/onyx/sandbox_proxy/ca_k8s.py @@ -13,16 +13,19 @@ import time from kubernetes import client from kubernetes.client.rest import ApiException -from onyx.sandbox_proxy.ca import CAStore -from onyx.sandbox_proxy.ca import CAStoreConflictError -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SANDBOX_PROXY_CA_CONFIGMAP -from onyx.server.features.build.configs import SANDBOX_PROXY_CA_SECRET -from onyx.server.features.build.configs import SANDBOX_PROXY_NAMESPACE +from onyx.sandbox_proxy.ca import CAStore, CAStoreConflictError +from onyx.server.features.build.configs import ( + SANDBOX_NAMESPACE, + SANDBOX_PROXY_CA_CONFIGMAP, + SANDBOX_PROXY_CA_SECRET, + SANDBOX_PROXY_NAMESPACE, +) from onyx.server.features.build.sandbox.kubernetes.k8s_client import build_core_v1_api -from onyx.server.features.build.sandbox.labels import LABEL_K8S_COMPONENT -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY_ONYX +from onyx.server.features.build.sandbox.labels import ( + LABEL_K8S_COMPONENT, + LABEL_K8S_MANAGED_BY, + LABEL_K8S_MANAGED_BY_ONYX, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/sandbox_proxy/identity.py b/backend/onyx/sandbox_proxy/identity.py index 37abf6d919..f6aa5a0023 100644 --- a/backend/onyx/sandbox_proxy/identity.py +++ b/backend/onyx/sandbox_proxy/identity.py @@ -20,8 +20,7 @@ from uuid import UUID from sqlalchemy import select from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox +from onyx.db.models import BuildSession, Sandbox from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/sandbox_proxy/identity_docker.py b/backend/onyx/sandbox_proxy/identity_docker.py index 499a528332..c4452a57bd 100644 --- a/backend/onyx/sandbox_proxy/identity_docker.py +++ b/backend/onyx/sandbox_proxy/identity_docker.py @@ -16,20 +16,22 @@ from typing import Any from uuid import UUID from docker import DockerClient -from docker.errors import APIError -from docker.errors import NotFound +from docker.errors import APIError, NotFound from docker.models.containers import Container from docker.types.daemon import CancellableStream from requests.exceptions import ConnectionError as RequestsConnectionError -from onyx.sandbox_proxy.identity import SandboxIdentity -from onyx.sandbox_proxy.identity import SandboxIPLookup -from onyx.server.features.build.configs import SANDBOX_DOCKER_NETWORK -from onyx.server.features.build.configs import SANDBOX_DOCKER_SOCKET -from onyx.server.features.build.sandbox.labels import LABEL_DOCKER_COMPONENT -from onyx.server.features.build.sandbox.labels import LABEL_DOCKER_COMPONENT_SANDBOX -from onyx.server.features.build.sandbox.labels import LABEL_SANDBOX_ID -from onyx.server.features.build.sandbox.labels import LABEL_TENANT_ID +from onyx.sandbox_proxy.identity import SandboxIdentity, SandboxIPLookup +from onyx.server.features.build.configs import ( + SANDBOX_DOCKER_NETWORK, + SANDBOX_DOCKER_SOCKET, +) +from onyx.server.features.build.sandbox.labels import ( + LABEL_DOCKER_COMPONENT, + LABEL_DOCKER_COMPONENT_SANDBOX, + LABEL_SANDBOX_ID, + LABEL_TENANT_ID, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/sandbox_proxy/identity_k8s.py b/backend/onyx/sandbox_proxy/identity_k8s.py index f10222451f..e0b1200fce 100644 --- a/backend/onyx/sandbox_proxy/identity_k8s.py +++ b/backend/onyx/sandbox_proxy/identity_k8s.py @@ -8,22 +8,21 @@ exponential backoff capped at `_RECONNECT_MAX_SECONDS`; on 410 Gone we relist. import threading from uuid import UUID -from kubernetes import client -from kubernetes import watch +from kubernetes import client, watch from kubernetes.client.rest import ApiException -from urllib3.exceptions import ProtocolError -from urllib3.exceptions import ReadTimeoutError +from urllib3.exceptions import ProtocolError, ReadTimeoutError -from onyx.sandbox_proxy.identity import SandboxIdentity -from onyx.sandbox_proxy.identity import SandboxIPLookup +from onyx.sandbox_proxy.identity import SandboxIdentity, SandboxIPLookup from onyx.server.features.build.configs import SANDBOX_NAMESPACE from onyx.server.features.build.sandbox.kubernetes.k8s_client import build_core_v1_api -from onyx.server.features.build.sandbox.labels import LABEL_K8S_COMPONENT -from onyx.server.features.build.sandbox.labels import LABEL_K8S_COMPONENT_SANDBOX -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY_ONYX -from onyx.server.features.build.sandbox.labels import LABEL_SANDBOX_ID -from onyx.server.features.build.sandbox.labels import LABEL_TENANT_ID +from onyx.server.features.build.sandbox.labels import ( + LABEL_K8S_COMPONENT, + LABEL_K8S_COMPONENT_SANDBOX, + LABEL_K8S_MANAGED_BY, + LABEL_K8S_MANAGED_BY_ONYX, + LABEL_SANDBOX_ID, + LABEL_TENANT_ID, +) from onyx.utils.logger import setup_logger _RECONNECT_INITIAL_SECONDS = 1.0 diff --git a/backend/onyx/sandbox_proxy/logging_utils.py b/backend/onyx/sandbox_proxy/logging_utils.py index e37974848b..7ec43a8515 100644 --- a/backend/onyx/sandbox_proxy/logging_utils.py +++ b/backend/onyx/sandbox_proxy/logging_utils.py @@ -7,8 +7,7 @@ from mitmproxy import http from onyx.db.enums import EndpointPolicy from onyx.external_apps.matching.engine import AllMatchedActions from onyx.sandbox_proxy.credential_injection import InjectionOutcome -from onyx.sandbox_proxy.identity import ResolvedSandbox -from onyx.sandbox_proxy.identity import SessionContext +from onyx.sandbox_proxy.identity import ResolvedSandbox, SessionContext _EGRESS_CONTEXT_FIELDS = "tenant=%s sandbox=%s" _EGRESS_SESSION_FIELDS = f"{_EGRESS_CONTEXT_FIELDS} session=%s" diff --git a/backend/onyx/sandbox_proxy/request_evaluator.py b/backend/onyx/sandbox_proxy/request_evaluator.py index b4cf35990c..a63b22840f 100644 --- a/backend/onyx/sandbox_proxy/request_evaluator.py +++ b/backend/onyx/sandbox_proxy/request_evaluator.py @@ -8,8 +8,7 @@ lockdown, not this heuristic. import json import re from collections.abc import Iterable -from typing import Any -from typing import Protocol +from typing import Any, Protocol from urllib.parse import parse_qs from uuid import UUID @@ -19,9 +18,11 @@ from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.external_app import get_external_apps from onyx.db.models import ExternalApp from onyx.external_apps.credentials import app_is_available -from onyx.external_apps.matching.engine import AllMatchedActions -from onyx.external_apps.matching.engine import apply_credential_gate -from onyx.external_apps.matching.engine import recognize_actions +from onyx.external_apps.matching.engine import ( + AllMatchedActions, + apply_credential_gate, + recognize_actions, +) from onyx.external_apps.matching.request import ProxiedRequest from onyx.utils.logger import setup_logger diff --git a/backend/onyx/sandbox_proxy/resolvers/external_app.py b/backend/onyx/sandbox_proxy/resolvers/external_app.py index ba613af74f..230f7bf7ec 100644 --- a/backend/onyx/sandbox_proxy/resolvers/external_app.py +++ b/backend/onyx/sandbox_proxy/resolvers/external_app.py @@ -13,9 +13,11 @@ from mitmproxy import http from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.external_apps.credentials import resolve_injection_headers from onyx.external_apps.token_refresh import ensure_fresh_credentials -from onyx.sandbox_proxy.credential_injection import CredentialResolver -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionContext +from onyx.sandbox_proxy.credential_injection import ( + CredentialResolver, + CredentialUnavailableError, + InjectionContext, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/sandbox_proxy/resolvers/llm_provider_key.py b/backend/onyx/sandbox_proxy/resolvers/llm_provider_key.py index 7597e8e2cd..43d3e077ed 100644 --- a/backend/onyx/sandbox_proxy/resolvers/llm_provider_key.py +++ b/backend/onyx/sandbox_proxy/resolvers/llm_provider_key.py @@ -11,14 +11,15 @@ from __future__ import annotations from mitmproxy import http -from onyx.auth.constants import API_KEY_HEADER_NAME -from onyx.auth.constants import BEARER_PREFIX +from onyx.auth.constants import API_KEY_HEADER_NAME, BEARER_PREFIX from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.llm import fetch_first_accessible_llm_provider_by_type from onyx.db.users import fetch_user_by_id -from onyx.sandbox_proxy.credential_injection import CredentialResolver -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionContext +from onyx.sandbox_proxy.credential_injection import ( + CredentialResolver, + CredentialUnavailableError, + InjectionContext, +) from onyx.sandbox_proxy.logging_utils import short_log_id from onyx.utils.credential_audit import emit_credential_access from onyx.utils.logger import setup_logger diff --git a/backend/onyx/sandbox_proxy/resolvers/onyx_pat.py b/backend/onyx/sandbox_proxy/resolvers/onyx_pat.py index 55dc22ca49..75c0aecd8b 100644 --- a/backend/onyx/sandbox_proxy/resolvers/onyx_pat.py +++ b/backend/onyx/sandbox_proxy/resolvers/onyx_pat.py @@ -12,15 +12,18 @@ from urllib.parse import urlparse from mitmproxy import http -from onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME -from onyx.auth.constants import API_KEY_HEADER_NAME -from onyx.auth.constants import BEARER_PREFIX +from onyx.auth.constants import ( + API_KEY_HEADER_ALTERNATIVE_NAME, + API_KEY_HEADER_NAME, + BEARER_PREFIX, +) from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.sandbox_proxy.credential_injection import CredentialResolver -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionContext -from onyx.sandbox_proxy.logging_utils import full_log_id -from onyx.sandbox_proxy.logging_utils import short_log_id +from onyx.sandbox_proxy.credential_injection import ( + CredentialResolver, + CredentialUnavailableError, + InjectionContext, +) +from onyx.sandbox_proxy.logging_utils import full_log_id, short_log_id from onyx.server.features.build.configs import SANDBOX_API_SERVER_URL from onyx.server.features.build.db.sandbox import get_sandbox_by_id from onyx.utils.logger import setup_logger diff --git a/backend/onyx/sandbox_proxy/server.py b/backend/onyx/sandbox_proxy/server.py index f2e1a56986..07e682a447 100644 --- a/backend/onyx/sandbox_proxy/server.py +++ b/backend/onyx/sandbox_proxy/server.py @@ -7,8 +7,7 @@ import sys import threading import uuid from collections.abc import Callable -from http.server import BaseHTTPRequestHandler -from http.server import HTTPServer +from http.server import BaseHTTPRequestHandler, HTTPServer from mitmproxy.options import Options from mitmproxy.tools.dump import DumpMaster @@ -17,21 +16,22 @@ from onyx.cache.factory import get_cache_backend from onyx.cache.interface import CacheBackend from onyx.db.engine.sql_engine import SqlEngine from onyx.sandbox_proxy.addons.gate import GateAddon -from onyx.sandbox_proxy.backend import build_ca_store -from onyx.sandbox_proxy.backend import build_ip_lookup -from onyx.sandbox_proxy.ca import CABootstrap -from onyx.sandbox_proxy.ca import MaterializedCA -from onyx.sandbox_proxy.credential_injection import CredentialInjectionDispatcher -from onyx.sandbox_proxy.credential_injection import CredentialResolver -from onyx.sandbox_proxy.identity import IdentityResolver -from onyx.sandbox_proxy.identity import SandboxIPLookup +from onyx.sandbox_proxy.backend import build_ca_store, build_ip_lookup +from onyx.sandbox_proxy.ca import CABootstrap, MaterializedCA +from onyx.sandbox_proxy.credential_injection import ( + CredentialInjectionDispatcher, + CredentialResolver, +) +from onyx.sandbox_proxy.identity import IdentityResolver, SandboxIPLookup from onyx.sandbox_proxy.request_evaluator import ExternalAppRequestEvaluator from onyx.sandbox_proxy.resolvers.external_app import ExternalAppResolver from onyx.sandbox_proxy.resolvers.llm_provider_key import LLMProviderKeyResolver from onyx.sandbox_proxy.resolvers.onyx_pat import OnyxPatResolver -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SANDBOX_PROXY_HEALTHZ_PORT -from onyx.server.features.build.configs import SANDBOX_PROXY_LISTEN_PORT +from onyx.server.features.build.configs import ( + SANDBOX_NAMESPACE, + SANDBOX_PROXY_HEALTHZ_PORT, + SANDBOX_PROXY_LISTEN_PORT, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable diff --git a/backend/onyx/secondary_llm_flows/chat_session_naming.py b/backend/onyx/secondary_llm_flows/chat_session_naming.py index c246f5766f..4ef7fb7845 100644 --- a/backend/onyx/secondary_llm_flows/chat_session_naming.py +++ b/backend/onyx/secondary_llm_flows/chat_session_naming.py @@ -4,11 +4,9 @@ from onyx.configs.constants import MessageType from onyx.llm.interfaces import LLM from onyx.llm.models import ReasoningEffort from onyx.llm.utils import llm_response_to_string -from onyx.prompts.chat_prompts import CHAT_NAMING_REMINDER -from onyx.prompts.chat_prompts import CHAT_NAMING_SYSTEM_PROMPT +from onyx.prompts.chat_prompts import CHAT_NAMING_REMINDER, CHAT_NAMING_SYSTEM_PROMPT from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/secondary_llm_flows/document_filter.py b/backend/onyx/secondary_llm_flows/document_filter.py index 54dbdfb682..21c35901b9 100644 --- a/backend/onyx/secondary_llm_flows/document_filter.py +++ b/backend/onyx/secondary_llm_flows/document_filter.py @@ -2,19 +2,21 @@ import json import re from onyx.configs.chat_configs import SECONDARY_LLM_FLOW_TIMEOUT_S -from onyx.context.search.models import ContextExpansionType -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection +from onyx.context.search.models import ( + ContextExpansionType, + InferenceChunk, + InferenceSection, +) from onyx.llm.interfaces import LLM -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import UserMessage -from onyx.prompts.search_prompts import DOCUMENT_CONTEXT_SELECTION_PROMPT -from onyx.prompts.search_prompts import DOCUMENT_SELECTION_PROMPT -from onyx.prompts.search_prompts import TRY_TO_FILL_TO_MAX_INSTRUCTIONS +from onyx.llm.models import ReasoningEffort, UserMessage +from onyx.prompts.search_prompts import ( + DOCUMENT_CONTEXT_SELECTION_PROMPT, + DOCUMENT_SELECTION_PROMPT, + TRY_TO_FILL_TO_MAX_INSTRUCTIONS, +) from onyx.tools.tool_implementations.search.constants import MAX_CHUNKS_FOR_RELEVANCE from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger from onyx.utils.timing import log_function_time diff --git a/backend/onyx/secondary_llm_flows/memory_update.py b/backend/onyx/secondary_llm_flows/memory_update.py index 94a67b3f0b..568bd3c609 100644 --- a/backend/onyx/secondary_llm_flows/memory_update.py +++ b/backend/onyx/secondary_llm_flows/memory_update.py @@ -1,12 +1,10 @@ from onyx.configs.constants import MessageType from onyx.llm.interfaces import LLM -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import UserMessage +from onyx.llm.models import ReasoningEffort, UserMessage from onyx.prompts.basic_memory import FULL_MEMORY_UPDATE_PROMPT from onyx.tools.models import ChatMinimalTextMessage from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger from onyx.utils.text_processing import parse_llm_json_response diff --git a/backend/onyx/secondary_llm_flows/query_expansion.py b/backend/onyx/secondary_llm_flows/query_expansion.py index 3acc401bbc..9043f88d89 100644 --- a/backend/onyx/secondary_llm_flows/query_expansion.py +++ b/backend/onyx/secondary_llm_flows/query_expansion.py @@ -1,20 +1,23 @@ from onyx.configs.constants import MessageType from onyx.llm.interfaces import LLM -from onyx.llm.models import AssistantMessage -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import SystemMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ( + AssistantMessage, + ChatCompletionMessage, + ReasoningEffort, + SystemMessage, + UserMessage, +) from onyx.prompts.prompt_utils import get_current_llm_day_time -from onyx.prompts.search_prompts import KEYWORD_REPHRASE_SYSTEM_PROMPT -from onyx.prompts.search_prompts import KEYWORD_REPHRASE_USER_PROMPT -from onyx.prompts.search_prompts import REPHRASE_CONTEXT_PROMPT -from onyx.prompts.search_prompts import SEMANTIC_QUERY_REPHRASE_SYSTEM_PROMPT -from onyx.prompts.search_prompts import SEMANTIC_QUERY_REPHRASE_USER_PROMPT +from onyx.prompts.search_prompts import ( + KEYWORD_REPHRASE_SYSTEM_PROMPT, + KEYWORD_REPHRASE_USER_PROMPT, + REPHRASE_CONTEXT_PROMPT, + SEMANTIC_QUERY_REPHRASE_SYSTEM_PROMPT, + SEMANTIC_QUERY_REPHRASE_USER_PROMPT, +) from onyx.tools.models import ChatMinimalTextMessage from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/secondary_llm_flows/source_filter.py b/backend/onyx/secondary_llm_flows/source_filter.py index 229ea51247..f9ce11fa6a 100644 --- a/backend/onyx/secondary_llm_flows/source_filter.py +++ b/backend/onyx/secondary_llm_flows/source_filter.py @@ -2,17 +2,13 @@ import json from pydantic import BaseModel -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MessageType +from onyx.configs.constants import DocumentSource, MessageType from onyx.llm.interfaces import LLM -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import UserMessage +from onyx.llm.models import ChatCompletionMessage, ReasoningEffort, UserMessage from onyx.prompts.filter_extration import SOURCE_SCOPE_DECISION_PROMPT from onyx.tools.models import ChatMinimalTextMessage from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger from onyx.utils.text_processing import parse_bracketed_list diff --git a/backend/onyx/secondary_llm_flows/time_filter.py b/backend/onyx/secondary_llm_flows/time_filter.py index 182b81f5ff..4f038d2246 100644 --- a/backend/onyx/secondary_llm_flows/time_filter.py +++ b/backend/onyx/secondary_llm_flows/time_filter.py @@ -1,25 +1,18 @@ import re -from datetime import datetime -from datetime import time -from datetime import timedelta -from datetime import timezone +from datetime import datetime, time, timedelta, timezone from enum import StrEnum from dateutil.relativedelta import relativedelta from pydantic import BaseModel from onyx.configs.constants import MessageType -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import TimeRange +from onyx.context.search.models import BaseFilters, TimeRange from onyx.llm.interfaces import LLM -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import UserMessage +from onyx.llm.models import ChatCompletionMessage, ReasoningEffort, UserMessage from onyx.prompts.filter_extration import TIME_SCOPE_DECISION_PROMPT from onyx.tools.models import ChatMinimalTextMessage from onyx.tracing.flows import LLMFlow -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/api_key/api.py b/backend/onyx/server/api_key/api.py index 148267bfcb..31f766ca1b 100644 --- a/backend/onyx/server/api_key/api.py +++ b/backend/onyx/server/api_key/api.py @@ -1,22 +1,25 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission -from onyx.db.api_key import ApiKeyDescriptor -from onyx.db.api_key import fetch_api_keys -from onyx.db.api_key import insert_api_key -from onyx.db.api_key import regenerate_api_key -from onyx.db.api_key import remove_api_key -from onyx.db.api_key import update_api_key +from onyx.db.api_key import ( + ApiKeyDescriptor, + fetch_api_keys, + insert_api_key, + regenerate_api_key, + remove_api_key, + update_api_key, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User from onyx.server.api_key.models import APIKeyArgs -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) router = APIRouter(prefix="/admin/api-key") diff --git a/backend/onyx/server/api_key_usage.py b/backend/onyx/server/api_key_usage.py index a35973bb87..1a90759c77 100644 --- a/backend/onyx/server/api_key_usage.py +++ b/backend/onyx/server/api_key_usage.py @@ -1,16 +1,13 @@ """API key and PAT usage tracking for cloud usage limits.""" -from fastapi import Depends -from fastapi import Request +from fastapi import Depends, Request from sqlalchemy.orm import Session from onyx.auth.api_key import get_hashed_api_key_from_request from onyx.auth.pat import get_hashed_pat_from_request from onyx.db.engine.sql_engine import get_session -from onyx.db.usage import increment_usage -from onyx.db.usage import UsageType -from onyx.server.usage_limits import check_usage_and_raise -from onyx.server.usage_limits import is_usage_limits_enabled +from onyx.db.usage import increment_usage, UsageType +from onyx.server.usage_limits import check_usage_and_raise, is_usage_limits_enabled from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/auth/captcha_api.py b/backend/onyx/server/auth/captcha_api.py index 12ca2b3897..38cbb950fc 100644 --- a/backend/onyx/server/auth/captcha_api.py +++ b/backend/onyx/server/auth/captcha_api.py @@ -14,26 +14,26 @@ Three entry points are gated: import hmac -from fastapi import APIRouter -from fastapi import Request -from fastapi import Response +from fastapi import APIRouter, Request, Response from pydantic import BaseModel -from starlette.middleware.base import BaseHTTPMiddleware -from starlette.middleware.base import RequestResponseEndpoint +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint -from onyx.auth.captcha import CAPTCHA_COOKIE_NAME -from onyx.auth.captcha import CaptchaAction -from onyx.auth.captcha import CaptchaVerificationError -from onyx.auth.captcha import is_captcha_enabled -from onyx.auth.captcha import issue_captcha_cookie_value -from onyx.auth.captcha import validate_captcha_cookie_value -from onyx.auth.captcha import verify_captcha_token -from onyx.configs.app_configs import CAPTCHA_COOKIE_TTL_SECONDS -from onyx.configs.app_configs import HEALTH_CHECK_BYPASS_TOKEN +from onyx.auth.captcha import ( + CAPTCHA_COOKIE_NAME, + CaptchaAction, + CaptchaVerificationError, + is_captcha_enabled, + issue_captcha_cookie_value, + validate_captcha_cookie_value, + verify_captcha_token, +) +from onyx.configs.app_configs import ( + CAPTCHA_COOKIE_TTL_SECONDS, + HEALTH_CHECK_BYPASS_TOKEN, +) from onyx.configs.constants import PUBLIC_API_TAGS from onyx.error_handling.error_codes import OnyxErrorCode -from onyx.error_handling.exceptions import onyx_error_to_json_response -from onyx.error_handling.exceptions import OnyxError +from onyx.error_handling.exceptions import onyx_error_to_json_response, OnyxError from onyx.utils.client_ip import get_client_ip from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/auth/mobile.py b/backend/onyx/server/auth/mobile.py index 35d72e1a12..9385c8c70b 100644 --- a/backend/onyx/server/auth/mobile.py +++ b/backend/onyx/server/auth/mobile.py @@ -17,8 +17,7 @@ from fastapi import APIRouter from pydantic import BaseModel from onyx.auth.mobile_sso.code_store import consume_sso_code -from onyx.auth.users import fastapi_users -from onyx.auth.users import mobile_auth_backend +from onyx.auth.users import fastapi_users, mobile_auth_backend from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/onyx/server/auth_check.py b/backend/onyx/server/auth_check.py index 436f49e92a..99f67083bd 100644 --- a/backend/onyx/server/auth_check.py +++ b/backend/onyx/server/auth_check.py @@ -4,12 +4,14 @@ from fastapi import FastAPI from fastapi.dependencies.models import Dependant from starlette.routing import BaseRoute -from onyx.auth.users import current_chat_accessible_user -from onyx.auth.users import current_curator_or_admin_user -from onyx.auth.users import current_limited_user -from onyx.auth.users import current_user -from onyx.auth.users import current_user_from_websocket -from onyx.auth.users import current_user_with_expired_token +from onyx.auth.users import ( + current_chat_accessible_user, + current_curator_or_admin_user, + current_limited_user, + current_user, + current_user_from_websocket, + current_user_with_expired_token, +) from onyx.configs.app_configs import APP_API_PREFIX from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop diff --git a/backend/onyx/server/documents/cc_pair.py b/backend/onyx/server/documents/cc_pair.py index 65ec5e2a1c..be8f1de010 100644 --- a/backend/onyx/server/documents/cc_pair.py +++ b/backend/onyx/server/documents/cc_pair.py @@ -1,10 +1,7 @@ from datetime import datetime from http import HTTPStatus -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query +from fastapi import APIRouter, Depends, HTTPException, Query from fastapi.responses import JSONResponse from sqlalchemy import select from sqlalchemy.exc import IntegrityError @@ -15,51 +12,50 @@ from onyx.auth.users import current_curator_or_admin_user from onyx.background.celery.tasks.pruning.tasks import try_creating_prune_generator_task from onyx.background.celery.versioned_apps.client import app as client_app from onyx.background.indexing.models import IndexAttemptErrorPydantic -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import PUBLIC_API_TAGS +from onyx.configs.constants import ( + DocumentSource, + OnyxCeleryPriority, + OnyxCeleryTask, + PUBLIC_API_TAGS, +) from onyx.connectors.exceptions import ValidationError -from onyx.connectors.factory import identify_connector_class -from onyx.connectors.factory import validate_ccpair_for_user +from onyx.connectors.factory import identify_connector_class, validate_ccpair_for_user from onyx.connectors.interfaces import Resolver from onyx.connectors.models import InputType from onyx.db.connector import delete_connector -from onyx.db.connector_credential_pair import add_credential_to_connector -from onyx.db.connector_credential_pair import get_connector_credential_pair_for_user from onyx.db.connector_credential_pair import ( + add_credential_to_connector, + get_connector_credential_pair_for_user, get_connector_credential_pair_from_id_for_user, + remove_credential_from_connector, + update_connector_credential_pair_from_id, + verify_user_has_access_to_cc_pair, ) -from onyx.db.connector_credential_pair import remove_credential_from_connector -from onyx.db.connector_credential_pair import update_connector_credential_pair_from_id -from onyx.db.connector_credential_pair import verify_user_has_access_to_cc_pair -from onyx.db.document import get_document_counts_for_cc_pairs -from onyx.db.document import get_documents_for_cc_pair +from onyx.db.document import get_document_counts_for_cc_pairs, get_documents_for_cc_pair from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingStatus -from onyx.db.enums import Permission -from onyx.db.enums import PermissionSyncStatus -from onyx.db.index_attempt import count_index_attempt_errors_for_cc_pair -from onyx.db.index_attempt import count_index_attempts_for_cc_pair -from onyx.db.index_attempt import get_error_counts_for_index_attempts -from onyx.db.index_attempt import get_index_attempt -from onyx.db.index_attempt import get_index_attempt_errors_for_cc_pair -from onyx.db.index_attempt import get_latest_index_attempt_for_cc_pair_id -from onyx.db.index_attempt import get_latest_successful_index_attempt_for_cc_pair_id -from onyx.db.index_attempt import get_paginated_index_attempts_for_cc_pair_id +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + IndexingStatus, + Permission, + PermissionSyncStatus, +) +from onyx.db.index_attempt import ( + count_index_attempt_errors_for_cc_pair, + count_index_attempts_for_cc_pair, + get_error_counts_for_index_attempts, + get_index_attempt, + get_index_attempt_errors_for_cc_pair, + get_latest_index_attempt_for_cc_pair_id, + get_latest_successful_index_attempt_for_cc_pair_id, + get_paginated_index_attempts_for_cc_pair_id, +) from onyx.db.index_attempt_metrics import get_stage_metrics_for_attempt from onyx.db.indexing_coordination import IndexingCoordination -from onyx.db.models import IndexAttempt -from onyx.db.models import User +from onyx.db.models import IndexAttempt, User from onyx.db.permission_sync_attempt import ( get_latest_doc_permission_sync_attempt_for_cc_pair, -) -from onyx.db.permission_sync_attempt import ( get_recent_doc_permission_sync_attempts_for_cc_pair, -) -from onyx.db.permission_sync_attempt import ( get_relevant_external_group_sync_attempts_for_cc_pair, ) from onyx.error_handling.error_codes import OnyxErrorCode @@ -68,26 +64,30 @@ from onyx.redis.redis_connector import RedisConnector from onyx.redis.redis_connector_utils import get_deletion_attempt_snapshot from onyx.redis.redis_pool import get_redis_client from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active -from onyx.server.documents.models import CCPairFullInfo -from onyx.server.documents.models import CCPairSyncAttemptsResponse -from onyx.server.documents.models import CCPropertyUpdateRequest -from onyx.server.documents.models import CCStatusUpdateRequest -from onyx.server.documents.models import ConnectorCredentialPairIdentifier -from onyx.server.documents.models import ConnectorCredentialPairMetadata -from onyx.server.documents.models import DocPermissionSyncAttemptSnapshot -from onyx.server.documents.models import DocumentSyncStatus -from onyx.server.documents.models import ExternalGroupSyncAttemptSnapshot -from onyx.server.documents.models import IndexAttemptSnapshot -from onyx.server.documents.models import IndexAttemptStageMetricSnapshot -from onyx.server.documents.models import IndexAttemptStageMetricsResponse -from onyx.server.documents.models import PaginatedReturn -from onyx.server.documents.models import synthesize_unaccounted +from onyx.server.documents.models import ( + CCPairFullInfo, + CCPairSyncAttemptsResponse, + CCPropertyUpdateRequest, + CCStatusUpdateRequest, + ConnectorCredentialPairIdentifier, + ConnectorCredentialPairMetadata, + DocPermissionSyncAttemptSnapshot, + DocumentSyncStatus, + ExternalGroupSyncAttemptSnapshot, + IndexAttemptSnapshot, + IndexAttemptStageMetricSnapshot, + IndexAttemptStageMetricsResponse, + PaginatedReturn, + synthesize_unaccounted, +) from onyx.server.models import StatusResponse from onyx.server.security.store import get_security_settings -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/documents/connector.py b/backend/onyx/server/documents/connector.py index ac233614d1..633e8f10f2 100644 --- a/backend/onyx/server/documents/connector.py +++ b/backend/onyx/server/documents/connector.py @@ -5,130 +5,146 @@ import os import zipfile from datetime import datetime from io import BytesIO -from typing import Any -from typing import cast +from typing import Any, cast -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import Form -from fastapi import HTTPException -from fastapi import Query -from fastapi import Request -from fastapi import Response -from fastapi import UploadFile +from fastapi import ( + APIRouter, + Depends, + File, + Form, + HTTPException, + Query, + Request, + Response, + UploadFile, +) from google.oauth2.credentials import Credentials from pydantic import BaseModel from sqlalchemy.orm import Session from onyx.auth.email_utils import send_email from onyx.auth.permissions import require_permission -from onyx.auth.users import current_chat_accessible_user -from onyx.auth.users import current_curator_or_admin_user +from onyx.auth.users import current_chat_accessible_user, current_curator_or_admin_user from onyx.background.celery.tasks.pruning.tasks import try_creating_prune_generator_task from onyx.background.celery.versioned_apps.client import app as client_app -from onyx.configs.app_configs import EMAIL_CONFIGURED -from onyx.configs.app_configs import ENABLED_CONNECTOR_TYPES -from onyx.configs.app_configs import MOCK_CONNECTOR_FILE_PATH -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import MilestoneRecordType -from onyx.configs.constants import ONYX_METADATA_FILENAME -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import PUBLIC_API_TAGS +from onyx.configs.app_configs import ( + EMAIL_CONFIGURED, + ENABLED_CONNECTOR_TYPES, + MOCK_CONNECTOR_FILE_PATH, +) +from onyx.configs.constants import ( + DocumentSource, + FileOrigin, + MilestoneRecordType, + ONYX_METADATA_FILENAME, + OnyxCeleryPriority, + OnyxCeleryTask, + PUBLIC_API_TAGS, +) from onyx.connectors.exceptions import ConnectorValidationError from onyx.connectors.factory import validate_ccpair_for_user from onyx.connectors.google_utils.google_auth import get_google_oauth_creds -from onyx.connectors.google_utils.google_kv import build_service_account_creds -from onyx.connectors.google_utils.google_kv import get_auth_url -from onyx.connectors.google_utils.google_kv import update_credential_access_tokens -from onyx.connectors.google_utils.google_kv import verify_csrf -from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY +from onyx.connectors.google_utils.google_kv import ( + build_service_account_creds, + get_auth_url, + update_credential_access_tokens, + verify_csrf, +) from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_DICT_TOKEN_KEY, GoogleOAuthAuthenticationMethod, ) -from onyx.db.connector import create_connector -from onyx.db.connector import delete_connector -from onyx.db.connector import fetch_connector_by_id -from onyx.db.connector import fetch_connectors -from onyx.db.connector import fetch_unique_document_sources -from onyx.db.connector import get_connector_credential_ids -from onyx.db.connector import mark_ccpair_with_indexing_trigger -from onyx.db.connector import update_connector -from onyx.db.connector_credential_pair import add_credential_to_connector +from onyx.db.connector import ( + create_connector, + delete_connector, + fetch_connector_by_id, + fetch_connectors, + fetch_unique_document_sources, + get_connector_credential_ids, + mark_ccpair_with_indexing_trigger, + update_connector, +) from onyx.db.connector_credential_pair import ( + add_credential_to_connector, fetch_connector_credential_pair_for_connector, -) -from onyx.db.connector_credential_pair import get_cc_pair_groups_for_ids -from onyx.db.connector_credential_pair import get_connector_credential_pair -from onyx.db.connector_credential_pair import get_connector_credential_pairs_for_user -from onyx.db.connector_credential_pair import ( + get_cc_pair_groups_for_ids, + get_connector_credential_pair, + get_connector_credential_pairs_for_user, get_connector_credential_pairs_for_user_parallel, + verify_user_has_access_to_cc_pair, ) -from onyx.db.connector_credential_pair import verify_user_has_access_to_cc_pair -from onyx.db.credentials import create_credential -from onyx.db.credentials import fetch_credential_by_id_for_user +from onyx.db.credentials import create_credential, fetch_credential_by_id_for_user from onyx.db.deletion_attempt import check_deletion_attempt_is_allowed from onyx.db.document import get_document_counts_for_all_cc_pairs from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingMode -from onyx.db.enums import Permission -from onyx.db.enums import ProcessingMode +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + IndexingMode, + Permission, + ProcessingMode, +) from onyx.db.federated import fetch_all_federated_connectors_parallel -from onyx.db.index_attempt import get_index_attempts_for_cc_pair -from onyx.db.index_attempt import get_latest_index_attempts_by_status -from onyx.db.index_attempt import get_latest_index_attempts_parallel -from onyx.db.index_attempt import get_latest_successful_index_attempts_parallel -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import FederatedConnector -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexingStatus -from onyx.db.models import User -from onyx.db.models import UserRole -from onyx.file_store.file_store import FileStore -from onyx.file_store.file_store import get_default_file_store +from onyx.db.index_attempt import ( + get_index_attempts_for_cc_pair, + get_latest_index_attempts_by_status, + get_latest_index_attempts_parallel, + get_latest_successful_index_attempts_parallel, +) +from onyx.db.models import ( + ConnectorCredentialPair, + FederatedConnector, + IndexAttempt, + IndexingStatus, + User, + UserRole, +) +from onyx.file_store.file_store import FileStore, get_default_file_store from onyx.redis.redis_pool import get_redis_client from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active -from onyx.server.documents.models import AuthStatus -from onyx.server.documents.models import AuthUrl -from onyx.server.documents.models import ConnectorBase -from onyx.server.documents.models import ConnectorCredentialPairIdentifier -from onyx.server.documents.models import ConnectorFileInfo -from onyx.server.documents.models import ConnectorFilesResponse -from onyx.server.documents.models import ConnectorIndexingStatusLite -from onyx.server.documents.models import ConnectorIndexingStatusLiteResponse -from onyx.server.documents.models import ConnectorRequestSubmission -from onyx.server.documents.models import ConnectorSnapshot -from onyx.server.documents.models import ConnectorStatus -from onyx.server.documents.models import ConnectorUpdateRequest -from onyx.server.documents.models import CredentialBase -from onyx.server.documents.models import CredentialSnapshot -from onyx.server.documents.models import DocsCountOperator -from onyx.server.documents.models import FailedConnectorIndexingStatus -from onyx.server.documents.models import FileUploadResponse -from onyx.server.documents.models import GDriveCallback -from onyx.server.documents.models import GmailCallback -from onyx.server.documents.models import GoogleServiceAccountCredentialRequest -from onyx.server.documents.models import IndexedSourcesResponse -from onyx.server.documents.models import IndexingStatusRequest -from onyx.server.documents.models import ObjectCreationIdResponse -from onyx.server.documents.models import RunConnectorRequest -from onyx.server.documents.models import SourceSummary +from onyx.server.documents.models import ( + AuthStatus, + AuthUrl, + ConnectorBase, + ConnectorCredentialPairIdentifier, + ConnectorFileInfo, + ConnectorFilesResponse, + ConnectorIndexingStatusLite, + ConnectorIndexingStatusLiteResponse, + ConnectorRequestSubmission, + ConnectorSnapshot, + ConnectorStatus, + ConnectorUpdateRequest, + CredentialBase, + CredentialSnapshot, + DocsCountOperator, + FailedConnectorIndexingStatus, + FileUploadResponse, + GDriveCallback, + GmailCallback, + GoogleServiceAccountCredentialRequest, + IndexedSourcesResponse, + IndexingStatusRequest, + ObjectCreationIdResponse, + RunConnectorRequest, + SourceSummary, +) from onyx.server.federated.models import FederatedConnectorStatus from onyx.server.models import StatusResponse from onyx.server.security.store import get_security_settings from onyx.server.utils_vector_db import require_vector_db -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.logger import setup_logger from onyx.utils.telemetry import mt_cloud_telemetry -from onyx.utils.threadpool_concurrency import CallableProtocol -from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel +from onyx.utils.threadpool_concurrency import ( + CallableProtocol, + run_functions_tuples_in_parallel, +) from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/documents/credential.py b/backend/onyx/server/documents/credential.py index 51c6a6585d..c310ccab3b 100644 --- a/backend/onyx/server/documents/credential.py +++ b/backend/onyx/server/documents/credential.py @@ -1,46 +1,47 @@ import json -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import Form -from fastapi import HTTPException -from fastapi import Query -from fastapi import UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.auth.users import current_curator_or_admin_user from onyx.configs.constants import PUBLIC_API_TAGS from onyx.connectors.factory import validate_ccpair_for_user -from onyx.db.credentials import alter_credential -from onyx.db.credentials import create_credential -from onyx.db.credentials import CREDENTIAL_PERMISSIONS_TO_IGNORE -from onyx.db.credentials import delete_credential -from onyx.db.credentials import delete_credential_for_user -from onyx.db.credentials import fetch_credential_by_id_for_user -from onyx.db.credentials import fetch_credentials_by_source_for_user -from onyx.db.credentials import fetch_credentials_for_user -from onyx.db.credentials import swap_credentials_connector -from onyx.db.credentials import update_credential +from onyx.db.credentials import ( + alter_credential, + create_credential, + CREDENTIAL_PERMISSIONS_TO_IGNORE, + delete_credential, + delete_credential_for_user, + fetch_credential_by_id_for_user, + fetch_credentials_by_source_for_user, + fetch_credentials_for_user, + swap_credentials_connector, + update_credential, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.models import DocumentSource -from onyx.db.models import User -from onyx.server.documents.models import CredentialBase -from onyx.server.documents.models import CredentialDataUpdateRequest -from onyx.server.documents.models import CredentialSnapshot -from onyx.server.documents.models import CredentialSwapRequest -from onyx.server.documents.models import ObjectCreationIdResponse -from onyx.server.documents.private_key_types import FILE_TYPE_TO_FILE_PROCESSOR -from onyx.server.documents.private_key_types import PrivateKeyFileTypes -from onyx.server.documents.private_key_types import ProcessPrivateKeyFileProtocol +from onyx.db.models import DocumentSource, User +from onyx.server.documents.models import ( + CredentialBase, + CredentialDataUpdateRequest, + CredentialSnapshot, + CredentialSwapRequest, + ObjectCreationIdResponse, +) +from onyx.server.documents.private_key_types import ( + FILE_TYPE_TO_FILE_PROCESSOR, + PrivateKeyFileTypes, + ProcessPrivateKeyFileProtocol, +) from onyx.server.models import StatusResponse from onyx.server.security.store import get_security_settings -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop diff --git a/backend/onyx/server/documents/document.py b/backend/onyx/server/documents/document.py index 085e35c395..ac96ca6ab6 100644 --- a/backend/onyx/server/documents/document.py +++ b/backend/onyx/server/documents/document.py @@ -1,7 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query +from fastapi import APIRouter, Depends, HTTPException, Query from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -17,8 +14,7 @@ from onyx.document_index.factory import get_default_document_index from onyx.document_index.interfaces_new import DocumentSectionRequest from onyx.natural_language_processing.utils import get_tokenizer from onyx.prompts.prompt_utils import build_doc_context_str -from onyx.server.documents.models import ChunkInfo -from onyx.server.documents.models import DocumentInfo +from onyx.server.documents.models import ChunkInfo, DocumentInfo from onyx.server.utils_vector_db import require_vector_db router = APIRouter(prefix="/document") diff --git a/backend/onyx/server/documents/models.py b/backend/onyx/server/documents/models.py index 456a19d29d..79fb74bf73 100644 --- a/backend/onyx/server/documents/models.py +++ b/backend/onyx/server/documents/models.py @@ -1,36 +1,36 @@ from collections.abc import Sequence -from datetime import datetime -from datetime import timezone -from datetime import UTC +from datetime import datetime, timezone, UTC from enum import Enum -from typing import Any -from typing import Generic -from typing import TypeVar +from typing import Any, Generic, TypeVar from uuid import UUID -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import PermissionSyncStatus -from onyx.db.enums import ProcessingMode -from onyx.db.index_attempt_metrics_models import IndexAttemptStage -from onyx.db.index_attempt_metrics_models import STAGE_SCOPE -from onyx.db.index_attempt_metrics_models import StageScope -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import DocPermissionSyncAttempt +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + PermissionSyncStatus, + ProcessingMode, +) +from onyx.db.index_attempt_metrics_models import ( + IndexAttemptStage, + STAGE_SCOPE, + StageScope, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + DocPermissionSyncAttempt, + ExternalGroupPermissionSyncAttempt, + IndexAttempt, + IndexAttemptStageMetric, + IndexingStatus, + TaskStatus, +) from onyx.db.models import Document as DbDocument -from onyx.db.models import ExternalGroupPermissionSyncAttempt -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptStageMetric -from onyx.db.models import IndexingStatus -from onyx.db.models import TaskStatus from onyx.server.federated.models import FederatedConnectorStatus from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop diff --git a/backend/onyx/server/documents/private_key_types.py b/backend/onyx/server/documents/private_key_types.py index aa8e26b325..5dd51c8095 100644 --- a/backend/onyx/server/documents/private_key_types.py +++ b/backend/onyx/server/documents/private_key_types.py @@ -2,8 +2,7 @@ import base64 from enum import Enum from typing import Protocol -from fastapi import HTTPException -from fastapi import UploadFile +from fastapi import HTTPException, UploadFile from onyx.server.documents.document_utils import validate_pkcs12_content diff --git a/backend/onyx/server/documents/standard_oauth.py b/backend/onyx/server/documents/standard_oauth.py index 8a78785857..ee12a5bfc4 100644 --- a/backend/onyx/server/documents/standard_oauth.py +++ b/backend/onyx/server/documents/standard_oauth.py @@ -1,15 +1,9 @@ import json import uuid -from typing import Annotated -from typing import cast +from typing import Annotated, cast -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query -from fastapi import Request -from pydantic import BaseModel -from pydantic import ValidationError +from fastapi import APIRouter, Depends, HTTPException, Query, Request +from pydantic import BaseModel, ValidationError from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission diff --git a/backend/onyx/server/documents/targeted_reindex.py b/backend/onyx/server/documents/targeted_reindex.py index 46dc51a952..4fcd4914e4 100644 --- a/backend/onyx/server/documents/targeted_reindex.py +++ b/backend/onyx/server/documents/targeted_reindex.py @@ -14,25 +14,24 @@ task with the pre-allocated task UUID. import datetime from typing import Any -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from pydantic import BaseModel from sqlalchemy.orm import Session from onyx.auth.users import current_curator_or_admin_user from onyx.background.celery.versioned_apps.client import app as client_app -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import OnyxCeleryPriority, OnyxCeleryQueues, OnyxCeleryTask from onyx.db.engine.sql_engine import get_session from onyx.db.enums import IndexingStatus from onyx.db.models import User -from onyx.db.targeted_reindex import count_targets_for_job -from onyx.db.targeted_reindex import create_targeted_reindex_job -from onyx.db.targeted_reindex import get_targeted_reindex_job -from onyx.db.targeted_reindex import MAX_TARGETS_PER_REQUEST -from onyx.db.targeted_reindex import resolve_error_ids_to_targets -from onyx.db.targeted_reindex import TargetSpec +from onyx.db.targeted_reindex import ( + count_targets_for_job, + create_targeted_reindex_job, + get_targeted_reindex_job, + MAX_TARGETS_PER_REQUEST, + resolve_error_ids_to_targets, + TargetSpec, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/admin_banner/api.py b/backend/onyx/server/features/admin_banner/api.py index 177f0f24de..c74614ef45 100644 --- a/backend/onyx/server/features/admin_banner/api.py +++ b/backend/onyx/server/features/admin_banner/api.py @@ -1,15 +1,15 @@ -from fastapi import APIRouter -from fastapi import Depends -from pydantic import BaseModel -from pydantic import Field +from fastapi import APIRouter, Depends +from pydantic import BaseModel, Field from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.configs.constants import NotificationType -from onyx.db.admin_banner import AdminBanner -from onyx.db.admin_banner import clear_admin_banner -from onyx.db.admin_banner import get_admin_banner -from onyx.db.admin_banner import set_admin_banner +from onyx.db.admin_banner import ( + AdminBanner, + clear_admin_banner, + get_admin_banner, + set_admin_banner, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User diff --git a/backend/onyx/server/features/build/api.py b/backend/onyx/server/features/build/api.py index d1a50295e6..6703ad9190 100644 --- a/backend/onyx/server/features/build/api.py +++ b/backend/onyx/server/features/build/api.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -19,8 +18,10 @@ from onyx.server.features.build.external_apps.oauth import ( ) from onyx.server.features.build.interactive_turns.api import router as turns_router from onyx.server.features.build.models import BaseInstructionsResponse -from onyx.server.features.build.rate_limit import get_user_rate_limit_status -from onyx.server.features.build.rate_limit import RateLimitResponse +from onyx.server.features.build.rate_limit import ( + get_user_rate_limit_status, + RateLimitResponse, +) from onyx.server.features.build.sandbox.util.agent_instructions import ( AGENT_INSTRUCTIONS_TEMPLATE_PATH, ) diff --git a/backend/onyx/server/features/build/approvals/api.py b/backend/onyx/server/features/build/approvals/api.py index a3b3bba465..e78875a98f 100644 --- a/backend/onyx/server/features/build/approvals/api.py +++ b/backend/onyx/server/features/build/approvals/api.py @@ -6,32 +6,23 @@ wakes the parked proxy via the `approval:wake:{id}` channel; a missed wake just falls back to the proxy's wait timeout. """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import Literal +from datetime import datetime, timedelta, timezone +from typing import Any, Literal from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from pydantic import BaseModel -from pydantic import computed_field -from pydantic import ConfigDict +from fastapi import APIRouter, Depends +from pydantic import BaseModel, computed_field, ConfigDict from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.cache.factory import get_cache_backend from onyx.cache.interface import CACHE_TRANSIENT_ERRORS from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import ApprovalDecidedVia -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import Permission +from onyx.db.enums import ApprovalDecidedVia, ApprovalDecision, Permission from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.matching.engine import actions_requiring_approval -from onyx.external_apps.matching.engine import MatchedAction +from onyx.external_apps.matching.engine import actions_requiring_approval, MatchedAction from onyx.external_apps.presentation.decode import decode_payload from onyx.sandbox_proxy import approval_cache from onyx.server.features.build.configs import SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS diff --git a/backend/onyx/server/features/build/connect_app.py b/backend/onyx/server/features/build/connect_app.py index c601ea656b..e64db28fb7 100644 --- a/backend/onyx/server/features/build/connect_app.py +++ b/backend/onyx/server/features/build/connect_app.py @@ -11,8 +11,7 @@ workers can be different processes. from enum import Enum -from pydantic import BaseModel -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from onyx.cache.interface import CacheBackend from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/db/action_approval.py b/backend/onyx/server/features/build/db/action_approval.py index 3aff2bb491..62377dae39 100644 --- a/backend/onyx/server/features/build/db/action_approval.py +++ b/backend/onyx/server/features/build/db/action_approval.py @@ -3,21 +3,20 @@ `try_record_decision`'s conditional UPDATE is the only race arbiter. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from uuid import UUID -from sqlalchemy import select -from sqlalchemy import update +from sqlalchemy import select, update from sqlalchemy.orm import Session -from onyx.db.enums import ApprovalDecidedVia -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import POLICY_SEVERITY -from onyx.db.models import ActionApproval -from onyx.db.models import BuildSession +from onyx.db.enums import ( + ApprovalDecidedVia, + ApprovalDecision, + EndpointPolicy, + POLICY_SEVERITY, +) +from onyx.db.models import ActionApproval, BuildSession from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/build/db/build_session.py b/backend/onyx/server/features/build/db/build_session.py index 81366edba9..c46be5c50f 100644 --- a/backend/onyx/server/features/build/db/build_session.py +++ b/backend/onyx/server/features/build/db/build_session.py @@ -1,35 +1,26 @@ """Database operations for Build Mode sessions.""" -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from uuid import UUID -from sqlalchemy import desc -from sqlalchemy import exists -from sqlalchemy import select +from sqlalchemy import desc, exists, select from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import selectinload -from sqlalchemy.orm import Session +from sqlalchemy.orm import selectinload, Session from onyx.auth.schemas import UserRole from onyx.configs.constants import MessageType -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SessionOrigin -from onyx.db.enums import SharingScope -from onyx.db.llm import can_user_access_llm_provider -from onyx.db.llm import fetch_user_group_ids -from onyx.db.models import Artifact -from onyx.db.models import BuildMessage -from onyx.db.models import BuildSession +from onyx.db.enums import BuildSessionStatus, SessionOrigin, SharingScope +from onyx.db.llm import can_user_access_llm_provider, fetch_user_group_ids +from onyx.db.models import Artifact, BuildMessage, BuildSession, Sandbox, User from onyx.db.models import LLMProvider as LLMProviderModel -from onyx.db.models import Sandbox -from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.build.configs import BUILD_MODE_ALLOWED_PROVIDER_TYPES -from onyx.server.features.build.configs import SANDBOX_NEXTJS_PORT_END -from onyx.server.features.build.configs import SANDBOX_NEXTJS_PORT_START +from onyx.server.features.build.configs import ( + BUILD_MODE_ALLOWED_PROVIDER_TYPES, + SANDBOX_NEXTJS_PORT_END, + SANDBOX_NEXTJS_PORT_START, +) from onyx.server.manage.llm.models import LLMProviderView from onyx.utils.logger import setup_logger from onyx.utils.postgres_sanitization import sanitize_json_like diff --git a/backend/onyx/server/features/build/db/rate_limit.py b/backend/onyx/server/features/build/db/rate_limit.py index 80bb3ccfac..176fc3a0f7 100644 --- a/backend/onyx/server/features/build/db/rate_limit.py +++ b/backend/onyx/server/features/build/db/rate_limit.py @@ -7,8 +7,7 @@ from sqlalchemy import func from sqlalchemy.orm import Session from onyx.configs.constants import MessageType -from onyx.db.models import BuildMessage -from onyx.db.models import BuildSession +from onyx.db.models import BuildMessage, BuildSession def count_user_messages_in_window( diff --git a/backend/onyx/server/features/build/db/sandbox.py b/backend/onyx/server/features/build/db/sandbox.py index 7be476d966..2f3218e7bc 100644 --- a/backend/onyx/server/features/build/db/sandbox.py +++ b/backend/onyx/server/features/build/db/sandbox.py @@ -3,23 +3,13 @@ import datetime from uuid import UUID -from sqlalchemy import func -from sqlalchemy import or_ -from sqlalchemy import select +from sqlalchemy import func, or_, select from sqlalchemy.orm import Session from onyx.auth.pat import hash_pat -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import PatType -from onyx.db.enums import Permission -from onyx.db.enums import SandboxStatus -from onyx.db.models import BuildSession -from onyx.db.models import PersonalAccessToken -from onyx.db.models import Sandbox -from onyx.db.models import Snapshot -from onyx.db.models import User -from onyx.db.pat import create_pat -from onyx.db.pat import revoke_pat +from onyx.db.enums import BuildSessionStatus, PatType, Permission, SandboxStatus +from onyx.db.models import BuildSession, PersonalAccessToken, Sandbox, Snapshot, User +from onyx.db.pat import create_pat, revoke_pat from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/build/db/user_library.py b/backend/onyx/server/features/build/db/user_library.py index 8226b5fdfd..fe0b189b12 100644 --- a/backend/onyx/server/features/build/db/user_library.py +++ b/backend/onyx/server/features/build/db/user_library.py @@ -8,44 +8,43 @@ import hashlib import io from uuid import UUID -from sqlalchemy import and_ -from sqlalchemy import cast -from sqlalchemy import func -from sqlalchemy import Integer -from sqlalchemy import select +from sqlalchemy import and_, cast, func, Integer, select from sqlalchemy.orm import Session -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.constants import DocumentSource, FileOrigin from onyx.connectors.models import InputType -from onyx.db.connector import create_connector -from onyx.db.connector import fetch_connectors -from onyx.db.connector_credential_pair import add_credential_to_connector -from onyx.db.connector_credential_pair import get_connector_credential_pairs_for_user -from onyx.db.credentials import create_credential -from onyx.db.credentials import fetch_credentials_for_user -from onyx.db.document import delete_document_by_id__no_commit -from onyx.db.document import get_document -from onyx.db.document import get_documents_by_source -from onyx.db.document import update_document_metadata__no_commit -from onyx.db.document import upsert_document_by_connector_credential_pair -from onyx.db.document import upsert_documents -from onyx.db.enums import AccessType -from onyx.db.enums import ProcessingMode -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair +from onyx.db.connector import create_connector, fetch_connectors +from onyx.db.connector_credential_pair import ( + add_credential_to_connector, + get_connector_credential_pairs_for_user, +) +from onyx.db.credentials import create_credential, fetch_credentials_for_user +from onyx.db.document import ( + delete_document_by_id__no_commit, + get_document, + get_documents_by_source, + update_document_metadata__no_commit, + upsert_document_by_connector_credential_pair, + upsert_documents, +) +from onyx.db.enums import AccessType, ProcessingMode +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + DocumentByConnectorCredentialPair, + User, +) from onyx.db.models import Document as DbDocument -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import User from onyx.document_index.document_metadata import DocumentMetadata from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import get_default_file_store -from onyx.server.documents.models import ConnectorBase -from onyx.server.documents.models import CredentialBase -from onyx.server.features.build.configs import USER_LIBRARY_CONNECTOR_NAME -from onyx.server.features.build.configs import USER_LIBRARY_CREDENTIAL_NAME -from onyx.server.features.build.configs import USER_LIBRARY_SOURCE_DIR +from onyx.server.documents.models import ConnectorBase, CredentialBase +from onyx.server.features.build.configs import ( + USER_LIBRARY_CONNECTOR_NAME, + USER_LIBRARY_CREDENTIAL_NAME, + USER_LIBRARY_SOURCE_DIR, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/build/debug.py b/backend/onyx/server/features/build/debug.py index d16b6324d0..6310b60a64 100644 --- a/backend/onyx/server/features/build/debug.py +++ b/backend/onyx/server/features/build/debug.py @@ -11,16 +11,13 @@ FE debug button to tail the user's sandbox pod logs in real time. import json from collections.abc import Generator -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import Permission -from onyx.db.enums import SandboxStatus +from onyx.db.enums import Permission, SandboxStatus from onyx.db.models import User from onyx.server.features.build.configs import ENABLE_OPENCODE_DEBUGGING from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id diff --git a/backend/onyx/server/features/build/external_apps/api.py b/backend/onyx/server/features/build/external_apps/api.py index 5aade79f5b..bb2206e370 100644 --- a/backend/onyx/server/features/build/external_apps/api.py +++ b/backend/onyx/server/features/build/external_apps/api.py @@ -1,63 +1,56 @@ from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import Form -from fastapi import UploadFile +from fastapi import APIRouter, Depends, File, Form, UploadFile from pydantic import TypeAdapter from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.cache.factory import get_cache_backend from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import ExternalAppType -from onyx.db.enums import Permission -from onyx.db.enums import SandboxStatus -from onyx.db.external_app import create_external_app -from onyx.db.external_app import delete_external_app -from onyx.db.external_app import get_external_app_by_id -from onyx.db.external_app import get_external_apps -from onyx.db.external_app import get_policies -from onyx.db.external_app import get_user_credentials_by_app_id -from onyx.db.external_app import required_user_credential_keys -from onyx.db.external_app import update_external_app -from onyx.db.external_app import upsert_external_app_user_credential -from onyx.db.external_app import validate_auth_template -from onyx.db.models import ExternalApp -from onyx.db.models import ExternalAppUserCredential -from onyx.db.models import User +from onyx.db.enums import ExternalAppType, Permission, SandboxStatus +from onyx.db.external_app import ( + create_external_app, + delete_external_app, + get_external_app_by_id, + get_external_apps, + get_policies, + get_user_credentials_by_app_id, + required_user_credential_keys, + update_external_app, + upsert_external_app_user_credential, + validate_auth_template, +) +from onyx.db.models import ExternalApp, ExternalAppUserCredential, User from onyx.db.skill import affected_user_ids_for_skill -from onyx.db.utils import none_as_unset -from onyx.db.utils import UNSET +from onyx.db.utils import none_as_unset, UNSET from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.external_apps.models import BuiltInExternalAppDescriptor from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.registry import action_policy_views -from onyx.external_apps.providers.registry import fetch_available_built_in_apps -from onyx.external_apps.providers.registry import get_onyx_managed_provider -from onyx.external_apps.providers.registry import get_provider_for_app -from onyx.external_apps.providers.registry import resolve_action_overrides +from onyx.external_apps.providers.registry import ( + action_policy_views, + fetch_available_built_in_apps, + get_onyx_managed_provider, + get_provider_for_app, + resolve_action_overrides, +) from onyx.external_apps.url_glob import UrlGlob from onyx.file_store.file_store import get_default_file_store from onyx.server.features.build import connect_app from onyx.server.features.build.db.build_session import get_build_session from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id -from onyx.server.features.build.external_apps.models import ConnectAppDecisionRequest from onyx.server.features.build.external_apps.models import ( + ConnectAppDecisionRequest, CreateBuiltInExternalAppRequest, + ExternalAppAdminResponse, + ExternalAppUserResponse, + UpdateExternalAppRequest, + UpsertUserCredentialsRequest, ) -from onyx.server.features.build.external_apps.models import ExternalAppAdminResponse -from onyx.server.features.build.external_apps.models import ExternalAppUserResponse -from onyx.server.features.build.external_apps.models import UpdateExternalAppRequest -from onyx.server.features.build.external_apps.models import UpsertUserCredentialsRequest from onyx.server.features.build.sandbox.factory import get_sandbox_manager from onyx.skills.bundle import read_bundle_file -from onyx.skills.ingest import delete_bundle_blob -from onyx.skills.ingest import ingested_skill_bundle -from onyx.skills.push import push_skill_to_affected_sandboxes -from onyx.skills.push import push_skills_for_users +from onyx.skills.ingest import delete_bundle_blob, ingested_skill_bundle +from onyx.skills.push import push_skill_to_affected_sandboxes, push_skills_for_users from onyx.utils.encryption import mask_string from onyx.utils.pydantic_util import parse_json_form_field from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/features/build/external_apps/models.py b/backend/onyx/server/features/build/external_apps/models.py index 9b306d98d1..22b1382390 100644 --- a/backend/onyx/server/features/build/external_apps/models.py +++ b/backend/onyx/server/features/build/external_apps/models.py @@ -1,10 +1,8 @@ from typing import Any -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.external_apps.models import ActionPolicyView from onyx.server.features.build.connect_app import ConnectAppDecision diff --git a/backend/onyx/server/features/build/external_apps/oauth.py b/backend/onyx/server/features/build/external_apps/oauth.py index 936d55c7d6..8ff85dabf2 100644 --- a/backend/onyx/server/features/build/external_apps/oauth.py +++ b/backend/onyx/server/features/build/external_apps/oauth.py @@ -1,12 +1,10 @@ import base64 import uuid -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from urllib.parse import urlencode import requests -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from pydantic import BaseModel from sqlalchemy.orm import Session @@ -14,19 +12,22 @@ from onyx.auth.permissions import require_permission from onyx.configs.app_configs import WEB_DOMAIN from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.external_app import get_external_app_by_id -from onyx.db.external_app import upsert_external_app_user_credential -from onyx.db.models import ExternalApp -from onyx.db.models import User +from onyx.db.external_app import ( + get_external_app_by_id, + upsert_external_app_user_credential, +) +from onyx.db.models import ExternalApp, User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.external_apps.providers.base import OAuthExternalAppProvider from onyx.external_apps.providers.registry import get_provider_or_raise from onyx.external_apps.token_utils import stamp_expires_at from onyx.redis.redis_pool import get_redis_client -from onyx.server.features.build.external_apps.models import OAuthCallbackRequest -from onyx.server.features.build.external_apps.models import OAuthCallbackResponse -from onyx.server.features.build.external_apps.models import OAuthStartResponse +from onyx.server.features.build.external_apps.models import ( + OAuthCallbackRequest, + OAuthCallbackResponse, + OAuthStartResponse, +) from onyx.skills.push import push_skills_for_users from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/build/interactive_turns/api.py b/backend/onyx/server/features/build/interactive_turns/api.py index b1714da889..f6ef8f8f92 100644 --- a/backend/onyx/server/features/build/interactive_turns/api.py +++ b/backend/onyx/server/features/build/interactive_turns/api.py @@ -5,12 +5,10 @@ from __future__ import annotations import json import time from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session @@ -18,8 +16,7 @@ from onyx.auth.permissions import require_permission from onyx.cache.factory import get_cache_backend from onyx.cache.interface import CacheBackend from onyx.configs.constants import PUBLIC_API_TAGS -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant from onyx.db.enums import Permission from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode @@ -29,9 +26,11 @@ from onyx.server.features.build.interactive_turns.executor import ( start_interactive_turn_runner, ) from onyx.server.features.build.interactive_turns.models import InteractiveTurnResponse -from onyx.server.features.build.interactive_turns.state import get_active_turn -from onyx.server.features.build.interactive_turns.state import get_turn -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_FAILED +from onyx.server.features.build.interactive_turns.state import ( + get_active_turn, + get_turn, + TURN_STATUS_FAILED, +) from onyx.server.features.build.session.manager import SessionManager from onyx.server.features.build.session.streaming import SSE_KEEPALIVE from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/interactive_turns/executor.py b/backend/onyx/server/features/build/interactive_turns/executor.py index 5465355bdc..343b54cece 100644 --- a/backend/onyx/server/features/build/interactive_turns/executor.py +++ b/backend/onyx/server/features/build/interactive_turns/executor.py @@ -5,43 +5,47 @@ from __future__ import annotations import threading import time from dataclasses import dataclass -from enum import auto -from enum import Enum +from enum import auto, Enum from uuid import UUID from onyx.cache.factory import get_cache_backend -from onyx.cache.interface import CACHE_TRANSIENT_ERRORS -from onyx.cache.interface import CacheBackend +from onyx.cache.interface import CACHE_TRANSIENT_ERRORS, CacheBackend from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.server.features.build.configs import ( OPENCODE_PROMPT_INACTIVITY_TIMEOUT_SECONDS, ) from onyx.server.features.build.db.build_session import update_session_activity -from onyx.server.features.build.interactive_turns.state import claim_turn_for_runner -from onyx.server.features.build.interactive_turns.state import finish_turn -from onyx.server.features.build.interactive_turns.state import get_active_turn -from onyx.server.features.build.interactive_turns.state import InteractiveTurn -from onyx.server.features.build.interactive_turns.state import touch_turn -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_CANCELLED -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_FAILED -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_SUCCEEDED -from onyx.server.features.build.sandbox.event_schema import ActivityTimeoutError +from onyx.server.features.build.interactive_turns.state import ( + claim_turn_for_runner, + finish_turn, + get_active_turn, + InteractiveTurn, + touch_turn, + TURN_STATUS_CANCELLED, + TURN_STATUS_FAILED, + TURN_STATUS_SUCCEEDED, +) +from onyx.server.features.build.sandbox.event_schema import ( + ActivityTimeoutError, + PromptResponse, +) from onyx.server.features.build.sandbox.event_schema import Error as SandboxError -from onyx.server.features.build.sandbox.event_schema import PromptResponse from onyx.server.features.build.sandbox.serve_transport import ( PROMPT_SLOT_FAST_FAIL_ACQUIRE_SECONDS, -) -from onyx.server.features.build.sandbox.serve_transport import ( PROMPT_SLOT_WAIT_OUT_ORPHAN_SECONDS, ) from onyx.server.features.build.sandbox.sse import SSEKeepalive -from onyx.server.features.build.session.interrupt_signal import clear_interrupt -from onyx.server.features.build.session.interrupt_signal import is_interrupt_requested +from onyx.server.features.build.session.interrupt_signal import ( + clear_interrupt, + is_interrupt_requested, +) from onyx.server.features.build.session.manager import SessionManager from onyx.server.features.build.session.streaming import BuildStreamingState from onyx.utils.logger import setup_logger -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) logger = setup_logger() diff --git a/backend/onyx/server/features/build/interactive_turns/state.py b/backend/onyx/server/features/build/interactive_turns/state.py index e1d9680135..daa8cbb751 100644 --- a/backend/onyx/server/features/build/interactive_turns/state.py +++ b/backend/onyx/server/features/build/interactive_turns/state.py @@ -3,16 +3,12 @@ from __future__ import annotations import json -from dataclasses import asdict -from dataclasses import dataclass -from datetime import datetime -from datetime import timezone +from dataclasses import asdict, dataclass +from datetime import datetime, timezone from enum import StrEnum -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheLock +from onyx.cache.interface import CacheBackend, CacheLock from onyx.utils.datetime import datetime_to_utc diff --git a/backend/onyx/server/features/build/packets.py b/backend/onyx/server/features/build/packets.py index 739a78aca9..c2f02c9851 100644 --- a/backend/onyx/server/features/build/packets.py +++ b/backend/onyx/server/features/build/packets.py @@ -23,13 +23,11 @@ Custom Onyx packets (defined here): - subagent_started: A child opencode session was created under a parent turn """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Literal from uuid import UUID -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field # ============================================================================= # Base Packet Type diff --git a/backend/onyx/server/features/build/rate_limit.py b/backend/onyx/server/features/build/rate_limit.py index 11f93a4939..4355b3a2cd 100644 --- a/backend/onyx/server/features/build/rate_limit.py +++ b/backend/onyx/server/features/build/rate_limit.py @@ -1,8 +1,6 @@ """Rate limiting logic for Build Mode.""" -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Literal from pydantic import BaseModel @@ -11,9 +9,11 @@ from sqlalchemy.orm import Session from onyx.db.models import User from onyx.feature_flags.factory import get_default_feature_flag_provider from onyx.server.features.build.configs import CRAFT_PAID_USER_RATE_LIMIT -from onyx.server.features.build.db.rate_limit import count_user_messages_in_window -from onyx.server.features.build.db.rate_limit import count_user_messages_total -from onyx.server.features.build.db.rate_limit import get_oldest_message_timestamp +from onyx.server.features.build.db.rate_limit import ( + count_user_messages_in_window, + count_user_messages_total, + get_oldest_message_timestamp, +) from onyx.server.features.build.subscription_check import is_user_subscribed from onyx.server.features.build.utils import CRAFT_HAS_USAGE_LIMITS from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/features/build/sandbox/base.py b/backend/onyx/server/features/build/sandbox/base.py index f4d95432c6..88f5b1022b 100644 --- a/backend/onyx/server/features/build/sandbox/base.py +++ b/backend/onyx/server/features/build/sandbox/base.py @@ -15,30 +15,32 @@ Architecture Note (User-Shared Sandbox Model): """ import time -from abc import ABC -from abc import abstractmethod -from collections.abc import Callable -from collections.abc import Generator +from abc import ABC, abstractmethod +from collections.abc import Callable, Generator from concurrent.futures import ThreadPoolExecutor from uuid import UUID -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import AgentPlanUpdate -from onyx.server.features.build.sandbox.event_schema import AgentThoughtChunk -from onyx.server.features.build.sandbox.event_schema import CurrentModeUpdate -from onyx.server.features.build.sandbox.event_schema import Error -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import ToolCallProgress -from onyx.server.features.build.sandbox.event_schema import ToolCallStart -from onyx.server.features.build.sandbox.models import FatalWriteError -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import FilesystemEntry -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import PushFailure -from onyx.server.features.build.sandbox.models import PushResult -from onyx.server.features.build.sandbox.models import RetriableWriteError -from onyx.server.features.build.sandbox.models import SandboxInfo -from onyx.server.features.build.sandbox.models import SnapshotResult +from onyx.server.features.build.sandbox.event_schema import ( + AgentMessageChunk, + AgentPlanUpdate, + AgentThoughtChunk, + CurrentModeUpdate, + Error, + PromptResponse, + ToolCallProgress, + ToolCallStart, +) +from onyx.server.features.build.sandbox.models import ( + FatalWriteError, + FileSet, + FilesystemEntry, + LLMProviderConfig, + PushFailure, + PushResult, + RetriableWriteError, + SandboxInfo, + SnapshotResult, +) from onyx.server.features.build.sandbox.serve_transport import _ServeMixin from onyx.server.features.build.sandbox.sse import SSEKeepalive from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/sandbox/docker/docker_sandbox_manager.py b/backend/onyx/server/features/build/sandbox/docker/docker_sandbox_manager.py index 91dc485753..faf9a26397 100644 --- a/backend/onyx/server/features/build/sandbox/docker/docker_sandbox_manager.py +++ b/backend/onyx/server/features/build/sandbox/docker/docker_sandbox_manager.py @@ -74,64 +74,63 @@ from typing import TypedDict from uuid import UUID from docker import DockerClient -from docker.errors import APIError -from docker.errors import NotFound +from docker.errors import APIError, NotFound from docker.models.containers import Container from onyx.configs.app_configs import DEV_MODE from onyx.db.enums import SandboxStatus from onyx.file_store.file_store import get_default_file_store -from onyx.server.features.build.configs import ATTACHMENTS_DIRECTORY -from onyx.server.features.build.configs import OPENCODE_DISABLED_TOOLS -from onyx.server.features.build.configs import OPENCODE_SERVE_PORT -from onyx.server.features.build.configs import OPENCODE_SERVER_PASSWORD -from onyx.server.features.build.configs import SANDBOX_API_SERVER_URL -from onyx.server.features.build.configs import SANDBOX_CONTAINER_IMAGE -from onyx.server.features.build.configs import SANDBOX_DOCKER_CPU_LIMIT -from onyx.server.features.build.configs import SANDBOX_DOCKER_MEMORY_LIMIT -from onyx.server.features.build.configs import SANDBOX_DOCKER_NETWORK -from onyx.server.features.build.configs import SANDBOX_DOCKER_SOCKET -from onyx.server.features.build.configs import SANDBOX_DOCKER_VOLUME_PREFIX -from onyx.server.features.build.configs import SANDBOX_PROXY_CA_VOLUME_NAME -from onyx.server.features.build.configs import SANDBOX_PROXY_HOST -from onyx.server.features.build.configs import SANDBOX_PROXY_INJECTED_PLACEHOLDER -from onyx.server.features.build.configs import SANDBOX_PROXY_PORT -from onyx.server.features.build.sandbox.base import BUN_CACHE_DIR -from onyx.server.features.build.sandbox.base import BUN_IMAGE_CACHE_DIR -from onyx.server.features.build.sandbox.base import SandboxManager +from onyx.server.features.build.configs import ( + ATTACHMENTS_DIRECTORY, + OPENCODE_DISABLED_TOOLS, + OPENCODE_SERVE_PORT, + OPENCODE_SERVER_PASSWORD, + SANDBOX_API_SERVER_URL, + SANDBOX_CONTAINER_IMAGE, + SANDBOX_DOCKER_CPU_LIMIT, + SANDBOX_DOCKER_MEMORY_LIMIT, + SANDBOX_DOCKER_NETWORK, + SANDBOX_DOCKER_SOCKET, + SANDBOX_DOCKER_VOLUME_PREFIX, + SANDBOX_PROXY_CA_VOLUME_NAME, + SANDBOX_PROXY_HOST, + SANDBOX_PROXY_INJECTED_PLACEHOLDER, + SANDBOX_PROXY_PORT, +) +from onyx.server.features.build.sandbox.base import ( + BUN_CACHE_DIR, + BUN_IMAGE_CACHE_DIR, + SandboxManager, +) from onyx.server.features.build.sandbox.docker.dev_mode_serve import ( opencode_serve_port_bindings, -) -from onyx.server.features.build.sandbox.docker.dev_mode_serve import ( published_opencode_serve_base_url, ) -from onyx.server.features.build.sandbox.docker.internal.exec_helpers import ExecError -from onyx.server.features.build.sandbox.docker.internal.exec_helpers import ExecResult from onyx.server.features.build.sandbox.docker.internal.exec_helpers import ( + ExecError, + ExecResult, run_in_container, -) -from onyx.server.features.build.sandbox.docker.internal.exec_helpers import ( stream_stdin_to_container, -) -from onyx.server.features.build.sandbox.docker.internal.exec_helpers import ( stream_stdout_from_container, ) -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY_ONYX -from onyx.server.features.build.sandbox.labels import LABEL_SANDBOX_ID -from onyx.server.features.build.sandbox.labels import LABEL_TENANT_ID -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import FilesystemEntry -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import SandboxInfo -from onyx.server.features.build.sandbox.models import SnapshotResult +from onyx.server.features.build.sandbox.labels import ( + LABEL_K8S_MANAGED_BY, + LABEL_K8S_MANAGED_BY_ONYX, + LABEL_SANDBOX_ID, + LABEL_TENANT_ID, +) +from onyx.server.features.build.sandbox.models import ( + FileSet, + FilesystemEntry, + LLMProviderConfig, + SandboxInfo, + SnapshotResult, +) from onyx.server.features.build.sandbox.nextjs_dev import build_nextjs_start_script from onyx.server.features.build.sandbox.serve_transport import ServeConnectionInfo from onyx.server.features.build.sandbox.snapshot_manager import SnapshotManager from onyx.server.features.build.sandbox.util.agent_instructions import ( ATTACHMENTS_SECTION_CONTENT, -) -from onyx.server.features.build.sandbox.util.agent_instructions import ( generate_agent_instructions, ) from onyx.server.features.build.sandbox.util.opencode_config import ( diff --git a/backend/onyx/server/features/build/sandbox/docker/internal/exec_helpers.py b/backend/onyx/server/features/build/sandbox/docker/internal/exec_helpers.py index 5c673a25c1..62e3b74612 100644 --- a/backend/onyx/server/features/build/sandbox/docker/internal/exec_helpers.py +++ b/backend/onyx/server/features/build/sandbox/docker/internal/exec_helpers.py @@ -16,13 +16,11 @@ from __future__ import annotations import socket import struct -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Generator, Iterator from contextlib import contextmanager from dataclasses import dataclass -from docker.errors import APIError -from docker.errors import NotFound +from docker.errors import APIError, NotFound from docker.models.containers import Container from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/sandbox/event_schema.py b/backend/onyx/server/features/build/sandbox/event_schema.py index 274617a18e..ad0d168791 100644 --- a/backend/onyx/server/features/build/sandbox/event_schema.py +++ b/backend/onyx/server/features/build/sandbox/event_schema.py @@ -7,15 +7,17 @@ a future inlining of these types (or swap to a different upstream) touches only this file. See docs/craft/drop-acp-layer.md. """ -from acp.schema import AgentMessageChunk -from acp.schema import AgentPlanUpdate -from acp.schema import AgentThoughtChunk -from acp.schema import CurrentModeUpdate -from acp.schema import Error -from acp.schema import PromptResponse -from acp.schema import RequestPermissionRequest -from acp.schema import ToolCallProgress -from acp.schema import ToolCallStart +from acp.schema import ( + AgentMessageChunk, + AgentPlanUpdate, + AgentThoughtChunk, + CurrentModeUpdate, + Error, + PromptResponse, + RequestPermissionRequest, + ToolCallProgress, + ToolCallStart, +) # Onyx-synthesized ``Error.code`` sentinels for turn-terminating failures. # Negative so they never collide with opencode / JSON-RPC error codes. diff --git a/backend/onyx/server/features/build/sandbox/factory.py b/backend/onyx/server/features/build/sandbox/factory.py index ea76778e58..19453aac2e 100644 --- a/backend/onyx/server/features/build/sandbox/factory.py +++ b/backend/onyx/server/features/build/sandbox/factory.py @@ -6,8 +6,7 @@ subclasses (which all import base). import threading -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from onyx.server.features.build.sandbox.base import SandboxManager from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/contract.py b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/contract.py index 711f1e1126..456f068ffd 100644 --- a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/contract.py +++ b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/contract.py @@ -8,8 +8,7 @@ api-server imports the full module path. from uuid import UUID -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict SIDECAR_HEALTH_PATH = "/health" SIDECAR_READY_PATH = "/ready" diff --git a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/extract.py b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/extract.py index b6b51bd7b8..d1d1f41649 100644 --- a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/extract.py +++ b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/extract.py @@ -3,8 +3,7 @@ import os import shutil import tarfile import threading -from datetime import datetime -from datetime import UTC +from datetime import datetime, UTC from pathlib import Path MAX_FILE_BYTES = 25 * 1024 * 1024 # 25 MiB per entry diff --git a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/filesystem.py b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/filesystem.py index 49a960e334..8b93760470 100644 --- a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/filesystem.py +++ b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/filesystem.py @@ -6,8 +6,7 @@ import mimetypes from pathlib import Path from uuid import UUID -from sandbox_daemon.contract import FilesystemListResponse -from sandbox_daemon.contract import SidecarFilesystemEntry +from sandbox_daemon.contract import FilesystemListResponse, SidecarFilesystemEntry from sandbox_daemon.snapshot import SESSIONS_ROOT _USER_LIBRARY_LINK_TARGET = Path("/workspace/managed/user_library") diff --git a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/opencode_history.py b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/opencode_history.py index 6d9bf472d4..7492fdb583 100644 --- a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/opencode_history.py +++ b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/opencode_history.py @@ -8,8 +8,7 @@ import tempfile import threading from pathlib import Path -from sandbox_daemon.snapshot import SESSIONS_ROOT -from sandbox_daemon.snapshot import SnapshotError +from sandbox_daemon.snapshot import SESSIONS_ROOT, SnapshotError # Opencode's data home. Defaults to the K8s shared-volume mount; the Docker # backend execs this module with OPENCODE_DATA_HOME set to its own data home. diff --git a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/server.py b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/server.py index c0fbb5ae02..cab6cf4086 100644 --- a/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/server.py +++ b/backend/onyx/server/features/build/sandbox/image/sandbox_daemon/server.py @@ -13,39 +13,38 @@ from uuid import UUID import uvicorn from cryptography.exceptions import InvalidSignature from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey -from fastapi import FastAPI -from fastapi import Header -from fastapi import HTTPException -from fastapi import Query -from fastapi import Request -from fastapi import Response +from fastapi import FastAPI, Header, HTTPException, Query, Request, Response from fastapi.responses import StreamingResponse -from sandbox_daemon.contract import FilesystemListRequest -from sandbox_daemon.contract import PUSH_DAEMON_PORT -from sandbox_daemon.contract import SIDECAR_FILESYSTEM_LIST_PATH -from sandbox_daemon.contract import SIDECAR_HEALTH_PATH -from sandbox_daemon.contract import SIDECAR_OPENCODE_HISTORY_CREATE_PATH -from sandbox_daemon.contract import SIDECAR_OPENCODE_HISTORY_MARK_RESTORED_PATH -from sandbox_daemon.contract import SIDECAR_OPENCODE_HISTORY_RESTORE_PATH -from sandbox_daemon.contract import SIDECAR_PUSH_PATH -from sandbox_daemon.contract import SIDECAR_PUSH_PUBLIC_KEY_ENV_VAR -from sandbox_daemon.contract import SIDECAR_READY_PATH -from sandbox_daemon.contract import SIDECAR_SNAPSHOT_CREATE_PATH -from sandbox_daemon.contract import sidecar_snapshot_restore_path -from sandbox_daemon.contract import SIDECAR_SNAPSHOT_RESTORE_ROUTE -from sandbox_daemon.contract import SnapshotCreateRequest -from sandbox_daemon.extract import MAX_BUNDLE_BYTES -from sandbox_daemon.extract import safe_extract_then_atomic_swap -from sandbox_daemon.filesystem import FilesystemPathError -from sandbox_daemon.filesystem import list_session_directory -from sandbox_daemon.opencode_history import create_opencode_history_archive_file -from sandbox_daemon.opencode_history import mark_opencode_history_restored -from sandbox_daemon.opencode_history import opencode_history_restored -from sandbox_daemon.opencode_history import restore_opencode_history_archive -from sandbox_daemon.snapshot import has_snapshot_content -from sandbox_daemon.snapshot import iter_snapshot_archive -from sandbox_daemon.snapshot import restore_snapshot -from sandbox_daemon.snapshot import SnapshotError +from sandbox_daemon.contract import ( + FilesystemListRequest, + PUSH_DAEMON_PORT, + SIDECAR_FILESYSTEM_LIST_PATH, + SIDECAR_HEALTH_PATH, + SIDECAR_OPENCODE_HISTORY_CREATE_PATH, + SIDECAR_OPENCODE_HISTORY_MARK_RESTORED_PATH, + SIDECAR_OPENCODE_HISTORY_RESTORE_PATH, + SIDECAR_PUSH_PATH, + SIDECAR_PUSH_PUBLIC_KEY_ENV_VAR, + SIDECAR_READY_PATH, + SIDECAR_SNAPSHOT_CREATE_PATH, + sidecar_snapshot_restore_path, + SIDECAR_SNAPSHOT_RESTORE_ROUTE, + SnapshotCreateRequest, +) +from sandbox_daemon.extract import MAX_BUNDLE_BYTES, safe_extract_then_atomic_swap +from sandbox_daemon.filesystem import FilesystemPathError, list_session_directory +from sandbox_daemon.opencode_history import ( + create_opencode_history_archive_file, + mark_opencode_history_restored, + opencode_history_restored, + restore_opencode_history_archive, +) +from sandbox_daemon.snapshot import ( + has_snapshot_content, + iter_snapshot_archive, + restore_snapshot, + SnapshotError, +) app = FastAPI(title="sandbox-sidecar", docs_url=None, redoc_url=None) diff --git a/backend/onyx/server/features/build/sandbox/kubernetes/k8s_client.py b/backend/onyx/server/features/build/sandbox/kubernetes/k8s_client.py index 277b8c2a79..9984bb30bf 100644 --- a/backend/onyx/server/features/build/sandbox/kubernetes/k8s_client.py +++ b/backend/onyx/server/features/build/sandbox/kubernetes/k8s_client.py @@ -1,8 +1,7 @@ import os from typing import Any -from kubernetes import client -from kubernetes import config +from kubernetes import client, config from urllib3.util.retry import Retry from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/sandbox/kubernetes/kubernetes_sandbox_manager.py b/backend/onyx/server/features/build/sandbox/kubernetes/kubernetes_sandbox_manager.py index fc762a2a0f..2b438ff008 100644 --- a/backend/onyx/server/features/build/sandbox/kubernetes/kubernetes_sandbox_manager.py +++ b/backend/onyx/server/features/build/sandbox/kubernetes/kubernetes_sandbox_manager.py @@ -55,8 +55,7 @@ from pathlib import Path from typing import cast from uuid import UUID -from kubernetes import client -from kubernetes import watch +from kubernetes import client, watch from kubernetes.client.rest import ApiException from kubernetes.stream import stream as k8s_stream @@ -64,79 +63,67 @@ from onyx.cache.factory import get_cache_backend from onyx.cache.interface import CACHE_TRANSIENT_ERRORS from onyx.db.enums import SandboxStatus from onyx.file_store.file_store import get_default_file_store -from onyx.server.features.build.configs import OPENCODE_DISABLED_TOOLS -from onyx.server.features.build.configs import OPENCODE_SERVE_PORT -from onyx.server.features.build.configs import OPENCODE_SERVER_PASSWORD -from onyx.server.features.build.configs import SANDBOX_API_SERVER_URL -from onyx.server.features.build.configs import SANDBOX_CONTAINER_IMAGE -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SANDBOX_NEXTJS_PORT_END -from onyx.server.features.build.configs import SANDBOX_NEXTJS_PORT_START -from onyx.server.features.build.configs import SANDBOX_PROXY_HOST -from onyx.server.features.build.configs import SANDBOX_PROXY_INJECTED_PLACEHOLDER -from onyx.server.features.build.configs import SANDBOX_PROXY_NAMESPACE -from onyx.server.features.build.configs import SANDBOX_SERVICE_ACCOUNT_NAME -from onyx.server.features.build.sandbox.base import BUN_CACHE_DIR -from onyx.server.features.build.sandbox.base import BUN_IMAGE_CACHE_DIR -from onyx.server.features.build.sandbox.base import SandboxManager +from onyx.server.features.build.configs import ( + OPENCODE_DISABLED_TOOLS, + OPENCODE_SERVE_PORT, + OPENCODE_SERVER_PASSWORD, + SANDBOX_API_SERVER_URL, + SANDBOX_CONTAINER_IMAGE, + SANDBOX_NAMESPACE, + SANDBOX_NEXTJS_PORT_END, + SANDBOX_NEXTJS_PORT_START, + SANDBOX_PROXY_HOST, + SANDBOX_PROXY_INJECTED_PLACEHOLDER, + SANDBOX_PROXY_NAMESPACE, + SANDBOX_SERVICE_ACCOUNT_NAME, +) +from onyx.server.features.build.sandbox.base import ( + BUN_CACHE_DIR, + BUN_IMAGE_CACHE_DIR, + SandboxManager, +) from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( PUSH_DAEMON_PORT, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_CREATE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_MARK_RESTORED_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_RESTORE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_PUSH_PUBLIC_KEY_ENV_VAR, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_SNAPSHOT_CREATE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( sidecar_snapshot_restore_path, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SnapshotCreateRequest, ) from onyx.server.features.build.sandbox.kubernetes.k8s_client import load_kube_config from onyx.server.features.build.sandbox.kubernetes.sidecar_client import ( get_push_key_pair, -) -from onyx.server.features.build.sandbox.kubernetes.sidecar_client import SidecarClient -from onyx.server.features.build.sandbox.kubernetes.sidecar_client import ( + SidecarClient, SidecarRequestError, -) -from onyx.server.features.build.sandbox.kubernetes.sidecar_client import ( SidecarStatusError, ) -from onyx.server.features.build.sandbox.labels import LABEL_K8S_COMPONENT -from onyx.server.features.build.sandbox.labels import LABEL_K8S_COMPONENT_SANDBOX -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY_ONYX -from onyx.server.features.build.sandbox.labels import LABEL_SANDBOX_ID -from onyx.server.features.build.sandbox.labels import LABEL_TENANT_ID -from onyx.server.features.build.sandbox.models import FatalWriteError -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import FilesystemEntry -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import RetriableWriteError -from onyx.server.features.build.sandbox.models import SandboxInfo -from onyx.server.features.build.sandbox.models import SnapshotResult +from onyx.server.features.build.sandbox.labels import ( + LABEL_K8S_COMPONENT, + LABEL_K8S_COMPONENT_SANDBOX, + LABEL_K8S_MANAGED_BY, + LABEL_K8S_MANAGED_BY_ONYX, + LABEL_SANDBOX_ID, + LABEL_TENANT_ID, +) +from onyx.server.features.build.sandbox.models import ( + FatalWriteError, + FileSet, + FilesystemEntry, + LLMProviderConfig, + RetriableWriteError, + SandboxInfo, + SnapshotResult, +) from onyx.server.features.build.sandbox.nextjs_dev import build_nextjs_start_script from onyx.server.features.build.sandbox.serve_transport import ( OPENCODE_SERVE_READY_TIMEOUT_SECONDS, + ServeConnectionInfo, ) -from onyx.server.features.build.sandbox.serve_transport import ServeConnectionInfo from onyx.server.features.build.sandbox.snapshot_manager import SnapshotManager from onyx.server.features.build.sandbox.util.agent_instructions import ( ATTACHMENTS_SECTION_CONTENT, -) -from onyx.server.features.build.sandbox.util.agent_instructions import ( generate_agent_instructions, ) from onyx.server.features.build.sandbox.util.opencode_config import ( diff --git a/backend/onyx/server/features/build/sandbox/kubernetes/sidecar_client.py b/backend/onyx/server/features/build/sandbox/kubernetes/sidecar_client.py index f6cf9f6e2f..ee87573663 100644 --- a/backend/onyx/server/features/build/sandbox/kubernetes/sidecar_client.py +++ b/backend/onyx/server/features/build/sandbox/kubernetes/sidecar_client.py @@ -6,35 +6,22 @@ import base64 import binascii import hashlib import time -from collections.abc import Callable -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager -from typing import cast -from typing import IO +from typing import cast, IO from uuid import UUID import httpx from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey -from cryptography.hazmat.primitives.serialization import Encoding -from cryptography.hazmat.primitives.serialization import PublicFormat +from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat from onyx.server.features.build.configs import SANDBOX_PUSH_PRIVATE_KEY from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( FilesystemListRequest, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( FilesystemListResponse, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( PUSH_DAEMON_PORT, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_FILESYSTEM_LIST_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_HEALTH_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_PUSH_PATH, ) from onyx.server.features.build.sandbox.models import FilesystemEntry diff --git a/backend/onyx/server/features/build/sandbox/opencode/event_bus.py b/backend/onyx/server/features/build/sandbox/opencode/event_bus.py index 7336d99309..696775d470 100644 --- a/backend/onyx/server/features/build/sandbox/opencode/event_bus.py +++ b/backend/onyx/server/features/build/sandbox/opencode/event_bus.py @@ -11,11 +11,8 @@ from __future__ import annotations import json import threading from collections.abc import Callable -from dataclasses import dataclass -from dataclasses import field -from queue import Empty -from queue import Full -from queue import Queue +from dataclasses import dataclass, field +from queue import Empty, Full, Queue from typing import Any import httpx diff --git a/backend/onyx/server/features/build/sandbox/opencode/serve_client.py b/backend/onyx/server/features/build/sandbox/opencode/serve_client.py index 110c6cfb6f..60397a210f 100644 --- a/backend/onyx/server/features/build/sandbox/opencode/serve_client.py +++ b/backend/onyx/server/features/build/sandbox/opencode/serve_client.py @@ -14,13 +14,9 @@ from __future__ import annotations import queue import time -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterable -from dataclasses import dataclass -from dataclasses import field -from typing import Any -from typing import cast +from collections.abc import Callable, Generator, Iterable +from dataclasses import dataclass, field +from typing import Any, cast from uuid import uuid4 import httpx @@ -29,29 +25,35 @@ from onyx.cache.factory import get_cache_backend from onyx.server.features.build import connect_app from onyx.server.features.build.configs import ( OPENCODE_PROMPT_INACTIVITY_TIMEOUT_SECONDS, + OPENCODE_SERVE_CONNECT_TIMEOUT, + OPENCODE_SERVE_EVENT_READ_TIMEOUT, + OPENCODE_SERVE_REQUEST_TIMEOUT, + OPENCODE_SERVER_USERNAME, + SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS, + SSE_KEEPALIVE_INTERVAL, +) +from onyx.server.features.build.packets import ( + CompactionPacket, + ContextUsagePacket, + SubagentStartedPacket, +) +from onyx.server.features.build.sandbox.event_schema import ( + ActivityTimeoutError, + AgentMessageChunk, + AgentThoughtChunk, + Error, + PromptResponse, + ToolCallProgress, + ToolCallStart, + TURN_ERROR_CODE_SESSION, + TURN_ERROR_CODE_TIMEOUT, + TURN_ERROR_CODE_TRANSPORT, +) +from onyx.server.features.build.sandbox.opencode.event_bus import ( + _Subscription, + BUS_CLOSED_SENTINEL, + PodEventBus, ) -from onyx.server.features.build.configs import OPENCODE_SERVE_CONNECT_TIMEOUT -from onyx.server.features.build.configs import OPENCODE_SERVE_EVENT_READ_TIMEOUT -from onyx.server.features.build.configs import OPENCODE_SERVE_REQUEST_TIMEOUT -from onyx.server.features.build.configs import OPENCODE_SERVER_USERNAME -from onyx.server.features.build.configs import SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS -from onyx.server.features.build.configs import SSE_KEEPALIVE_INTERVAL -from onyx.server.features.build.packets import CompactionPacket -from onyx.server.features.build.packets import ContextUsagePacket -from onyx.server.features.build.packets import SubagentStartedPacket -from onyx.server.features.build.sandbox.event_schema import ActivityTimeoutError -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import AgentThoughtChunk -from onyx.server.features.build.sandbox.event_schema import Error -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import ToolCallProgress -from onyx.server.features.build.sandbox.event_schema import ToolCallStart -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_SESSION -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TIMEOUT -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TRANSPORT -from onyx.server.features.build.sandbox.opencode.event_bus import _Subscription -from onyx.server.features.build.sandbox.opencode.event_bus import BUS_CLOSED_SENTINEL -from onyx.server.features.build.sandbox.opencode.event_bus import PodEventBus from onyx.server.features.build.sandbox.sse import SSEKeepalive from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/build/sandbox/serve_transport.py b/backend/onyx/server/features/build/sandbox/serve_transport.py index e66930fe04..e3ace8eadd 100644 --- a/backend/onyx/server/features/build/sandbox/serve_transport.py +++ b/backend/onyx/server/features/build/sandbox/serve_transport.py @@ -14,8 +14,7 @@ import queue import threading import time from abc import abstractmethod -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from dataclasses import dataclass from typing import Any from uuid import UUID @@ -23,26 +22,28 @@ from uuid import UUID import httpx from onyx.cache.factory import get_cache_backend -from onyx.cache.interface import CACHE_TRANSIENT_ERRORS -from onyx.cache.interface import CacheLock -from onyx.cache.interface import CacheLockLostError +from onyx.cache.interface import CACHE_TRANSIENT_ERRORS, CacheLock, CacheLockLostError from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.server.features.build.configs import ( OPENCODE_PROMPT_INACTIVITY_TIMEOUT_SECONDS, + OPENCODE_SERVE_EVENT_READ_TIMEOUT, + OPENCODE_SERVER_USERNAME, + PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS, + PROMPT_SLOT_LEASE_SECONDS, ) -from onyx.server.features.build.configs import OPENCODE_SERVE_EVENT_READ_TIMEOUT -from onyx.server.features.build.configs import OPENCODE_SERVER_USERNAME -from onyx.server.features.build.configs import PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS -from onyx.server.features.build.configs import PROMPT_SLOT_LEASE_SECONDS from onyx.server.features.build.db.sandbox import get_sandbox_by_id -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import AgentThoughtChunk -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.opencode.event_bus import BUS_CLOSED_SENTINEL -from onyx.server.features.build.sandbox.opencode.event_bus import PodEventBus -from onyx.server.features.build.sandbox.opencode.serve_client import _TurnState -from onyx.server.features.build.sandbox.opencode.serve_client import OpencodeServeClient +from onyx.server.features.build.sandbox.event_schema import ( + AgentMessageChunk, + AgentThoughtChunk, + PromptResponse, +) +from onyx.server.features.build.sandbox.opencode.event_bus import ( + BUS_CLOSED_SENTINEL, + PodEventBus, +) from onyx.server.features.build.sandbox.opencode.serve_client import ( + _TurnState, + OpencodeServeClient, translate_opencode_event, ) from onyx.server.features.build.sandbox.sse import SSEKeepalive diff --git a/backend/onyx/server/features/build/sandbox/user_library.py b/backend/onyx/server/features/build/sandbox/user_library.py index 0c0bd8afb5..a4b144319e 100644 --- a/backend/onyx/server/features/build/sandbox/user_library.py +++ b/backend/onyx/server/features/build/sandbox/user_library.py @@ -9,8 +9,7 @@ from onyx.server.features.build.db.sandbox import get_sandbox_user_map from onyx.server.features.build.db.user_library import list_user_files from onyx.server.features.build.sandbox.base import SandboxManager from onyx.server.features.build.sandbox.factory import get_sandbox_manager -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import PushResult +from onyx.server.features.build.sandbox.models import FileSet, PushResult from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/build/scheduled_tasks/api.py b/backend/onyx/server/features/build/scheduled_tasks/api.py index e5e006eaa9..e4019521ea 100644 --- a/backend/onyx/server/features/build/scheduled_tasks/api.py +++ b/backend/onyx/server/features/build/scheduled_tasks/api.py @@ -13,50 +13,45 @@ The router is mounted under the existing ``/build`` prefix (see from __future__ import annotations -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Query -from fastapi import Response -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import model_validator +from fastapi import APIRouter, Depends, Query, Response +from pydantic import BaseModel, ConfigDict, Field, model_validator from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.background.celery.versioned_apps.client import app as celery_app -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import OnyxCeleryPriority, OnyxCeleryQueues, OnyxCeleryTask from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import Permission -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import ScheduledTaskStatus -from onyx.db.enums import ScheduledTaskTriggerSource +from onyx.db.enums import ( + Permission, + ScheduledTaskRunStatus, + ScheduledTaskStatus, + ScheduledTaskTriggerSource, +) from onyx.db.external_app import get_external_apps -from onyx.db.models import ScheduledTask -from onyx.db.models import ScheduledTaskRun -from onyx.db.models import User -from onyx.db.scheduled_task import create_scheduled_task -from onyx.db.scheduled_task import get_scheduled_task -from onyx.db.scheduled_task import insert_run -from onyx.db.scheduled_task import list_runs_for_task -from onyx.db.scheduled_task import list_scheduled_tasks_for_user -from onyx.db.scheduled_task import soft_delete_scheduled_task -from onyx.db.scheduled_task import update_scheduled_task +from onyx.db.models import ScheduledTask, ScheduledTaskRun, User +from onyx.db.scheduled_task import ( + create_scheduled_task, + get_scheduled_task, + insert_run, + list_runs_for_task, + list_scheduled_tasks_for_user, + soft_delete_scheduled_task, + update_scheduled_task, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.build.scheduled_tasks.schedule import compile_to_cron -from onyx.server.features.build.scheduled_tasks.schedule import EDITOR_PAYLOAD_MODELS -from onyx.server.features.build.scheduled_tasks.schedule import EditorMode -from onyx.server.features.build.scheduled_tasks.schedule import EditorPayload -from onyx.server.features.build.scheduled_tasks.schedule import human_readable -from onyx.server.features.build.scheduled_tasks.schedule import next_n_fires +from onyx.server.features.build.scheduled_tasks.schedule import ( + compile_to_cron, + EDITOR_PAYLOAD_MODELS, + EditorMode, + EditorPayload, + human_readable, + next_n_fires, +) from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/build/scheduled_tasks/executor.py b/backend/onyx/server/features/build/scheduled_tasks/executor.py index 5b617563c9..9c2d4efd52 100644 --- a/backend/onyx/server/features/build/scheduled_tasks/executor.py +++ b/backend/onyx/server/features/build/scheduled_tasks/executor.py @@ -41,21 +41,21 @@ import time from typing import Any from uuid import UUID -from onyx.configs.constants import MessageType -from onyx.configs.constants import NotificationType +from onyx.configs.constants import MessageType, NotificationType from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import ScheduledTaskErrorClass -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import SessionOrigin +from onyx.db.enums import ScheduledTaskErrorClass, ScheduledTaskRunStatus, SessionOrigin from onyx.db.notification import create_notification -from onyx.db.scheduled_task import get_run -from onyx.db.scheduled_task import mark_run_status -from onyx.server.features.build.db.build_session import create_message -from onyx.server.features.build.db.build_session import get_session_messages -from onyx.server.features.build.sandbox.event_schema import Error -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import RequestPermissionRequest -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TIMEOUT +from onyx.db.scheduled_task import get_run, mark_run_status +from onyx.server.features.build.db.build_session import ( + create_message, + get_session_messages, +) +from onyx.server.features.build.sandbox.event_schema import ( + Error, + PromptResponse, + RequestPermissionRequest, + TURN_ERROR_CODE_TIMEOUT, +) from onyx.server.features.build.session.manager import SessionManager from onyx.server.features.build.session.streaming import BuildStreamingState from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/scheduled_tasks/schedule.py b/backend/onyx/server/features/build/scheduled_tasks/schedule.py index 866a247340..f36ed40239 100644 --- a/backend/onyx/server/features/build/scheduled_tasks/schedule.py +++ b/backend/onyx/server/features/build/scheduled_tasks/schedule.py @@ -25,11 +25,7 @@ from typing import Literal from cron_descriptor import ExpressionDescriptor from croniter import croniter -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_validator -from pydantic import model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/onyx/server/features/build/session/api.py b/backend/onyx/server/features/build/session/api.py index 3c7aaf556d..cf4000edc4 100644 --- a/backend/onyx/server/features/build/session/api.py +++ b/backend/onyx/server/features/build/session/api.py @@ -3,81 +3,72 @@ import json import time from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import HTTPException -from fastapi import Response -from fastapi import UploadFile -from fastapi.responses import JSONResponse -from fastapi.responses import StreamingResponse +from fastapi import APIRouter, Depends, File, HTTPException, Response, UploadFile +from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel from sqlalchemy import exists from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import Permission -from onyx.db.enums import SandboxStatus -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.models import BuildMessage -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant +from onyx.db.enums import ( + BuildSessionStatus, + Permission, + SandboxStatus, + ScheduledTaskRunStatus, +) +from onyx.db.models import BuildMessage, Sandbox, User from onyx.db.scheduled_task import get_scheduled_run_context from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.redis.redis_pool import get_redis_client -from onyx.server.features.build.db.build_session import allocate_nextjs_port -from onyx.server.features.build.db.build_session import get_build_session -from onyx.server.features.build.db.build_session import set_build_session_sharing_scope -from onyx.server.features.build.db.build_session import skills_are_stale -from onyx.server.features.build.db.sandbox import get_latest_snapshot_for_session -from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id -from onyx.server.features.build.db.sandbox import update_sandbox_heartbeat +from onyx.server.features.build.db.build_session import ( + allocate_nextjs_port, + get_build_session, + set_build_session_sharing_scope, + skills_are_stale, +) +from onyx.server.features.build.db.sandbox import ( + get_latest_snapshot_for_session, + get_sandbox_by_user_id, + update_sandbox_heartbeat, +) from onyx.server.features.build.models import UploadResponse from onyx.server.features.build.sandbox.factory import get_sandbox_manager from onyx.server.features.build.sandbox.models import DirectoryListing from onyx.server.features.build.session.errors import UploadLimitExceededError from onyx.server.features.build.session.manager import SessionManager -from onyx.server.features.build.session.models import ArtifactResponse -from onyx.server.features.build.session.models import DetailedSessionResponse -from onyx.server.features.build.session.models import OpencodeHistorySnapshotResponse -from onyx.server.features.build.session.models import PptxPreviewResponse -from onyx.server.features.build.session.models import PreProvisionedCheckResponse -from onyx.server.features.build.session.models import SandboxStatusResponse -from onyx.server.features.build.session.models import SessionCreateRequest -from onyx.server.features.build.session.models import SessionListResponse -from onyx.server.features.build.session.models import SessionNameGenerateResponse -from onyx.server.features.build.session.models import SessionResponse -from onyx.server.features.build.session.models import SessionSkillsStateResponse -from onyx.server.features.build.session.models import SessionUpdateRequest -from onyx.server.features.build.session.models import SetSessionSharingRequest -from onyx.server.features.build.session.models import SetSessionSharingResponse -from onyx.server.features.build.session.models import SnapshotResponse -from onyx.server.features.build.session.models import WebappInfo +from onyx.server.features.build.session.models import ( + ArtifactResponse, + DetailedSessionResponse, + OpencodeHistorySnapshotResponse, + PptxPreviewResponse, + PreProvisionedCheckResponse, + SandboxStatusResponse, + SessionCreateRequest, + SessionListResponse, + SessionNameGenerateResponse, + SessionResponse, + SessionSkillsStateResponse, + SessionUpdateRequest, + SetSessionSharingRequest, + SetSessionSharingResponse, + SnapshotResponse, + WebappInfo, +) from onyx.server.features.build.session.sandbox_lifecycle import ( create_session_snapshot_keep_latest, -) -from onyx.server.features.build.session.sandbox_lifecycle import hydrate_managed_content -from onyx.server.features.build.session.sandbox_lifecycle import ( + hydrate_managed_content, mark_sandbox_provisioning, -) -from onyx.server.features.build.session.sandbox_lifecycle import provision_sandbox -from onyx.server.features.build.session.sandbox_lifecycle import ( + provision_sandbox, recover_unhealthy_sandbox, -) -from onyx.server.features.build.session.sandbox_lifecycle import ( rollback_failed_provisioning, ) from onyx.server.features.build.session.streaming import SSE_KEEPALIVE -from onyx.server.features.build.utils import sanitize_filename -from onyx.server.features.build.utils import validate_file +from onyx.server.features.build.utils import sanitize_filename, validate_file from onyx.skills.push import build_user_skills_payload from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/build/session/llm_config.py b/backend/onyx/server/features/build/session/llm_config.py index 801c200ca0..12668d20b9 100644 --- a/backend/onyx/server/features/build/session/llm_config.py +++ b/backend/onyx/server/features/build/session/llm_config.py @@ -11,8 +11,10 @@ from onyx.error_handling.exceptions import OnyxError from onyx.llm.well_known_providers.llm_provider_options import ( fetch_default_model_for_provider, ) -from onyx.server.features.build.configs import BUILD_MODE_ALLOWED_PROVIDER_TYPES -from onyx.server.features.build.configs import BUILD_MODE_NOT_CONFIGURED_API_KEY +from onyx.server.features.build.configs import ( + BUILD_MODE_ALLOWED_PROVIDER_TYPES, + BUILD_MODE_NOT_CONFIGURED_API_KEY, +) from onyx.server.features.build.sandbox.models import LLMProviderConfig from onyx.server.manage.llm.models import LLMProviderView from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/session/manager.py b/backend/onyx/server/features/build/session/manager.py index faec668192..e0c4398f03 100644 --- a/backend/onyx/server/features/build/session/manager.py +++ b/backend/onyx/server/features/build/session/manager.py @@ -11,12 +11,9 @@ import mimetypes import threading import uuid import zipfile -from collections.abc import Callable -from collections.abc import Generator -from contextlib import AbstractContextManager -from contextlib import nullcontext -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Generator +from contextlib import AbstractContextManager, nullcontext +from datetime import datetime, timezone from pathlib import Path from typing import Any from uuid import UUID @@ -26,59 +23,67 @@ from sqlalchemy.orm import Session as DBSession from onyx.cache.factory import get_cache_backend from onyx.configs.app_configs import WEB_DOMAIN -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SessionOrigin +from onyx.db.enums import SandboxStatus, SessionOrigin from onyx.db.external_app import get_connectable_apps_for_user -from onyx.db.models import BuildMessage -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import BuildMessage, BuildSession, Sandbox, User from onyx.db.users import fetch_user_by_id from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import get_default_file_store -from onyx.server.features.build.configs import MAX_TOTAL_UPLOAD_SIZE_BYTES -from onyx.server.features.build.configs import MAX_UPLOAD_FILES_PER_SESSION -from onyx.server.features.build.configs import PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS -from onyx.server.features.build.db.build_session import allocate_nextjs_port -from onyx.server.features.build.db.build_session import create_build_session__no_commit -from onyx.server.features.build.db.build_session import delete_build_session__no_commit +from onyx.server.features.build.configs import ( + MAX_TOTAL_UPLOAD_SIZE_BYTES, + MAX_UPLOAD_FILES_PER_SESSION, + PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS, +) from onyx.server.features.build.db.build_session import ( + allocate_nextjs_port, + create_build_session__no_commit, + delete_build_session__no_commit, fetch_all_supported_build_llm_providers, + get_build_session, + get_empty_session_for_user, + get_session_messages, + get_user_build_sessions, + skills_are_stale, + update_session_activity, +) +from onyx.server.features.build.db.sandbox import ( + get_sandbox_by_user_id, + get_snapshots_for_session, + update_sandbox_heartbeat, ) -from onyx.server.features.build.db.build_session import get_build_session -from onyx.server.features.build.db.build_session import get_empty_session_for_user -from onyx.server.features.build.db.build_session import get_session_messages -from onyx.server.features.build.db.build_session import get_user_build_sessions -from onyx.server.features.build.db.build_session import skills_are_stale -from onyx.server.features.build.db.build_session import update_session_activity -from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id -from onyx.server.features.build.db.sandbox import get_snapshots_for_session -from onyx.server.features.build.db.sandbox import update_sandbox_heartbeat from onyx.server.features.build.rate_limit import get_user_rate_limit_status from onyx.server.features.build.sandbox.factory import get_sandbox_manager -from onyx.server.features.build.sandbox.models import DirectoryListing -from onyx.server.features.build.sandbox.models import FilesystemEntry -from onyx.server.features.build.sandbox.models import LLMProviderConfig +from onyx.server.features.build.sandbox.models import ( + DirectoryListing, + FilesystemEntry, + LLMProviderConfig, +) from onyx.server.features.build.sandbox.serve_transport import ( PROMPT_SLOT_FAST_FAIL_ACQUIRE_SECONDS, + PromptSlot, ) -from onyx.server.features.build.sandbox.serve_transport import PromptSlot from onyx.server.features.build.sandbox.snapshot_manager import SnapshotManager from onyx.server.features.build.sandbox.util.agent_instructions import ( build_connectable_apps_list, ) from onyx.server.features.build.session import streaming as _streaming -from onyx.server.features.build.session.errors import RateLimitError -from onyx.server.features.build.session.errors import UploadLimitExceededError +from onyx.server.features.build.session.errors import ( + RateLimitError, + UploadLimitExceededError, +) from onyx.server.features.build.session.interrupt_signal import request_interrupt -from onyx.server.features.build.session.llm_config import get_all_build_mode_llm_configs -from onyx.server.features.build.session.llm_config import select_default_llm_config +from onyx.server.features.build.session.llm_config import ( + get_all_build_mode_llm_configs, + select_default_llm_config, +) from onyx.server.features.build.session.md_to_docx import markdown_to_docx_bytes from onyx.server.features.build.session.naming import generate_session_name -from onyx.server.features.build.session.sandbox_lifecycle import ensure_sandbox_ready -from onyx.server.features.build.session.sandbox_lifecycle import hydrate_managed_content -from onyx.server.features.build.session.sandbox_lifecycle import ProvisioningPolicy +from onyx.server.features.build.session.sandbox_lifecycle import ( + ensure_sandbox_ready, + hydrate_managed_content, + ProvisioningPolicy, +) from onyx.server.features.build.session.streaming import BuildStreamingState from onyx.skills.push import build_user_skills_payload from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/build/session/md_to_docx.py b/backend/onyx/server/features/build/session/md_to_docx.py index 7e7941192f..c654f3c2da 100644 --- a/backend/onyx/server/features/build/session/md_to_docx.py +++ b/backend/onyx/server/features/build/session/md_to_docx.py @@ -13,12 +13,10 @@ start values), blockquotes, fenced code blocks, GFM tables, hyperlinks dropped rather than shown as literal markup. """ -from dataclasses import dataclass -from dataclasses import replace +from dataclasses import dataclass, replace from html import unescape from io import BytesIO -from typing import Any -from typing import cast +from typing import Any, cast import mistune from docx import Document @@ -29,13 +27,9 @@ from docx.opc.constants import RELATIONSHIP_TYPE from docx.opc.package import OpcPackage from docx.opc.packuri import PackURI from docx.opc.part import XmlPart -from docx.oxml import OxmlElement -from docx.oxml import parse_xml +from docx.oxml import OxmlElement, parse_xml from docx.oxml.ns import qn -from docx.shared import Inches -from docx.shared import Pt -from docx.shared import RGBColor -from docx.shared import Twips +from docx.shared import Inches, Pt, RGBColor, Twips from docx.styles.style import ParagraphStyle from docx.table import _Cell from docx.text.paragraph import Paragraph diff --git a/backend/onyx/server/features/build/session/messages.py b/backend/onyx/server/features/build/session/messages.py index 535a9d0a0d..38007d30b7 100644 --- a/backend/onyx/server/features/build/session/messages.py +++ b/backend/onyx/server/features/build/session/messages.py @@ -1,48 +1,51 @@ """API endpoints for Build Mode message management.""" from collections.abc import Generator -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.cache.factory import get_cache_backend -from onyx.configs.constants import MessageType -from onyx.configs.constants import PUBLIC_API_TAGS -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.configs.constants import MessageType, PUBLIC_API_TAGS +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant from onyx.db.enums import Permission from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.build.db.build_session import count_user_messages -from onyx.server.features.build.db.build_session import create_message -from onyx.server.features.build.db.build_session import get_build_session -from onyx.server.features.build.db.build_session import skills_are_stale -from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id -from onyx.server.features.build.db.sandbox import update_sandbox_heartbeat +from onyx.server.features.build.db.build_session import ( + count_user_messages, + create_message, + get_build_session, + skills_are_stale, +) +from onyx.server.features.build.db.sandbox import ( + get_sandbox_by_user_id, + update_sandbox_heartbeat, +) from onyx.server.features.build.interactive_turns.executor import ( start_interactive_turn_runner, ) from onyx.server.features.build.interactive_turns.models import InteractiveTurnResponse -from onyx.server.features.build.interactive_turns.state import acquire_active_turn_lock -from onyx.server.features.build.interactive_turns.state import create_interactive_turn -from onyx.server.features.build.interactive_turns.state import finish_turn -from onyx.server.features.build.interactive_turns.state import get_active_turn -from onyx.server.features.build.interactive_turns.state import get_turn_for_request -from onyx.server.features.build.interactive_turns.state import InteractiveTurnLockError -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_FAILED +from onyx.server.features.build.interactive_turns.state import ( + acquire_active_turn_lock, + create_interactive_turn, + finish_turn, + get_active_turn, + get_turn_for_request, + InteractiveTurnLockError, + TURN_STATUS_FAILED, +) from onyx.server.features.build.session.errors import RateLimitError from onyx.server.features.build.session.manager import SessionManager -from onyx.server.features.build.session.models import MessageInterruptResponse -from onyx.server.features.build.session.models import MessageListResponse -from onyx.server.features.build.session.models import MessageRequest -from onyx.server.features.build.session.models import MessageResponse +from onyx.server.features.build.session.models import ( + MessageInterruptResponse, + MessageListResponse, + MessageRequest, + MessageResponse, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/build/session/models.py b/backend/onyx/server/features/build/session/models.py index 3559927f08..717a2ecd65 100644 --- a/backend/onyx/server/features/build/session/models.py +++ b/backend/onyx/server/features/build/session/models.py @@ -1,21 +1,20 @@ from datetime import datetime -from typing import Any -from typing import TYPE_CHECKING -from typing import Union +from typing import Any, TYPE_CHECKING, Union from pydantic import BaseModel from onyx.configs.constants import MessageType -from onyx.db.enums import ArtifactType -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SessionOrigin -from onyx.db.enums import SharingScope +from onyx.db.enums import ( + ArtifactType, + BuildSessionStatus, + SandboxStatus, + SessionOrigin, + SharingScope, +) from onyx.server.features.build.db.build_session import skills_are_stale if TYPE_CHECKING: - from onyx.db.models import BuildSession - from onyx.db.models import Sandbox + from onyx.db.models import BuildSession, Sandbox # ===== Session Models ===== diff --git a/backend/onyx/server/features/build/session/naming.py b/backend/onyx/server/features/build/session/naming.py index 955ae184dd..3412ff22c1 100644 --- a/backend/onyx/server/features/build/session/naming.py +++ b/backend/onyx/server/features/build/session/naming.py @@ -12,16 +12,17 @@ from sqlalchemy.orm import Session as DBSession from onyx.configs.constants import MessageType from onyx.llm.factory import get_default_llm -from onyx.llm.models import LanguageModelInput -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import SystemMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ( + LanguageModelInput, + ReasoningEffort, + SystemMessage, + UserMessage, +) from onyx.llm.utils import llm_response_to_string from onyx.server.features.build.db.build_session import get_session_messages from onyx.tracing.flows import LLMFlow from onyx.tracing.framework.create import ensure_trace -from onyx.tracing.llm_utils import llm_generation_span -from onyx.tracing.llm_utils import record_llm_response +from onyx.tracing.llm_utils import llm_generation_span, record_llm_response from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/build/session/sandbox_lifecycle.py b/backend/onyx/server/features/build/session/sandbox_lifecycle.py index 7d4d033170..5dc0ecf21c 100644 --- a/backend/onyx/server/features/build/session/sandbox_lifecycle.py +++ b/backend/onyx/server/features/build/session/sandbox_lifecycle.py @@ -3,44 +3,49 @@ interactive (``create_session__no_commit``) and headless (``ensure_sandbox_running``) flows.""" import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from enum import Enum from uuid import UUID from sqlalchemy.orm import Session as DBSession from onyx.db.enums import SandboxStatus -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import Sandbox, User from onyx.db.users import fetch_user_by_id from onyx.file_store.file_store import get_default_file_store -from onyx.server.features.build.configs import SANDBOX_IDLE_TIMEOUT_SECONDS -from onyx.server.features.build.configs import SANDBOX_MAX_CONCURRENT_PER_ORG -from onyx.server.features.build.db.build_session import clear_nextjs_ports_for_user +from onyx.server.features.build.configs import ( + SANDBOX_IDLE_TIMEOUT_SECONDS, + SANDBOX_MAX_CONCURRENT_PER_ORG, +) from onyx.server.features.build.db.build_session import ( + clear_nextjs_ports_for_user, mark_user_sessions_idle__no_commit, ) -from onyx.server.features.build.db.sandbox import create_sandbox__no_commit -from onyx.server.features.build.db.sandbox import create_snapshot__no_commit -from onyx.server.features.build.db.sandbox import delete_snapshot__no_commit -from onyx.server.features.build.db.sandbox import ensure_sandbox_pat -from onyx.server.features.build.db.sandbox import get_running_sandbox_count -from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id -from onyx.server.features.build.db.sandbox import get_snapshots_for_session -from onyx.server.features.build.db.sandbox import set_sandbox_skills_hashes__no_commit -from onyx.server.features.build.db.sandbox import update_sandbox_status__no_commit +from onyx.server.features.build.db.sandbox import ( + create_sandbox__no_commit, + create_snapshot__no_commit, + delete_snapshot__no_commit, + ensure_sandbox_pat, + get_running_sandbox_count, + get_sandbox_by_user_id, + get_snapshots_for_session, + set_sandbox_skills_hashes__no_commit, + update_sandbox_status__no_commit, +) from onyx.server.features.build.sandbox.base import SandboxManager -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import SnapshotResult +from onyx.server.features.build.sandbox.models import ( + FileSet, + LLMProviderConfig, + SnapshotResult, +) from onyx.server.features.build.sandbox.snapshot_manager import SnapshotManager from onyx.server.features.build.sandbox.user_library import hydrate_user_library from onyx.server.features.build.session.errors import SandboxProvisioningError -from onyx.skills.push import build_skills_fileset_for_user -from onyx.skills.push import compute_skills_hash -from onyx.skills.push import hydrate_sandbox_skills +from onyx.skills.push import ( + build_skills_fileset_for_user, + compute_skills_hash, + hydrate_sandbox_skills, +) from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/build/session/streaming.py b/backend/onyx/server/features/build/session/streaming.py index d1336e1ee7..1cb154fbc4 100644 --- a/backend/onyx/server/features/build/session/streaming.py +++ b/backend/onyx/server/features/build/session/streaming.py @@ -17,10 +17,8 @@ import json import queue as queue_lib import threading import time -from collections.abc import Callable -from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Generator +from datetime import datetime, timezone from typing import Any from uuid import UUID @@ -33,32 +31,40 @@ from onyx.db.enums import SandboxStatus from onyx.db.models import BuildSession from onyx.sandbox_proxy import approval_cache from onyx.server.features.build import connect_app -from onyx.server.features.build.configs import PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS from onyx.server.features.build.configs import ( + PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS, SANDBOX_HEARTBEAT_REFRESH_INTERVAL_SECONDS, ) -from onyx.server.features.build.db.build_session import create_message -from onyx.server.features.build.db.build_session import get_build_session -from onyx.server.features.build.db.build_session import update_session_activity -from onyx.server.features.build.db.build_session import upsert_agent_plan -from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id -from onyx.server.features.build.db.sandbox import update_sandbox_heartbeat -from onyx.server.features.build.packets import ApprovalRequestedPacket -from onyx.server.features.build.packets import BuildPacket -from onyx.server.features.build.packets import CompactionPacket -from onyx.server.features.build.packets import ConnectAppRequestPacket -from onyx.server.features.build.packets import ContextUsagePacket -from onyx.server.features.build.packets import ErrorPacket -from onyx.server.features.build.packets import SubagentStartedPacket +from onyx.server.features.build.db.build_session import ( + create_message, + get_build_session, + update_session_activity, + upsert_agent_plan, +) +from onyx.server.features.build.db.sandbox import ( + get_sandbox_by_user_id, + update_sandbox_heartbeat, +) +from onyx.server.features.build.packets import ( + ApprovalRequestedPacket, + BuildPacket, + CompactionPacket, + ConnectAppRequestPacket, + ContextUsagePacket, + ErrorPacket, + SubagentStartedPacket, +) from onyx.server.features.build.sandbox.base import SandboxManager -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import AgentPlanUpdate -from onyx.server.features.build.sandbox.event_schema import AgentThoughtChunk -from onyx.server.features.build.sandbox.event_schema import CurrentModeUpdate +from onyx.server.features.build.sandbox.event_schema import ( + AgentMessageChunk, + AgentPlanUpdate, + AgentThoughtChunk, + CurrentModeUpdate, + PromptResponse, + ToolCallProgress, + ToolCallStart, +) from onyx.server.features.build.sandbox.event_schema import Error as SandboxError -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import ToolCallProgress -from onyx.server.features.build.sandbox.event_schema import ToolCallStart from onyx.server.features.build.sandbox.opencode.serve_client import _merge_field_meta from onyx.server.features.build.sandbox.serve_transport import PromptSlot from onyx.server.features.build.sandbox.sse import SSEKeepalive diff --git a/backend/onyx/server/features/build/user_library/api.py b/backend/onyx/server/features/build/user_library/api.py index a4c4fb2cca..eea8b34057 100644 --- a/backend/onyx/server/features/build/user_library/api.py +++ b/backend/onyx/server/features/build/user_library/api.py @@ -3,40 +3,40 @@ import mimetypes import re import zipfile -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from io import BytesIO -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import Form -from fastapi import UploadFile +from fastapi import APIRouter, Depends, File, Form, UploadFile from pydantic import BaseModel from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission -from onyx.configs.app_configs import MAX_EMBEDDED_IMAGES_PER_FILE -from onyx.configs.app_configs import MAX_EMBEDDED_IMAGES_PER_UPLOAD +from onyx.configs.app_configs import ( + MAX_EMBEDDED_IMAGES_PER_FILE, + MAX_EMBEDDED_IMAGES_PER_UPLOAD, +) from onyx.db.connector_credential_pair import update_connector_credential_pair from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import Permission +from onyx.db.enums import ConnectorCredentialPairStatus, Permission from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_processing.extract_file_text import count_pdf_embedded_images -from onyx.server.features.build.configs import USER_LIBRARY_MAX_FILE_SIZE_BYTES -from onyx.server.features.build.configs import USER_LIBRARY_MAX_FILES_PER_UPLOAD -from onyx.server.features.build.configs import USER_LIBRARY_MAX_TOTAL_SIZE_BYTES -from onyx.server.features.build.db.user_library import cleanup_old_blobs -from onyx.server.features.build.db.user_library import create_directory_record -from onyx.server.features.build.db.user_library import delete_user_file -from onyx.server.features.build.db.user_library import fetch_user_file_for_user -from onyx.server.features.build.db.user_library import get_or_create_craft_connector -from onyx.server.features.build.db.user_library import get_user_storage_bytes -from onyx.server.features.build.db.user_library import list_user_files -from onyx.server.features.build.db.user_library import store_user_file +from onyx.server.features.build.configs import ( + USER_LIBRARY_MAX_FILE_SIZE_BYTES, + USER_LIBRARY_MAX_FILES_PER_UPLOAD, + USER_LIBRARY_MAX_TOTAL_SIZE_BYTES, +) +from onyx.server.features.build.db.user_library import ( + cleanup_old_blobs, + create_directory_record, + delete_user_file, + fetch_user_file_for_user, + get_or_create_craft_connector, + get_user_storage_bytes, + list_user_files, + store_user_file, +) from onyx.server.features.build.sandbox.user_library import ( sync_user_library_to_active_sandboxes, ) diff --git a/backend/onyx/server/features/build/utils.py b/backend/onyx/server/features/build/utils.py index b11667b7a8..41662c31f1 100644 --- a/backend/onyx/server/features/build/utils.py +++ b/backend/onyx/server/features/build/utils.py @@ -11,8 +11,7 @@ from onyx.db.models import User from onyx.db.notification import create_notification from onyx.feature_flags.factory import get_default_feature_flag_provider from onyx.feature_flags.interface import NoOpFeatureFlagProvider -from onyx.server.features.build.configs import ENABLE_CRAFT -from onyx.server.features.build.configs import MAX_UPLOAD_FILE_SIZE_BYTES +from onyx.server.features.build.configs import ENABLE_CRAFT, MAX_UPLOAD_FILE_SIZE_BYTES from onyx.server.settings.store import load_settings from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/build/webapp_proxy.py b/backend/onyx/server/features/build/webapp_proxy.py index a4a3312046..75bd7480da 100644 --- a/backend/onyx/server/features/build/webapp_proxy.py +++ b/backend/onyx/server/features/build/webapp_proxy.py @@ -1,21 +1,21 @@ import asyncio from collections.abc import AsyncGenerator from pathlib import Path -from urllib.parse import parse_qsl -from urllib.parse import urlencode +from urllib.parse import parse_qsl, urlencode from uuid import UUID import httpx -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Request -from fastapi import Response -from fastapi import WebSocket -from fastapi import WebSocketDisconnect -from fastapi import WebSocketException -from fastapi.responses import RedirectResponse -from fastapi.responses import StreamingResponse +from fastapi import ( + APIRouter, + Depends, + HTTPException, + Request, + Response, + WebSocket, + WebSocketDisconnect, + WebSocketException, +) +from fastapi.responses import RedirectResponse, StreamingResponse from fastapi_users.authentication.strategy.base import Strategy from fastapi_users.manager import BaseUserManager from starlette.websockets import WebSocketState @@ -24,19 +24,22 @@ from websockets.asyncio.client import connect as websocket_connect from websockets.exceptions import ConnectionClosed from onyx.auth.permissions import get_effective_permissions -from onyx.auth.users import auth_backend -from onyx.auth.users import get_user_manager -from onyx.auth.users import is_same_origin -from onyx.auth.users import optional_user +from onyx.auth.users import ( + auth_backend, + get_user_manager, + is_same_origin, + optional_user, +) from onyx.cache.factory import get_cache_backend from onyx.configs.app_configs import WEB_DOMAIN from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.db.engine.async_sql_engine import get_async_session_context_manager -from onyx.db.enums import Permission -from onyx.db.enums import SharingScope +from onyx.db.enums import Permission, SharingScope from onyx.db.models import User -from onyx.server.features.build.db.build_session import get_webapp_access_async -from onyx.server.features.build.db.build_session import get_webapp_target_async +from onyx.server.features.build.db.build_session import ( + get_webapp_access_async, + get_webapp_target_async, +) from onyx.server.features.build.sandbox.factory import get_sandbox_manager from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/default_assistant/api.py b/backend/onyx/server/features/default_assistant/api.py index 483ebc58ff..5cd757bf0a 100644 --- a/backend/onyx/server/features/default_assistant/api.py +++ b/backend/onyx/server/features/default_assistant/api.py @@ -1,19 +1,21 @@ """API endpoints for default assistant configuration.""" -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.persona import get_default_assistant -from onyx.db.persona import update_default_assistant_configuration +from onyx.db.persona import ( + get_default_assistant, + update_default_assistant_configuration, +) from onyx.prompts.chat_prompts import DEFAULT_SYSTEM_PROMPT -from onyx.server.features.default_assistant.models import DefaultAssistantConfiguration -from onyx.server.features.default_assistant.models import DefaultAssistantUpdateRequest +from onyx.server.features.default_assistant.models import ( + DefaultAssistantConfiguration, + DefaultAssistantUpdateRequest, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/default_assistant/models.py b/backend/onyx/server/features/default_assistant/models.py index e5abd2dc7b..3770a57372 100644 --- a/backend/onyx/server/features/default_assistant/models.py +++ b/backend/onyx/server/features/default_assistant/models.py @@ -1,7 +1,6 @@ """Models for default assistant configuration API.""" -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field class DefaultAssistantConfiguration(BaseModel): diff --git a/backend/onyx/server/features/document_set/api.py b/backend/onyx/server/features/document_set/api.py index aeea8631ef..a95db23fb2 100644 --- a/backend/onyx/server/features/document_set/api.py +++ b/backend/onyx/server/features/document_set/api.py @@ -1,30 +1,30 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query +from fastapi import APIRouter, Depends, HTTPException, Query from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.auth.users import current_curator_or_admin_user from onyx.background.celery.versioned_apps.client import app as client_app from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryTask -from onyx.db.document_set import check_document_sets_are_public +from onyx.configs.constants import OnyxCeleryPriority, OnyxCeleryTask +from onyx.db.document_set import ( + check_document_sets_are_public, + fetch_all_document_sets_for_user, + get_document_set_by_id, + insert_document_set, + mark_document_set_as_to_be_deleted, + update_document_set, +) from onyx.db.document_set import delete_document_set as db_delete_document_set -from onyx.db.document_set import fetch_all_document_sets_for_user -from onyx.db.document_set import get_document_set_by_id -from onyx.db.document_set import insert_document_set -from onyx.db.document_set import mark_document_set_as_to_be_deleted -from onyx.db.document_set import update_document_set from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.server.features.document_set.models import CheckDocSetPublicRequest -from onyx.server.features.document_set.models import CheckDocSetPublicResponse -from onyx.server.features.document_set.models import DocumentSetCreationRequest -from onyx.server.features.document_set.models import DocumentSetSummary -from onyx.server.features.document_set.models import DocumentSetUpdateRequest +from onyx.server.features.document_set.models import ( + CheckDocSetPublicRequest, + CheckDocSetPublicResponse, + DocumentSetCreationRequest, + DocumentSetSummary, + DocumentSetUpdateRequest, +) from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/document_set/models.py b/backend/onyx/server/features/document_set/models.py index c4135df752..1eaf897ceb 100644 --- a/backend/onyx/server/features/document_set/models.py +++ b/backend/onyx/server/features/document_set/models.py @@ -1,15 +1,16 @@ from typing import Any from uuid import UUID -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.db.models import DocumentSet as DocumentSetDBModel from onyx.db.models import FederatedConnector__DocumentSet -from onyx.server.documents.models import CCPairSummary -from onyx.server.documents.models import ConnectorCredentialPairDescriptor -from onyx.server.documents.models import ConnectorSnapshot -from onyx.server.documents.models import CredentialSnapshot +from onyx.server.documents.models import ( + CCPairSummary, + ConnectorCredentialPairDescriptor, + ConnectorSnapshot, + CredentialSnapshot, +) from onyx.server.federated.models import FederatedConnectorSummary diff --git a/backend/onyx/server/features/hierarchy/api.py b/backend/onyx/server/features/hierarchy/api.py index d781da2572..275a8249fc 100644 --- a/backend/onyx/server/features/hierarchy/api.py +++ b/backend/onyx/server/features/hierarchy/api.py @@ -1,7 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query +from fastapi import APIRouter, Depends, HTTPException, Query from sqlalchemy.orm import Session from onyx.access.hierarchy_access import get_user_external_group_ids @@ -11,26 +8,32 @@ from onyx.configs.constants import DocumentSource from onyx.db.document import get_accessible_documents_for_hierarchy_node_paginated from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.hierarchy import get_accessible_hierarchy_nodes_for_source -from onyx.db.hierarchy import search_accessible_hierarchy_nodes +from onyx.db.hierarchy import ( + get_accessible_hierarchy_nodes_for_source, + search_accessible_hierarchy_nodes, +) from onyx.db.models import User from onyx.db.opensearch_migration import get_opensearch_retrieval_state -from onyx.server.features.hierarchy.constants import DOCUMENT_PAGE_SIZE -from onyx.server.features.hierarchy.constants import HIERARCHY_NODE_DOCUMENTS_PATH -from onyx.server.features.hierarchy.constants import HIERARCHY_NODE_SEARCH_LIMIT -from onyx.server.features.hierarchy.constants import HIERARCHY_NODES_LIST_PATH -from onyx.server.features.hierarchy.constants import HIERARCHY_NODES_PREFIX -from onyx.server.features.hierarchy.constants import HIERARCHY_NODES_SEARCH_PATH -from onyx.server.features.hierarchy.models import DocumentPageCursor -from onyx.server.features.hierarchy.models import DocumentSortDirection -from onyx.server.features.hierarchy.models import DocumentSortField -from onyx.server.features.hierarchy.models import DocumentSummary -from onyx.server.features.hierarchy.models import HierarchyNodeDocumentsRequest -from onyx.server.features.hierarchy.models import HierarchyNodeDocumentsResponse -from onyx.server.features.hierarchy.models import HierarchyNodeSearchResponse -from onyx.server.features.hierarchy.models import HierarchyNodeSearchSummary -from onyx.server.features.hierarchy.models import HierarchyNodesResponse -from onyx.server.features.hierarchy.models import HierarchyNodeSummary +from onyx.server.features.hierarchy.constants import ( + DOCUMENT_PAGE_SIZE, + HIERARCHY_NODE_DOCUMENTS_PATH, + HIERARCHY_NODE_SEARCH_LIMIT, + HIERARCHY_NODES_LIST_PATH, + HIERARCHY_NODES_PREFIX, + HIERARCHY_NODES_SEARCH_PATH, +) +from onyx.server.features.hierarchy.models import ( + DocumentPageCursor, + DocumentSortDirection, + DocumentSortField, + DocumentSummary, + HierarchyNodeDocumentsRequest, + HierarchyNodeDocumentsResponse, + HierarchyNodeSearchResponse, + HierarchyNodeSearchSummary, + HierarchyNodesResponse, + HierarchyNodeSummary, +) OPENSEARCH_NOT_ENABLED_MESSAGE = "Per-source knowledge selection is coming soon in v3.0! OpenSearch indexing must be enabled to use this feature." diff --git a/backend/onyx/server/features/image_generation/api.py b/backend/onyx/server/features/image_generation/api.py index 50baece93c..11835f1f64 100644 --- a/backend/onyx/server/features/image_generation/api.py +++ b/backend/onyx/server/features/image_generation/api.py @@ -1,8 +1,7 @@ import base64 import binascii -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from pydantic import BaseModel from onyx.auth.permissions import require_permission @@ -12,8 +11,7 @@ from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.image_gen.exceptions import ImageGenerationNotConfiguredError from onyx.image_gen.generation import generate_images_with_default_config -from onyx.image_gen.interfaces import ImageShape -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.interfaces import ImageShape, ReferenceImage from onyx.utils.b64 import get_image_type from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/input_prompt/api.py b/backend/onyx/server/features/input_prompt/api.py index a8ac00ca9f..52008bd971 100644 --- a/backend/onyx/server/features/input_prompt/api.py +++ b/backend/onyx/server/features/input_prompt/api.py @@ -1,23 +1,24 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.input_prompt import disable_input_prompt_for_user -from onyx.db.input_prompt import fetch_input_prompt_by_id -from onyx.db.input_prompt import fetch_input_prompts_by_user -from onyx.db.input_prompt import insert_input_prompt -from onyx.db.input_prompt import remove_input_prompt -from onyx.db.input_prompt import remove_public_input_prompt -from onyx.db.input_prompt import update_input_prompt -from onyx.db.models import InputPrompt__User -from onyx.db.models import User -from onyx.server.features.input_prompt.models import CreateInputPromptRequest -from onyx.server.features.input_prompt.models import InputPromptSnapshot -from onyx.server.features.input_prompt.models import UpdateInputPromptRequest +from onyx.db.input_prompt import ( + disable_input_prompt_for_user, + fetch_input_prompt_by_id, + fetch_input_prompts_by_user, + insert_input_prompt, + remove_input_prompt, + remove_public_input_prompt, + update_input_prompt, +) +from onyx.db.models import InputPrompt__User, User +from onyx.server.features.input_prompt.models import ( + CreateInputPromptRequest, + InputPromptSnapshot, + UpdateInputPromptRequest, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/mcp/api.py b/backend/onyx/server/features/mcp/api.py index be233b7b44..2588d340e7 100644 --- a/backend/onyx/server/features/mcp/api.py +++ b/backend/onyx/server/features/mcp/api.py @@ -12,102 +12,109 @@ from urllib.parse import urlparse from uuid import UUID import requests -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Request +from fastapi import APIRouter, Depends, HTTPException, Request from mcp.client.auth import OAuthClientProvider -from mcp.shared.auth import OAuthClientInformationFull -from mcp.shared.auth import OAuthToken +from mcp.shared.auth import OAuthClientInformationFull, OAuthToken from mcp.types import InitializeResult from mcp.types import Tool as MCPLibTool -from pydantic import AnyUrl -from pydantic import BaseModel +from pydantic import AnyUrl, BaseModel from sqlalchemy.orm import Session -from onyx.auth.oauth_token_manager import build_oauth_authorization_url -from onyx.auth.oauth_token_manager import exchange_oauth_code_for_token -from onyx.auth.oauth_token_manager import OAuthFlowParams -from onyx.auth.oauth_token_manager import validate_oauth_endpoint_url +from onyx.auth.oauth_token_manager import ( + build_oauth_authorization_url, + exchange_oauth_code_for_token, + OAuthFlowParams, + validate_oauth_endpoint_url, +) from onyx.auth.permissions import require_permission from onyx.auth.schemas import UserRole from onyx.auth.users import current_curator_or_admin_user from onyx.configs.app_configs import WEB_DOMAIN from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPOAuthProviderMode -from onyx.db.enums import MCPServerStatus -from onyx.db.enums import MCPTransport -from onyx.db.enums import Permission -from onyx.db.mcp import create_connection_config -from onyx.db.mcp import create_mcp_server__no_commit -from onyx.db.mcp import delete_all_user_connection_configs_for_server_no_commit -from onyx.db.mcp import delete_connection_config -from onyx.db.mcp import delete_mcp_server -from onyx.db.mcp import extract_connection_data -from onyx.db.mcp import get_all_mcp_servers -from onyx.db.mcp import get_mcp_server_by_id -from onyx.db.mcp import get_mcp_servers_accessible_to_user -from onyx.db.mcp import get_mcp_servers_for_persona -from onyx.db.mcp import get_server_auth_template -from onyx.db.mcp import get_user_connection_config -from onyx.db.mcp import update_connection_config -from onyx.db.mcp import update_mcp_server__no_commit -from onyx.db.mcp import upsert_user_connection_config -from onyx.db.mcp import user_can_access_mcp_server -from onyx.db.models import MCPConnectionConfig +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPOAuthProviderMode, + MCPServerStatus, + MCPTransport, + Permission, +) +from onyx.db.mcp import ( + create_connection_config, + create_mcp_server__no_commit, + delete_all_user_connection_configs_for_server_no_commit, + delete_connection_config, + delete_mcp_server, + extract_connection_data, + get_all_mcp_servers, + get_mcp_server_by_id, + get_mcp_servers_accessible_to_user, + get_mcp_servers_for_persona, + get_server_auth_template, + get_user_connection_config, + update_connection_config, + update_mcp_server__no_commit, + upsert_user_connection_config, + user_can_access_mcp_server, +) +from onyx.db.models import MCPConnectionConfig, Tool, User from onyx.db.models import MCPServer as DbMCPServer -from onyx.db.models import Tool -from onyx.db.models import User -from onyx.db.tools import create_tool__no_commit -from onyx.db.tools import delete_tool__no_commit -from onyx.db.tools import get_tools_by_mcp_server_id +from onyx.db.tools import ( + create_tool__no_commit, + delete_tool__no_commit, + get_tools_by_mcp_server_id, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.redis.redis_pool import get_redis_client -from onyx.server.features.mcp.client import discover_mcp_tools -from onyx.server.features.mcp.client import initialize_mcp_client -from onyx.server.features.mcp.client import log_exception_group -from onyx.server.features.mcp.models import apply_auto_substitutions -from onyx.server.features.mcp.models import MCPApiKeyResponse -from onyx.server.features.mcp.models import MCPAuthTemplate -from onyx.server.features.mcp.models import MCPConnectionData -from onyx.server.features.mcp.models import MCPOAuthCallbackResponse -from onyx.server.features.mcp.models import MCPOAuthKeys -from onyx.server.features.mcp.models import MCPServer -from onyx.server.features.mcp.models import MCPServerCreateResponse -from onyx.server.features.mcp.models import MCPServerSimpleCreateRequest -from onyx.server.features.mcp.models import MCPServerSimpleUpdateRequest -from onyx.server.features.mcp.models import MCPServersResponse -from onyx.server.features.mcp.models import MCPServerUpdateResponse -from onyx.server.features.mcp.models import MCPToolCreateRequest -from onyx.server.features.mcp.models import MCPToolListResponse -from onyx.server.features.mcp.models import MCPToolUpdateRequest -from onyx.server.features.mcp.models import MCPUserCredentialsRequest -from onyx.server.features.mcp.models import MCPUserOAuthConnectRequest -from onyx.server.features.mcp.models import MCPUserOAuthConnectResponse -from onyx.server.features.mcp.oauth import _absolute_token_expiry -from onyx.server.features.mcp.oauth import key_auth_url -from onyx.server.features.mcp.oauth import key_code -from onyx.server.features.mcp.oauth import key_state -from onyx.server.features.mcp.oauth import key_tokens -from onyx.server.features.mcp.oauth import make_oauth_provider -from onyx.server.features.mcp.oauth import MCPOauthState -from onyx.server.features.mcp.oauth import OAUTH_WAIT_SECONDS -from onyx.server.features.mcp.oauth import REQUESTED_SCOPE -from onyx.server.features.mcp.oauth import STATE_TTL_SECONDS -from onyx.server.features.mcp.oauth import UNUSED_RETURN_PATH +from onyx.server.features.mcp.client import ( + discover_mcp_tools, + initialize_mcp_client, + log_exception_group, +) +from onyx.server.features.mcp.models import ( + apply_auto_substitutions, + MCPApiKeyResponse, + MCPAuthTemplate, + MCPConnectionData, + MCPOAuthCallbackResponse, + MCPOAuthKeys, + MCPServer, + MCPServerCreateResponse, + MCPServerSimpleCreateRequest, + MCPServerSimpleUpdateRequest, + MCPServersResponse, + MCPServerUpdateResponse, + MCPToolCreateRequest, + MCPToolListResponse, + MCPToolUpdateRequest, + MCPUserCredentialsRequest, + MCPUserOAuthConnectRequest, + MCPUserOAuthConnectResponse, +) +from onyx.server.features.mcp.oauth import ( + _absolute_token_expiry, + key_auth_url, + key_code, + key_state, + key_tokens, + make_oauth_provider, + MCPOauthState, + OAUTH_WAIT_SECONDS, + REQUESTED_SCOPE, + STATE_TTL_SECONDS, + UNUSED_RETURN_PATH, +) from onyx.server.features.mcp.ssrf import validate_mcp_outbound_url from onyx.server.features.tool.models import ToolSnapshot -from onyx.utils.encryption import mask_string -from onyx.utils.encryption import reject_masked_credentials +from onyx.utils.encryption import mask_string, reject_masked_credentials from onyx.utils.logger import setup_logger -from onyx.utils.url import BLOCKED_HOSTNAMES -from onyx.utils.url import SSRFException -from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.url import BLOCKED_HOSTNAMES, SSRFException +from onyx.utils.variable_functionality import ( + fetch_ee_implementation_or_noop, + fetch_versioned_implementation, + global_version, +) from shared_configs.contextvars import get_current_tenant_id logger = setup_logger() diff --git a/backend/onyx/server/features/mcp/client.py b/backend/onyx/server/features/mcp/client.py index 1e4b1a32f9..1db0d740af 100644 --- a/backend/onyx/server/features/mcp/client.py +++ b/backend/onyx/server/features/mcp/client.py @@ -5,21 +5,20 @@ This module provides a proper MCP client that follows the JSON-RPC 2.0 specifica and handles connection initialization, session management, and protocol communication. """ -from collections.abc import Callable -from collections.abc import Coroutine +from collections.abc import Callable, Coroutine from enum import Enum -from typing import Any -from typing import Dict -from typing import TypeVar +from typing import Any, Dict, TypeVar from mcp import ClientSession from mcp.client.auth import OAuthClientProvider from mcp.client.sse import sse_client from mcp.client.streamable_http import streamablehttp_client # or use stdio_client -from mcp.types import CallToolResult -from mcp.types import InitializeResult -from mcp.types import ListResourcesResult -from mcp.types import TextResourceContents +from mcp.types import ( + CallToolResult, + InitializeResult, + ListResourcesResult, + TextResourceContents, +) from mcp.types import Tool as MCPLibTool from pydantic import BaseModel diff --git a/backend/onyx/server/features/mcp/models.py b/backend/onyx/server/features/mcp/models.py index d03d5f5f37..29cd139334 100644 --- a/backend/onyx/server/features/mcp/models.py +++ b/backend/onyx/server/features/mcp/models.py @@ -1,23 +1,19 @@ import datetime import re from enum import Enum -from typing import Any -from typing import List -from typing import NotRequired -from typing import Optional -from typing import TypedDict +from typing import Any, List, NotRequired, Optional, TypedDict from uuid import UUID from mcp.types import Tool as MCPLibTool -from pydantic import BaseModel -from pydantic import Field -from pydantic import model_validator +from pydantic import BaseModel, Field, model_validator -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPOAuthProviderMode -from onyx.db.enums import MCPServerStatus -from onyx.db.enums import MCPTransport +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPOAuthProviderMode, + MCPServerStatus, + MCPTransport, +) # Matches `{placeholder_name}` inside header value templates. _PLACEHOLDER_RE = re.compile(r"\{([^}]+)\}") diff --git a/backend/onyx/server/features/mcp/oauth.py b/backend/onyx/server/features/mcp/oauth.py index e16908ed14..9d04202b16 100644 --- a/backend/onyx/server/features/mcp/oauth.py +++ b/backend/onyx/server/features/mcp/oauth.py @@ -13,15 +13,15 @@ from typing import Any from urllib.parse import urlparse import httpx -from mcp.client.auth import OAuthClientProvider -from mcp.client.auth import TokenStorage +from mcp.client.auth import OAuthClientProvider, TokenStorage from mcp.client.auth.oauth2 import OAuthContext -from mcp.shared.auth import OAuthClientInformationFull -from mcp.shared.auth import OAuthClientMetadata -from mcp.shared.auth import OAuthMetadata -from mcp.shared.auth import OAuthToken -from pydantic import AnyUrl -from pydantic import BaseModel +from mcp.shared.auth import ( + OAuthClientInformationFull, + OAuthClientMetadata, + OAuthMetadata, + OAuthToken, +) +from pydantic import AnyUrl, BaseModel from sqlalchemy.orm import Session from onyx.cache.interface import CacheLockAcquisitionError @@ -29,16 +29,17 @@ from onyx.cache.locks import cache_shared_lock from onyx.configs.app_configs import WEB_DOMAIN from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import MCPOAuthProviderMode -from onyx.db.mcp import extract_connection_data -from onyx.db.mcp import get_connection_config_by_id -from onyx.db.mcp import update_connection_config +from onyx.db.mcp import ( + extract_connection_data, + get_connection_config_by_id, + update_connection_config, +) from onyx.db.models import MCPConnectionConfig from onyx.db.models import MCPServer as DbMCPServer from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.redis.redis_pool import get_redis_client -from onyx.server.features.mcp.models import MCPConnectionData -from onyx.server.features.mcp.models import MCPOAuthKeys +from onyx.server.features.mcp.models import MCPConnectionData, MCPOAuthKeys from onyx.server.features.mcp.ssrf import mcp_ssrf_httpx_client_factory from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_async_sync_no_cancel diff --git a/backend/onyx/server/features/notifications/api.py b/backend/onyx/server/features/notifications/api.py index b581ab19ee..b4ebafd94f 100644 --- a/backend/onyx/server/features/notifications/api.py +++ b/backend/onyx/server/features/notifications/api.py @@ -1,6 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Query +from fastapi import APIRouter, Depends, Query from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -8,21 +6,23 @@ from onyx.configs.constants import NotificationType from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.notification import count_notifications -from onyx.db.notification import dismiss_notification -from onyx.db.notification import dismiss_user_notifications -from onyx.db.notification import get_notification_by_id -from onyx.db.notification import get_notifications +from onyx.db.notification import ( + count_notifications, + dismiss_notification, + dismiss_user_notifications, + get_notification_by_id, + get_notifications, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.server.features.build.utils import ensure_build_mode_intro_notification -from onyx.server.features.notifications.models import NotificationResponse -from onyx.server.features.notifications.models import NotificationSummary -from onyx.server.features.notifications.models import PaginatedNotifications -from onyx.server.features.notifications.utils import ( - ensure_permissions_migration_notification, +from onyx.server.features.notifications.models import ( + NotificationResponse, + NotificationSummary, + PaginatedNotifications, ) from onyx.server.features.notifications.utils import ( + ensure_permissions_migration_notification, ensure_system_announcement_notification, ) from onyx.server.features.release_notes.utils import ( diff --git a/backend/onyx/server/features/notifications/models.py b/backend/onyx/server/features/notifications/models.py index 8d3b16c2e4..f4c92c0ab6 100644 --- a/backend/onyx/server/features/notifications/models.py +++ b/backend/onyx/server/features/notifications/models.py @@ -1,7 +1,6 @@ from datetime import datetime -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from onyx.configs.constants import NotificationType diff --git a/backend/onyx/server/features/oauth_config/__init__.py b/backend/onyx/server/features/oauth_config/__init__.py index c01f10b171..f6f3628042 100644 --- a/backend/onyx/server/features/oauth_config/__init__.py +++ b/backend/onyx/server/features/oauth_config/__init__.py @@ -1,6 +1,5 @@ """OAuth configuration feature module.""" -from onyx.server.features.oauth_config.api import admin_router -from onyx.server.features.oauth_config.api import router +from onyx.server.features.oauth_config.api import admin_router, router __all__ = ["admin_router", "router"] diff --git a/backend/onyx/server/features/oauth_config/api.py b/backend/onyx/server/features/oauth_config/api.py index f2df6a4248..6b01508728 100644 --- a/backend/onyx/server/features/oauth_config/api.py +++ b/backend/onyx/server/features/oauth_config/api.py @@ -1,8 +1,6 @@ """API endpoints for OAuth configuration management.""" -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from onyx.auth.oauth_token_manager import OAuthTokenManager @@ -11,24 +9,29 @@ from onyx.auth.users import current_curator_or_admin_user from onyx.configs.app_configs import WEB_DOMAIN from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.models import OAuthConfig -from onyx.db.models import User -from onyx.db.oauth_config import create_oauth_config -from onyx.db.oauth_config import delete_oauth_config -from onyx.db.oauth_config import delete_user_oauth_token -from onyx.db.oauth_config import get_oauth_config -from onyx.db.oauth_config import get_oauth_configs -from onyx.db.oauth_config import get_tools_by_oauth_config -from onyx.db.oauth_config import update_oauth_config -from onyx.db.oauth_config import upsert_user_oauth_token -from onyx.federated_connectors.oauth_utils import generate_oauth_state -from onyx.federated_connectors.oauth_utils import verify_oauth_state -from onyx.server.features.oauth_config.models import OAuthCallbackResponse -from onyx.server.features.oauth_config.models import OAuthConfigCreate -from onyx.server.features.oauth_config.models import OAuthConfigSnapshot -from onyx.server.features.oauth_config.models import OAuthConfigUpdate -from onyx.server.features.oauth_config.models import OAuthInitiateRequest -from onyx.server.features.oauth_config.models import OAuthInitiateResponse +from onyx.db.models import OAuthConfig, User +from onyx.db.oauth_config import ( + create_oauth_config, + delete_oauth_config, + delete_user_oauth_token, + get_oauth_config, + get_oauth_configs, + get_tools_by_oauth_config, + update_oauth_config, + upsert_user_oauth_token, +) +from onyx.federated_connectors.oauth_utils import ( + generate_oauth_state, + verify_oauth_state, +) +from onyx.server.features.oauth_config.models import ( + OAuthCallbackResponse, + OAuthConfigCreate, + OAuthConfigSnapshot, + OAuthConfigUpdate, + OAuthInitiateRequest, + OAuthInitiateResponse, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/password/api.py b/backend/onyx/server/features/password/api.py index eff5423aca..8c8ddec7e3 100644 --- a/backend/onyx/server/features/password/api.py +++ b/backend/onyx/server/features/password/api.py @@ -1,19 +1,17 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from fastapi_users.exceptions import InvalidPasswordException from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission -from onyx.auth.users import get_user_manager -from onyx.auth.users import User -from onyx.auth.users import UserManager +from onyx.auth.users import get_user_manager, User, UserManager from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.users import get_user_by_email -from onyx.server.features.password.models import ChangePasswordRequest -from onyx.server.features.password.models import UserResetRequest -from onyx.server.features.password.models import UserResetResponse +from onyx.server.features.password.models import ( + ChangePasswordRequest, + UserResetRequest, + UserResetResponse, +) router = APIRouter(prefix="/password") diff --git a/backend/onyx/server/features/persona/api.py b/backend/onyx/server/features/persona/api.py index cd5ce3be38..fcd7a2dc11 100644 --- a/backend/onyx/server/features/persona/api.py +++ b/backend/onyx/server/features/persona/api.py @@ -1,67 +1,75 @@ from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query -from fastapi import Request -from fastapi import Response -from fastapi import UploadFile +from fastapi import ( + APIRouter, + Depends, + HTTPException, + Query, + Request, + Response, + UploadFile, +) from fastapi.responses import StreamingResponse -from pydantic import BaseModel -from pydantic import model_validator +from pydantic import BaseModel, model_validator from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission -from onyx.auth.users import current_chat_accessible_user -from onyx.auth.users import current_curator_or_admin_user -from onyx.auth.users import current_limited_user +from onyx.auth.users import ( + current_chat_accessible_user, + current_curator_or_admin_user, + current_limited_user, +) from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import MilestoneRecordType -from onyx.configs.constants import PUBLIC_API_TAGS +from onyx.configs.constants import FileOrigin, MilestoneRecordType, PUBLIC_API_TAGS from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import Permission -from onyx.db.enums import PersonaSharePermission +from onyx.db.enums import Permission, PersonaSharePermission from onyx.db.file_record import get_filerecord_by_file_id_optional from onyx.db.models import User -from onyx.db.persona import create_assistant_label -from onyx.db.persona import create_update_persona -from onyx.db.persona import delete_persona_label -from onyx.db.persona import get_assistant_labels -from onyx.db.persona import get_minimal_persona_snapshots_for_user -from onyx.db.persona import get_minimal_persona_snapshots_paginated -from onyx.db.persona import get_persona_by_id -from onyx.db.persona import get_persona_count_for_user -from onyx.db.persona import get_persona_snapshots_for_user -from onyx.db.persona import get_persona_snapshots_paginated -from onyx.db.persona import mark_persona_as_deleted -from onyx.db.persona import mark_persona_as_not_deleted -from onyx.db.persona import remove_user_from_persona_shares -from onyx.db.persona import update_persona_featured -from onyx.db.persona import update_persona_label -from onyx.db.persona import update_persona_public_status -from onyx.db.persona import update_persona_shared -from onyx.db.persona import update_persona_visibility -from onyx.db.persona import update_personas_display_priority -from onyx.db.persona_sharing import get_persona_access_level -from onyx.db.persona_sharing import get_user_group_ids_for_user -from onyx.db.persona_sharing import persona_ownership_is_vacant +from onyx.db.persona import ( + create_assistant_label, + create_update_persona, + delete_persona_label, + get_assistant_labels, + get_minimal_persona_snapshots_for_user, + get_minimal_persona_snapshots_paginated, + get_persona_by_id, + get_persona_count_for_user, + get_persona_snapshots_for_user, + get_persona_snapshots_paginated, + mark_persona_as_deleted, + mark_persona_as_not_deleted, + remove_user_from_persona_shares, + update_persona_featured, + update_persona_label, + update_persona_public_status, + update_persona_shared, + update_persona_visibility, + update_personas_display_priority, +) +from onyx.db.persona_sharing import ( + get_persona_access_level, + get_user_group_ids_for_user, + persona_ownership_is_vacant, +) from onyx.db.users import get_active_admin_count from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import get_default_file_store from onyx.file_store.models import ChatFileType from onyx.server.documents.models import PaginatedReturn -from onyx.server.features.persona.constants import ADMIN_AGENTS_RESOURCE -from onyx.server.features.persona.constants import AGENTS_RESOURCE -from onyx.server.features.persona.models import FullPersonaSnapshot -from onyx.server.features.persona.models import MinimalPersonaSnapshot -from onyx.server.features.persona.models import PersonaLabelCreate -from onyx.server.features.persona.models import PersonaLabelResponse -from onyx.server.features.persona.models import PersonaSnapshot -from onyx.server.features.persona.models import PersonaUpsertRequest +from onyx.server.features.persona.constants import ( + ADMIN_AGENTS_RESOURCE, + AGENTS_RESOURCE, +) +from onyx.server.features.persona.models import ( + FullPersonaSnapshot, + MinimalPersonaSnapshot, + PersonaLabelCreate, + PersonaLabelResponse, + PersonaSnapshot, + PersonaUpsertRequest, +) from onyx.server.manage.llm.api import get_valid_model_configuration_ids_for_persona from onyx.server.models import DisplayPriorityRequest from onyx.server.settings.store import load_settings diff --git a/backend/onyx/server/features/persona/models.py b/backend/onyx/server/features/persona/models.py index 5783766014..eef5f04a95 100644 --- a/backend/onyx/server/features/persona/models.py +++ b/backend/onyx/server/features/persona/models.py @@ -1,19 +1,22 @@ from datetime import datetime from uuid import UUID -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.configs.constants import DocumentSource -from onyx.db.enums import HierarchyNodeType -from onyx.db.enums import PersonaAccessLevel -from onyx.db.enums import PersonaSharePermission -from onyx.db.enums import PersonaSharingStatus -from onyx.db.models import Document -from onyx.db.models import HierarchyNode -from onyx.db.models import Persona -from onyx.db.models import PersonaLabel -from onyx.db.models import StarterMessage +from onyx.db.enums import ( + HierarchyNodeType, + PersonaAccessLevel, + PersonaSharePermission, + PersonaSharingStatus, +) +from onyx.db.models import ( + Document, + HierarchyNode, + Persona, + PersonaLabel, + StarterMessage, +) from onyx.db.persona_sharing import derive_persona_sharing_status from onyx.server.features.document_set.models import DocumentSetSummary from onyx.server.features.tool.models import ToolSnapshot diff --git a/backend/onyx/server/features/projects/api.py b/backend/onyx/server/features/projects/api.py index d7bbb22348..c3d5afec72 100644 --- a/backend/onyx/server/features/projects/api.py +++ b/backend/onyx/server/features/projects/api.py @@ -1,40 +1,43 @@ import json from uuid import UUID -from fastapi import APIRouter -from fastapi import BackgroundTasks -from fastapi import Depends -from fastapi import File -from fastapi import Form -from fastapi import HTTPException -from fastapi import Response -from fastapi import UploadFile +from fastapi import ( + APIRouter, + BackgroundTasks, + Depends, + File, + Form, + HTTPException, + Response, + UploadFile, +) from pydantic import BaseModel from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import PUBLIC_API_TAGS -from onyx.configs.constants import USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH +from onyx.configs.constants import ( + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + PUBLIC_API_TAGS, + USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH, +) from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import Permission -from onyx.db.enums import UserFileStatus -from onyx.db.models import ChatSession -from onyx.db.models import Project__UserFile -from onyx.db.models import User -from onyx.db.models import UserFile -from onyx.db.models import UserProject +from onyx.db.enums import Permission, UserFileStatus +from onyx.db.models import ChatSession, Project__UserFile, User, UserFile, UserProject from onyx.db.persona import get_personas_by_ids -from onyx.db.projects import get_project_token_count -from onyx.db.projects import upload_files_to_user_files_with_indexing -from onyx.server.features.projects.models import CategorizedFilesSnapshot -from onyx.server.features.projects.models import ChatSessionRequest -from onyx.server.features.projects.models import TokenCountResponse -from onyx.server.features.projects.models import UserFileSnapshot -from onyx.server.features.projects.models import UserProjectSnapshot +from onyx.db.projects import ( + get_project_token_count, + upload_files_to_user_files_with_indexing, +) +from onyx.server.features.projects.models import ( + CategorizedFilesSnapshot, + ChatSessionRequest, + TokenCountResponse, + UserFileSnapshot, + UserProjectSnapshot, +) from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id @@ -64,8 +67,6 @@ def _trigger_user_file_project_sync( from onyx.background.celery.tasks.user_file_processing.tasks import ( enqueue_user_file_project_sync_task, - ) - from onyx.background.celery.tasks.user_file_processing.tasks import ( get_user_file_project_sync_queue_depth, ) from onyx.background.celery.versioned_apps.client import app as client_app diff --git a/backend/onyx/server/features/projects/models.py b/backend/onyx/server/features/projects/models.py index 94209ccb1b..4c3dba4072 100644 --- a/backend/onyx/server/features/projects/models.py +++ b/backend/onyx/server/features/projects/models.py @@ -4,8 +4,7 @@ from uuid import UUID from pydantic import BaseModel from onyx.db.enums import UserFileStatus -from onyx.db.models import UserFile -from onyx.db.models import UserProject +from onyx.db.models import UserFile, UserProject from onyx.db.projects import CategorizedFilesResult from onyx.file_store.models import ChatFileType from onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type diff --git a/backend/onyx/server/features/projects/projects_file_utils.py b/backend/onyx/server/features/projects/projects_file_utils.py index aa8d43fd5d..8bf18e30ac 100644 --- a/backend/onyx/server/features/projects/projects_file_utils.py +++ b/backend/onyx/server/features/projects/projects_file_utils.py @@ -1,26 +1,25 @@ from math import ceil from fastapi import UploadFile -from PIL import Image -from PIL import ImageOps -from PIL import UnidentifiedImageError -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field +from PIL import Image, ImageOps, UnidentifiedImageError +from pydantic import BaseModel, ConfigDict, Field from sqlalchemy.orm import Session -from onyx.configs.app_configs import MAX_EMBEDDED_IMAGES_PER_FILE -from onyx.configs.app_configs import MAX_EMBEDDED_IMAGES_PER_UPLOAD +from onyx.configs.app_configs import ( + MAX_EMBEDDED_IMAGES_PER_FILE, + MAX_EMBEDDED_IMAGES_PER_UPLOAD, +) from onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled from onyx.db.llm import fetch_default_llm_model -from onyx.file_processing.extract_file_text import count_docx_embedded_images -from onyx.file_processing.extract_file_text import count_pdf_embedded_images -from onyx.file_processing.extract_file_text import extract_file_text -from onyx.file_processing.extract_file_text import get_file_ext +from onyx.file_processing.extract_file_text import ( + count_docx_embedded_images, + count_pdf_embedded_images, + extract_file_text, + get_file_ext, +) from onyx.file_processing.file_types import OnyxFileExtensions from onyx.file_processing.password_validation import is_file_password_protected -from onyx.natural_language_processing.utils import count_tokens -from onyx.natural_language_processing.utils import get_tokenizer +from onyx.natural_language_processing.utils import count_tokens, get_tokenizer from onyx.server.settings.store import load_settings from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/release_notes/utils.py b/backend/onyx/server/features/release_notes/utils.py index 937906b8d5..c96bab3b2b 100644 --- a/backend/onyx/server/features/release_notes/utils.py +++ b/backend/onyx/server/features/release_notes/utils.py @@ -1,8 +1,7 @@ """Utility functions for release notes parsing and caching.""" import re -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import httpx from sqlalchemy.orm import Session @@ -12,12 +11,14 @@ from onyx.cache.factory import get_shared_cache_backend from onyx.configs.app_configs import INSTANCE_TYPE from onyx.configs.constants import OnyxRedisLocks from onyx.db.release_notes import create_release_notifications_for_versions -from onyx.server.features.release_notes.constants import AUTO_REFRESH_THRESHOLD_SECONDS -from onyx.server.features.release_notes.constants import FETCH_TIMEOUT -from onyx.server.features.release_notes.constants import GITHUB_CHANGELOG_RAW_URL -from onyx.server.features.release_notes.constants import REDIS_CACHE_TTL -from onyx.server.features.release_notes.constants import REDIS_KEY_ETAG -from onyx.server.features.release_notes.constants import REDIS_KEY_FETCHED_AT +from onyx.server.features.release_notes.constants import ( + AUTO_REFRESH_THRESHOLD_SECONDS, + FETCH_TIMEOUT, + GITHUB_CHANGELOG_RAW_URL, + REDIS_CACHE_TTL, + REDIS_KEY_ETAG, + REDIS_KEY_FETCHED_AT, +) from onyx.server.features.release_notes.models import ReleaseNoteEntry from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/features/search/api.py b/backend/onyx/server/features/search/api.py index a039712a07..7f28657a9b 100644 --- a/backend/onyx/server/features/search/api.py +++ b/backend/onyx/server/features/search/api.py @@ -13,22 +13,21 @@ directly — a lighter-weight flow with optional query expansion. import json from typing import cast -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.auth.schemas import UserRole from onyx.chat.emitter import NullEmitter from onyx.configs.constants import MessageType -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import PersonaSearchInfo -from onyx.context.search.models import TimeRange +from onyx.context.search.models import BaseFilters, PersonaSearchInfo, TimeRange from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.llm import can_user_access_llm_provider -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import fetch_user_group_ids +from onyx.db.llm import ( + can_user_access_llm_provider, + fetch_existing_llm_provider, + fetch_user_group_ids, +) from onyx.db.models import User from onyx.db.persona import get_persona_by_id from onyx.db.search_settings import get_current_search_settings @@ -36,20 +35,19 @@ from onyx.db.tools import get_tools from onyx.document_index.factory import get_default_document_index from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.llm.factory import get_default_llm -from onyx.llm.factory import get_llm_for_persona -from onyx.llm.factory import llm_from_provider -from onyx.server.features.search.models import SearchRequest -from onyx.server.features.search.models import SearchResponse -from onyx.server.features.search.models import SearchResult +from onyx.llm.factory import get_default_llm, get_llm_for_persona, llm_from_provider +from onyx.server.features.search.models import ( + SearchRequest, + SearchResponse, + SearchResult, +) from onyx.server.manage.llm.models import LLMProviderView from onyx.server.query_and_chat.placement import Placement from onyx.server.settings.store import load_settings from onyx.server.usage_limits import check_llm_cost_limit_for_provider from onyx.server.utils_vector_db import require_vector_db from onyx.tools.constants import SEARCH_TOOL_ID -from onyx.tools.models import ChatMinimalTextMessage -from onyx.tools.models import SearchToolOverrideKwargs +from onyx.tools.models import ChatMinimalTextMessage, SearchToolOverrideKwargs from onyx.tools.tool_implementations.search.search_tool import SearchTool from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/features/search/models.py b/backend/onyx/server/features/search/models.py index 3ea3ceb1fa..8e36a5748e 100644 --- a/backend/onyx/server/features/search/models.py +++ b/backend/onyx/server/features/search/models.py @@ -1,8 +1,6 @@ from datetime import datetime -from pydantic import BaseModel -from pydantic import Field -from pydantic import model_validator +from pydantic import BaseModel, Field, model_validator from onyx.configs.constants import DocumentSource from onyx.context.search.models import Tag diff --git a/backend/onyx/server/features/skill/api.py b/backend/onyx/server/features/skill/api.py index eddeedaeaa..aced3fbb94 100644 --- a/backend/onyx/server/features/skill/api.py +++ b/backend/onyx/server/features/skill/api.py @@ -3,68 +3,70 @@ import zipfile from typing import Annotated from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import Form -from fastapi import UploadFile +from fastapi import APIRouter, Depends, File, Form, UploadFile from pydantic import ValidationError from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.auth.schemas import UserRole from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import AccountType -from onyx.db.enums import Permission -from onyx.db.enums import SkillSharePermission -from onyx.db.models import Skill -from onyx.db.models import User -from onyx.db.skill import affected_user_ids_for_skill -from onyx.db.skill import create_skill__no_commit -from onyx.db.skill import delete_skill -from onyx.db.skill import enable_skill_for_user_if_unset__no_commit -from onyx.db.skill import fetch_skill -from onyx.db.skill import list_skills -from onyx.db.skill import replace_skill_bundle -from onyx.db.skill import replace_skill_shares -from onyx.db.skill import set_skill_enabled_for_user -from onyx.db.skill import set_skill_public_permission -from onyx.db.skill import SkillAccessPolicy -from onyx.db.skill import transfer_skill_ownership +from onyx.db.enums import AccountType, Permission, SkillSharePermission +from onyx.db.models import Skill, User +from onyx.db.skill import ( + affected_user_ids_for_skill, + create_skill__no_commit, + delete_skill, + enable_skill_for_user_if_unset__no_commit, + fetch_skill, + list_skills, + replace_skill_bundle, + replace_skill_shares, + set_skill_enabled_for_user, + set_skill_public_permission, + SkillAccessPolicy, + transfer_skill_ownership, +) from onyx.db.users import fetch_user_by_id from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import get_default_file_store -from onyx.server.features.skill.models import SkillBundleInspectResponse -from onyx.server.features.skill.models import SkillCreateRequest -from onyx.server.features.skill.models import SkillEditableDetailResponse -from onyx.server.features.skill.models import SkillEnableRequest -from onyx.server.features.skill.models import SkillPatchRequest -from onyx.server.features.skill.models import SkillPreviewResponse -from onyx.server.features.skill.models import SkillResponse -from onyx.server.features.skill.models import SkillShareRequest -from onyx.server.features.skill.models import SkillsList -from onyx.server.features.skill.models import TransferSkillOwnershipRequest -from onyx.server.features.skill.response_helpers import skill_preview_response -from onyx.server.features.skill.response_helpers import skill_response_for_user -from onyx.server.features.skill.response_helpers import skills_list_response_for_user -from onyx.skills.bundle import build_single_file_bundle -from onyx.skills.bundle import build_skill_md -from onyx.skills.bundle import compute_bundle_sha256 -from onyx.skills.bundle import inspect_custom_bundle -from onyx.skills.bundle import normalize_custom_bundle -from onyx.skills.bundle import read_bundle_file -from onyx.skills.bundle import read_custom_bundle_instructions -from onyx.skills.bundle import rewrite_custom_bundle_skill_md -from onyx.skills.bundle import SKILL_MD_NAME -from onyx.skills.bundle import update_custom_bundle_files +from onyx.server.features.skill.models import ( + SkillBundleInspectResponse, + SkillCreateRequest, + SkillEditableDetailResponse, + SkillEnableRequest, + SkillPatchRequest, + SkillPreviewResponse, + SkillResponse, + SkillShareRequest, + SkillsList, + TransferSkillOwnershipRequest, +) +from onyx.server.features.skill.response_helpers import ( + skill_preview_response, + skill_response_for_user, + skills_list_response_for_user, +) +from onyx.skills.bundle import ( + build_single_file_bundle, + build_skill_md, + compute_bundle_sha256, + inspect_custom_bundle, + normalize_custom_bundle, + read_bundle_file, + read_custom_bundle_instructions, + rewrite_custom_bundle_skill_md, + SKILL_MD_NAME, + update_custom_bundle_files, +) from onyx.skills.content import read_custom_skill_bundle_bytes -from onyx.skills.ingest import delete_bundle_blob -from onyx.skills.ingest import ingested_skill_bundle -from onyx.skills.ingest import save_skill_bundle_bytes +from onyx.skills.ingest import ( + delete_bundle_blob, + ingested_skill_bundle, + save_skill_bundle_bytes, +) from onyx.skills.metadata import parse_skill_document -from onyx.skills.push import push_skill_to_affected_sandboxes -from onyx.skills.push import push_skills_for_users +from onyx.skills.push import push_skill_to_affected_sandboxes, push_skills_for_users user_router = APIRouter(prefix="/skills") diff --git a/backend/onyx/server/features/skill/models.py b/backend/onyx/server/features/skill/models.py index f2f2b421be..333a9ec079 100644 --- a/backend/onyx/server/features/skill/models.py +++ b/backend/onyx/server/features/skill/models.py @@ -1,18 +1,13 @@ """Pydantic request and response models for the skills API.""" import datetime -from typing import Any -from typing import Literal +from typing import Any, Literal from uuid import UUID -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from sqlalchemy.orm import Session -from onyx.db.enums import SkillAccessLevel -from onyx.db.enums import SkillSharePermission +from onyx.db.enums import SkillAccessLevel, SkillSharePermission from onyx.db.models import Skill from onyx.server.models import MinimalUserSnapshot from onyx.skills.built_in import BuiltInSkillDefinition diff --git a/backend/onyx/server/features/skill/response_helpers.py b/backend/onyx/server/features/skill/response_helpers.py index dfea484525..49ae6886e3 100644 --- a/backend/onyx/server/features/skill/response_helpers.py +++ b/backend/onyx/server/features/skill/response_helpers.py @@ -1,22 +1,25 @@ from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.enums import SkillAccessLevel -from onyx.db.enums import SkillSharePermission -from onyx.db.models import Skill -from onyx.db.models import User -from onyx.db.persona_sharing import get_curated_user_group_ids_for_user -from onyx.db.persona_sharing import get_user_group_ids_for_user -from onyx.db.skill import skill_user_states -from onyx.db.skill import SkillUserState +from onyx.db.enums import SkillAccessLevel, SkillSharePermission +from onyx.db.models import Skill, User +from onyx.db.persona_sharing import ( + get_curated_user_group_ids_for_user, + get_user_group_ids_for_user, +) +from onyx.db.skill import skill_user_states, SkillUserState from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.skill.models import SkillPreviewResponse -from onyx.server.features.skill.models import SkillResponse -from onyx.server.features.skill.models import SkillsList +from onyx.server.features.skill.models import ( + SkillPreviewResponse, + SkillResponse, + SkillsList, +) from onyx.skills.built_in import BUILT_IN_SKILLS -from onyx.skills.content import read_builtin_skill_instructions -from onyx.skills.content import read_custom_skill_bundle_instructions +from onyx.skills.content import ( + read_builtin_skill_instructions, + read_custom_skill_bundle_instructions, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/features/tool/api.py b/backend/onyx/server/features/tool/api.py index 4583a90376..9848ff4972 100644 --- a/backend/onyx/server/features/tool/api.py +++ b/backend/onyx/server/features/tool/api.py @@ -1,8 +1,6 @@ from typing import Any -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from sqlalchemy.orm import Session @@ -13,24 +11,25 @@ from onyx.configs.constants import PUBLIC_API_TAGS from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.mcp import get_mcp_servers_accessible_to_user -from onyx.db.models import Tool -from onyx.db.models import User -from onyx.db.tools import create_tool__no_commit -from onyx.db.tools import delete_tool__no_commit -from onyx.db.tools import get_tool_by_id -from onyx.db.tools import get_tools -from onyx.db.tools import get_tools_by_ids -from onyx.db.tools import update_tool -from onyx.server.features.tool.models import CustomToolCreate -from onyx.server.features.tool.models import CustomToolUpdate -from onyx.server.features.tool.models import ToolSnapshot +from onyx.db.models import Tool, User +from onyx.db.tools import ( + create_tool__no_commit, + delete_tool__no_commit, + get_tool_by_id, + get_tools, + get_tools_by_ids, + update_tool, +) +from onyx.server.features.tool.models import ( + CustomToolCreate, + CustomToolUpdate, + ToolSnapshot, +) from onyx.server.features.tool.tool_visibility import should_expose_tool_to_fe from onyx.tools.built_in_tools import get_built_in_tool_by_id -from onyx.tools.tool_implementations.custom.openapi_parsing import MethodSpec from onyx.tools.tool_implementations.custom.openapi_parsing import ( + MethodSpec, openapi_to_method_specs, -) -from onyx.tools.tool_implementations.custom.openapi_parsing import ( validate_openapi_schema, ) diff --git a/backend/onyx/server/features/tool/tool_visibility.py b/backend/onyx/server/features/tool/tool_visibility.py index febb58eb31..5d4ee25de6 100644 --- a/backend/onyx/server/features/tool/tool_visibility.py +++ b/backend/onyx/server/features/tool/tool_visibility.py @@ -3,8 +3,7 @@ from pydantic import BaseModel from onyx.db.models import Tool -from onyx.tools.constants import MEMORY_TOOL_ID -from onyx.tools.constants import OPEN_URL_TOOL_ID +from onyx.tools.constants import MEMORY_TOOL_ID, OPEN_URL_TOOL_ID # Tool class name constant for OktaProfileTool (not in main constants.py as it's hidden) OKTA_PROFILE_TOOL_ID = "OktaProfileTool" diff --git a/backend/onyx/server/features/user_oauth_token/api.py b/backend/onyx/server/features/user_oauth_token/api.py index dc5edd4029..88f7387a97 100644 --- a/backend/onyx/server/features/user_oauth_token/api.py +++ b/backend/onyx/server/features/user_oauth_token/api.py @@ -1,7 +1,6 @@ """API endpoints for user OAuth token management.""" -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from pydantic import BaseModel from sqlalchemy.orm import Session diff --git a/backend/onyx/server/features/web_search/api.py b/backend/onyx/server/features/web_search/api.py index a9ff07d8f9..0422cab44b 100644 --- a/backend/onyx/server/features/web_search/api.py +++ b/backend/onyx/server/features/web_search/api.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -7,47 +6,47 @@ from onyx.configs.constants import PUBLIC_API_TAGS from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.web_search import fetch_active_web_content_provider -from onyx.db.web_search import fetch_active_web_search_provider +from onyx.db.web_search import ( + fetch_active_web_content_provider, + fetch_active_web_search_provider, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.web_search.models import OpenUrlsToolRequest -from onyx.server.features.web_search.models import OpenUrlsToolResponse -from onyx.server.features.web_search.models import WebSearchToolRequest -from onyx.server.features.web_search.models import WebSearchToolResponse -from onyx.server.features.web_search.models import WebSearchWithContentResponse -from onyx.server.manage.web_search.models import WebContentProviderView -from onyx.server.manage.web_search.models import WebSearchProviderView -from onyx.tools.models import LlmOpenUrlResult -from onyx.tools.models import LlmWebSearchResult +from onyx.server.features.web_search.models import ( + OpenUrlsToolRequest, + OpenUrlsToolResponse, + WebSearchToolRequest, + WebSearchToolResponse, + WebSearchWithContentResponse, +) +from onyx.server.manage.web_search.models import ( + WebContentProviderView, + WebSearchProviderView, +) +from onyx.tools.models import LlmOpenUrlResult, LlmWebSearchResult from onyx.tools.tool_implementations.open_url.models import WebContentProvider from onyx.tools.tool_implementations.open_url.onyx_web_crawler import ( DEFAULT_MAX_HTML_SIZE_BYTES, -) -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import ( DEFAULT_MAX_PDF_SIZE_BYTES, + OnyxWebCrawler, ) -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import OnyxWebCrawler from onyx.tools.tool_implementations.open_url.utils import ( filter_web_contents_with_no_title_or_content, ) -from onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.providers import ( - build_content_provider_from_config, +from onyx.tools.tool_implementations.web_search.models import ( + WebContentProviderConfig, + WebSearchProvider, ) from onyx.tools.tool_implementations.web_search.providers import ( + build_content_provider_from_config, build_search_provider_from_config, ) from onyx.tools.tool_implementations.web_search.utils import ( filter_web_search_results_with_no_title_or_snippet, -) -from onyx.tools.tool_implementations.web_search.utils import ( truncate_search_result_content, ) from onyx.utils.logger import setup_logger -from shared_configs.enums import WebContentProviderType -from shared_configs.enums import WebSearchProviderType +from shared_configs.enums import WebContentProviderType, WebSearchProviderType router = APIRouter(prefix="/web-search", tags=PUBLIC_API_TAGS) logger = setup_logger() diff --git a/backend/onyx/server/features/web_search/models.py b/backend/onyx/server/features/web_search/models.py index c434ed9b0b..2066ab860d 100644 --- a/backend/onyx/server/features/web_search/models.py +++ b/backend/onyx/server/features/web_search/models.py @@ -1,11 +1,7 @@ -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator +from pydantic import BaseModel, Field, field_validator -from onyx.tools.models import LlmOpenUrlResult -from onyx.tools.models import LlmWebSearchResult -from shared_configs.enums import WebContentProviderType -from shared_configs.enums import WebSearchProviderType +from onyx.tools.models import LlmOpenUrlResult, LlmWebSearchResult +from shared_configs.enums import WebContentProviderType, WebSearchProviderType class WebSearchToolRequest(BaseModel): diff --git a/backend/onyx/server/federated/api.py b/backend/onyx/server/federated/api.py index d5cffe1536..2159db34c2 100644 --- a/backend/onyx/server/federated/api.py +++ b/backend/onyx/server/federated/api.py @@ -2,11 +2,7 @@ import json from typing import Any from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Request -from fastapi import Response +from fastapi import APIRouter, Depends, HTTPException, Request, Response from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -17,32 +13,40 @@ from onyx.db.enums import Permission from onyx.db.federated import ( create_federated_connector as db_create_federated_connector, ) -from onyx.db.federated import delete_federated_connector -from onyx.db.federated import fetch_all_federated_connectors -from onyx.db.federated import fetch_federated_connector_by_id -from onyx.db.federated import update_federated_connector -from onyx.db.federated import update_federated_connector_oauth_token -from onyx.db.federated import validate_federated_connector_credentials +from onyx.db.federated import ( + delete_federated_connector, + fetch_all_federated_connectors, + fetch_federated_connector_by_id, + update_federated_connector, + update_federated_connector_oauth_token, + validate_federated_connector_credentials, +) from onyx.db.models import User -from onyx.federated_connectors.factory import get_federated_connector -from onyx.federated_connectors.factory import get_federated_connector_cls +from onyx.federated_connectors.factory import ( + get_federated_connector, + get_federated_connector_cls, +) from onyx.federated_connectors.interfaces import FederatedConnector -from onyx.federated_connectors.oauth_utils import add_state_to_oauth_url -from onyx.federated_connectors.oauth_utils import generate_oauth_state -from onyx.federated_connectors.oauth_utils import get_oauth_callback_uri -from onyx.federated_connectors.oauth_utils import verify_oauth_state -from onyx.server.federated.models import AuthorizeUrlResponse -from onyx.server.federated.models import ConfigurationSchemaResponse -from onyx.server.federated.models import CredentialSchemaResponse -from onyx.server.federated.models import EntitySpecResponse -from onyx.server.federated.models import FederatedConnectorCredentials -from onyx.server.federated.models import FederatedConnectorDetail -from onyx.server.federated.models import FederatedConnectorRequest -from onyx.server.federated.models import FederatedConnectorResponse -from onyx.server.federated.models import FederatedConnectorStatus -from onyx.server.federated.models import FederatedConnectorUpdateRequest -from onyx.server.federated.models import OAuthCallbackResult -from onyx.server.federated.models import UserOAuthStatus +from onyx.federated_connectors.oauth_utils import ( + add_state_to_oauth_url, + generate_oauth_state, + get_oauth_callback_uri, + verify_oauth_state, +) +from onyx.server.federated.models import ( + AuthorizeUrlResponse, + ConfigurationSchemaResponse, + CredentialSchemaResponse, + EntitySpecResponse, + FederatedConnectorCredentials, + FederatedConnectorDetail, + FederatedConnectorRequest, + FederatedConnectorResponse, + FederatedConnectorStatus, + FederatedConnectorUpdateRequest, + OAuthCallbackResult, + UserOAuthStatus, +) from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/federated/models.py b/backend/onyx/server/federated/models.py index 804388986d..bdd6e40f07 100644 --- a/backend/onyx/server/federated/models.py +++ b/backend/onyx/server/federated/models.py @@ -1,8 +1,7 @@ from datetime import datetime from typing import Any -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.configs.constants import FederatedConnectorSource diff --git a/backend/onyx/server/kg/api.py b/backend/onyx/server/kg/api.py index 3e1fcda370..9bd58bdfad 100644 --- a/backend/onyx/server/kg/api.py +++ b/backend/onyx/server/kg/api.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -7,33 +6,43 @@ from onyx.configs.constants import TMP_DRALPHA_PERSONA_NAME from onyx.configs.kg_configs import KG_BETA_ASSISTANT_DESCRIPTION from onyx.db.engine.sql_engine import get_session from onyx.db.entities import get_entity_stats_by_grounded_source_name -from onyx.db.entity_type import get_configured_entity_types -from onyx.db.entity_type import update_entity_types_and_related_connectors__commit +from onyx.db.entity_type import ( + get_configured_entity_types, + update_entity_types_and_related_connectors__commit, +) from onyx.db.enums import Permission -from onyx.db.kg_config import disable_kg -from onyx.db.kg_config import enable_kg -from onyx.db.kg_config import get_kg_config_settings -from onyx.db.kg_config import set_kg_config_settings +from onyx.db.kg_config import ( + disable_kg, + enable_kg, + get_kg_config_settings, + set_kg_config_settings, +) from onyx.db.models import User -from onyx.db.persona import create_update_persona -from onyx.db.persona import get_persona_by_id -from onyx.db.persona import mark_persona_as_deleted -from onyx.db.persona import mark_persona_as_not_deleted +from onyx.db.persona import ( + create_update_persona, + get_persona_by_id, + mark_persona_as_deleted, + mark_persona_as_not_deleted, +) from onyx.db.tools import get_builtin_tool from onyx.kg.resets.reset_index import reset_full_kg_index__commit from onyx.kg.setup.kg_default_entity_definitions import ( populate_missing_default_entity_types__commit, ) -from onyx.prompts.kg_prompts import KG_BETA_ASSISTANT_SYSTEM_PROMPT -from onyx.prompts.kg_prompts import KG_BETA_ASSISTANT_TASK_PROMPT +from onyx.prompts.kg_prompts import ( + KG_BETA_ASSISTANT_SYSTEM_PROMPT, + KG_BETA_ASSISTANT_TASK_PROMPT, +) from onyx.server.features.persona.models import PersonaUpsertRequest -from onyx.server.kg.models import DisableKGConfigRequest -from onyx.server.kg.models import EnableKGConfigRequest -from onyx.server.kg.models import EntityType -from onyx.server.kg.models import KGConfig +from onyx.server.kg.models import ( + DisableKGConfigRequest, + EnableKGConfigRequest, + EntityType, + KGConfig, + SourceAndEntityTypeView, + SourceStatistics, +) from onyx.server.kg.models import KGConfig as KGConfigAPIModel -from onyx.server.kg.models import SourceAndEntityTypeView -from onyx.server.kg.models import SourceStatistics from onyx.tools.tool_implementations.knowledge_graph.knowledge_graph_tool import ( KnowledgeGraphTool, ) diff --git a/backend/onyx/server/manage/administrative.py b/backend/onyx/server/manage/administrative.py index dc7d73e483..84e2ab4f61 100644 --- a/backend/onyx/server/manage/administrative.py +++ b/backend/onyx/server/manage/administrative.py @@ -1,11 +1,7 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import cast -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -13,32 +9,39 @@ from onyx.auth.users import current_curator_or_admin_user from onyx.background.celery.versioned_apps.client import app as client_app from onyx.background.indexing.models import IndexAttemptErrorPydantic from onyx.configs.app_configs import GENERATIVE_MODEL_ACCESS_CHECK_FREQ -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import KV_GEN_AI_KEY_CHECK_TIME -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import PUBLIC_API_TAGS -from onyx.db.connector_credential_pair import get_connector_credential_pair_for_user -from onyx.db.connector_credential_pair import update_connector_credential_pair_from_id +from onyx.configs.constants import ( + DocumentSource, + KV_GEN_AI_KEY_CHECK_TIME, + OnyxCeleryPriority, + OnyxCeleryTask, + PUBLIC_API_TAGS, +) +from onyx.db.connector_credential_pair import ( + get_connector_credential_pair_for_user, + update_connector_credential_pair_from_id, +) from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import Permission -from onyx.db.feedback import fetch_docs_ranked_by_boost_for_user -from onyx.db.feedback import update_document_boost_for_user -from onyx.db.feedback import update_document_hidden_for_user -from onyx.db.index_attempt import cancel_indexing_attempts_for_ccpair -from onyx.db.index_attempt import get_index_attempt_errors_across_connectors +from onyx.db.enums import ConnectorCredentialPairStatus, Permission +from onyx.db.feedback import ( + fetch_docs_ranked_by_boost_for_user, + update_document_boost_for_user, + update_document_hidden_for_user, +) +from onyx.db.index_attempt import ( + cancel_indexing_attempts_for_ccpair, + get_index_attempt_errors_across_connectors, +) from onyx.db.models import User from onyx.file_store.file_store import get_default_file_store from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError from onyx.llm.factory import get_default_llm from onyx.llm.utils import test_llm -from onyx.server.documents.models import ConnectorCredentialPairIdentifier -from onyx.server.documents.models import PaginatedReturn -from onyx.server.manage.models import BoostDoc -from onyx.server.manage.models import BoostUpdateRequest -from onyx.server.manage.models import HiddenUpdateRequest +from onyx.server.documents.models import ( + ConnectorCredentialPairIdentifier, + PaginatedReturn, +) +from onyx.server.manage.models import BoostDoc, BoostUpdateRequest, HiddenUpdateRequest from onyx.server.models import StatusResponse from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/manage/code_interpreter/api.py b/backend/onyx/server/manage/code_interpreter/api.py index 85f5ed1a74..93b8a269c2 100644 --- a/backend/onyx/server/manage/code_interpreter/api.py +++ b/backend/onyx/server/manage/code_interpreter/api.py @@ -1,15 +1,18 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission -from onyx.db.code_interpreter import fetch_code_interpreter_server -from onyx.db.code_interpreter import update_code_interpreter_server_enabled +from onyx.db.code_interpreter import ( + fetch_code_interpreter_server, + update_code_interpreter_server_enabled, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.server.manage.code_interpreter.models import CodeInterpreterServer -from onyx.server.manage.code_interpreter.models import CodeInterpreterServerHealth +from onyx.server.manage.code_interpreter.models import ( + CodeInterpreterServer, + CodeInterpreterServerHealth, +) from onyx.tools.tool_implementations.python.code_interpreter_client import ( CodeInterpreterClient, ) diff --git a/backend/onyx/server/manage/discord_bot/api.py b/backend/onyx/server/manage/discord_bot/api.py index 680bf3ffe0..99d6350408 100644 --- a/backend/onyx/server/manage/discord_bot/api.py +++ b/backend/onyx/server/manage/discord_bot/api.py @@ -1,35 +1,36 @@ """Discord bot admin API endpoints.""" -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import status +from fastapi import APIRouter, Depends, HTTPException, status from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.configs.app_configs import DISCORD_BOT_TOKEN -from onyx.db.discord_bot import create_discord_bot_config -from onyx.db.discord_bot import create_guild_config -from onyx.db.discord_bot import delete_discord_bot_config -from onyx.db.discord_bot import delete_discord_service_api_key -from onyx.db.discord_bot import delete_guild_config -from onyx.db.discord_bot import get_channel_config_by_internal_ids -from onyx.db.discord_bot import get_channel_configs -from onyx.db.discord_bot import get_discord_bot_config -from onyx.db.discord_bot import get_guild_config_by_internal_id -from onyx.db.discord_bot import get_guild_configs -from onyx.db.discord_bot import update_discord_channel_config -from onyx.db.discord_bot import update_guild_config +from onyx.db.discord_bot import ( + create_discord_bot_config, + create_guild_config, + delete_discord_bot_config, + delete_discord_service_api_key, + delete_guild_config, + get_channel_config_by_internal_ids, + get_channel_configs, + get_discord_bot_config, + get_guild_config_by_internal_id, + get_guild_configs, + update_discord_channel_config, + update_guild_config, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.server.manage.discord_bot.models import DiscordBotConfigCreateRequest -from onyx.server.manage.discord_bot.models import DiscordBotConfigResponse -from onyx.server.manage.discord_bot.models import DiscordChannelConfigResponse -from onyx.server.manage.discord_bot.models import DiscordChannelConfigUpdateRequest -from onyx.server.manage.discord_bot.models import DiscordGuildConfigCreateResponse -from onyx.server.manage.discord_bot.models import DiscordGuildConfigResponse -from onyx.server.manage.discord_bot.models import DiscordGuildConfigUpdateRequest +from onyx.server.manage.discord_bot.models import ( + DiscordBotConfigCreateRequest, + DiscordBotConfigResponse, + DiscordChannelConfigResponse, + DiscordChannelConfigUpdateRequest, + DiscordGuildConfigCreateResponse, + DiscordGuildConfigResponse, + DiscordGuildConfigUpdateRequest, +) from onyx.server.manage.discord_bot.utils import generate_discord_registration_key from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/manage/discord_bot/utils.py b/backend/onyx/server/manage/discord_bot/utils.py index c4c5f2cc0a..fe11ddba5f 100644 --- a/backend/onyx/server/manage/discord_bot/utils.py +++ b/backend/onyx/server/manage/discord_bot/utils.py @@ -1,8 +1,7 @@ """Discord registration key generation and parsing.""" import secrets -from urllib.parse import quote -from urllib.parse import unquote +from urllib.parse import quote, unquote from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/manage/embedding/api.py b/backend/onyx/server/manage/embedding/api.py index e304f56101..0a09bd3cf7 100644 --- a/backend/onyx/server/manage/embedding/api.py +++ b/backend/onyx/server/manage/embedding/api.py @@ -1,28 +1,31 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.llm import fetch_existing_embedding_providers -from onyx.db.llm import remove_embedding_provider -from onyx.db.llm import upsert_cloud_embedding_provider +from onyx.db.llm import ( + fetch_existing_embedding_providers, + remove_embedding_provider, + upsert_cloud_embedding_provider, +) from onyx.db.models import User -from onyx.db.search_settings import get_all_search_settings -from onyx.db.search_settings import get_current_db_embedding_provider +from onyx.db.search_settings import ( + get_all_search_settings, + get_current_db_embedding_provider, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.indexing.models import EmbeddingModelDetail from onyx.natural_language_processing.search_nlp_models import EmbeddingModel -from onyx.server.manage.embedding.models import CloudEmbeddingProvider -from onyx.server.manage.embedding.models import CloudEmbeddingProviderCreationRequest -from onyx.server.manage.embedding.models import TestEmbeddingRequest +from onyx.server.manage.embedding.models import ( + CloudEmbeddingProvider, + CloudEmbeddingProviderCreationRequest, + TestEmbeddingRequest, +) from onyx.utils.logger import setup_logger -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType +from shared_configs.configs import MODEL_SERVER_HOST, MODEL_SERVER_PORT +from shared_configs.enums import EmbeddingProvider, EmbedTextType logger = setup_logger() diff --git a/backend/onyx/server/manage/get_state.py b/backend/onyx/server/manage/get_state.py index ba766dce53..b180f4d892 100644 --- a/backend/onyx/server/manage/get_state.py +++ b/backend/onyx/server/manage/get_state.py @@ -4,27 +4,28 @@ import threading import requests from cachetools import TTLCache -from fastapi import APIRouter -from fastapi import HTTPException -from fastapi import Response +from fastapi import APIRouter, HTTPException, Response from fastapi.concurrency import run_in_threadpool from onyx import __version__ -from onyx.auth.users import anonymous_user_enabled -from onyx.auth.users import user_needs_to_be_verified +from onyx.auth.users import anonymous_user_enabled, user_needs_to_be_verified from onyx.configs.app_configs import OAUTH_ENABLED -from onyx.configs.constants import DEV_VERSION_PATTERN -from onyx.configs.constants import PUBLIC_API_TAGS -from onyx.configs.constants import STABLE_VERSION_PATTERN +from onyx.configs.constants import ( + DEV_VERSION_PATTERN, + PUBLIC_API_TAGS, + STABLE_VERSION_PATTERN, +) from onyx.db.auth import get_user_count from onyx.db.engine.sql_engine import get_session_with_shared_schema from onyx.db.enums import SSOProviderType from onyx.db.sso_provider import fetch_sso_providers -from onyx.server.manage.models import AllVersions -from onyx.server.manage.models import AuthConfigResponse -from onyx.server.manage.models import ContainerVersions -from onyx.server.manage.models import SSOProviderOption -from onyx.server.manage.models import VersionResponse +from onyx.server.manage.models import ( + AllVersions, + AuthConfigResponse, + ContainerVersions, + SSOProviderOption, + VersionResponse, +) from onyx.server.models import StatusResponse from onyx.server.security.store import get_security_settings from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/manage/image_generation/api.py b/backend/onyx/server/manage/image_generation/api.py index 23b351b810..9f0cb24f33 100644 --- a/backend/onyx/server/manage/image_generation/api.py +++ b/backend/onyx/server/manage/image_generation/api.py @@ -1,35 +1,39 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.image_generation import create_image_generation_config__no_commit -from onyx.db.image_generation import delete_image_generation_config__no_commit -from onyx.db.image_generation import get_all_image_generation_configs -from onyx.db.image_generation import get_image_generation_config -from onyx.db.image_generation import set_default_image_generation_config -from onyx.db.image_generation import unset_default_image_generation_config +from onyx.db.image_generation import ( + create_image_generation_config__no_commit, + delete_image_generation_config__no_commit, + get_all_image_generation_configs, + get_image_generation_config, + set_default_image_generation_config, + unset_default_image_generation_config, +) from onyx.db.llm import remove_llm_provider from onyx.db.models import LLMProvider as LLMProviderModel -from onyx.db.models import ModelConfiguration -from onyx.db.models import User +from onyx.db.models import ModelConfiguration, User from onyx.image_gen.exceptions import ImageProviderCredentialsError -from onyx.image_gen.factory import get_image_generation_provider -from onyx.image_gen.factory import validate_credentials +from onyx.image_gen.factory import get_image_generation_provider, validate_credentials from onyx.image_gen.interfaces import ImageGenerationProviderCredentials from onyx.llm.model_capabilities import get_max_input_tokens -from onyx.server.manage.image_generation.models import ImageGenerationConfigCreate -from onyx.server.manage.image_generation.models import ImageGenerationConfigUpdate -from onyx.server.manage.image_generation.models import ImageGenerationConfigView -from onyx.server.manage.image_generation.models import ImageGenerationCredentials -from onyx.server.manage.image_generation.models import TestImageGenerationRequest -from onyx.server.manage.llm.api import _validate_and_normalize_vertex_auth -from onyx.server.manage.llm.api import _validate_llm_provider_change -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.image_generation.models import ( + ImageGenerationConfigCreate, + ImageGenerationConfigUpdate, + ImageGenerationConfigView, + ImageGenerationCredentials, + TestImageGenerationRequest, +) +from onyx.server.manage.llm.api import ( + _validate_and_normalize_vertex_auth, + _validate_llm_provider_change, +) +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) from onyx.server.manage.llm.provider_cache import invalidate_provider_listing_cache from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/manage/llm/api.py b/backend/onyx/server/manage/llm/api.py index 7d5577ebd1..f4605d35dd 100644 --- a/backend/onyx/server/manage/llm/api.py +++ b/backend/onyx/server/manage/llm/api.py @@ -1,20 +1,14 @@ from collections import defaultdict -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import boto3 import botocore.session import httpx from botocore.config import Config -from botocore.exceptions import BotoCoreError -from botocore.exceptions import ClientError -from botocore.exceptions import NoCredentialsError -from botocore.tokens import FrozenAuthToken -from botocore.tokens import TokenProviderChain -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Query +from botocore.exceptions import BotoCoreError, ClientError, NoCredentialsError +from botocore.tokens import FrozenAuthToken, TokenProviderChain +from fastapi import APIRouter, Depends, Query from pydantic import ValidationError from sqlalchemy.orm import Session @@ -22,99 +16,115 @@ from onyx.auth.permissions import require_permission from onyx.auth.schemas import UserRole from onyx.auth.users import current_chat_accessible_user from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import LLMModelFlowType -from onyx.db.enums import Permission -from onyx.db.llm import can_user_access_llm_provider -from onyx.db.llm import fetch_default_llm_model -from onyx.db.llm import fetch_default_vision_model -from onyx.db.llm import fetch_existing_llm_provider_by_id -from onyx.db.llm import fetch_existing_llm_providers -from onyx.db.llm import fetch_existing_models -from onyx.db.llm import fetch_model_configuration_by_id -from onyx.db.llm import fetch_persona_with_groups -from onyx.db.llm import fetch_user_group_ids -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import sync_model_configurations -from onyx.db.llm import update_default_provider -from onyx.db.llm import update_default_vision_provider -from onyx.db.llm import upsert_llm_provider -from onyx.db.llm import validate_persona_ids_exist -from onyx.db.models import Persona -from onyx.db.models import User +from onyx.db.enums import LLMModelFlowType, Permission +from onyx.db.llm import ( + can_user_access_llm_provider, + fetch_default_llm_model, + fetch_default_vision_model, + fetch_existing_llm_provider_by_id, + fetch_existing_llm_providers, + fetch_existing_models, + fetch_model_configuration_by_id, + fetch_persona_with_groups, + fetch_user_group_ids, + remove_llm_provider, + sync_model_configurations, + update_default_provider, + update_default_vision_provider, + upsert_llm_provider, + validate_persona_ids_exist, +) +from onyx.db.models import Persona, User from onyx.db.persona import user_can_access_persona from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.llm.constants import LlmProviderNames -from onyx.llm.constants import PROVIDER_DISPLAY_NAMES -from onyx.llm.constants import WELL_KNOWN_PROVIDER_NAMES -from onyx.llm.factory import get_default_llm -from onyx.llm.factory import get_llm -from onyx.llm.factory import get_max_input_tokens_from_llm_provider -from onyx.llm.model_capabilities import get_bedrock_token_limit -from onyx.llm.model_capabilities import litellm_thinks_model_supports_image_input -from onyx.llm.model_capabilities import model_is_reasoning_model -from onyx.llm.utils import get_llm_contextual_cost -from onyx.llm.utils import is_sensitive_custom_config_key -from onyx.llm.utils import test_llm +from onyx.llm.constants import ( + LlmProviderNames, + PROVIDER_DISPLAY_NAMES, + WELL_KNOWN_PROVIDER_NAMES, +) +from onyx.llm.factory import ( + get_default_llm, + get_llm, + get_max_input_tokens_from_llm_provider, +) +from onyx.llm.model_capabilities import ( + get_bedrock_token_limit, + litellm_thinks_model_supports_image_input, + model_is_reasoning_model, +) +from onyx.llm.utils import ( + get_llm_contextual_cost, + is_sensitive_custom_config_key, + test_llm, +) from onyx.llm.well_known_providers.auto_update_service import ( fetch_llm_recommendations_from_github, ) -from onyx.llm.well_known_providers.constants import LM_STUDIO_API_KEY_CONFIG_KEY -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_SERVICE_ACCOUNT -from onyx.llm.well_known_providers.constants import VERTEX_AUTH_METHOD_WORKLOAD_IDENTITY -from onyx.llm.well_known_providers.constants import VERTEX_CREDENTIALS_FILE_KWARG -from onyx.llm.well_known_providers.constants import VERTEX_PROJECT_KWARG +from onyx.llm.well_known_providers.constants import ( + LM_STUDIO_API_KEY_CONFIG_KEY, + VERTEX_AUTH_METHOD_KWARG, + VERTEX_AUTH_METHOD_SERVICE_ACCOUNT, + VERTEX_AUTH_METHOD_WORKLOAD_IDENTITY, + VERTEX_CREDENTIALS_FILE_KWARG, + VERTEX_PROJECT_KWARG, +) from onyx.llm.well_known_providers.llm_provider_options import ( fetch_available_well_known_llms, -) -from onyx.llm.well_known_providers.llm_provider_options import ( WellKnownLLMProviderDescriptor, ) -from onyx.server.manage.llm.models import BedrockFinalModelResponse -from onyx.server.manage.llm.models import BedrockModelsRequest -from onyx.server.manage.llm.models import BifrostFinalModelResponse -from onyx.server.manage.llm.models import BifrostModelsRequest -from onyx.server.manage.llm.models import CustomProviderOption -from onyx.server.manage.llm.models import DefaultModel -from onyx.server.manage.llm.models import LitellmFinalModelResponse -from onyx.server.manage.llm.models import LitellmModelDetails -from onyx.server.manage.llm.models import LitellmModelsRequest -from onyx.server.manage.llm.models import LLMCost -from onyx.server.manage.llm.models import LLMProviderDescriptor -from onyx.server.manage.llm.models import LLMProviderResponse -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import LMStudioFinalModelResponse -from onyx.server.manage.llm.models import LMStudioModelsRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest -from onyx.server.manage.llm.models import NebiusTokenfactoryFinalModelResponse -from onyx.server.manage.llm.models import NebiusTokenfactoryModelsRequest -from onyx.server.manage.llm.models import OllamaFinalModelResponse -from onyx.server.manage.llm.models import OllamaModelDetails -from onyx.server.manage.llm.models import OllamaModelsRequest -from onyx.server.manage.llm.models import OpenAICompatibleFinalModelResponse -from onyx.server.manage.llm.models import OpenAICompatibleModelsRequest -from onyx.server.manage.llm.models import OpenRouterFinalModelResponse -from onyx.server.manage.llm.models import OpenRouterModelDetails -from onyx.server.manage.llm.models import OpenRouterModelsRequest -from onyx.server.manage.llm.models import SyncModelEntry -from onyx.server.manage.llm.models import TestLLMRequest -from onyx.server.manage.llm.models import VisionProviderResponse -from onyx.server.manage.llm.provider_cache import cache_provider_listing -from onyx.server.manage.llm.provider_cache import get_cached_provider_listing -from onyx.server.manage.llm.provider_cache import invalidate_provider_listing_cache -from onyx.server.manage.llm.utils import generate_bedrock_display_name -from onyx.server.manage.llm.utils import generate_ollama_display_name -from onyx.server.manage.llm.utils import is_embedding_model -from onyx.server.manage.llm.utils import is_reasoning_model -from onyx.server.manage.llm.utils import is_valid_bedrock_model -from onyx.server.manage.llm.utils import ModelMetadata -from onyx.server.manage.llm.utils import strip_openrouter_vendor_prefix -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.server.manage.llm.models import ( + BedrockFinalModelResponse, + BedrockModelsRequest, + BifrostFinalModelResponse, + BifrostModelsRequest, + CustomProviderOption, + DefaultModel, + LitellmFinalModelResponse, + LitellmModelDetails, + LitellmModelsRequest, + LLMCost, + LLMProviderDescriptor, + LLMProviderResponse, + LLMProviderUpsertRequest, + LLMProviderView, + LMStudioFinalModelResponse, + LMStudioModelsRequest, + ModelConfigurationUpsertRequest, + NebiusTokenfactoryFinalModelResponse, + NebiusTokenfactoryModelsRequest, + OllamaFinalModelResponse, + OllamaModelDetails, + OllamaModelsRequest, + OpenAICompatibleFinalModelResponse, + OpenAICompatibleModelsRequest, + OpenRouterFinalModelResponse, + OpenRouterModelDetails, + OpenRouterModelsRequest, + SyncModelEntry, + TestLLMRequest, + VisionProviderResponse, +) +from onyx.server.manage.llm.provider_cache import ( + cache_provider_listing, + get_cached_provider_listing, + invalidate_provider_listing_cache, +) +from onyx.server.manage.llm.utils import ( + generate_bedrock_display_name, + generate_ollama_display_name, + is_embedding_model, + is_reasoning_model, + is_valid_bedrock_model, + ModelMetadata, + strip_openrouter_vendor_prefix, +) +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.encryption import mask_string as mask_with_ellipsis from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/manage/llm/models.py b/backend/onyx/server/manage/llm/models.py index ac7a1cf48e..3276ce9320 100644 --- a/backend/onyx/server/manage/llm/models.py +++ b/backend/onyx/server/manage/llm/models.py @@ -1,23 +1,22 @@ from __future__ import annotations from functools import cached_property -from typing import Any -from typing import Generic -from typing import TYPE_CHECKING -from typing import TypeVar +from typing import Any, Generic, TYPE_CHECKING, TypeVar -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator +from pydantic import BaseModel, Field, field_validator from onyx.db.enums import LLMModelFlowType from onyx.llm.constants import DYNAMIC_LLM_PROVIDERS -from onyx.llm.model_capabilities import get_max_input_tokens -from onyx.llm.model_capabilities import litellm_thinks_model_supports_image_input -from onyx.llm.model_capabilities import model_is_reasoning_model -from onyx.server.manage.llm.utils import extract_vendor_from_model_name -from onyx.server.manage.llm.utils import filter_model_configurations -from onyx.server.manage.llm.utils import is_reasoning_model +from onyx.llm.model_capabilities import ( + get_max_input_tokens, + litellm_thinks_model_supports_image_input, + model_is_reasoning_model, +) +from onyx.server.manage.llm.utils import ( + extract_vendor_from_model_name, + filter_model_configurations, + is_reasoning_model, +) if TYPE_CHECKING: from onyx.db.models import LLMProvider as LLMProviderModel @@ -74,8 +73,6 @@ class LLMProviderDescriptor(BaseModel): ) -> "LLMProviderDescriptor": from onyx.llm.well_known_providers.llm_provider_options import ( fetch_default_model_for_provider, - ) - from onyx.llm.well_known_providers.llm_provider_options import ( get_provider_display_name, ) diff --git a/backend/onyx/server/manage/llm/provider_cache.py b/backend/onyx/server/manage/llm/provider_cache.py index 3fa9735c21..147bebb976 100644 --- a/backend/onyx/server/manage/llm/provider_cache.py +++ b/backend/onyx/server/manage/llm/provider_cache.py @@ -18,14 +18,11 @@ Cache failures are non-fatal: readers fall through to Postgres. import hashlib import uuid -from pydantic import BaseModel -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from onyx.cache.factory import get_cache_backend -from onyx.cache.interface import CACHE_TRANSIENT_ERRORS -from onyx.cache.interface import CacheBackend -from onyx.server.manage.llm.models import LLMProviderDescriptor -from onyx.server.manage.llm.models import LLMProviderResponse +from onyx.cache.interface import CACHE_TRANSIENT_ERRORS, CacheBackend +from onyx.server.manage.llm.models import LLMProviderDescriptor, LLMProviderResponse from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/manage/llm/utils.py b/backend/onyx/server/manage/llm/utils.py index 758de41d92..8040317dd0 100644 --- a/backend/onyx/server/manage/llm/utils.py +++ b/backend/onyx/server/manage/llm/utils.py @@ -10,12 +10,14 @@ Utilities for dynamic LLM providers (Bedrock, Ollama, OpenRouter): import re from typing import TypedDict -from onyx.llm.constants import BEDROCK_MODEL_NAME_MAPPINGS -from onyx.llm.constants import LlmProviderNames -from onyx.llm.constants import MODEL_PREFIX_TO_VENDOR -from onyx.llm.constants import OLLAMA_MODEL_NAME_MAPPINGS -from onyx.llm.constants import OLLAMA_MODEL_TO_VENDOR -from onyx.llm.constants import PROVIDER_DISPLAY_NAMES +from onyx.llm.constants import ( + BEDROCK_MODEL_NAME_MAPPINGS, + LlmProviderNames, + MODEL_PREFIX_TO_VENDOR, + OLLAMA_MODEL_NAME_MAPPINGS, + OLLAMA_MODEL_TO_VENDOR, + PROVIDER_DISPLAY_NAMES, +) class ModelMetadata(TypedDict): diff --git a/backend/onyx/server/manage/models.py b/backend/onyx/server/manage/models.py index f8aa23bae0..7298a436a7 100644 --- a/backend/onyx/server/manage/models.py +++ b/backend/onyx/server/manage/models.py @@ -1,34 +1,27 @@ import re from datetime import datetime from enum import Enum -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_validator -from pydantic import model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from onyx.auth.schemas import UserRole from onyx.context.search.models import SavedSearchSettings -from onyx.db.enums import DefaultAppMode -from onyx.db.enums import SSOProviderType -from onyx.db.enums import SupportedLanguage -from onyx.db.enums import ThemePreference +from onyx.db.enums import ( + DefaultAppMode, + SSOProviderType, + SupportedLanguage, + ThemePreference, +) from onyx.db.memory import MAX_MEMORIES_PER_USER -from onyx.db.models import AllowedAnswerFilters -from onyx.db.models import ChannelConfig +from onyx.db.models import AllowedAnswerFilters, ChannelConfig, User from onyx.db.models import SlackBot as SlackAppModel from onyx.db.models import SlackChannelConfig as SlackChannelConfigModel from onyx.db.models import StandardAnswer as StandardAnswerModel from onyx.db.models import StandardAnswerCategory as StandardAnswerCategoryModel -from onyx.db.models import User from onyx.onyxbot.slack.config import VALID_SLACK_FILTERS -from onyx.server.features.persona.models import FullPersonaSnapshot -from onyx.server.features.persona.models import PersonaSnapshot -from onyx.server.models import FullUserSnapshot -from onyx.server.models import InvitedUserSnapshot +from onyx.server.features.persona.models import FullPersonaSnapshot, PersonaSnapshot +from onyx.server.models import FullUserSnapshot, InvitedUserSnapshot if TYPE_CHECKING: pass diff --git a/backend/onyx/server/manage/oauth_test.py b/backend/onyx/server/manage/oauth_test.py index 4dcc17cba9..ed6a97b0f9 100644 --- a/backend/onyx/server/manage/oauth_test.py +++ b/backend/onyx/server/manage/oauth_test.py @@ -9,12 +9,13 @@ flow again. from typing import Any -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from pydantic import BaseModel -from onyx.auth.oauth_claims_capture import get_captured_oauth_claims -from onyx.auth.oauth_claims_capture import get_idp_profile_fields +from onyx.auth.oauth_claims_capture import ( + get_captured_oauth_claims, + get_idp_profile_fields, +) from onyx.auth.permissions import require_permission from onyx.configs.app_configs import IDP_PROFILE_ENRICHMENT_ENABLED from onyx.db.enums import Permission diff --git a/backend/onyx/server/manage/opensearch_migration/api.py b/backend/onyx/server/manage/opensearch_migration/api.py index 5772a940db..78457dd3b4 100644 --- a/backend/onyx/server/manage/opensearch_migration/api.py +++ b/backend/onyx/server/manage/opensearch_migration/api.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -7,16 +6,14 @@ from onyx.configs.app_configs import ONYX_DISABLE_VESPA from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.opensearch_migration import get_opensearch_migration_state -from onyx.db.opensearch_migration import get_opensearch_retrieval_state -from onyx.db.opensearch_migration import set_enable_opensearch_retrieval_with_commit +from onyx.db.opensearch_migration import ( + get_opensearch_migration_state, + get_opensearch_retrieval_state, + set_enable_opensearch_retrieval_with_commit, +) from onyx.server.manage.opensearch_migration.models import ( OpenSearchMigrationStatusResponse, -) -from onyx.server.manage.opensearch_migration.models import ( OpenSearchRetrievalStatusRequest, -) -from onyx.server.manage.opensearch_migration.models import ( OpenSearchRetrievalStatusResponse, ) diff --git a/backend/onyx/server/manage/search_settings.py b/backend/onyx/server/manage/search_settings.py index 715b47ad45..d433c1bf91 100644 --- a/backend/onyx/server/manage/search_settings.py +++ b/backend/onyx/server/manage/search_settings.py @@ -1,53 +1,58 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import status +from fastapi import APIRouter, Depends, HTTPException, status from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP -from onyx.context.search.models import SavedSearchSettings -from onyx.context.search.models import SearchSettingsCreationRequest +from onyx.context.search.models import ( + SavedSearchSettings, + SearchSettingsCreationRequest, +) from onyx.db.connector_credential_pair import ( fetch_indexable_standard_connector_credential_pair_ids, + get_connector_credential_pairs, + get_last_successful_attempt_poll_range_end, + resync_cc_pair, ) -from onyx.db.connector_credential_pair import get_connector_credential_pairs -from onyx.db.connector_credential_pair import get_last_successful_attempt_poll_range_end -from onyx.db.connector_credential_pair import resync_cc_pair from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import Permission -from onyx.db.enums import SwitchoverType -from onyx.db.index_attempt import create_synthetic_seed_attempt -from onyx.db.index_attempt import expire_index_attempts -from onyx.db.llm import fetch_default_contextual_rag_model -from onyx.db.llm import update_default_contextual_model -from onyx.db.llm import update_no_default_contextual_rag_provider -from onyx.db.models import IndexModelStatus -from onyx.db.models import User -from onyx.db.port_attempt import cancel_active_port_attempts -from onyx.db.port_attempt import port_backfill_has_pending_work -from onyx.db.search_settings import create_search_settings -from onyx.db.search_settings import delete_search_settings -from onyx.db.search_settings import get_current_search_settings -from onyx.db.search_settings import get_embedding_provider_from_provider_type -from onyx.db.search_settings import get_secondary_search_settings -from onyx.db.search_settings import update_current_search_settings -from onyx.db.search_settings import update_search_settings_status -from onyx.document_index.factory import get_all_document_indices -from onyx.document_index.factory import get_default_document_index +from onyx.db.enums import Permission, SwitchoverType +from onyx.db.index_attempt import create_synthetic_seed_attempt, expire_index_attempts +from onyx.db.llm import ( + fetch_default_contextual_rag_model, + update_default_contextual_model, + update_no_default_contextual_rag_provider, +) +from onyx.db.models import IndexModelStatus, User +from onyx.db.port_attempt import ( + cancel_active_port_attempts, + port_backfill_has_pending_work, +) +from onyx.db.search_settings import ( + create_search_settings, + delete_search_settings, + get_current_search_settings, + get_embedding_provider_from_provider_type, + get_secondary_search_settings, + update_current_search_settings, + update_search_settings_status, +) +from onyx.document_index.factory import ( + get_all_document_indices, + get_default_document_index, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.file_processing.unstructured import delete_unstructured_api_key -from onyx.file_processing.unstructured import get_unstructured_api_key -from onyx.file_processing.unstructured import update_unstructured_api_key +from onyx.file_processing.unstructured import ( + delete_unstructured_api_key, + get_unstructured_api_key, + update_unstructured_api_key, +) from onyx.natural_language_processing.search_nlp_models import clean_model_name from onyx.server.manage.embedding.models import SearchSettingsDeleteRequest from onyx.server.manage.models import FullModelVersionResponse from onyx.server.models import IdReturn from onyx.server.utils_vector_db import require_vector_db from onyx.utils.logger import setup_logger -from shared_configs.configs import ALT_INDEX_SUFFIX -from shared_configs.configs import MULTI_TENANT +from shared_configs.configs import ALT_INDEX_SUFFIX, MULTI_TENANT router = APIRouter(prefix="/search-settings") logger = setup_logger() diff --git a/backend/onyx/server/manage/slack_bot.py b/backend/onyx/server/manage/slack_bot.py index cae386a59b..6bb86c7c4f 100644 --- a/backend/onyx/server/manage/slack_bot.py +++ b/backend/onyx/server/manage/slack_bot.py @@ -1,6 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -8,28 +6,35 @@ from onyx.configs.constants import MilestoneRecordType from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.models import ChannelConfig -from onyx.db.models import User +from onyx.db.models import ChannelConfig, User from onyx.db.persona import get_persona_by_id -from onyx.db.slack_bot import fetch_slack_bot -from onyx.db.slack_bot import fetch_slack_bots -from onyx.db.slack_bot import insert_slack_bot -from onyx.db.slack_bot import remove_slack_bot -from onyx.db.slack_bot import update_slack_bot -from onyx.db.slack_channel_config import create_slack_channel_persona -from onyx.db.slack_channel_config import fetch_slack_channel_config -from onyx.db.slack_channel_config import fetch_slack_channel_configs -from onyx.db.slack_channel_config import insert_slack_channel_config -from onyx.db.slack_channel_config import remove_slack_channel_config -from onyx.db.slack_channel_config import update_slack_channel_config +from onyx.db.slack_bot import ( + fetch_slack_bot, + fetch_slack_bots, + insert_slack_bot, + remove_slack_bot, + update_slack_bot, +) +from onyx.db.slack_channel_config import ( + create_slack_channel_persona, + fetch_slack_channel_config, + fetch_slack_channel_configs, + insert_slack_channel_config, + remove_slack_channel_config, + update_slack_channel_config, +) from onyx.onyxbot.slack.config import validate_channel_name -from onyx.server.manage.models import SlackBot -from onyx.server.manage.models import SlackBotCreationRequest -from onyx.server.manage.models import SlackChannelConfig -from onyx.server.manage.models import SlackChannelConfigCreationRequest -from onyx.server.manage.validate_tokens import validate_app_token -from onyx.server.manage.validate_tokens import validate_bot_token -from onyx.server.manage.validate_tokens import validate_user_token +from onyx.server.manage.models import ( + SlackBot, + SlackBotCreationRequest, + SlackChannelConfig, + SlackChannelConfigCreationRequest, +) +from onyx.server.manage.validate_tokens import ( + validate_app_token, + validate_bot_token, + validate_user_token, +) from onyx.utils.logger import setup_logger from onyx.utils.telemetry import mt_cloud_telemetry from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/manage/sso/api.py b/backend/onyx/server/manage/sso/api.py index 8ded2df9a7..f1c6b17d7e 100644 --- a/backend/onyx/server/manage/sso/api.py +++ b/backend/onyx/server/manage/sso/api.py @@ -1,7 +1,6 @@ from typing import Any -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session @@ -9,21 +8,23 @@ from onyx.auth.permissions import require_permission from onyx.configs.app_configs import WEB_DOMAIN from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.models import SSOProvider -from onyx.db.models import User -from onyx.db.sso_provider import create_sso_provider -from onyx.db.sso_provider import fetch_sso_providers -from onyx.db.sso_provider import set_sso_provider_enabled -from onyx.db.sso_provider import update_sso_provider +from onyx.db.models import SSOProvider, User +from onyx.db.sso_provider import ( + create_sso_provider, + fetch_sso_providers, + set_sso_provider_enabled, + update_sso_provider, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.server.manage.get_state import invalidate_sso_provider_options_cache -from onyx.server.manage.sso.models import SSOProviderCreateRequest -from onyx.server.manage.sso.models import SSOProviderEnabledRequest -from onyx.server.manage.sso.models import SSOProviderResponse -from onyx.server.manage.sso.models import SSOProviderUpdateRequest -from onyx.utils.encryption import reject_masked_credentials -from onyx.utils.encryption import restore_masked_credentials +from onyx.server.manage.sso.models import ( + SSOProviderCreateRequest, + SSOProviderEnabledRequest, + SSOProviderResponse, + SSOProviderUpdateRequest, +) +from onyx.utils.encryption import reject_masked_credentials, restore_masked_credentials from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop admin_router = APIRouter(prefix="/admin/sso") diff --git a/backend/onyx/server/manage/tracing/api.py b/backend/onyx/server/manage/tracing/api.py index b138ea9113..a93726f7de 100644 --- a/backend/onyx/server/manage/tracing/api.py +++ b/backend/onyx/server/manage/tracing/api.py @@ -1,26 +1,30 @@ from __future__ import annotations -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.models import TracingProviderConfig -from onyx.db.models import User -from onyx.db.tracing import delete_tracing_provider -from onyx.db.tracing import fetch_tracing_provider -from onyx.db.tracing import upsert_tracing_provider +from onyx.db.models import TracingProviderConfig, User +from onyx.db.tracing import ( + delete_tracing_provider, + fetch_tracing_provider, + upsert_tracing_provider, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.manage.tracing.models import TracingProviderTestRequest -from onyx.server.manage.tracing.models import TracingProviderUpsertRequest -from onyx.server.manage.tracing.models import TracingProviderView -from onyx.tracing.provider_config import BraintrustConfig -from onyx.tracing.provider_config import EffectiveTracingConfig -from onyx.tracing.provider_config import LangfuseConfig -from onyx.tracing.provider_config import resolve_effective_tracing_config +from onyx.server.manage.tracing.models import ( + TracingProviderTestRequest, + TracingProviderUpsertRequest, + TracingProviderView, +) +from onyx.tracing.provider_config import ( + BraintrustConfig, + EffectiveTracingConfig, + LangfuseConfig, + resolve_effective_tracing_config, +) from onyx.tracing.validation import validate_tracing_credentials from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/manage/tracing/models.py b/backend/onyx/server/manage/tracing/models.py index 261abacf19..404e286a78 100644 --- a/backend/onyx/server/manage/tracing/models.py +++ b/backend/onyx/server/manage/tracing/models.py @@ -1,7 +1,6 @@ from typing import Literal -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from shared_configs.enums import TracingProviderType diff --git a/backend/onyx/server/manage/users.py b/backend/onyx/server/manage/users.py index d4ddec77c1..2c74e5b14d 100644 --- a/backend/onyx/server/manage/users.py +++ b/backend/onyx/server/manage/users.py @@ -1,22 +1,13 @@ import csv import io import re -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import cast from uuid import UUID import jwt -from email_validator import EmailNotValidError -from email_validator import EmailUndeliverableError -from email_validator import validate_email -from fastapi import APIRouter -from fastapi import Body -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query -from fastapi import Request +from email_validator import EmailNotValidError, EmailUndeliverableError, validate_email +from fastapi import APIRouter, Body, Depends, HTTPException, Query, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel from sqlalchemy import select @@ -24,115 +15,127 @@ from sqlalchemy.orm import Session from onyx.auth.anonymous_user import fetch_anonymous_user_info from onyx.auth.email_utils import send_user_email_invite -from onyx.auth.invited_users import get_invited_users -from onyx.auth.invited_users import remove_user_from_invited_users -from onyx.auth.invited_users import write_invited_users -from onyx.auth.permissions import get_effective_permissions -from onyx.auth.permissions import require_permission +from onyx.auth.invited_users import ( + get_invited_users, + remove_user_from_invited_users, + write_invited_users, +) +from onyx.auth.permissions import get_effective_permissions, require_permission from onyx.auth.schemas import UserRole -from onyx.auth.session_tokens import build_session_rejection_error -from onyx.auth.session_tokens import classify_session_token_value -from onyx.auth.session_tokens import SessionRejection -from onyx.auth.users import anonymous_user_enabled -from onyx.auth.users import current_curator_or_admin_user -from onyx.auth.users import enforce_seat_limit_locked -from onyx.auth.users import optional_user -from onyx.auth.users import scope_exempt -from onyx.configs.app_configs import AUTH_BACKEND -from onyx.configs.app_configs import AuthBackend -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import EMAIL_CONFIGURED -from onyx.configs.app_configs import ENABLE_EMAIL_INVITES -from onyx.configs.app_configs import NUM_FREE_TRIAL_USER_INVITES -from onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX -from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS -from onyx.configs.app_configs import USER_AUTH_SECRET -from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME -from onyx.configs.constants import PUBLIC_API_TAGS +from onyx.auth.session_tokens import ( + build_session_rejection_error, + classify_session_token_value, + SessionRejection, +) +from onyx.auth.users import ( + anonymous_user_enabled, + current_curator_or_admin_user, + enforce_seat_limit_locked, + optional_user, + scope_exempt, +) +from onyx.configs.app_configs import ( + AUTH_BACKEND, + AuthBackend, + DEV_MODE, + EMAIL_CONFIGURED, + ENABLE_EMAIL_INVITES, + NUM_FREE_TRIAL_USER_INVITES, + REDIS_AUTH_KEY_PREFIX, + SESSION_EXPIRE_TIME_SECONDS, + USER_AUTH_SECRET, +) +from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME, PUBLIC_API_TAGS from onyx.db.api_key import is_api_key_email_address from onyx.db.auth import get_live_users_count -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_shared_schema -from onyx.db.enums import AccountType -from onyx.db.enums import Permission -from onyx.db.enums import UserFileStatus -from onyx.db.models import User -from onyx.db.models import UserFile -from onyx.db.tenant_invite_counter import release_trial_invites -from onyx.db.tenant_invite_counter import reserve_trial_invites -from onyx.db.user_preferences import activate_user -from onyx.db.user_preferences import deactivate_user -from onyx.db.user_preferences import get_all_user_assistant_specific_configs -from onyx.db.user_preferences import get_latest_access_token_for_user -from onyx.db.user_preferences import get_memories_for_user -from onyx.db.user_preferences import update_assistant_preferences -from onyx.db.user_preferences import update_user_assistant_visibility -from onyx.db.user_preferences import update_user_auto_scroll -from onyx.db.user_preferences import update_user_chat_background -from onyx.db.user_preferences import update_user_default_app_mode -from onyx.db.user_preferences import update_user_default_model -from onyx.db.user_preferences import update_user_language -from onyx.db.user_preferences import update_user_paste_as_tile -from onyx.db.user_preferences import update_user_personalization -from onyx.db.user_preferences import update_user_pinned_assistants -from onyx.db.user_preferences import update_user_role -from onyx.db.user_preferences import update_user_shortcut_enabled -from onyx.db.user_preferences import update_user_temperature_override_enabled -from onyx.db.user_preferences import update_user_theme_preference -from onyx.db.user_preferences import update_users_craft_enabled -from onyx.db.users import batch_get_user_groups -from onyx.db.users import delete_user_from_db -from onyx.db.users import get_all_accepted_users -from onyx.db.users import get_all_users -from onyx.db.users import get_page_of_filtered_users -from onyx.db.users import get_total_filtered_users_count -from onyx.db.users import get_user_by_email -from onyx.db.users import get_user_counts_by_role_and_status -from onyx.db.users import validate_user_role_update +from onyx.db.engine.sql_engine import get_session, get_session_with_shared_schema +from onyx.db.enums import AccountType, Permission, UserFileStatus +from onyx.db.models import User, UserFile +from onyx.db.tenant_invite_counter import release_trial_invites, reserve_trial_invites +from onyx.db.user_preferences import ( + activate_user, + deactivate_user, + get_all_user_assistant_specific_configs, + get_latest_access_token_for_user, + get_memories_for_user, + update_assistant_preferences, + update_user_assistant_visibility, + update_user_auto_scroll, + update_user_chat_background, + update_user_default_app_mode, + update_user_default_model, + update_user_language, + update_user_paste_as_tile, + update_user_personalization, + update_user_pinned_assistants, + update_user_role, + update_user_shortcut_enabled, + update_user_temperature_override_enabled, + update_user_theme_preference, + update_users_craft_enabled, +) +from onyx.db.users import ( + batch_get_user_groups, + delete_user_from_db, + get_all_accepted_users, + get_all_users, + get_page_of_filtered_users, + get_total_filtered_users_count, + get_user_by_email, + get_user_counts_by_role_and_status, + validate_user_role_update, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.key_value_store.factory import get_kv_store -from onyx.redis.redis_pool import get_raw_redis_client -from onyx.redis.redis_pool import get_redis_client +from onyx.redis.redis_pool import get_raw_redis_client, get_redis_client from onyx.server.documents.models import PaginatedReturn from onyx.server.features.projects.models import UserFileSnapshot -from onyx.server.manage.invite_rate_limit import enforce_invite_rate_limit -from onyx.server.manage.invite_rate_limit import enforce_remove_invited_rate_limit -from onyx.server.manage.models import AllUsersResponse -from onyx.server.manage.models import AutoScrollRequest -from onyx.server.manage.models import BulkInviteResponse -from onyx.server.manage.models import ChatBackgroundRequest -from onyx.server.manage.models import DefaultAppModeRequest -from onyx.server.manage.models import EmailInviteStatus -from onyx.server.manage.models import LanguageRequest -from onyx.server.manage.models import MemoryItem -from onyx.server.manage.models import PersonalizationUpdateRequest -from onyx.server.manage.models import TenantInfo -from onyx.server.manage.models import TenantSnapshot -from onyx.server.manage.models import ThemePreferenceRequest -from onyx.server.manage.models import UserByEmail -from onyx.server.manage.models import UserCraftAccessUpdateRequest -from onyx.server.manage.models import UserInfo -from onyx.server.manage.models import UserPreferences -from onyx.server.manage.models import UserRoleResponse -from onyx.server.manage.models import UserRoleUpdateRequest -from onyx.server.manage.models import UserSpecificAssistantPreference -from onyx.server.manage.models import UserSpecificAssistantPreferences -from onyx.server.models import FullUserSnapshot -from onyx.server.models import InvitedUserSnapshot -from onyx.server.models import MinimalUserSnapshot -from onyx.server.models import UserGroupInfo +from onyx.server.manage.invite_rate_limit import ( + enforce_invite_rate_limit, + enforce_remove_invited_rate_limit, +) +from onyx.server.manage.models import ( + AllUsersResponse, + AutoScrollRequest, + BulkInviteResponse, + ChatBackgroundRequest, + DefaultAppModeRequest, + EmailInviteStatus, + LanguageRequest, + MemoryItem, + PersonalizationUpdateRequest, + TenantInfo, + TenantSnapshot, + ThemePreferenceRequest, + UserByEmail, + UserCraftAccessUpdateRequest, + UserInfo, + UserPreferences, + UserRoleResponse, + UserRoleUpdateRequest, + UserSpecificAssistantPreference, + UserSpecificAssistantPreferences, +) +from onyx.server.models import ( + FullUserSnapshot, + InvitedUserSnapshot, + MinimalUserSnapshot, + UserGroupInfo, +) from onyx.server.security.store import get_security_settings from onyx.server.usage_limits import is_tenant_on_trial_fn from onyx.server.utils import BasicAuthenticationError -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.csv_utils import sanitize_csv_cell from onyx.utils.logger import setup_logger -from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop from onyx.utils.variable_functionality import ( + fetch_ee_implementation_or_noop, fetch_versioned_implementation_with_fallback, ) from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/manage/voice/api.py b/backend/onyx/server/manage/voice/api.py index 717ec9f9e8..a0b065f389 100644 --- a/backend/onyx/server/manage/voice/api.py +++ b/backend/onyx/server/manage/voice/api.py @@ -1,34 +1,34 @@ -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Response +from fastapi import APIRouter, Depends, Response from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import LLMProvider as LLMProviderModel -from onyx.db.models import User -from onyx.db.models import VoiceProvider -from onyx.db.voice import deactivate_stt_provider -from onyx.db.voice import deactivate_tts_provider -from onyx.db.voice import delete_voice_provider -from onyx.db.voice import fetch_voice_provider_by_id -from onyx.db.voice import fetch_voice_provider_by_type -from onyx.db.voice import fetch_voice_providers -from onyx.db.voice import set_default_stt_provider -from onyx.db.voice import set_default_tts_provider -from onyx.db.voice import upsert_voice_provider +from onyx.db.models import User, VoiceProvider +from onyx.db.voice import ( + deactivate_stt_provider, + deactivate_tts_provider, + delete_voice_provider, + fetch_voice_provider_by_id, + fetch_voice_provider_by_type, + fetch_voice_providers, + set_default_stt_provider, + set_default_tts_provider, + upsert_voice_provider, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.manage.voice.models import VoiceOption -from onyx.server.manage.voice.models import VoiceProviderTestRequest -from onyx.server.manage.voice.models import VoiceProviderUpdateSuccess -from onyx.server.manage.voice.models import VoiceProviderUpsertRequest -from onyx.server.manage.voice.models import VoiceProviderView +from onyx.server.manage.voice.models import ( + VoiceOption, + VoiceProviderTestRequest, + VoiceProviderUpdateSuccess, + VoiceProviderUpsertRequest, + VoiceProviderView, +) from onyx.utils.encryption import mask_string from onyx.utils.logger import setup_logger -from onyx.utils.url import SSRFException -from onyx.utils.url import validate_outbound_http_url +from onyx.utils.url import SSRFException, validate_outbound_http_url from onyx.voice.factory import get_voice_provider logger = setup_logger() diff --git a/backend/onyx/server/manage/voice/models.py b/backend/onyx/server/manage/voice/models.py index 9c1682ac22..d49227aef0 100644 --- a/backend/onyx/server/manage/voice/models.py +++ b/backend/onyx/server/manage/voice/models.py @@ -1,7 +1,6 @@ from typing import Any -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field class VoiceProviderView(BaseModel): diff --git a/backend/onyx/server/manage/voice/user_api.py b/backend/onyx/server/manage/voice/user_api.py index 96f85aaccc..c6fd8b39e6 100644 --- a/backend/onyx/server/manage/voice/user_api.py +++ b/backend/onyx/server/manage/voice/user_api.py @@ -2,27 +2,23 @@ import json import secrets from collections.abc import AsyncIterator -from fastapi import APIRouter -from fastapi import Depends -from fastapi import File -from fastapi import UploadFile +from fastapi import APIRouter, Depends, File, UploadFile from fastapi.responses import StreamingResponse -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.voice import fetch_default_stt_provider -from onyx.db.voice import fetch_default_tts_provider -from onyx.db.voice import update_user_voice_settings +from onyx.db.voice import ( + fetch_default_stt_provider, + fetch_default_tts_provider, + update_user_voice_settings, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.redis.redis_pool import store_ws_token -from onyx.redis.redis_pool import WsTokenRateLimitExceeded +from onyx.redis.redis_pool import store_ws_token, WsTokenRateLimitExceeded from onyx.server.manage.models import VoiceSettingsUpdateRequest from onyx.server.manage.voice.text_utils import strip_markdown_for_tts from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/manage/voice/websocket_api.py b/backend/onyx/server/manage/voice/websocket_api.py index 7652f1ca3f..a8dfb1e059 100644 --- a/backend/onyx/server/manage/voice/websocket_api.py +++ b/backend/onyx/server/manage/voice/websocket_api.py @@ -8,23 +8,21 @@ from collections.abc import MutableMapping from typing import Any import numpy as np -from fastapi import APIRouter -from fastapi import Depends -from fastapi import WebSocket -from fastapi import WebSocketDisconnect +from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect from sqlalchemy.orm import Session from onyx.auth.users import current_user_from_websocket from onyx.db.engine.sql_engine import get_sqlalchemy_engine from onyx.db.models import User -from onyx.db.voice import fetch_default_stt_provider -from onyx.db.voice import fetch_default_tts_provider +from onyx.db.voice import fetch_default_stt_provider, fetch_default_tts_provider from onyx.server.manage.voice.text_utils import strip_markdown_for_tts from onyx.utils.logger import setup_logger from onyx.voice.factory import get_voice_provider -from onyx.voice.interface import StreamingSynthesizerProtocol -from onyx.voice.interface import StreamingTranscriberProtocol -from onyx.voice.interface import TranscriptResult +from onyx.voice.interface import ( + StreamingSynthesizerProtocol, + StreamingTranscriberProtocol, + TranscriptResult, +) logger = setup_logger() diff --git a/backend/onyx/server/manage/web_search/api.py b/backend/onyx/server/manage/web_search/api.py index 693be38d0c..5e6814e709 100644 --- a/backend/onyx/server/manage/web_search/api.py +++ b/backend/onyx/server/manage/web_search/api.py @@ -1,55 +1,49 @@ from __future__ import annotations -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Response +from fastapi import APIRouter, Depends, HTTPException, Response from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission -from onyx.db.models import InternetContentProvider -from onyx.db.models import InternetSearchProvider -from onyx.db.models import User -from onyx.db.web_search import deactivate_web_content_provider -from onyx.db.web_search import deactivate_web_search_provider -from onyx.db.web_search import delete_web_content_provider -from onyx.db.web_search import delete_web_search_provider -from onyx.db.web_search import fetch_web_content_provider_by_name -from onyx.db.web_search import fetch_web_content_provider_by_type -from onyx.db.web_search import fetch_web_content_providers -from onyx.db.web_search import fetch_web_search_provider_by_name -from onyx.db.web_search import fetch_web_search_provider_by_type -from onyx.db.web_search import fetch_web_search_providers -from onyx.db.web_search import set_active_web_content_provider -from onyx.db.web_search import set_active_web_search_provider -from onyx.db.web_search import upsert_web_content_provider -from onyx.db.web_search import upsert_web_search_provider -from onyx.server.manage.web_search.models import WebContentProviderTestRequest -from onyx.server.manage.web_search.models import WebContentProviderUpsertRequest -from onyx.server.manage.web_search.models import WebContentProviderView -from onyx.server.manage.web_search.models import WebSearchProviderTestRequest -from onyx.server.manage.web_search.models import WebSearchProviderUpsertRequest -from onyx.server.manage.web_search.models import WebSearchProviderView +from onyx.db.models import InternetContentProvider, InternetSearchProvider, User +from onyx.db.web_search import ( + deactivate_web_content_provider, + deactivate_web_search_provider, + delete_web_content_provider, + delete_web_search_provider, + fetch_web_content_provider_by_name, + fetch_web_content_provider_by_type, + fetch_web_content_providers, + fetch_web_search_provider_by_name, + fetch_web_search_provider_by_type, + fetch_web_search_providers, + set_active_web_content_provider, + set_active_web_search_provider, + upsert_web_content_provider, + upsert_web_search_provider, +) +from onyx.server.manage.web_search.models import ( + WebContentProviderTestRequest, + WebContentProviderUpsertRequest, + WebContentProviderView, + WebSearchProviderTestRequest, + WebSearchProviderUpsertRequest, + WebSearchProviderView, +) from onyx.tools.tool_implementations.open_url.utils import ( filter_web_contents_with_no_title_or_content, ) from onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig from onyx.tools.tool_implementations.web_search.providers import ( build_content_provider_from_config, -) -from onyx.tools.tool_implementations.web_search.providers import ( build_search_provider_from_config, -) -from onyx.tools.tool_implementations.web_search.providers import ( provider_requires_api_key, ) from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT -from shared_configs.enums import WebContentProviderType -from shared_configs.enums import WebSearchProviderType +from shared_configs.enums import WebContentProviderType, WebSearchProviderType logger = setup_logger() diff --git a/backend/onyx/server/manage/web_search/models.py b/backend/onyx/server/manage/web_search/models.py index 1fa6968fe9..fb79913e7d 100644 --- a/backend/onyx/server/manage/web_search/models.py +++ b/backend/onyx/server/manage/web_search/models.py @@ -1,11 +1,9 @@ from typing import Any -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig -from shared_configs.enums import WebContentProviderType -from shared_configs.enums import WebSearchProviderType +from shared_configs.enums import WebContentProviderType, WebSearchProviderType class WebSearchProviderView(BaseModel): diff --git a/backend/onyx/server/metrics/celery_task_metrics.py b/backend/onyx/server/metrics/celery_task_metrics.py index e6caf6eb07..6c13df3369 100644 --- a/backend/onyx/server/metrics/celery_task_metrics.py +++ b/backend/onyx/server/metrics/celery_task_metrics.py @@ -21,9 +21,7 @@ import threading import time from celery import Task -from prometheus_client import Counter -from prometheus_client import Gauge -from prometheus_client import Histogram +from prometheus_client import Counter, Gauge, Histogram from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/metrics/connector_health_metrics.py b/backend/onyx/server/metrics/connector_health_metrics.py index 463d3ab34b..ca4e23b516 100644 --- a/backend/onyx/server/metrics/connector_health_metrics.py +++ b/backend/onyx/server/metrics/connector_health_metrics.py @@ -15,8 +15,7 @@ of connectors (one series per connector), not by the number of task executions. """ -from prometheus_client import Counter -from prometheus_client import Gauge +from prometheus_client import Counter, Gauge from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/metrics/connector_state_metrics.py b/backend/onyx/server/metrics/connector_state_metrics.py index 5e2a5f54cc..8d9d032a56 100644 --- a/backend/onyx/server/metrics/connector_state_metrics.py +++ b/backend/onyx/server/metrics/connector_state_metrics.py @@ -6,17 +6,12 @@ from collections.abc import Iterator from datetime import datetime from typing import NamedTuple -from prometheus_client.core import GaugeMetricFamily -from prometheus_client.core import InfoMetricFamily -from prometheus_client.core import Metric -from prometheus_client.core import REGISTRY +from prometheus_client.core import GaugeMetricFamily, InfoMetricFamily, Metric, REGISTRY from prometheus_client.registry import Collector from onyx.db.connector_credential_pair import get_connector_state_snapshots from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingMode +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, IndexingMode from onyx.utils.datetime import datetime_to_utc from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/metrics/deletion_metrics.py b/backend/onyx/server/metrics/deletion_metrics.py index 1bda4c5a40..45ce0ffb11 100644 --- a/backend/onyx/server/metrics/deletion_metrics.py +++ b/backend/onyx/server/metrics/deletion_metrics.py @@ -24,8 +24,7 @@ Usage: ) """ -from prometheus_client import Counter -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/metrics/embedding.py b/backend/onyx/server/metrics/embedding.py index eeb321331f..20c4bd6433 100644 --- a/backend/onyx/server/metrics/embedding.py +++ b/backend/onyx/server/metrics/embedding.py @@ -9,12 +9,9 @@ from collections.abc import Generator from contextlib import contextmanager from enum import Enum -from prometheus_client import Counter -from prometheus_client import Gauge -from prometheus_client import Histogram +from prometheus_client import Counter, Gauge, Histogram -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType +from shared_configs.enums import EmbeddingProvider, EmbedTextType class QueryEmbeddingCacheLookupOutcome(str, Enum): diff --git a/backend/onyx/server/metrics/image_processing.py b/backend/onyx/server/metrics/image_processing.py index 2fcb822802..00881649f1 100644 --- a/backend/onyx/server/metrics/image_processing.py +++ b/backend/onyx/server/metrics/image_processing.py @@ -5,11 +5,9 @@ import io import logging import time from collections.abc import Callable -from typing import ParamSpec -from typing import TypeVar +from typing import ParamSpec, TypeVar -from prometheus_client import Counter -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram logger = logging.getLogger(__name__) diff --git a/backend/onyx/server/metrics/indexing_pipeline.py b/backend/onyx/server/metrics/indexing_pipeline.py index 2810054f7b..84f4443124 100644 --- a/backend/onyx/server/metrics/indexing_pipeline.py +++ b/backend/onyx/server/metrics/indexing_pipeline.py @@ -24,9 +24,11 @@ from prometheus_client.core import GaugeMetricFamily from prometheus_client.registry import Collector from redis import Redis -from onyx.background.celery.celery_redis import celery_get_broker_client -from onyx.background.celery.celery_redis import celery_get_queue_length -from onyx.background.celery.celery_redis import celery_get_unacked_task_ids +from onyx.background.celery.celery_redis import ( + celery_get_broker_client, + celery_get_queue_length, + celery_get_unacked_task_ids, +) from onyx.configs.constants import OnyxCeleryQueues from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/metrics/indexing_pipeline_setup.py b/backend/onyx/server/metrics/indexing_pipeline_setup.py index b4bd152ea8..df2fc177ca 100644 --- a/backend/onyx/server/metrics/indexing_pipeline_setup.py +++ b/backend/onyx/server/metrics/indexing_pipeline_setup.py @@ -6,10 +6,12 @@ Called once by the monitoring celery worker after Redis and DB are ready. from celery import Celery from prometheus_client.registry import REGISTRY -from onyx.server.metrics.indexing_pipeline import QueueDepthCollector -from onyx.server.metrics.indexing_pipeline import RedisHealthCollector -from onyx.server.metrics.indexing_pipeline import WorkerHealthCollector -from onyx.server.metrics.indexing_pipeline import WorkerHeartbeatMonitor +from onyx.server.metrics.indexing_pipeline import ( + QueueDepthCollector, + RedisHealthCollector, + WorkerHealthCollector, + WorkerHeartbeatMonitor, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/metrics/indexing_task_metrics.py b/backend/onyx/server/metrics/indexing_task_metrics.py index 1af3a0f770..9bb3d3f911 100644 --- a/backend/onyx/server/metrics/indexing_task_metrics.py +++ b/backend/onyx/server/metrics/indexing_task_metrics.py @@ -26,8 +26,7 @@ import time from dataclasses import dataclass from celery import Task -from prometheus_client import Counter -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram from onyx.configs.constants import OnyxCeleryTask from onyx.server.metrics.celery_task_metrics import _MAX_START_TIME_AGE_SECONDS diff --git a/backend/onyx/server/metrics/mcp_client.py b/backend/onyx/server/metrics/mcp_client.py index c8e49a6951..ef17690ee5 100644 --- a/backend/onyx/server/metrics/mcp_client.py +++ b/backend/onyx/server/metrics/mcp_client.py @@ -1,7 +1,6 @@ import time -from prometheus_client import Counter -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram from onyx.server.metrics.mcp_common import MCPToolCallStatus from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/metrics/mcp_server.py b/backend/onyx/server/metrics/mcp_server.py index 6cb253d909..b4c63e6fcc 100644 --- a/backend/onyx/server/metrics/mcp_server.py +++ b/backend/onyx/server/metrics/mcp_server.py @@ -1,8 +1,7 @@ import time from enum import Enum -from prometheus_client import Counter -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram from onyx.server.metrics.mcp_common import MCPToolCallStatus from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/metrics/metrics_auth.py b/backend/onyx/server/metrics/metrics_auth.py index 8263f5714a..bfe699f9af 100644 --- a/backend/onyx/server/metrics/metrics_auth.py +++ b/backend/onyx/server/metrics/metrics_auth.py @@ -14,8 +14,7 @@ import secrets from fastapi import Request from onyx.auth.constants import BEARER_PREFIX -from onyx.configs.app_configs import DISABLE_METRICS_AUTH -from onyx.configs.app_configs import METRICS_AUTH_TOKEN +from onyx.configs.app_configs import DISABLE_METRICS_AUTH, METRICS_AUTH_TOKEN from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/onyx/server/metrics/opensearch_search.py b/backend/onyx/server/metrics/opensearch_search.py index 5bf0dd7d5a..30b215d8fb 100644 --- a/backend/onyx/server/metrics/opensearch_search.py +++ b/backend/onyx/server/metrics/opensearch_search.py @@ -14,9 +14,7 @@ import logging from collections.abc import Generator from contextlib import contextmanager -from prometheus_client import Counter -from prometheus_client import Gauge -from prometheus_client import Histogram +from prometheus_client import Counter, Gauge, Histogram from onyx.document_index.opensearch.constants import OpenSearchSearchType diff --git a/backend/onyx/server/metrics/perm_sync_metrics.py b/backend/onyx/server/metrics/perm_sync_metrics.py index b33d9ca194..0bbb19634f 100644 --- a/backend/onyx/server/metrics/perm_sync_metrics.py +++ b/backend/onyx/server/metrics/perm_sync_metrics.py @@ -33,8 +33,7 @@ Usage: ) """ -from prometheus_client import Counter -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/metrics/postgres_connection_pool.py b/backend/onyx/server/metrics/postgres_connection_pool.py index 47b44f2257..95ec81f426 100644 --- a/backend/onyx/server/metrics/postgres_connection_pool.py +++ b/backend/onyx/server/metrics/postgres_connection_pool.py @@ -16,23 +16,20 @@ import time from fastapi import Request from fastapi.responses import JSONResponse -from prometheus_client import Counter -from prometheus_client import Gauge -from prometheus_client import Histogram +from prometheus_client import Counter, Gauge, Histogram from prometheus_client.core import GaugeMetricFamily -from prometheus_client.registry import Collector -from prometheus_client.registry import REGISTRY +from prometheus_client.registry import Collector, REGISTRY from sqlalchemy import event from sqlalchemy.engine import Engine from sqlalchemy.engine.interfaces import DBAPIConnection from sqlalchemy.ext.asyncio import AsyncEngine -from sqlalchemy.pool import ConnectionPoolEntry -from sqlalchemy.pool import PoolProxiedConnection -from sqlalchemy.pool import QueuePool +from sqlalchemy.pool import ConnectionPoolEntry, PoolProxiedConnection, QueuePool from onyx.utils.logger import setup_logger -from shared_configs.contextvars import CURRENT_ENDPOINT_CONTEXTVAR -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR +from shared_configs.contextvars import ( + CURRENT_ENDPOINT_CONTEXTVAR, + CURRENT_TENANT_ID_CONTEXTVAR, +) logger = setup_logger() diff --git a/backend/onyx/server/metrics/pruning_metrics.py b/backend/onyx/server/metrics/pruning_metrics.py index 949f074915..135a2abb30 100644 --- a/backend/onyx/server/metrics/pruning_metrics.py +++ b/backend/onyx/server/metrics/pruning_metrics.py @@ -17,8 +17,7 @@ Usage: ) """ -from prometheus_client import Counter -from prometheus_client import Histogram +from prometheus_client import Counter, Histogram from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/middleware/latency_logging.py b/backend/onyx/server/middleware/latency_logging.py index b061a370c8..2bba484441 100644 --- a/backend/onyx/server/middleware/latency_logging.py +++ b/backend/onyx/server/middleware/latency_logging.py @@ -1,11 +1,8 @@ import logging import time -from collections.abc import Awaitable -from collections.abc import Callable +from collections.abc import Awaitable, Callable -from fastapi import FastAPI -from fastapi import Request -from fastapi import Response +from fastapi import FastAPI, Request, Response def add_latency_logging_middleware(app: FastAPI, logger: logging.LoggerAdapter) -> None: diff --git a/backend/onyx/server/middleware/rate_limiting.py b/backend/onyx/server/middleware/rate_limiting.py index 5bc0e9847d..c81abbb5f6 100644 --- a/backend/onyx/server/middleware/rate_limiting.py +++ b/backend/onyx/server/middleware/rate_limiting.py @@ -1,17 +1,16 @@ -from fastapi import Depends -from fastapi import params -from fastapi import Request -from fastapi import Response +from fastapi import Depends, params, Request, Response from fastapi_limiter import FastAPILimiter from fastapi_limiter.depends import RateLimiter from onyx.auth.users import current_chat_accessible_user -from onyx.configs.app_configs import AUTH_RATE_LIMITING_ENABLED -from onyx.configs.app_configs import FEEDBACK_RATE_LIMIT_MAX_REQUESTS -from onyx.configs.app_configs import FEEDBACK_RATE_LIMIT_WINDOW_SECONDS -from onyx.configs.app_configs import FEEDBACK_RATE_LIMITING_ENABLED -from onyx.configs.app_configs import RATE_LIMIT_MAX_REQUESTS -from onyx.configs.app_configs import RATE_LIMIT_WINDOW_SECONDS +from onyx.configs.app_configs import ( + AUTH_RATE_LIMITING_ENABLED, + FEEDBACK_RATE_LIMIT_MAX_REQUESTS, + FEEDBACK_RATE_LIMIT_WINDOW_SECONDS, + FEEDBACK_RATE_LIMITING_ENABLED, + RATE_LIMIT_MAX_REQUESTS, + RATE_LIMIT_WINDOW_SECONDS, +) from onyx.db.enums import AccountType from onyx.db.models import User from onyx.redis.redis_pool import get_async_redis_connection diff --git a/backend/onyx/server/models.py b/backend/onyx/server/models.py index 2cb02b3fde..40dfd541c9 100644 --- a/backend/onyx/server/models.py +++ b/backend/onyx/server/models.py @@ -1,7 +1,5 @@ import datetime -from typing import Generic -from typing import Optional -from typing import TypeVar +from typing import Generic, Optional, TypeVar from uuid import UUID from pydantic import BaseModel diff --git a/backend/onyx/server/oidc_multi.py b/backend/onyx/server/oidc_multi.py index 77ececd873..ae063c7eae 100644 --- a/backend/onyx/server/oidc_multi.py +++ b/backend/onyx/server/oidc_multi.py @@ -11,46 +11,47 @@ import uuid from typing import Any from cachetools import TTLCache -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Request -from fastapi import Response +from fastapi import APIRouter, Depends, Request, Response from fastapi.concurrency import run_in_threadpool from fastapi.responses import JSONResponse from fastapi_users.authentication import Strategy from httpx_oauth.clients.google import GoogleOAuth2 from httpx_oauth.clients.openid import BASE_SCOPES -from httpx_oauth.oauth2 import BaseOAuth2 -from httpx_oauth.oauth2 import GetAccessTokenError +from httpx_oauth.oauth2 import BaseOAuth2, GetAccessTokenError from sqlalchemy.orm import Session from onyx.auth.oidc_client import VerifiedEmailOpenID -from onyx.auth.users import auth_backend -from onyx.auth.users import complete_login_flow -from onyx.auth.users import CSRF_TOKEN_COOKIE_NAME -from onyx.auth.users import CSRF_TOKEN_KEY -from onyx.auth.users import decode_and_validate_oauth_state -from onyx.auth.users import generate_csrf_token -from onyx.auth.users import generate_pkce_pair -from onyx.auth.users import generate_state_token -from onyx.auth.users import get_pkce_cookie_name -from onyx.auth.users import get_user_manager -from onyx.auth.users import OAuth2AuthorizeResponse -from onyx.auth.users import STATE_TOKEN_LIFETIME_SECONDS -from onyx.auth.users import UserManager -from onyx.configs.app_configs import GOOGLE_LOGIN_BASE_SCOPES -from onyx.configs.app_configs import GOOGLE_OAUTH_SCOPE_OVERRIDE -from onyx.configs.app_configs import OIDC_PKCE_ENABLED -from onyx.configs.app_configs import OIDC_SCOPE_OVERRIDE -from onyx.configs.app_configs import USER_AUTH_SECRET -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.auth.users import ( + auth_backend, + complete_login_flow, + CSRF_TOKEN_COOKIE_NAME, + CSRF_TOKEN_KEY, + decode_and_validate_oauth_state, + generate_csrf_token, + generate_pkce_pair, + generate_state_token, + get_pkce_cookie_name, + get_user_manager, + OAuth2AuthorizeResponse, + STATE_TOKEN_LIFETIME_SECONDS, + UserManager, +) +from onyx.configs.app_configs import ( + GOOGLE_LOGIN_BASE_SCOPES, + GOOGLE_OAUTH_SCOPE_OVERRIDE, + OIDC_PKCE_ENABLED, + OIDC_SCOPE_OVERRIDE, + USER_AUTH_SECRET, + WEB_DOMAIN, +) from onyx.db.engine.sql_engine import get_session from onyx.db.enums import SSOProviderType -from onyx.db.models import SSOProvider -from onyx.db.models import User -from onyx.db.sso_provider import fetch_sso_provider_by_name -from onyx.db.sso_provider import sso_login_callback_uri -from onyx.db.sso_provider import validate_sso_config +from onyx.db.models import SSOProvider, User +from onyx.db.sso_provider import ( + fetch_sso_provider_by_name, + sso_login_callback_uri, + validate_sso_config, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.utils.url import sanitize_next_url diff --git a/backend/onyx/server/onyx_api/ingestion.py b/backend/onyx/server/onyx_api/ingestion.py index 16cb47e3fb..c95a09b5d8 100644 --- a/backend/onyx/server/onyx_api/ingestion.py +++ b/backend/onyx/server/onyx_api/ingestion.py @@ -1,35 +1,36 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from onyx.auth.users import current_curator_or_admin_user -from onyx.configs.constants import DEFAULT_CC_PAIR_ID -from onyx.configs.constants import PUBLIC_API_TAGS -from onyx.connectors.models import Document -from onyx.connectors.models import IndexAttemptMetadata +from onyx.configs.constants import DEFAULT_CC_PAIR_ID, PUBLIC_API_TAGS +from onyx.connectors.models import Document, IndexAttemptMetadata from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.document import delete_documents_complete -from onyx.db.document import get_document -from onyx.db.document import get_documents_by_cc_pair -from onyx.db.document import get_ingestion_documents +from onyx.db.document import ( + delete_documents_complete, + get_document, + get_documents_by_cc_pair, + get_ingestion_documents, +) from onyx.db.engine.sql_engine import get_session from onyx.db.models import User -from onyx.db.search_settings import get_active_search_settings -from onyx.db.search_settings import get_current_search_settings -from onyx.db.search_settings import get_secondary_search_settings +from onyx.db.search_settings import ( + get_active_search_settings, + get_current_search_settings, + get_secondary_search_settings, +) 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 -from onyx.server.onyx_api.models import DocMinimalInfo -from onyx.server.onyx_api.models import IngestionDocument -from onyx.server.onyx_api.models import IngestionResult +from onyx.server.onyx_api.models import ( + DocMinimalInfo, + IngestionDocument, + IngestionResult, +) from onyx.server.utils_vector_db import require_vector_db from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/pat/api.py b/backend/onyx/server/pat/api.py index 8555989c3a..4419cbacaa 100644 --- a/backend/onyx/server/pat/api.py +++ b/backend/onyx/server/pat/api.py @@ -1,24 +1,22 @@ """API endpoints for Personal Access Tokens.""" -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.db.engine.sql_engine import get_session -from onyx.db.enums import PatType -from onyx.db.enums import Permission +from onyx.db.enums import PatType, Permission from onyx.db.models import User -from onyx.db.pat import create_pat -from onyx.db.pat import list_user_pats -from onyx.db.pat import revoke_pat +from onyx.db.pat import create_pat, list_user_pats, revoke_pat from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.pat.models import CreatedTokenResponse -from onyx.server.pat.models import CreateTokenRequest -from onyx.server.pat.models import PatScopeOption -from onyx.server.pat.models import SELECTABLE_PAT_SCOPES -from onyx.server.pat.models import TokenResponse +from onyx.server.pat.models import ( + CreatedTokenResponse, + CreateTokenRequest, + PatScopeOption, + SELECTABLE_PAT_SCOPES, + TokenResponse, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/pat/models.py b/backend/onyx/server/pat/models.py index 3a65650e2b..0e80b1fc78 100644 --- a/backend/onyx/server/pat/models.py +++ b/backend/onyx/server/pat/models.py @@ -2,11 +2,7 @@ from datetime import datetime -from pydantic import BaseModel -from pydantic import computed_field -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_validator +from pydantic import BaseModel, computed_field, ConfigDict, Field, field_validator from onyx.auth.permissions import resolve_effective_permissions from onyx.db.enums import Permission diff --git a/backend/onyx/server/query_and_chat/chat_backend.py b/backend/onyx/server/query_and_chat/chat_backend.py index ec6376d924..70fd9464ed 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -5,14 +5,8 @@ from collections.abc import Generator from datetime import timedelta from uuid import UUID -from fastapi import APIRouter -from fastapi import Depends -from fastapi import HTTPException -from fastapi import Query -from fastapi import Request -from fastapi import Response -from fastapi.responses import JSONResponse -from fastapi.responses import StreamingResponse +from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response +from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel from sqlalchemy.orm import Session @@ -22,94 +16,98 @@ from onyx.auth.pat import get_hashed_pat_from_request from onyx.auth.permissions import require_permission from onyx.auth.users import current_chat_accessible_user from onyx.cache.factory import get_cache_backend -from onyx.chat.chat_processing_checker import get_processing_run_id -from onyx.chat.chat_processing_checker import is_chat_session_processing +from onyx.chat.chat_processing_checker import ( + get_processing_run_id, + is_chat_session_processing, +) from onyx.chat.chat_state import ChatStateContainer -from onyx.chat.chat_utils import convert_chat_history_basic -from onyx.chat.chat_utils import create_chat_history_chain -from onyx.chat.chat_utils import create_chat_session_from_request -from onyx.chat.chat_utils import extract_headers -from onyx.chat.models import ChatFullResponse -from onyx.chat.models import CreateChatSessionID -from onyx.chat.process_message import gather_stream_full -from onyx.chat.process_message import handle_multi_model_stream -from onyx.chat.process_message import handle_stream_message_objects +from onyx.chat.chat_utils import ( + convert_chat_history_basic, + create_chat_history_chain, + create_chat_session_from_request, + extract_headers, +) +from onyx.chat.models import ChatFullResponse, CreateChatSessionID +from onyx.chat.process_message import ( + gather_stream_full, + handle_multi_model_stream, + handle_stream_message_objects, +) from onyx.chat.prompt_utils import get_default_base_system_prompt from onyx.chat.stop_signal_checker import set_fence -from onyx.chat.stream_buffer import has_stream_buffer -from onyx.chat.stream_buffer import read_stream_chunks +from onyx.chat.stream_buffer import has_stream_buffer, read_stream_chunks from onyx.configs.app_configs import WEB_DOMAIN -from onyx.configs.chat_configs import CHAT_HEARTBEAT_INTERVAL_S -from onyx.configs.chat_configs import CHAT_RESUME_POLL_INTERVAL_S -from onyx.configs.chat_configs import HARD_DELETE_CHATS -from onyx.configs.constants import MessageType -from onyx.configs.constants import MilestoneRecordType -from onyx.configs.constants import PUBLIC_API_TAGS +from onyx.configs.chat_configs import ( + CHAT_HEARTBEAT_INTERVAL_S, + CHAT_RESUME_POLL_INTERVAL_S, + HARD_DELETE_CHATS, +) +from onyx.configs.constants import MessageType, MilestoneRecordType, PUBLIC_API_TAGS from onyx.configs.model_configs import LITELLM_PASS_THROUGH_HEADERS -from onyx.db.chat import add_chats_to_session_from_slack_thread -from onyx.db.chat import delete_all_chat_sessions_for_user -from onyx.db.chat import delete_chat_session -from onyx.db.chat import duplicate_chat_session_for_user_from_slack -from onyx.db.chat import get_chat_message -from onyx.db.chat import get_chat_messages_by_session -from onyx.db.chat import get_chat_session_by_id -from onyx.db.chat import get_chat_sessions_by_user -from onyx.db.chat import set_as_latest_chat_message -from onyx.db.chat import set_preferred_response -from onyx.db.chat import translate_db_message_to_chat_message_detail -from onyx.db.chat import update_chat_session +from onyx.db.chat import ( + add_chats_to_session_from_slack_thread, + delete_all_chat_sessions_for_user, + delete_chat_session, + duplicate_chat_session_for_user_from_slack, + get_chat_message, + get_chat_messages_by_session, + get_chat_session_by_id, + get_chat_sessions_by_user, + set_as_latest_chat_message, + set_preferred_response, + translate_db_message_to_chat_message_detail, + update_chat_session, +) from onyx.db.chat_search import search_chat_sessions -from onyx.db.engine.sql_engine import get_session -from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant from onyx.db.enums import Permission -from onyx.db.feedback import create_chat_message_feedback -from onyx.db.feedback import remove_chat_message_feedback -from onyx.db.models import ChatSessionSharedStatus -from onyx.db.models import Persona -from onyx.db.models import User +from onyx.db.feedback import create_chat_message_feedback, remove_chat_message_feedback +from onyx.db.models import ChatSessionSharedStatus, Persona, User from onyx.db.persona import get_persona_by_id -from onyx.db.usage import increment_usage -from onyx.db.usage import UsageType +from onyx.db.usage import increment_usage, UsageType from onyx.db.user_file import get_file_id_by_user_file_id from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import get_default_file_store from onyx.llm.constants import LlmProviderNames -from onyx.llm.factory import get_default_llm -from onyx.llm.factory import get_llm_for_persona -from onyx.llm.factory import get_llm_token_counter +from onyx.llm.factory import get_default_llm, get_llm_for_persona, get_llm_token_counter from onyx.secondary_llm_flows.chat_session_naming import generate_chat_session_name from onyx.server.api_key_usage import check_api_key_usage from onyx.server.middleware.rate_limiting import get_feedback_rate_limiters -from onyx.server.query_and_chat.chat_utils import is_spreadsheet_mime_type -from onyx.server.query_and_chat.chat_utils import parse_spreadsheet_for_preview -from onyx.server.query_and_chat.models import ChatFeedbackRequest -from onyx.server.query_and_chat.models import ChatMessageIdentifier -from onyx.server.query_and_chat.models import ChatRenameRequest -from onyx.server.query_and_chat.models import ChatSearchResponse -from onyx.server.query_and_chat.models import ChatSessionCreationRequest -from onyx.server.query_and_chat.models import ChatSessionDetailResponse -from onyx.server.query_and_chat.models import ChatSessionDetails -from onyx.server.query_and_chat.models import ChatSessionGroup -from onyx.server.query_and_chat.models import ChatSessionsResponse -from onyx.server.query_and_chat.models import ChatSessionSummary -from onyx.server.query_and_chat.models import ChatSessionUpdateRequest -from onyx.server.query_and_chat.models import CurrentRunInfo -from onyx.server.query_and_chat.models import MessageOrigin -from onyx.server.query_and_chat.models import RenameChatSessionResponse -from onyx.server.query_and_chat.models import SendMessageRequest -from onyx.server.query_and_chat.models import SetPreferredResponseRequest -from onyx.server.query_and_chat.models import UpdateChatSessionTemperatureRequest -from onyx.server.query_and_chat.models import UpdateChatSessionThreadRequest +from onyx.server.query_and_chat.chat_utils import ( + is_spreadsheet_mime_type, + parse_spreadsheet_for_preview, +) +from onyx.server.query_and_chat.models import ( + ChatFeedbackRequest, + ChatMessageIdentifier, + ChatRenameRequest, + ChatSearchResponse, + ChatSessionCreationRequest, + ChatSessionDetailResponse, + ChatSessionDetails, + ChatSessionGroup, + ChatSessionsResponse, + ChatSessionSummary, + ChatSessionUpdateRequest, + CurrentRunInfo, + MessageOrigin, + RenameChatSessionResponse, + SendMessageRequest, + SetPreferredResponseRequest, + UpdateChatSessionTemperatureRequest, + UpdateChatSessionThreadRequest, +) from onyx.server.query_and_chat.session_loading import ( translate_assistant_message_to_packets, ) -from onyx.server.query_and_chat.streaming_models import heartbeat_packet -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import heartbeat_packet, Packet from onyx.server.query_and_chat.token_limit import check_token_rate_limits -from onyx.server.usage_limits import check_llm_cost_limit_for_provider -from onyx.server.usage_limits import check_usage_and_raise -from onyx.server.usage_limits import is_usage_limits_enabled +from onyx.server.usage_limits import ( + check_llm_cost_limit_for_provider, + check_usage_and_raise, + is_usage_limits_enabled, +) from onyx.server.utils import get_json_line from onyx.tracing.framework.create import ensure_trace from onyx.utils.headers import get_custom_tool_additional_request_headers diff --git a/backend/onyx/server/query_and_chat/chat_utils.py b/backend/onyx/server/query_and_chat/chat_utils.py index 9f5fa35e14..8c1e3c37a4 100644 --- a/backend/onyx/server/query_and_chat/chat_utils.py +++ b/backend/onyx/server/query_and_chat/chat_utils.py @@ -1,11 +1,9 @@ -from typing import Any -from typing import IO +from typing import Any, IO from pydantic import BaseModel from onyx.file_processing.extract_file_text import xlsx_sheet_extraction -from onyx.file_processing.file_types import OnyxMimeTypes -from onyx.file_processing.file_types import SPREADSHEET_MIME_TYPES +from onyx.file_processing.file_types import OnyxMimeTypes, SPREADSHEET_MIME_TYPES from onyx.file_store.models import ChatFileType # Per-sheet cap on CSV text returned for in-chat spreadsheet previews. Sheets are diff --git a/backend/onyx/server/query_and_chat/models.py b/backend/onyx/server/query_and_chat/models.py index 45fcb91ac1..7f8ea21fc7 100644 --- a/backend/onyx/server/query_and_chat/models.py +++ b/backend/onyx/server/query_and_chat/models.py @@ -3,16 +3,10 @@ from enum import Enum from typing import Any from uuid import UUID -from pydantic import BaseModel -from pydantic import model_validator +from pydantic import BaseModel, model_validator -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MessageType -from onyx.configs.constants import SessionType -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import SavedSearchDoc -from onyx.context.search.models import SearchDoc -from onyx.context.search.models import Tag +from onyx.configs.constants import DocumentSource, MessageType, SessionType +from onyx.context.search.models import BaseFilters, SavedSearchDoc, SearchDoc, Tag from onyx.db.enums import ChatSessionSharedStatus from onyx.db.models import ChatSession from onyx.file_store.models import FileDescriptor diff --git a/backend/onyx/server/query_and_chat/query_backend.py b/backend/onyx/server/query_and_chat/query_backend.py index 6bb237ea2f..a5635ba912 100644 --- a/backend/onyx/server/query_and_chat/query_backend.py +++ b/backend/onyx/server/query_and_chat/query_backend.py @@ -1,13 +1,11 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission from onyx.auth.users import current_curator_or_admin_user from onyx.configs.chat_configs import NUM_RETURNED_HITS from onyx.configs.constants import DocumentSource -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import SearchDoc +from onyx.context.search.models import IndexFilters, SearchDoc from onyx.context.search.preprocessing.access_filters import ( build_access_filters_for_user, ) @@ -17,10 +15,12 @@ from onyx.db.models import User from onyx.db.search_settings import get_current_search_settings from onyx.db.tag import find_tags from onyx.document_index.factory import get_default_document_index -from onyx.server.query_and_chat.models import AdminSearchRequest -from onyx.server.query_and_chat.models import AdminSearchResponse -from onyx.server.query_and_chat.models import SourceTag -from onyx.server.query_and_chat.models import TagResponse +from onyx.server.query_and_chat.models import ( + AdminSearchRequest, + AdminSearchResponse, + SourceTag, + TagResponse, +) from onyx.server.utils_vector_db import require_vector_db from onyx.utils.logger import setup_logger from shared_configs.contextvars import get_current_tenant_id diff --git a/backend/onyx/server/query_and_chat/session_loading.py b/backend/onyx/server/query_and_chat/session_loading.py index 661b5e397c..e9128fb8a8 100644 --- a/backend/onyx/server/query_and_chat/session_loading.py +++ b/backend/onyx/server/query_and_chat/session_loading.py @@ -1,59 +1,61 @@ from __future__ import annotations import json -from typing import Any -from typing import cast -from typing import Literal +from typing import Any, cast, Literal from pydantic import ValidationError from sqlalchemy.orm import Session from onyx.chat.citation_utils import extract_citation_order_from_text -from onyx.coding_agent.mock_tools import CODING_AGENT_QUERY_KEY -from onyx.coding_agent.mock_tools import CODING_AGENT_REPO_KEY +from onyx.coding_agent.mock_tools import CODING_AGENT_QUERY_KEY, CODING_AGENT_REPO_KEY from onyx.configs.constants import MessageType -from onyx.context.search.models import SavedSearchDoc -from onyx.context.search.models import SearchDoc -from onyx.db.chat import get_db_search_doc_by_id -from onyx.db.chat import translate_db_search_doc_to_saved_search_doc +from onyx.context.search.models import SavedSearchDoc, SearchDoc +from onyx.db.chat import ( + get_db_search_doc_by_id, + translate_db_search_doc_to_saved_search_doc, +) from onyx.db.models import ChatMessage from onyx.db.tools import get_tool_by_id -from onyx.deep_research.dr_mock_tools import RESEARCH_AGENT_IN_CODE_ID -from onyx.deep_research.dr_mock_tools import RESEARCH_AGENT_TASK_KEY +from onyx.deep_research.dr_mock_tools import ( + RESEARCH_AGENT_IN_CODE_ID, + RESEARCH_AGENT_TASK_KEY, +) from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import CitationInfo -from onyx.server.query_and_chat.streaming_models import CodingAgentFinal -from onyx.server.query_and_chat.streaming_models import CodingAgentStart -from onyx.server.query_and_chat.streaming_models import CustomToolArgs -from onyx.server.query_and_chat.streaming_models import CustomToolDelta -from onyx.server.query_and_chat.streaming_models import CustomToolErrorInfo -from onyx.server.query_and_chat.streaming_models import CustomToolStart -from onyx.server.query_and_chat.streaming_models import FileReaderResult -from onyx.server.query_and_chat.streaming_models import FileReaderStart -from onyx.server.query_and_chat.streaming_models import GeneratedImage -from onyx.server.query_and_chat.streaming_models import ImageGenerationFinal -from onyx.server.query_and_chat.streaming_models import ImageGenerationToolStart -from onyx.server.query_and_chat.streaming_models import IntermediateReportDelta -from onyx.server.query_and_chat.streaming_models import IntermediateReportStart -from onyx.server.query_and_chat.streaming_models import MemoryToolDelta -from onyx.server.query_and_chat.streaming_models import MemoryToolStart -from onyx.server.query_and_chat.streaming_models import OpenUrlDocuments -from onyx.server.query_and_chat.streaming_models import OpenUrlStart -from onyx.server.query_and_chat.streaming_models import OpenUrlUrls -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import PythonToolDelta -from onyx.server.query_and_chat.streaming_models import PythonToolStart -from onyx.server.query_and_chat.streaming_models import ReasoningDelta -from onyx.server.query_and_chat.streaming_models import ReasoningStart -from onyx.server.query_and_chat.streaming_models import ResearchAgentStart -from onyx.server.query_and_chat.streaming_models import SearchToolDocumentsDelta -from onyx.server.query_and_chat.streaming_models import SearchToolQueriesDelta -from onyx.server.query_and_chat.streaming_models import SearchToolStart -from onyx.server.query_and_chat.streaming_models import SectionEnd -from onyx.server.query_and_chat.streaming_models import TopLevelBranching +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, + CitationInfo, + CodingAgentFinal, + CodingAgentStart, + CustomToolArgs, + CustomToolDelta, + CustomToolErrorInfo, + CustomToolStart, + FileReaderResult, + FileReaderStart, + GeneratedImage, + ImageGenerationFinal, + ImageGenerationToolStart, + IntermediateReportDelta, + IntermediateReportStart, + MemoryToolDelta, + MemoryToolStart, + OpenUrlDocuments, + OpenUrlStart, + OpenUrlUrls, + OverallStop, + Packet, + PythonToolDelta, + PythonToolStart, + ReasoningDelta, + ReasoningStart, + ResearchAgentStart, + SearchToolDocumentsDelta, + SearchToolQueriesDelta, + SearchToolStart, + SectionEnd, + TopLevelBranching, +) from onyx.tools.tool_implementations.coding_agent.coding_agent_tool import ( CodingAgentTool, ) diff --git a/backend/onyx/server/query_and_chat/streaming_models.py b/backend/onyx/server/query_and_chat/streaming_models.py index 123418f3fe..b8ad771270 100644 --- a/backend/onyx/server/query_and_chat/streaming_models.py +++ b/backend/onyx/server/query_and_chat/streaming_models.py @@ -1,12 +1,8 @@ from datetime import datetime from enum import Enum -from typing import Annotated -from typing import Any -from typing import Literal -from typing import Union +from typing import Annotated, Any, Literal, Union -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.context.search.models import SearchDoc from onyx.server.query_and_chat.placement import Placement diff --git a/backend/onyx/server/query_and_chat/token_limit.py b/backend/onyx/server/query_and_chat/token_limit.py index 632be5e361..35677c407b 100644 --- a/backend/onyx/server/query_and_chat/token_limit.py +++ b/backend/onyx/server/query_and_chat/token_limit.py @@ -1,23 +1,16 @@ from collections.abc import Sequence -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from threading import RLock from cachetools import TTLCache from dateutil import tz -from fastapi import Depends -from fastapi import HTTPException -from sqlalchemy import func -from sqlalchemy import select +from fastapi import Depends, HTTPException +from sqlalchemy import func, select from sqlalchemy.orm import Session from onyx.auth.users import current_chat_accessible_user from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession -from onyx.db.models import TokenRateLimit -from onyx.db.models import User +from onyx.db.models import ChatMessage, ChatSession, TokenRateLimit, User from onyx.db.token_limit import fetch_all_global_token_rate_limits from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_versioned_implementation diff --git a/backend/onyx/server/runtime/onyx_runtime.py b/backend/onyx/server/runtime/onyx_runtime.py index d576fe9a8b..b11654f896 100644 --- a/backend/onyx/server/runtime/onyx_runtime.py +++ b/backend/onyx/server/runtime/onyx_runtime.py @@ -2,21 +2,24 @@ import io from PIL import Image -from onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT from onyx.background.celery.tasks.beat_schedule import ( + CLOUD_BEAT_MULTIPLIER_DEFAULT, CLOUD_DOC_PERMISSION_SYNC_MULTIPLIER_DEFAULT, ) -from onyx.configs.app_configs import ENABLE_TENANT_WORK_GATING -from onyx.configs.app_configs import TENANT_WORK_GATING_FULL_FANOUT_INTERVAL_SECONDS -from onyx.configs.app_configs import TENANT_WORK_GATING_TTL_SECONDS -from onyx.configs.constants import CLOUD_BUILD_FENCE_LOOKUP_TABLE_INTERVAL_DEFAULT -from onyx.configs.constants import ONYX_CLOUD_REDIS_RUNTIME -from onyx.configs.constants import ONYX_CLOUD_TENANT_ID -from onyx.configs.constants import ONYX_EMAILABLE_LOGO_MAX_DIM +from onyx.configs.app_configs import ( + ENABLE_TENANT_WORK_GATING, + TENANT_WORK_GATING_FULL_FANOUT_INTERVAL_SECONDS, + TENANT_WORK_GATING_TTL_SECONDS, +) +from onyx.configs.constants import ( + CLOUD_BUILD_FENCE_LOOKUP_TABLE_INTERVAL_DEFAULT, + ONYX_CLOUD_REDIS_RUNTIME, + ONYX_CLOUD_TENANT_ID, + ONYX_EMAILABLE_LOGO_MAX_DIM, +) from onyx.file_store.file_store import get_default_file_store from onyx.redis.redis_pool import get_redis_replica_client -from onyx.utils.file import FileWithMimeType -from onyx.utils.file import OnyxStaticFileManager +from onyx.utils.file import FileWithMimeType, OnyxStaticFileManager from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop diff --git a/backend/onyx/server/saml.py b/backend/onyx/server/saml.py index 9e3853310c..0beeee87a8 100644 --- a/backend/onyx/server/saml.py +++ b/backend/onyx/server/saml.py @@ -8,12 +8,10 @@ from fastapi import Request from fastapi_users import exceptions from pydantic import BaseModel -from onyx.auth.schemas import UserCreate -from onyx.auth.schemas import UserRole +from onyx.auth.schemas import UserCreate, UserRole from onyx.auth.users import get_user_manager from onyx.configs.app_configs import WEB_DOMAIN -from onyx.db.auth import get_user_count -from onyx.db.auth import get_user_db +from onyx.db.auth import get_user_count, get_user_db from onyx.db.engine.async_sql_engine import get_async_session_context_manager from onyx.db.models import User from onyx.utils.logger import setup_logger diff --git a/backend/onyx/server/saml_multi.py b/backend/onyx/server/saml_multi.py index 40a3ff82ce..b1fd1920f3 100644 --- a/backend/onyx/server/saml_multi.py +++ b/backend/onyx/server/saml_multi.py @@ -1,40 +1,34 @@ import base64 import uuid -from typing import Any -from typing import NoReturn -from typing import TypedDict +from typing import Any, NoReturn, TypedDict -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Request -from fastapi import Response +from fastapi import APIRouter, Depends, Request, Response from fastapi_users.authentication import Strategy from onelogin.saml2.auth import OneLogin_Saml2_Auth from onelogin.saml2.xml_utils import OneLogin_Saml2_XML from pydantic import ValidationError from sqlalchemy.orm import Session -from onyx.auth.users import auth_backend -from onyx.auth.users import fastapi_users -from onyx.auth.users import get_user_manager -from onyx.auth.users import UserManager -from onyx.configs.app_configs import REQUIRE_EMAIL_VERIFICATION -from onyx.configs.app_configs import WEB_DOMAIN +from onyx.auth.users import auth_backend, fastapi_users, get_user_manager, UserManager +from onyx.configs.app_configs import REQUIRE_EMAIL_VERIFICATION, WEB_DOMAIN from onyx.db.engine.sql_engine import get_session from onyx.db.enums import SSOProviderType -from onyx.db.models import SSOProvider -from onyx.db.models import User -from onyx.db.sso_provider import fetch_sso_provider_by_name -from onyx.db.sso_provider import fetch_sso_providers -from onyx.db.sso_provider import SAMLProviderConfig +from onyx.db.models import SSOProvider, User +from onyx.db.sso_provider import ( + fetch_sso_provider_by_name, + fetch_sso_providers, + SAMLProviderConfig, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.saml import _sanitize_relay_state -from onyx.server.saml import EMAIL_ATTRIBUTE_KEYS -from onyx.server.saml import EMAIL_ATTRIBUTE_KEYS_LOWER -from onyx.server.saml import prepare_from_fastapi_request -from onyx.server.saml import SAMLAuthorizeResponse -from onyx.server.saml import upsert_saml_user +from onyx.server.saml import ( + _sanitize_relay_state, + EMAIL_ATTRIBUTE_KEYS, + EMAIL_ATTRIBUTE_KEYS_LOWER, + prepare_from_fastapi_request, + SAMLAuthorizeResponse, + upsert_saml_user, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/security/api.py b/backend/onyx/server/security/api.py index fb50afad57..584ad0fc9b 100644 --- a/backend/onyx/server/security/api.py +++ b/backend/onyx/server/security/api.py @@ -1,9 +1,7 @@ import json from typing import Any -from fastapi import APIRouter -from fastapi import Depends -from fastapi import Request +from fastapi import APIRouter, Depends, Request from fastapi.concurrency import run_in_threadpool from pydantic import ValidationError @@ -12,11 +10,12 @@ from onyx.db.enums import Permission from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.security.models import OPERATOR_LOCKED_FIELDS -from onyx.server.security.models import SecuritySettings -from onyx.server.security.models import SecuritySettingsOverrides -from onyx.server.security.store import apply_patch -from onyx.server.security.store import get_security_settings +from onyx.server.security.models import ( + OPERATOR_LOCKED_FIELDS, + SecuritySettings, + SecuritySettingsOverrides, +) +from onyx.server.security.store import apply_patch, get_security_settings from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT diff --git a/backend/onyx/server/security/models.py b/backend/onyx/server/security/models.py index c8d036121a..c2f82356fc 100644 --- a/backend/onyx/server/security/models.py +++ b/backend/onyx/server/security/models.py @@ -1,11 +1,7 @@ from enum import Enum from typing import NamedTuple -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_validator -from pydantic import model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Self diff --git a/backend/onyx/server/security/store.py b/backend/onyx/server/security/store.py index 39e8c9871a..c9922a169a 100644 --- a/backend/onyx/server/security/store.py +++ b/backend/onyx/server/security/store.py @@ -14,13 +14,14 @@ from onyx.db.security_settings import load_overrides as _db_load_overrides from onyx.db.security_settings import upsert_overrides as _db_upsert_overrides from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.security.models import OPERATOR_LOCKED_FIELDS -from onyx.server.security.models import SecuritySettings -from onyx.server.security.models import SecuritySettingsOverrides -from onyx.server.security.models import SSRFProtectionLevel +from onyx.server.security.models import ( + OPERATOR_LOCKED_FIELDS, + SecuritySettings, + SecuritySettingsOverrides, + SSRFProtectionLevel, +) from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR logger = setup_logger() diff --git a/backend/onyx/server/settings/api.py b/backend/onyx/server/settings/api.py index d6e973e254..6e8cad3ac6 100644 --- a/backend/onyx/server/settings/api.py +++ b/backend/onyx/server/settings/api.py @@ -1,53 +1,58 @@ from typing import cast -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session from onyx import __version__ as onyx_version from onyx.auth.permissions import require_permission from onyx.auth.users import is_user_admin -from onyx.configs.app_configs import DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import MAX_ALLOWED_UPLOAD_SIZE_MB -from onyx.configs.app_configs import POSTHOG_API_KEY -from onyx.configs.app_configs import POSTHOG_HOST -from onyx.configs.constants import KV_REINDEX_KEY -from onyx.configs.constants import NotificationType +from onyx.configs.app_configs import ( + DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB, + DISABLE_VECTOR_DB, + MAX_ALLOWED_UPLOAD_SIZE_MB, + POSTHOG_API_KEY, + POSTHOG_HOST, +) +from onyx.configs.constants import KV_REINDEX_KEY, NotificationType from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.notification import dismiss_all_notifications -from onyx.db.notification import get_notifications -from onyx.db.notification import update_notification_last_shown +from onyx.db.notification import ( + dismiss_all_notifications, + get_notifications, + update_notification_last_shown, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError -from onyx.server.features.build.utils import is_craft_available_for_deployment -from onyx.server.features.build.utils import is_craft_enabled_for_user +from onyx.server.features.build.utils import ( + is_craft_available_for_deployment, + is_craft_enabled_for_user, +) from onyx.server.features.notifications.models import NotificationResponse from onyx.server.settings.models import ( DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_NO_VECTOR_DB, + DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_VECTOR_DB, + Settings, + Tier, + UserSettings, ) -from onyx.server.settings.models import DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_VECTOR_DB -from onyx.server.settings.models import Settings -from onyx.server.settings.models import Tier -from onyx.server.settings.models import UserSettings -from onyx.server.settings.store import load_settings -from onyx.server.settings.store import store_settings +from onyx.server.settings.store import load_settings, store_settings from onyx.server.settings.tier_order import tier_at_least -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event +from onyx.utils.audit import ( + actor_from_user, + AuditAction, + AuditOutcome, + emit_audit_event, +) from onyx.utils.logger import setup_logger from onyx.utils.platform_utils import is_running_in_container from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, + global_version, ) -from onyx.utils.variable_functionality import global_version from shared_configs.configs import MULTI_TENANT logger = setup_logger() diff --git a/backend/onyx/server/settings/models.py b/backend/onyx/server/settings/models.py index 03550aff1e..ae91330376 100644 --- a/backend/onyx/server/settings/models.py +++ b/backend/onyx/server/settings/models.py @@ -1,12 +1,13 @@ from enum import Enum -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field -from onyx.configs.app_configs import DEFAULT_PRUNING_FREQ -from onyx.configs.app_configs import DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import MAX_ALLOWED_UPLOAD_SIZE_MB +from onyx.configs.app_configs import ( + DEFAULT_PRUNING_FREQ, + DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB, + DISABLE_VECTOR_DB, + MAX_ALLOWED_UPLOAD_SIZE_MB, +) from onyx.configs.constants import QueryHistoryType from onyx.server.features.notifications.models import NotificationResponse from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA diff --git a/backend/onyx/server/settings/store.py b/backend/onyx/server/settings/store.py index 205450ca88..d5e243bf29 100644 --- a/backend/onyx/server/settings/store.py +++ b/backend/onyx/server/settings/store.py @@ -1,21 +1,22 @@ from onyx.cache.factory import get_cache_backend -from onyx.configs.app_configs import DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB -from onyx.configs.app_configs import DISABLE_USER_KNOWLEDGE -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX -from onyx.configs.app_configs import HIDE_QUERY_HISTORY_FROM_ADMIN_PANEL -from onyx.configs.app_configs import MAX_ALLOWED_UPLOAD_SIZE_MB -from onyx.configs.app_configs import ONYX_QUERY_HISTORY_TYPE -from onyx.configs.app_configs import SHOW_EXTRA_CONNECTORS -from onyx.configs.constants import KV_SETTINGS_KEY -from onyx.configs.constants import OnyxRedisLocks +from onyx.configs.app_configs import ( + DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB, + DISABLE_USER_KNOWLEDGE, + DISABLE_VECTOR_DB, + ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + HIDE_QUERY_HISTORY_FROM_ADMIN_PANEL, + MAX_ALLOWED_UPLOAD_SIZE_MB, + ONYX_QUERY_HISTORY_TYPE, + SHOW_EXTRA_CONNECTORS, +) +from onyx.configs.constants import KV_SETTINGS_KEY, OnyxRedisLocks from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError from onyx.server.settings.models import ( DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_NO_VECTOR_DB, + DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_VECTOR_DB, + Settings, ) -from onyx.server.settings.models import DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_VECTOR_DB -from onyx.server.settings.models import Settings from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/server/token_rate_limits/api.py b/backend/onyx/server/token_rate_limits/api.py index 0233f87f6e..2f8b189649 100644 --- a/backend/onyx/server/token_rate_limits/api.py +++ b/backend/onyx/server/token_rate_limits/api.py @@ -1,5 +1,4 @@ -from fastapi import APIRouter -from fastapi import Depends +from fastapi import APIRouter, Depends from sqlalchemy.orm import Session from onyx.auth.permissions import require_permission @@ -7,15 +6,19 @@ from onyx.configs.constants import PUBLIC_API_TAGS from onyx.db.engine.sql_engine import get_session from onyx.db.enums import Permission from onyx.db.models import User -from onyx.db.token_limit import delete_token_rate_limit -from onyx.db.token_limit import fetch_all_global_token_rate_limits -from onyx.db.token_limit import insert_global_token_rate_limit -from onyx.db.token_limit import update_token_rate_limit +from onyx.db.token_limit import ( + delete_token_rate_limit, + fetch_all_global_token_rate_limits, + insert_global_token_rate_limit, + update_token_rate_limit, +) from onyx.server.query_and_chat.token_limit import ( invalidate_any_rate_limit_exists_cache, ) -from onyx.server.token_rate_limits.models import TokenRateLimitArgs -from onyx.server.token_rate_limits.models import TokenRateLimitDisplay +from onyx.server.token_rate_limits.models import ( + TokenRateLimitArgs, + TokenRateLimitDisplay, +) router = APIRouter(prefix="/admin/token-rate-limits", tags=PUBLIC_API_TAGS) diff --git a/backend/onyx/server/token_rate_limits/models.py b/backend/onyx/server/token_rate_limits/models.py index 032dd57c44..42f4619c2e 100644 --- a/backend/onyx/server/token_rate_limits/models.py +++ b/backend/onyx/server/token_rate_limits/models.py @@ -1,5 +1,4 @@ -from pydantic import BaseModel -from pydantic import model_validator +from pydantic import BaseModel, model_validator from onyx.db.models import TokenRateLimit diff --git a/backend/onyx/server/usage_limits.py b/backend/onyx/server/usage_limits.py index e4dbc3a57e..5acc552cd0 100644 --- a/backend/onyx/server/usage_limits.py +++ b/backend/onyx/server/usage_limits.py @@ -4,28 +4,32 @@ from collections.abc import Callable from sqlalchemy.orm import Session -from onyx.configs.app_configs import ANTHROPIC_DEFAULT_API_KEY -from onyx.configs.app_configs import COHERE_DEFAULT_API_KEY -from onyx.configs.app_configs import OPENAI_DEFAULT_API_KEY -from onyx.configs.app_configs import OPENROUTER_DEFAULT_API_KEY -from onyx.db.usage import check_usage_limit -from onyx.db.usage import UsageLimitExceededError -from onyx.db.usage import UsageType +from onyx.configs.app_configs import ( + ANTHROPIC_DEFAULT_API_KEY, + COHERE_DEFAULT_API_KEY, + OPENAI_DEFAULT_API_KEY, + OPENROUTER_DEFAULT_API_KEY, +) +from onyx.db.usage import check_usage_limit, UsageLimitExceededError, UsageType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.tenant_usage_limits import TenantUsageLimitKeys -from onyx.server.tenant_usage_limits import TenantUsageLimitOverrides +from onyx.server.tenant_usage_limits import ( + TenantUsageLimitKeys, + TenantUsageLimitOverrides, +) from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import fetch_versioned_implementation -from shared_configs.configs import USAGE_LIMIT_API_CALLS_PAID -from shared_configs.configs import USAGE_LIMIT_API_CALLS_TRIAL -from shared_configs.configs import USAGE_LIMIT_CHUNKS_INDEXED_PAID -from shared_configs.configs import USAGE_LIMIT_CHUNKS_INDEXED_TRIAL -from shared_configs.configs import USAGE_LIMIT_LLM_COST_CENTS_PAID -from shared_configs.configs import USAGE_LIMIT_LLM_COST_CENTS_TRIAL -from shared_configs.configs import USAGE_LIMIT_NON_STREAMING_CALLS_PAID -from shared_configs.configs import USAGE_LIMIT_NON_STREAMING_CALLS_TRIAL -from shared_configs.configs import USAGE_LIMITS_ENABLED +from shared_configs.configs import ( + USAGE_LIMIT_API_CALLS_PAID, + USAGE_LIMIT_API_CALLS_TRIAL, + USAGE_LIMIT_CHUNKS_INDEXED_PAID, + USAGE_LIMIT_CHUNKS_INDEXED_TRIAL, + USAGE_LIMIT_LLM_COST_CENTS_PAID, + USAGE_LIMIT_LLM_COST_CENTS_TRIAL, + USAGE_LIMIT_NON_STREAMING_CALLS_PAID, + USAGE_LIMIT_NON_STREAMING_CALLS_TRIAL, + USAGE_LIMITS_ENABLED, +) logger = setup_logger() diff --git a/backend/onyx/server/utils.py b/backend/onyx/server/utils.py index 6b08a322d4..e1702e1136 100644 --- a/backend/onyx/server/utils.py +++ b/backend/onyx/server/utils.py @@ -5,8 +5,7 @@ from datetime import datetime from typing import Any from uuid import UUID -from fastapi import HTTPException -from fastapi import status +from fastapi import HTTPException, status class BasicAuthenticationError(HTTPException): diff --git a/backend/onyx/setup.py b/backend/onyx/setup.py index 4b3ffc509a..cd9102b654 100644 --- a/backend/onyx/setup.py +++ b/backend/onyx/setup.py @@ -2,41 +2,53 @@ import time from sqlalchemy.orm import Session -from onyx.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP -from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX -from onyx.configs.app_configs import INTEGRATION_TESTS_MODE -from onyx.configs.app_configs import MANAGED_VESPA -from onyx.configs.app_configs import ONYX_DISABLE_VESPA -from onyx.configs.app_configs import VESPA_NUM_ATTEMPTS_ON_STARTUP +from onyx.configs.app_configs import ( + DISABLE_INDEX_UPDATE_ON_SWAP, + DISABLE_VECTOR_DB, + ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + INTEGRATION_TESTS_MODE, + MANAGED_VESPA, + ONYX_DISABLE_VESPA, + VESPA_NUM_ATTEMPTS_ON_STARTUP, +) from onyx.configs.constants import KV_REINDEX_KEY -from onyx.configs.embedding_configs import SUPPORTED_EMBEDDING_MODELS -from onyx.configs.embedding_configs import SupportedEmbeddingModel -from onyx.configs.model_configs import GEN_AI_API_KEY -from onyx.configs.model_configs import GEN_AI_MODEL_VERSION +from onyx.configs.embedding_configs import ( + SUPPORTED_EMBEDDING_MODELS, + SupportedEmbeddingModel, +) +from onyx.configs.model_configs import GEN_AI_API_KEY, GEN_AI_MODEL_VERSION from onyx.context.search.models import SavedSearchSettings -from onyx.db.connector import check_connectors_exist -from onyx.db.connector import create_initial_default_connector -from onyx.db.connector_credential_pair import associate_default_cc_pair -from onyx.db.connector_credential_pair import get_connector_credential_pairs -from onyx.db.connector_credential_pair import resync_cc_pair +from onyx.db.connector import check_connectors_exist, create_initial_default_connector +from onyx.db.connector_credential_pair import ( + associate_default_cc_pair, + get_connector_credential_pairs, + resync_cc_pair, +) from onyx.db.credentials import create_initial_public_credential from onyx.db.document import check_docs_exist from onyx.db.enums import EmbeddingPrecision -from onyx.db.index_attempt import cancel_indexing_attempts_past_model -from onyx.db.index_attempt import expire_index_attempts -from onyx.db.llm import fetch_default_llm_model -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider -from onyx.db.search_settings import get_active_search_settings -from onyx.db.search_settings import get_current_search_settings -from onyx.db.search_settings import update_current_search_settings +from onyx.db.index_attempt import ( + cancel_indexing_attempts_past_model, + expire_index_attempts, +) +from onyx.db.llm import ( + fetch_default_llm_model, + fetch_existing_llm_provider, + update_default_provider, + upsert_llm_provider, +) +from onyx.db.search_settings import ( + get_active_search_settings, + get_current_search_settings, + update_current_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.document_index.interfaces_new import DocumentIndex -from onyx.document_index.opensearch.client import OpenSearchClient -from onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout +from onyx.document_index.opensearch.client import ( + OpenSearchClient, + wait_for_opensearch_with_timeout, +) from onyx.document_index.opensearch.opensearch_document_index import set_cluster_state from onyx.document_index.vespa.vespa_document_index import ( register_multitenant_vespa_indices, @@ -46,18 +58,23 @@ from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError from onyx.llm.constants import LlmProviderNames from onyx.llm.well_known_providers.llm_provider_options import get_openai_model_names -from onyx.natural_language_processing.search_nlp_models import EmbeddingModel -from onyx.natural_language_processing.search_nlp_models import warm_up_bi_encoder -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest -from onyx.server.settings.store import load_settings -from onyx.server.settings.store import store_settings +from onyx.natural_language_processing.search_nlp_models import ( + EmbeddingModel, + warm_up_bi_encoder, +) +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) +from onyx.server.settings.store import load_settings, store_settings from onyx.utils.gpu_utils import gpu_status_request from onyx.utils.logger import setup_logger -from shared_configs.configs import ALT_INDEX_SUFFIX -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT -from shared_configs.configs import MULTI_TENANT +from shared_configs.configs import ( + ALT_INDEX_SUFFIX, + MODEL_SERVER_HOST, + MODEL_SERVER_PORT, + MULTI_TENANT, +) logger = setup_logger() diff --git a/backend/onyx/skills/built_in.py b/backend/onyx/skills/built_in.py index 9001ff591d..928220ba79 100644 --- a/backend/onyx/skills/built_in.py +++ b/backend/onyx/skills/built_in.py @@ -20,17 +20,11 @@ or external app is a single ``_REGISTRY`` entry plus its on-disk ``builtin//`` directory. """ -from collections.abc import Callable -from collections.abc import Iterable +from collections.abc import Callable, Iterable from pathlib import Path -from typing import Final -from typing import Literal +from typing import Final, Literal -from pydantic import BaseModel -from pydantic import computed_field -from pydantic import ConfigDict -from pydantic import Field -from pydantic import model_validator +from pydantic import BaseModel, computed_field, ConfigDict, Field, model_validator from sqlalchemy.orm import Session from onyx.db.enums import ExternalAppType diff --git a/backend/onyx/skills/builtin/pptx/scripts/lint.py b/backend/onyx/skills/builtin/pptx/scripts/lint.py index f00e01ff59..ec565aca49 100644 --- a/backend/onyx/skills/builtin/pptx/scripts/lint.py +++ b/backend/onyx/skills/builtin/pptx/scripts/lint.py @@ -36,8 +36,7 @@ from pathlib import Path from PIL import ImageFont from pptx import Presentation -from pptx.enum.dml import MSO_COLOR_TYPE -from pptx.enum.dml import MSO_FILL +from pptx.enum.dml import MSO_COLOR_TYPE, MSO_FILL from pptx.enum.shapes import MSO_SHAPE_TYPE from pptx.enum.text import MSO_AUTO_SIZE from pptx.shapes.base import BaseShape diff --git a/backend/onyx/skills/builtin/pptx/scripts/office/pack.py b/backend/onyx/skills/builtin/pptx/scripts/office/pack.py index 541b6045ac..944d226403 100755 --- a/backend/onyx/skills/builtin/pptx/scripts/office/pack.py +++ b/backend/onyx/skills/builtin/pptx/scripts/office/pack.py @@ -18,9 +18,7 @@ import zipfile from pathlib import Path import defusedxml.minidom -from validators import DOCXSchemaValidator -from validators import PPTXSchemaValidator -from validators import RedliningValidator +from validators import DOCXSchemaValidator, PPTXSchemaValidator, RedliningValidator def pack( diff --git a/backend/onyx/skills/builtin/pptx/scripts/office/validate.py b/backend/onyx/skills/builtin/pptx/scripts/office/validate.py index 2efe8009be..0d33edba90 100755 --- a/backend/onyx/skills/builtin/pptx/scripts/office/validate.py +++ b/backend/onyx/skills/builtin/pptx/scripts/office/validate.py @@ -19,9 +19,7 @@ import tempfile import zipfile from pathlib import Path -from validators import DOCXSchemaValidator -from validators import PPTXSchemaValidator -from validators import RedliningValidator +from validators import DOCXSchemaValidator, PPTXSchemaValidator, RedliningValidator def main(): diff --git a/backend/onyx/skills/builtin/pptx/scripts/thumbnail.py b/backend/onyx/skills/builtin/pptx/scripts/thumbnail.py index 253db294e8..edcbdc0f81 100755 --- a/backend/onyx/skills/builtin/pptx/scripts/thumbnail.py +++ b/backend/onyx/skills/builtin/pptx/scripts/thumbnail.py @@ -24,9 +24,7 @@ from pathlib import Path import defusedxml.minidom from office.soffice import get_soffice_env -from PIL import Image -from PIL import ImageDraw -from PIL import ImageFont +from PIL import Image, ImageDraw, ImageFont THUMBNAIL_WIDTH = 300 CONVERSION_DPI = 100 diff --git a/backend/onyx/skills/bundle.py b/backend/onyx/skills/bundle.py index 4c6c126593..9129e5eca1 100644 --- a/backend/onyx/skills/bundle.py +++ b/backend/onyx/skills/bundle.py @@ -11,18 +11,18 @@ import zipfile from contextlib import ExitStack from copy import copy from dataclasses import dataclass -from typing import Final -from typing import IO +from typing import Final, IO from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.skills.built_in import BUILT_IN_SKILLS -from onyx.skills.metadata import parse_skill_document -from onyx.skills.metadata import parse_skill_md_frontmatter -from onyx.skills.metadata import serialize_skill_md -from onyx.skills.metadata import split_skill_md -from onyx.skills.models import CustomSkillBundleContents -from onyx.skills.models import SkillBundleFile +from onyx.skills.metadata import ( + parse_skill_document, + parse_skill_md_frontmatter, + serialize_skill_md, + split_skill_md, +) +from onyx.skills.models import CustomSkillBundleContents, SkillBundleFile DEFAULT_PER_FILE_MAX_BYTES: Final[int] = int( os.environ.get("SKILL_BUNDLE_PER_FILE_MAX_BYTES") or 25 * 1024 * 1024 diff --git a/backend/onyx/skills/content.py b/backend/onyx/skills/content.py index 5037a8fc2a..a7ff0dd077 100644 --- a/backend/onyx/skills/content.py +++ b/backend/onyx/skills/content.py @@ -3,13 +3,14 @@ from onyx.db.models import Skill from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.file_store.file_store import FileStore -from onyx.file_store.file_store import get_default_file_store +from onyx.file_store.file_store import FileStore, get_default_file_store from onyx.skills.built_in import BuiltInSkillDefinition -from onyx.skills.bundle import read_custom_bundle_instructions -from onyx.skills.bundle import SKILL_MD_NAME -from onyx.skills.bundle import strip_skill_md_frontmatter -from onyx.skills.bundle import TEMPLATE_SUFFIX +from onyx.skills.bundle import ( + read_custom_bundle_instructions, + SKILL_MD_NAME, + strip_skill_md_frontmatter, + TEMPLATE_SUFFIX, +) def read_builtin_skill_instructions(definition: BuiltInSkillDefinition) -> str: diff --git a/backend/onyx/skills/ingest.py b/backend/onyx/skills/ingest.py index 336fbec5da..9439a675e0 100644 --- a/backend/onyx/skills/ingest.py +++ b/backend/onyx/skills/ingest.py @@ -9,10 +9,12 @@ from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import FileStore from onyx.skills.built_in import BUILT_IN_SKILLS -from onyx.skills.bundle import build_single_file_bundle -from onyx.skills.bundle import compute_bundle_sha256 -from onyx.skills.bundle import normalize_custom_bundle -from onyx.skills.bundle import SKILL_MD_NAME +from onyx.skills.bundle import ( + build_single_file_bundle, + compute_bundle_sha256, + normalize_custom_bundle, + SKILL_MD_NAME, +) from onyx.skills.metadata import parse_skill_document from onyx.utils.logger import setup_logger diff --git a/backend/onyx/skills/metadata.py b/backend/onyx/skills/metadata.py index 5f55f723a4..c98afe5dbf 100644 --- a/backend/onyx/skills/metadata.py +++ b/backend/onyx/skills/metadata.py @@ -4,8 +4,7 @@ from __future__ import annotations import re from collections.abc import Mapping -from typing import Any -from typing import Final +from typing import Any, Final import yaml from pydantic import ValidationError @@ -14,8 +13,7 @@ from yaml.nodes import MappingNode from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.skills.models import SkillDocument -from onyx.skills.models import SkillMetadata +from onyx.skills.models import SkillDocument, SkillMetadata _FRONTMATTER_REGEX: Final[re.Pattern[str]] = re.compile( r"\A---[ \t]*\r?\n(?P.*?)(?:\r?\n)---[ \t]*(?:\r?\n|\Z)", diff --git a/backend/onyx/skills/models.py b/backend/onyx/skills/models.py index db9f72a6df..a42313adc8 100644 --- a/backend/onyx/skills/models.py +++ b/backend/onyx/skills/models.py @@ -1,10 +1,7 @@ import re from typing import Final -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import Field -from pydantic import field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator SKILL_NAME_PATTERN: Final[re.Pattern[str]] = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$") diff --git a/backend/onyx/skills/push.py b/backend/onyx/skills/push.py index 102ef0829f..aba2fa50dd 100644 --- a/backend/onyx/skills/push.py +++ b/backend/onyx/skills/push.py @@ -9,34 +9,41 @@ from uuid import UUID from sqlalchemy.orm import Session -from onyx.db.external_app import get_connectable_apps_for_user -from onyx.db.external_app import get_external_app_by_skill_id -from onyx.db.models import Skill -from onyx.db.models import User -from onyx.db.skill import affected_user_ids_for_skill -from onyx.db.skill import list_skills -from onyx.db.skill import persist_skill_validity -from onyx.db.skill import SkillAccessPolicy -from onyx.db.skill import SkillValidityUpdate +from onyx.db.external_app import ( + get_connectable_apps_for_user, + get_external_app_by_skill_id, +) +from onyx.db.models import Skill, User +from onyx.db.skill import ( + affected_user_ids_for_skill, + list_skills, + persist_skill_validity, + SkillAccessPolicy, + SkillValidityUpdate, +) from onyx.file_store.file_store import get_default_file_store -from onyx.server.features.build.db.sandbox import get_sandbox_user_map -from onyx.server.features.build.db.sandbox import lock_sandbox_skills_hashes -from onyx.server.features.build.db.sandbox import set_sandbox_skills_hashes__no_commit +from onyx.server.features.build.db.sandbox import ( + get_sandbox_user_map, + lock_sandbox_skills_hashes, + set_sandbox_skills_hashes__no_commit, +) from onyx.server.features.build.sandbox.base import SandboxManager from onyx.server.features.build.sandbox.factory import get_sandbox_manager -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import PushResult +from onyx.server.features.build.sandbox.models import FileSet, PushResult from onyx.server.features.build.sandbox.util.agent_instructions import ( build_connectable_apps_list, ) -from onyx.skills.built_in import BUILT_IN_SKILLS -from onyx.skills.built_in import BuiltInSkillDefinition -from onyx.skills.built_in import COMPANY_SEARCH -from onyx.skills.built_in import EXTERNAL_APP_SKILL_ID_TO_APP_TYPE -from onyx.skills.rendering import render_company_search_skill -from onyx.skills.rendering import render_external_app_skill -from onyx.skills.validation import load_stored_custom_skill_bundle -from onyx.skills.validation import validate_stored_custom_skill +from onyx.skills.built_in import ( + BUILT_IN_SKILLS, + BuiltInSkillDefinition, + COMPANY_SEARCH, + EXTERNAL_APP_SKILL_ID_TO_APP_TYPE, +) +from onyx.skills.rendering import render_company_search_skill, render_external_app_skill +from onyx.skills.validation import ( + load_stored_custom_skill_bundle, + validate_stored_custom_skill, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/skills/rendering.py b/backend/onyx/skills/rendering.py index 60fd99a219..c5a5103dcd 100644 --- a/backend/onyx/skills/rendering.py +++ b/backend/onyx/skills/rendering.py @@ -4,15 +4,12 @@ from pathlib import Path from sqlalchemy.orm import Session -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import DocumentSourceDescription +from onyx.configs.constants import DocumentSource, DocumentSourceDescription from onyx.db.connector import _INTERNAL_ONLY_SOURCES from onyx.db.connector_credential_pair import get_connector_credential_pairs_for_user -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.db.external_app import get_policies -from onyx.db.models import ExternalApp -from onyx.db.models import User +from onyx.db.models import ExternalApp, User from onyx.external_apps.providers.registry import action_policy_views from onyx.utils.logger import setup_logger diff --git a/backend/onyx/skills/validation.py b/backend/onyx/skills/validation.py index 796a9cac67..ade21a7846 100644 --- a/backend/onyx/skills/validation.py +++ b/backend/onyx/skills/validation.py @@ -11,10 +11,12 @@ from onyx.db.models import Skill from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import FileStore from onyx.skills.built_in import BUILT_IN_SKILLS -from onyx.skills.bundle import normalize_custom_bundle -from onyx.skills.bundle import NormalizedSkillBundle -from onyx.skills.bundle import read_bundle_file -from onyx.skills.bundle import SKILL_MD_NAME +from onyx.skills.bundle import ( + normalize_custom_bundle, + NormalizedSkillBundle, + read_bundle_file, + SKILL_MD_NAME, +) from onyx.skills.metadata import parse_skill_document from onyx.skills.models import SKILL_NAME_PATTERN diff --git a/backend/onyx/tools/built_in_tools.py b/backend/onyx/tools/built_in_tools.py index decd1b8895..6ced3fa15f 100644 --- a/backend/onyx/tools/built_in_tools.py +++ b/backend/onyx/tools/built_in_tools.py @@ -1,5 +1,4 @@ -from typing import Type -from typing import Union +from typing import Type, Union from onyx.tools.tool_implementations.coding_agent.coding_agent_tool import ( CodingAgentTool, diff --git a/backend/onyx/tools/fake_tools/coding_agent.py b/backend/onyx/tools/fake_tools/coding_agent.py index df0be40927..627d90e1d2 100644 --- a/backend/onyx/tools/fake_tools/coding_agent.py +++ b/backend/onyx/tools/fake_tools/coding_agent.py @@ -6,43 +6,49 @@ from typing import Callable from onyx.chat.emitter import Emitter from onyx.chat.llm_loop import construct_message_history from onyx.chat.llm_step import run_llm_step_pkt_generator -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import ToolCallSimple -from onyx.coding_agent.mock_tools import BASH_TOOL_CMD_KEY -from onyx.coding_agent.mock_tools import BASH_TOOL_NAME -from onyx.coding_agent.mock_tools import CODING_AGENT_QUERY_KEY -from onyx.coding_agent.mock_tools import CODING_AGENT_REPO_KEY -from onyx.coding_agent.mock_tools import GENERATE_ANSWER_TOOL_NAME -from onyx.coding_agent.mock_tools import get_coding_agent_tool_definitions -from onyx.coding_agent.models import CodingAgentCallResult -from onyx.coding_agent.models import CodingAgentSpecialToolCalls +from onyx.chat.models import ChatMessageSimple, ToolCallSimple +from onyx.coding_agent.mock_tools import ( + BASH_TOOL_CMD_KEY, + BASH_TOOL_NAME, + CODING_AGENT_QUERY_KEY, + CODING_AGENT_REPO_KEY, + GENERATE_ANSWER_TOOL_NAME, + get_coding_agent_tool_definitions, +) +from onyx.coding_agent.models import CodingAgentCallResult, CodingAgentSpecialToolCalls from onyx.configs.constants import MessageType -from onyx.deep_research.dr_mock_tools import THINK_TOOL_NAME -from onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_MESSAGE -from onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_TOKEN_COUNT +from onyx.deep_research.dr_mock_tools import ( + THINK_TOOL_NAME, + THINK_TOOL_RESPONSE_MESSAGE, + THINK_TOOL_RESPONSE_TOKEN_COUNT, +) from onyx.deep_research.utils import create_think_tool_token_processor -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMUserIdentity +from onyx.llm.interfaces import LLM, LLMUserIdentity from onyx.llm.model_capabilities import model_is_reasoning_model -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import ToolChoiceOptions -from onyx.prompts.coding_agent.coding_agent import CODING_AGENT_FINAL_ANSWER_PROMPT -from onyx.prompts.coding_agent.coding_agent import CODING_AGENT_PROMPT -from onyx.prompts.coding_agent.coding_agent import CODING_AGENT_PROMPT_REASONING -from onyx.prompts.coding_agent.coding_agent import MAX_CODING_AGENT_CYCLES -from onyx.prompts.coding_agent.coding_agent import USER_FINAL_ANSWER_QUERY +from onyx.llm.models import ReasoningEffort, ToolChoiceOptions +from onyx.prompts.coding_agent.coding_agent import ( + CODING_AGENT_FINAL_ANSWER_PROMPT, + CODING_AGENT_PROMPT, + CODING_AGENT_PROMPT_REASONING, + MAX_CODING_AGENT_CYCLES, + USER_FINAL_ANSWER_QUERY, +) from onyx.prompts.prompt_utils import get_current_llm_day_time from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import CodingAgentFinal -from onyx.server.query_and_chat.streaming_models import CodingAgentThinkingDelta -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import PacketException -from onyx.server.query_and_chat.streaming_models import StreamingType +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, + CodingAgentFinal, + CodingAgentThinkingDelta, + Packet, + PacketException, + StreamingType, +) from onyx.tools.models import ToolCallKickoff -from onyx.tools.tool_implementations.bash.bash_tool import BashTool -from onyx.tools.tool_implementations.bash.bash_tool import BashToolOverrideKwargs +from onyx.tools.tool_implementations.bash.bash_tool import ( + BashTool, + BashToolOverrideKwargs, +) from onyx.tools.tool_implementations.python.code_interpreter_client import ( CodeInterpreterClient, ) diff --git a/backend/onyx/tools/fake_tools/research_agent.py b/backend/onyx/tools/fake_tools/research_agent.py index 19c383d9c7..205666d707 100644 --- a/backend/onyx/tools/fake_tools/research_agent.py +++ b/backend/onyx/tools/fake_tools/research_agent.py @@ -1,76 +1,82 @@ import queue import time from collections.abc import Callable -from typing import Any -from typing import cast +from typing import Any, cast from onyx.chat.chat_state import ChatStateContainer from onyx.chat.chat_utils import create_tool_call_failure_messages -from onyx.chat.citation_processor import CitationMapping -from onyx.chat.citation_processor import CitationMode -from onyx.chat.citation_processor import DynamicCitationProcessor -from onyx.chat.citation_utils import collapse_citations -from onyx.chat.citation_utils import update_citation_processor_from_tool_response +from onyx.chat.citation_processor import ( + CitationMapping, + CitationMode, + DynamicCitationProcessor, +) +from onyx.chat.citation_utils import ( + collapse_citations, + update_citation_processor_from_tool_response, +) from onyx.chat.emitter import Emitter from onyx.chat.llm_loop import construct_message_history -from onyx.chat.llm_step import run_llm_step -from onyx.chat.llm_step import run_llm_step_pkt_generator -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import LlmStepResult -from onyx.chat.models import ToolCallSimple +from onyx.chat.llm_step import run_llm_step, run_llm_step_pkt_generator +from onyx.chat.models import ChatMessageSimple, LlmStepResult, ToolCallSimple from onyx.configs.chat_configs import DR_REPORT_LLM_TIMEOUT_S from onyx.configs.constants import MessageType from onyx.context.search.models import SearchDocsResponse from onyx.deep_research.dr_mock_tools import ( get_research_agent_additional_tool_definitions, + RESEARCH_AGENT_TASK_KEY, + THINK_TOOL_RESPONSE_MESSAGE, + THINK_TOOL_RESPONSE_TOKEN_COUNT, ) -from onyx.deep_research.dr_mock_tools import RESEARCH_AGENT_TASK_KEY -from onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_MESSAGE -from onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_TOKEN_COUNT -from onyx.deep_research.models import CombinedResearchAgentCallResult -from onyx.deep_research.models import ResearchAgentCallResult -from onyx.deep_research.utils import check_special_tool_calls -from onyx.deep_research.utils import create_think_tool_token_processor -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMUserIdentity -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import ToolChoiceOptions -from onyx.prompts.deep_research.dr_tool_prompts import OPEN_URLS_TOOL_DESCRIPTION +from onyx.deep_research.models import ( + CombinedResearchAgentCallResult, + ResearchAgentCallResult, +) +from onyx.deep_research.utils import ( + check_special_tool_calls, + create_think_tool_token_processor, +) +from onyx.llm.interfaces import LLM, LLMUserIdentity +from onyx.llm.models import ReasoningEffort, ToolChoiceOptions from onyx.prompts.deep_research.dr_tool_prompts import ( + OPEN_URLS_TOOL_DESCRIPTION, OPEN_URLS_TOOL_DESCRIPTION_REASONING, + WEB_SEARCH_TOOL_DESCRIPTION, +) +from onyx.prompts.deep_research.research_agent import ( + MAX_RESEARCH_CYCLES, + OPEN_URL_REMINDER_RESEARCH_AGENT, + RESEARCH_AGENT_PROMPT, + RESEARCH_AGENT_PROMPT_REASONING, + RESEARCH_REPORT_PROMPT, + USER_REPORT_QUERY, ) -from onyx.prompts.deep_research.dr_tool_prompts import WEB_SEARCH_TOOL_DESCRIPTION -from onyx.prompts.deep_research.research_agent import MAX_RESEARCH_CYCLES -from onyx.prompts.deep_research.research_agent import OPEN_URL_REMINDER_RESEARCH_AGENT -from onyx.prompts.deep_research.research_agent import RESEARCH_AGENT_PROMPT -from onyx.prompts.deep_research.research_agent import RESEARCH_AGENT_PROMPT_REASONING -from onyx.prompts.deep_research.research_agent import RESEARCH_REPORT_PROMPT -from onyx.prompts.deep_research.research_agent import USER_REPORT_QUERY from onyx.prompts.prompt_utils import get_current_llm_day_time from onyx.prompts.tool_prompts import INTERNAL_SEARCH_GUIDANCE from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import IntermediateReportCitedDocs -from onyx.server.query_and_chat.streaming_models import IntermediateReportDelta -from onyx.server.query_and_chat.streaming_models import IntermediateReportStart -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import PacketException -from onyx.server.query_and_chat.streaming_models import ResearchAgentStart -from onyx.server.query_and_chat.streaming_models import SectionEnd -from onyx.server.query_and_chat.streaming_models import StreamingType +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, + IntermediateReportCitedDocs, + IntermediateReportDelta, + IntermediateReportStart, + Packet, + PacketException, + ResearchAgentStart, + SectionEnd, + StreamingType, +) from onyx.tools.interface import Tool -from onyx.tools.models import ToolCallInfo -from onyx.tools.models import ToolCallKickoff -from onyx.tools.models import ToolResponse +from onyx.tools.models import ToolCallInfo, ToolCallKickoff, ToolResponse from onyx.tools.tool_implementations.open_url.open_url_tool import OpenURLTool from onyx.tools.tool_implementations.search.search_tool import SearchTool from onyx.tools.tool_implementations.web_search.utils import extract_url_snippet_map from onyx.tools.tool_implementations.web_search.web_search_tool import WebSearchTool from onyx.tools.tool_runner import run_tool_calls -from onyx.tools.utils import compute_all_tool_tokens -from onyx.tools.utils import compute_tool_definition_tokens -from onyx.tools.utils import generate_tools_description +from onyx.tools.utils import ( + compute_all_tool_tokens, + compute_tool_definition_tokens, + generate_tools_description, +) from onyx.tracing.framework.create import function_span from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel @@ -724,12 +730,10 @@ if __name__ == "__main__": from uuid import uuid4 from onyx.chat.chat_state import ChatStateContainer - from onyx.db.engine.sql_engine import get_session_with_current_tenant - from onyx.db.engine.sql_engine import SqlEngine + from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from onyx.db.models import User from onyx.db.persona import get_default_behavior_persona - from onyx.llm.factory import get_default_llm - from onyx.llm.factory import get_llm_token_counter + from onyx.llm.factory import get_default_llm, get_llm_token_counter from onyx.llm.model_capabilities import model_is_reasoning_model from onyx.server.query_and_chat.placement import Placement from onyx.tools.models import ToolCallKickoff diff --git a/backend/onyx/tools/interface.py b/backend/onyx/tools/interface.py index 0b054bfa2b..55267b282b 100644 --- a/backend/onyx/tools/interface.py +++ b/backend/onyx/tools/interface.py @@ -1,9 +1,7 @@ from __future__ import annotations import abc -from typing import Any -from typing import Generic -from typing import TypeVar +from typing import Any, Generic, TypeVar from sqlalchemy.orm import Session diff --git a/backend/onyx/tools/models.py b/backend/onyx/tools/models.py index cb7fec2d1d..a3bd2da77d 100644 --- a/backend/onyx/tools/models.py +++ b/backend/onyx/tools/models.py @@ -2,27 +2,25 @@ from __future__ import annotations import json from enum import Enum -from typing import Any -from typing import Callable -from typing import Literal +from typing import Any, Callable, Literal from uuid import UUID -from pydantic import BaseModel -from pydantic import ConfigDict -from pydantic import model_validator +from pydantic import BaseModel, ConfigDict, model_validator from onyx.chat.emitter import Emitter -from onyx.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT -from onyx.configs.chat_configs import NUM_RETURNED_HITS +from onyx.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT, NUM_RETURNED_HITS from onyx.configs.constants import MessageType -from onyx.context.search.models import SearchDoc -from onyx.context.search.models import SearchDocsResponse +from onyx.context.search.models import SearchDoc, SearchDocsResponse from onyx.db.memory import UserMemoryContext -from onyx.file_store.models import install_lazy_content_loader -from onyx.file_store.models import maybe_materialize_lazy_content +from onyx.file_store.models import ( + install_lazy_content_loader, + maybe_materialize_lazy_content, +) from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import CustomToolErrorInfo -from onyx.server.query_and_chat.streaming_models import GeneratedImage +from onyx.server.query_and_chat.streaming_models import ( + CustomToolErrorInfo, + GeneratedImage, +) from onyx.tools.tool_implementations.images.models import FinalImageGenerationResponse from onyx.tools.tool_implementations.memory.models import MemoryToolResponse diff --git a/backend/onyx/tools/tool_constructor.py b/backend/onyx/tools/tool_constructor.py index cd8004e65b..030a2578d8 100644 --- a/backend/onyx/tools/tool_constructor.py +++ b/backend/onyx/tools/tool_constructor.py @@ -9,27 +9,25 @@ from onyx.auth.oauth_token_manager import OAuthTokenManager from onyx.chat.emitter import Emitter from onyx.configs.app_configs import DISABLE_VECTOR_DB from onyx.configs.model_configs import GEN_AI_TEMPERATURE -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import PersonaSearchInfo +from onyx.context.search.models import BaseFilters, PersonaSearchInfo from onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none -from onyx.db.mcp import get_all_mcp_tools_for_server -from onyx.db.mcp import get_mcp_server_by_id -from onyx.db.mcp import MCPCredentialsError -from onyx.db.mcp import resolve_mcp_credentials -from onyx.db.models import Persona -from onyx.db.models import User +from onyx.db.mcp import ( + get_all_mcp_tools_for_server, + get_mcp_server_by_id, + MCPCredentialsError, + resolve_mcp_credentials, +) +from onyx.db.models import Persona, User from onyx.db.oauth_config import get_oauth_config from onyx.db.search_settings import get_current_search_settings from onyx.db.tools import get_builtin_tool from onyx.document_index.factory import get_default_document_index from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMConfig +from onyx.llm.interfaces import LLM, LLMConfig from onyx.onyxbot.slack.models import SlackContext from onyx.tools.built_in_tools import get_built_in_tool_by_id from onyx.tools.interface import Tool -from onyx.tools.models import DynamicSchemaInfo -from onyx.tools.models import SearchToolUsage +from onyx.tools.models import DynamicSchemaInfo, SearchToolUsage from onyx.tools.tool_implementations.coding_agent.coding_agent_tool import ( CodingAgentTool, ) diff --git a/backend/onyx/tools/tool_implementations/bash/bash_tool.py b/backend/onyx/tools/tool_implementations/bash/bash_tool.py index 87561d2766..768c70006c 100644 --- a/backend/onyx/tools/tool_implementations/bash/bash_tool.py +++ b/backend/onyx/tools/tool_implementations/bash/bash_tool.py @@ -1,22 +1,24 @@ from typing import Any -from pydantic import BaseModel -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from sqlalchemy.orm import Session from typing_extensions import override from onyx.chat.emitter import Emitter -from onyx.configs.app_configs import CODE_INTERPRETER_BASE_URL -from onyx.configs.app_configs import CODE_INTERPRETER_DEFAULT_TIMEOUT_MS -from onyx.configs.app_configs import CODE_INTERPRETER_MAX_OUTPUT_LENGTH +from onyx.configs.app_configs import ( + CODE_INTERPRETER_BASE_URL, + CODE_INTERPRETER_DEFAULT_TIMEOUT_MS, + CODE_INTERPRETER_MAX_OUTPUT_LENGTH, +) from onyx.db.code_interpreter import fetch_code_interpreter_server from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import BashToolDelta -from onyx.server.query_and_chat.streaming_models import BashToolStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + BashToolDelta, + BashToolStart, + Packet, +) from onyx.tools.interface import Tool -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse +from onyx.tools.models import ToolCallException, ToolResponse from onyx.tools.tool_implementations.python.code_interpreter_client import ( CodeInterpreterClient, ) diff --git a/backend/onyx/tools/tool_implementations/coding_agent/coding_agent_tool.py b/backend/onyx/tools/tool_implementations/coding_agent/coding_agent_tool.py index 277647f6c5..94aa9a7b9e 100644 --- a/backend/onyx/tools/tool_implementations/coding_agent/coding_agent_tool.py +++ b/backend/onyx/tools/tool_implementations/coding_agent/coding_agent_tool.py @@ -1,5 +1,4 @@ -from typing import Any -from typing import cast +from typing import Any, cast from uuid import uuid4 from pydantic import BaseModel @@ -7,18 +6,17 @@ from sqlalchemy.orm import Session from typing_extensions import override from onyx.chat.emitter import Emitter -from onyx.coding_agent.mock_tools import CODING_AGENT_QUERY_KEY -from onyx.coding_agent.mock_tools import CODING_AGENT_REPO_KEY -from onyx.coding_agent.mock_tools import CODING_AGENT_TOOL_NAME +from onyx.coding_agent.mock_tools import ( + CODING_AGENT_QUERY_KEY, + CODING_AGENT_REPO_KEY, + CODING_AGENT_TOOL_NAME, +) from onyx.llm.factory import get_llm_token_counter from onyx.llm.interfaces import LLM from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import CodingAgentStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import CodingAgentStart, Packet from onyx.tools.interface import Tool -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolCallKickoff -from onyx.tools.models import ToolResponse +from onyx.tools.models import ToolCallException, ToolCallKickoff, ToolResponse from onyx.tools.tool_implementations.bash.bash_tool import BashTool from onyx.utils.logger import setup_logger diff --git a/backend/onyx/tools/tool_implementations/custom/custom_tool.py b/backend/onyx/tools/tool_implementations/custom/custom_tool.py index 854f918418..805e8e6363 100644 --- a/backend/onyx/tools/tool_implementations/custom/custom_tool.py +++ b/backend/onyx/tools/tool_implementations/custom/custom_tool.py @@ -2,11 +2,8 @@ import csv import json import queue import uuid -from io import BytesIO -from io import StringIO -from typing import Any -from typing import Dict -from typing import List +from io import BytesIO, StringIO +from typing import Any, Dict, List import requests from requests import JSONDecodeError @@ -15,32 +12,33 @@ from onyx.chat.emitter import Emitter from onyx.configs.constants import FileOrigin from onyx.file_store.file_store import get_default_file_store from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import CustomToolArgs -from onyx.server.query_and_chat.streaming_models import CustomToolDelta -from onyx.server.query_and_chat.streaming_models import CustomToolErrorInfo -from onyx.server.query_and_chat.streaming_models import CustomToolStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + CustomToolArgs, + CustomToolDelta, + CustomToolErrorInfo, + CustomToolStart, + Packet, +) from onyx.tools.interface import Tool -from onyx.tools.models import CHAT_SESSION_ID_PLACEHOLDER -from onyx.tools.models import CustomToolCallSummary -from onyx.tools.models import CustomToolUserFileSnapshot -from onyx.tools.models import DynamicSchemaInfo -from onyx.tools.models import MESSAGE_ID_PLACEHOLDER -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse -from onyx.tools.models import USER_EMAIL_PLACEHOLDER -from onyx.tools.models import USER_ID_PLACEHOLDER -from onyx.tools.tool_implementations.custom.openapi_parsing import MethodSpec -from onyx.tools.tool_implementations.custom.openapi_parsing import ( - openapi_to_method_specs, +from onyx.tools.models import ( + CHAT_SESSION_ID_PLACEHOLDER, + CustomToolCallSummary, + CustomToolUserFileSnapshot, + DynamicSchemaInfo, + MESSAGE_ID_PLACEHOLDER, + ToolCallException, + ToolResponse, + USER_EMAIL_PLACEHOLDER, + USER_ID_PLACEHOLDER, ) -from onyx.tools.tool_implementations.custom.openapi_parsing import openapi_to_url -from onyx.tools.tool_implementations.custom.openapi_parsing import REQUEST_BODY from onyx.tools.tool_implementations.custom.openapi_parsing import ( + MethodSpec, + openapi_to_method_specs, + openapi_to_url, + REQUEST_BODY, validate_openapi_schema, ) -from onyx.utils.headers import header_list_to_header_dict -from onyx.utils.headers import HeaderItemDict +from onyx.utils.headers import header_list_to_header_dict, HeaderItemDict from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/custom/openapi_parsing.py b/backend/onyx/tools/tool_implementations/custom/openapi_parsing.py index 87cce8741a..84fc573890 100644 --- a/backend/onyx/tools/tool_implementations/custom/openapi_parsing.py +++ b/backend/onyx/tools/tool_implementations/custom/openapi_parsing.py @@ -1,5 +1,4 @@ -from typing import Any -from typing import cast +from typing import Any, cast from pydantic import BaseModel diff --git a/backend/onyx/tools/tool_implementations/file_reader/file_reader_tool.py b/backend/onyx/tools/tool_implementations/file_reader/file_reader_tool.py index c992119d1c..996c1ff761 100644 --- a/backend/onyx/tools/tool_implementations/file_reader/file_reader_tool.py +++ b/backend/onyx/tools/tool_implementations/file_reader/file_reader_tool.py @@ -1,7 +1,6 @@ import io import json -from typing import Any -from typing import cast +from typing import Any, cast from uuid import UUID from sqlalchemy.orm import Session @@ -11,17 +10,16 @@ from onyx.chat.emitter import Emitter from onyx.configs.app_configs import DISABLE_VECTOR_DB from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.file_processing.extract_file_text import extract_file_text -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import InMemoryChatFile -from onyx.file_store.utils import load_chat_file_by_id -from onyx.file_store.utils import load_user_file +from onyx.file_store.models import ChatFileType, InMemoryChatFile +from onyx.file_store.utils import load_chat_file_by_id, load_user_file from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import FileReaderResult -from onyx.server.query_and_chat.streaming_models import FileReaderStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + FileReaderResult, + FileReaderStart, + Packet, +) from onyx.tools.interface import Tool -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse +from onyx.tools.models import ToolCallException, ToolResponse from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/images/image_generation_tool.py b/backend/onyx/tools/tool_implementations/images/image_generation_tool.py index 6f7c7948f0..51265c5509 100644 --- a/backend/onyx/tools/tool_implementations/images/image_generation_tool.py +++ b/backend/onyx/tools/tool_implementations/images/image_generation_tool.py @@ -1,38 +1,44 @@ import json import threading -from typing import Any -from typing import cast +from typing import Any, cast import requests from sqlalchemy.orm import Session from typing_extensions import override from onyx.chat.emitter import Emitter -from onyx.configs.app_configs import IMAGE_MODEL_NAME -from onyx.configs.app_configs import IMAGE_MODEL_PROVIDER +from onyx.configs.app_configs import IMAGE_MODEL_NAME, IMAGE_MODEL_PROVIDER from onyx.file_store.models import ChatFileType -from onyx.file_store.utils import build_frontend_file_url -from onyx.file_store.utils import load_chat_file_by_id -from onyx.file_store.utils import save_files +from onyx.file_store.utils import ( + build_frontend_file_url, + load_chat_file_by_id, + save_files, +) from onyx.image_gen.factory import get_image_generation_provider -from onyx.image_gen.generation import generate_images_with_provider -from onyx.image_gen.generation import is_image_generation_configured -from onyx.image_gen.generation import resolve_image_size -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.image_gen.interfaces import ImageShape -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.generation import ( + generate_images_with_provider, + is_image_generation_configured, + resolve_image_size, +) +from onyx.image_gen.interfaces import ( + ImageGenerationProviderCredentials, + ImageShape, + ReferenceImage, +) from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import GeneratedImage -from onyx.server.query_and_chat.streaming_models import ImageGenerationFinal -from onyx.server.query_and_chat.streaming_models import ImageGenerationToolHeartbeat -from onyx.server.query_and_chat.streaming_models import ImageGenerationToolStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + GeneratedImage, + ImageGenerationFinal, + ImageGenerationToolHeartbeat, + ImageGenerationToolStart, + Packet, +) from onyx.tools.interface import Tool -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolExecutionException -from onyx.tools.models import ToolResponse -from onyx.tools.tool_implementations.images.models import FinalImageGenerationResponse -from onyx.tools.tool_implementations.images.models import ImageGenerationResponse +from onyx.tools.models import ToolCallException, ToolExecutionException, ToolResponse +from onyx.tools.tool_implementations.images.models import ( + FinalImageGenerationResponse, + ImageGenerationResponse, +) from onyx.utils.b64 import get_image_type_from_bytes from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel diff --git a/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py b/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py index 87f020135c..6f4ce4845a 100644 --- a/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py +++ b/backend/onyx/tools/tool_implementations/mcp/mcp_tool.py @@ -5,25 +5,26 @@ from typing import Any from mcp.client.auth import OAuthClientProvider from onyx.chat.emitter import Emitter -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPTransport +from onyx.db.enums import MCPAuthenticationType, MCPTransport from onyx.db.mcp import ResolvedMCPCredentials -from onyx.db.models import MCPConnectionConfig -from onyx.db.models import MCPServer +from onyx.db.models import MCPConnectionConfig, MCPServer from onyx.server.features.mcp.client import call_mcp_tool from onyx.server.features.mcp.models import DENYLISTED_MCP_HEADERS -from onyx.server.features.mcp.oauth import make_oauth_provider -from onyx.server.features.mcp.oauth import refresh_mcp_oauth_token_if_expired -from onyx.server.features.mcp.oauth import UNUSED_RETURN_PATH +from onyx.server.features.mcp.oauth import ( + make_oauth_provider, + refresh_mcp_oauth_token_if_expired, + UNUSED_RETURN_PATH, +) from onyx.server.metrics.mcp_client import record_mcp_client_tool_outcome from onyx.server.metrics.mcp_common import MCPToolCallStatus from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import CustomToolDelta -from onyx.server.query_and_chat.streaming_models import CustomToolStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + CustomToolDelta, + CustomToolStart, + Packet, +) from onyx.tools.interface import Tool -from onyx.tools.models import CustomToolCallSummary -from onyx.tools.models import ToolResponse +from onyx.tools.models import CustomToolCallSummary, ToolResponse from onyx.tools.tool_name import sanitize_tool_name from onyx.utils.logger import setup_logger diff --git a/backend/onyx/tools/tool_implementations/memory/memory_tool.py b/backend/onyx/tools/tool_implementations/memory/memory_tool.py index 0b8af0c6c9..df37d47206 100644 --- a/backend/onyx/tools/tool_implementations/memory/memory_tool.py +++ b/backend/onyx/tools/tool_implementations/memory/memory_tool.py @@ -6,9 +6,7 @@ The memories are passed in via override_kwargs which contains the current list o memories that exist for the user. """ -from typing import Any -from typing import cast -from typing import Literal +from typing import Any, cast, Literal from pydantic import BaseModel from typing_extensions import override @@ -17,13 +15,13 @@ from onyx.chat.emitter import Emitter from onyx.llm.interfaces import LLM from onyx.secondary_llm_flows.memory_update import process_memory_update from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import MemoryToolDelta -from onyx.server.query_and_chat.streaming_models import MemoryToolStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + MemoryToolDelta, + MemoryToolStart, + Packet, +) from onyx.tools.interface import Tool -from onyx.tools.models import ChatMinimalTextMessage -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse +from onyx.tools.models import ChatMinimalTextMessage, ToolCallException, ToolResponse from onyx.tools.tool_implementations.memory.models import MemoryToolResponse from onyx.utils.logger import setup_logger diff --git a/backend/onyx/tools/tool_implementations/open_url/firecrawl.py b/backend/onyx/tools/tool_implementations/open_url/firecrawl.py index ae07292a45..4bebb5e888 100644 --- a/backend/onyx/tools/tool_implementations/open_url/firecrawl.py +++ b/backend/onyx/tools/tool_implementations/open_url/firecrawl.py @@ -9,8 +9,10 @@ from typing import Any import requests from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.tools.tool_implementations.open_url.models import WebContent -from onyx.tools.tool_implementations.open_url.models import WebContentProvider +from onyx.tools.tool_implementations.open_url.models import ( + WebContent, + WebContentProvider, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/open_url/models.py b/backend/onyx/tools/tool_implementations/open_url/models.py index 13692f32b4..dbdde336ea 100644 --- a/backend/onyx/tools/tool_implementations/open_url/models.py +++ b/backend/onyx/tools/tool_implementations/open_url/models.py @@ -1,10 +1,8 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import Sequence from datetime import datetime -from pydantic import BaseModel -from pydantic import field_validator +from pydantic import BaseModel, field_validator from onyx.utils.url import normalize_url diff --git a/backend/onyx/tools/tool_implementations/open_url/onyx_web_crawler.py b/backend/onyx/tools/tool_implementations/open_url/onyx_web_crawler.py index 0fd910390c..e61d0d5834 100644 --- a/backend/onyx/tools/tool_implementations/open_url/onyx_web_crawler.py +++ b/backend/onyx/tools/tool_implementations/open_url/onyx_web_crawler.py @@ -6,23 +6,27 @@ from concurrent.futures import ThreadPoolExecutor import requests from onyx.configs.app_configs import OPEN_URL_PLAYWRIGHT_FALLBACK_ENABLED -from onyx.file_processing.html_utils import ParsedHTML -from onyx.file_processing.html_utils import web_html_cleanup +from onyx.file_processing.html_utils import ParsedHTML, web_html_cleanup from onyx.server.security.models import outbound_allow_private_network from onyx.server.security.store import get_security_settings -from onyx.tools.tool_implementations.open_url.models import WebContent -from onyx.tools.tool_implementations.open_url.models import WebContentProvider +from onyx.tools.tool_implementations.open_url.models import ( + WebContent, + WebContentProvider, +) from onyx.utils.logger import setup_logger -from onyx.utils.playwright_fetch import fetch_rendered_html -from onyx.utils.playwright_fetch import looks_like_cloudflare_challenge -from onyx.utils.playwright_fetch import RenderedPage -from onyx.utils.url import ssrf_safe_get -from onyx.utils.url import SSRFException -from onyx.utils.web_content import decode_html_bytes -from onyx.utils.web_content import extract_pdf_text -from onyx.utils.web_content import is_pdf_resource -from onyx.utils.web_content import title_from_pdf_metadata -from onyx.utils.web_content import title_from_url +from onyx.utils.playwright_fetch import ( + fetch_rendered_html, + looks_like_cloudflare_challenge, + RenderedPage, +) +from onyx.utils.url import ssrf_safe_get, SSRFException +from onyx.utils.web_content import ( + decode_html_bytes, + extract_pdf_text, + is_pdf_resource, + title_from_pdf_metadata, + title_from_url, +) logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/open_url/open_url_tool.py b/backend/onyx/tools/tool_implementations/open_url/open_url_tool.py index 8ea1a4405f..62c8b28f22 100644 --- a/backend/onyx/tools/tool_implementations/open_url/open_url_tool.py +++ b/backend/onyx/tools/tool_implementations/open_url/open_url_tool.py @@ -8,35 +8,37 @@ from typing_extensions import override from onyx.chat.emitter import Emitter from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceSection -from onyx.context.search.models import SearchDocsResponse +from onyx.context.search.models import ( + IndexFilters, + InferenceSection, + SearchDocsResponse, +) from onyx.context.search.preprocessing.access_filters import ( build_access_filters_for_user, ) -from onyx.context.search.utils import convert_inference_sections_to_search_docs -from onyx.context.search.utils import inference_section_from_chunks -from onyx.db.document import fetch_document_ids_by_links -from onyx.db.document import filter_existing_document_ids +from onyx.context.search.utils import ( + convert_inference_sections_to_search_docs, + inference_section_from_chunks, +) +from onyx.db.document import fetch_document_ids_by_links, filter_existing_document_ids from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import User -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentSectionRequest +from onyx.document_index.interfaces_new import DocumentIndex, DocumentSectionRequest from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import OpenUrlDocuments -from onyx.server.query_and_chat.streaming_models import OpenUrlStart -from onyx.server.query_and_chat.streaming_models import OpenUrlUrls -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.streaming_models import ( + OpenUrlDocuments, + OpenUrlStart, + OpenUrlUrls, + Packet, +) from onyx.tools.interface import Tool -from onyx.tools.models import OpenURLToolOverrideKwargs -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse -from onyx.tools.tool_implementations.open_url.models import FailedFetch -from onyx.tools.tool_implementations.open_url.models import WebContentProvider -from onyx.tools.tool_implementations.open_url.url_normalization import ( - _default_url_normalizer, +from onyx.tools.models import OpenURLToolOverrideKwargs, ToolCallException, ToolResponse +from onyx.tools.tool_implementations.open_url.models import ( + FailedFetch, + WebContentProvider, ) from onyx.tools.tool_implementations.open_url.url_normalization import ( + _default_url_normalizer, normalize_url_candidates, ) from onyx.tools.tool_implementations.open_url.utils import ( @@ -47,8 +49,8 @@ from onyx.tools.tool_implementations.web_search.providers import ( ) from onyx.tools.tool_implementations.web_search.utils import ( inference_section_from_internet_page_scrape, + MAX_CHARS_PER_URL, ) -from onyx.tools.tool_implementations.web_search.utils import MAX_CHARS_PER_URL from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel from onyx.utils.url import normalize_url as normalize_web_content_url diff --git a/backend/onyx/tools/tool_implementations/open_url/snippet_matcher.py b/backend/onyx/tools/tool_implementations/open_url/snippet_matcher.py index 0a320fbd25..4d9e38444d 100644 --- a/backend/onyx/tools/tool_implementations/open_url/snippet_matcher.py +++ b/backend/onyx/tools/tool_implementations/open_url/snippet_matcher.py @@ -1,11 +1,9 @@ import unicodedata from pydantic import BaseModel -from rapidfuzz import fuzz -from rapidfuzz import utils +from rapidfuzz import fuzz, utils -from onyx.utils.text_processing import is_zero_width_char -from onyx.utils.text_processing import normalize_char +from onyx.utils.text_processing import is_zero_width_char, normalize_char class SnippetMatchResult(BaseModel): diff --git a/backend/onyx/tools/tool_implementations/open_url/tavily.py b/backend/onyx/tools/tool_implementations/open_url/tavily.py index fa20443d9a..2754bd6102 100644 --- a/backend/onyx/tools/tool_implementations/open_url/tavily.py +++ b/backend/onyx/tools/tool_implementations/open_url/tavily.py @@ -5,8 +5,10 @@ from typing import Any import requests -from onyx.tools.tool_implementations.open_url.models import WebContent -from onyx.tools.tool_implementations.open_url.models import WebContentProvider +from onyx.tools.tool_implementations.open_url.models import ( + WebContent, + WebContentProvider, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/open_url/url_normalization.py b/backend/onyx/tools/tool_implementations/open_url/url_normalization.py index 0ba150f3ca..7a6600197d 100644 --- a/backend/onyx/tools/tool_implementations/open_url/url_normalization.py +++ b/backend/onyx/tools/tool_implementations/open_url/url_normalization.py @@ -9,8 +9,7 @@ Usage: # Returns: "https://docs.google.com/document/d/123" """ -from urllib.parse import urlparse -from urllib.parse import urlunparse +from urllib.parse import urlparse, urlunparse from onyx.configs.constants import DocumentSource from onyx.connectors.factory import identify_connector_class diff --git a/backend/onyx/tools/tool_implementations/python/code_interpreter_client.py b/backend/onyx/tools/tool_implementations/python/code_interpreter_client.py index 30fdea6ab6..f2cfa74be6 100644 --- a/backend/onyx/tools/tool_implementations/python/code_interpreter_client.py +++ b/backend/onyx/tools/tool_implementations/python/code_interpreter_client.py @@ -3,16 +3,9 @@ from __future__ import annotations import json import re import time -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from functools import wraps -from typing import Any -from typing import Concatenate -from typing import Literal -from typing import ParamSpec -from typing import TypedDict -from typing import TypeVar -from typing import Union +from typing import Any, Concatenate, Literal, ParamSpec, TypedDict, TypeVar, Union import requests from pydantic import BaseModel diff --git a/backend/onyx/tools/tool_implementations/python/python_tool.py b/backend/onyx/tools/tool_implementations/python/python_tool.py index 98c3af71ce..1198c3adbd 100644 --- a/backend/onyx/tools/tool_implementations/python/python_tool.py +++ b/backend/onyx/tools/tool_implementations/python/python_tool.py @@ -3,48 +3,45 @@ import mimetypes import os import re from io import BytesIO -from typing import Any -from typing import cast +from typing import Any, cast -from pydantic import BaseModel -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from sqlalchemy.orm import Session from typing_extensions import override from onyx.chat.emitter import Emitter -from onyx.configs.app_configs import CODE_INTERPRETER_BASE_URL -from onyx.configs.app_configs import CODE_INTERPRETER_DEFAULT_TIMEOUT_MS -from onyx.configs.app_configs import CODE_INTERPRETER_MAX_OUTPUT_LENGTH -from onyx.configs.app_configs import CODE_INTERPRETER_MAX_STAGED_BYTES -from onyx.configs.app_configs import CODE_INTERPRETER_MAX_STAGED_FILES -from onyx.configs.app_configs import CODE_INTERPRETER_STAGING_CONCURRENCY +from onyx.configs.app_configs import ( + CODE_INTERPRETER_BASE_URL, + CODE_INTERPRETER_DEFAULT_TIMEOUT_MS, + CODE_INTERPRETER_MAX_OUTPUT_LENGTH, + CODE_INTERPRETER_MAX_STAGED_BYTES, + CODE_INTERPRETER_MAX_STAGED_FILES, + CODE_INTERPRETER_STAGING_CONCURRENCY, +) from onyx.configs.constants import FileOrigin from onyx.db.code_interpreter import fetch_code_interpreter_server -from onyx.file_store.utils import build_full_frontend_file_url -from onyx.file_store.utils import get_default_file_store +from onyx.file_store.utils import build_full_frontend_file_url, get_default_file_store from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import PythonToolDelta -from onyx.server.query_and_chat.streaming_models import PythonToolStart +from onyx.server.query_and_chat.streaming_models import ( + Packet, + PythonToolDelta, + PythonToolStart, +) from onyx.tools.interface import Tool -from onyx.tools.models import ChatFile -from onyx.tools.models import LlmPythonExecutionResult -from onyx.tools.models import PythonExecutionFile -from onyx.tools.models import PythonToolOverrideKwargs -from onyx.tools.models import PythonToolRichResponse -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse +from onyx.tools.models import ( + ChatFile, + LlmPythonExecutionResult, + PythonExecutionFile, + PythonToolOverrideKwargs, + PythonToolRichResponse, + ToolCallException, + ToolResponse, +) from onyx.tools.tool_implementations.python.code_interpreter_client import ( CodeInterpreterClient, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import FileInput -from onyx.tools.tool_implementations.python.code_interpreter_client import ( + FileInput, StreamErrorEvent, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import ( StreamOutputEvent, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import ( StreamResultEvent, ) from onyx.tools.tool_implementations.utils import truncate_output as _truncate_output diff --git a/backend/onyx/tools/tool_implementations/search/search_tool.py b/backend/onyx/tools/tool_implementations/search/search_tool.py index 2a4870d555..e42cd40838 100644 --- a/backend/onyx/tools/tool_implementations/search/search_tool.py +++ b/backend/onyx/tools/tool_implementations/search/search_tool.py @@ -36,88 +36,94 @@ refer to by using matching keywords to other parts of the prompt and reminders. import time from collections.abc import Callable -from typing import Any -from typing import cast +from typing import Any, cast from pydantic import BaseModel from sqlalchemy.orm import Session from onyx.chat.emitter import Emitter from onyx.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FederatedConnectorSource +from onyx.configs.constants import DocumentSource, FederatedConnectorSource from onyx.context.search.federated.slack_search import slack_retrieval -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import ChunkIndexRequest -from onyx.context.search.models import ChunkSearchRequest -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection -from onyx.context.search.models import PersonaSearchInfo -from onyx.context.search.models import SearchDocsResponse -from onyx.context.search.pipeline import merge_individual_chunks -from onyx.context.search.pipeline import search_pipeline +from onyx.context.search.models import ( + BaseFilters, + ChunkIndexRequest, + ChunkSearchRequest, + IndexFilters, + InferenceChunk, + InferenceSection, + PersonaSearchInfo, + SearchDocsResponse, +) +from onyx.context.search.pipeline import merge_individual_chunks, search_pipeline from onyx.context.search.preprocessing.access_filters import ( build_access_filters_for_user, ) -from onyx.context.search.utils import convert_inference_sections_to_search_docs -from onyx.context.search.utils import populate_file_ids_on_sections -from onyx.db.connector import check_connectors_exist -from onyx.db.connector import check_federated_connectors_exist -from onyx.db.connector import fetch_unique_document_sources +from onyx.context.search.utils import ( + convert_inference_sections_to_search_docs, + populate_file_ids_on_sections, +) +from onyx.db.connector import ( + check_connectors_exist, + check_federated_connectors_exist, + fetch_unique_document_sources, +) from onyx.db.document_set import filter_document_set_names_by_user_access from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.federated import ( get_federated_connector_document_set_mappings_by_document_set_names, + list_federated_connector_oauth_tokens, ) -from onyx.db.federated import list_federated_connector_oauth_tokens -from onyx.db.models import SearchSettings -from onyx.db.models import User +from onyx.db.models import SearchSettings, User from onyx.db.search_settings import get_current_search_settings from onyx.db.slack_bot import fetch_slack_bots from onyx.document_index.interfaces_new import DocumentIndex from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.federated_connectors.federated_retrieval import FederatedRetrievalInfo from onyx.federated_connectors.federated_retrieval import ( + FederatedRetrievalInfo, get_federated_retrieval_functions, ) from onyx.llm.factory import get_llm_token_counter from onyx.llm.interfaces import LLM from onyx.natural_language_processing.search_nlp_models import EmbeddingModel from onyx.onyxbot.slack.models import SlackContext -from onyx.secondary_llm_flows.document_filter import select_chunks_for_relevance -from onyx.secondary_llm_flows.document_filter import select_sections_for_expansion -from onyx.secondary_llm_flows.query_expansion import keyword_query_expansion -from onyx.secondary_llm_flows.query_expansion import semantic_query_rephrase -from onyx.secondary_llm_flows.source_filter import decide_search_scope -from onyx.secondary_llm_flows.source_filter import SearchCycle -from onyx.secondary_llm_flows.time_filter import decide_time_filter -from onyx.secondary_llm_flows.time_filter import TimeFilter +from onyx.secondary_llm_flows.document_filter import ( + select_chunks_for_relevance, + select_sections_for_expansion, +) +from onyx.secondary_llm_flows.query_expansion import ( + keyword_query_expansion, + semantic_query_rephrase, +) +from onyx.secondary_llm_flows.source_filter import decide_search_scope, SearchCycle +from onyx.secondary_llm_flows.time_filter import decide_time_filter, TimeFilter from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import SearchToolDocumentsDelta -from onyx.server.query_and_chat.streaming_models import SearchToolFilterDelta -from onyx.server.query_and_chat.streaming_models import SearchToolQueriesDelta -from onyx.server.query_and_chat.streaming_models import SearchToolStart +from onyx.server.query_and_chat.streaming_models import ( + Packet, + SearchToolDocumentsDelta, + SearchToolFilterDelta, + SearchToolQueriesDelta, + SearchToolStart, +) from onyx.tools.interface import Tool -from onyx.tools.models import ChatMinimalTextMessage -from onyx.tools.models import SearchToolOverrideKwargs -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse -from onyx.tools.tool_implementations.search.constants import KEYWORD_QUERY_HYBRID_ALPHA -from onyx.tools.tool_implementations.search.constants import LLM_KEYWORD_QUERY_WEIGHT -from onyx.tools.tool_implementations.search.constants import LLM_NON_CUSTOM_QUERY_WEIGHT -from onyx.tools.tool_implementations.search.constants import LLM_SEMANTIC_QUERY_WEIGHT -from onyx.tools.tool_implementations.search.constants import MAX_CHUNKS_FOR_RELEVANCE -from onyx.tools.tool_implementations.search.constants import ORIGINAL_QUERY_WEIGHT +from onyx.tools.models import ( + ChatMinimalTextMessage, + SearchToolOverrideKwargs, + ToolCallException, + ToolResponse, +) +from onyx.tools.tool_implementations.search.constants import ( + KEYWORD_QUERY_HYBRID_ALPHA, + LLM_KEYWORD_QUERY_WEIGHT, + LLM_NON_CUSTOM_QUERY_WEIGHT, + LLM_SEMANTIC_QUERY_WEIGHT, + MAX_CHUNKS_FOR_RELEVANCE, + ORIGINAL_QUERY_WEIGHT, +) from onyx.tools.tool_implementations.search.search_utils import ( expand_section_with_context, -) -from onyx.tools.tool_implementations.search.search_utils import ( merge_overlapping_sections, -) -from onyx.tools.tool_implementations.search.search_utils import ( weighted_reciprocal_rank_fusion, ) from onyx.tools.tool_implementations.utils import ( @@ -126,9 +132,11 @@ from onyx.tools.tool_implementations.utils import ( from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel from onyx.utils.timing import log_function_time -from shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT +from shared_configs.configs import ( + DOC_EMBEDDING_CONTEXT_SIZE, + MODEL_SERVER_HOST, + MODEL_SERVER_PORT, +) logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/search/search_utils.py b/backend/onyx/tools/tool_implementations/search/search_utils.py index b91b9503f3..b1df713336 100644 --- a/backend/onyx/tools/tool_implementations/search/search_utils.py +++ b/backend/onyx/tools/tool_implementations/search/search_utils.py @@ -2,21 +2,24 @@ from collections import defaultdict from collections.abc import Callable from typing import TypeVar -from onyx.context.search.models import ContextExpansionType -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection +from onyx.context.search.models import ( + ContextExpansionType, + IndexFilters, + InferenceChunk, + InferenceSection, +) from onyx.context.search.utils import inference_section_from_chunks -from onyx.document_index.interfaces_new import DocumentIndex -from onyx.document_index.interfaces_new import DocumentSectionRequest +from onyx.document_index.interfaces_new import DocumentIndex, DocumentSectionRequest from onyx.document_index.vespa.shared_utils.utils import ( replace_invalid_doc_id_characters, ) from onyx.llm.interfaces import LLM from onyx.prompts.prompt_utils import clean_up_source from onyx.secondary_llm_flows.document_filter import classify_section_relevance -from onyx.tools.tool_implementations.search.constants import FULL_DOC_NUM_CHUNKS_AROUND -from onyx.tools.tool_implementations.search.constants import RRF_K_VALUE +from onyx.tools.tool_implementations.search.constants import ( + FULL_DOC_NUM_CHUNKS_AROUND, + RRF_K_VALUE, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/search_like_tool_utils.py b/backend/onyx/tools/tool_implementations/search_like_tool_utils.py index 76c31d48ea..ec0e05ef76 100644 --- a/backend/onyx/tools/tool_implementations/search_like_tool_utils.py +++ b/backend/onyx/tools/tool_implementations/search_like_tool_utils.py @@ -1,6 +1,4 @@ -from onyx.connectors.models import Document -from onyx.connectors.models import IndexingDocument -from onyx.connectors.models import Section +from onyx.connectors.models import Document, IndexingDocument, Section FINAL_CONTEXT_DOCUMENTS_ID = "final_context_documents" FINAL_SEARCH_QUERIES_ID = "final_search_queries" diff --git a/backend/onyx/tools/tool_implementations/web_search/clients/brave_client.py b/backend/onyx/tools/tool_implementations/web_search/clients/brave_client.py index d377624d9d..7976526585 100644 --- a/backend/onyx/tools/tool_implementations/web_search/clients/brave_client.py +++ b/backend/onyx/tools/tool_implementations/web_search/clients/brave_client.py @@ -5,8 +5,10 @@ from typing import Any import requests from fastapi import HTTPException -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.web_search.models import ( + WebSearchProvider, + WebSearchResult, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/tools/tool_implementations/web_search/clients/exa_client.py b/backend/onyx/tools/tool_implementations/web_search/clients/exa_client.py index 4cd02b56cc..fdcf8990ff 100644 --- a/backend/onyx/tools/tool_implementations/web_search/clients/exa_client.py +++ b/backend/onyx/tools/tool_implementations/web_search/clients/exa_client.py @@ -8,10 +8,14 @@ from exa_py.api import HighlightsContentsOptions from fastapi import HTTPException from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.tools.tool_implementations.open_url.models import WebContent -from onyx.tools.tool_implementations.open_url.models import WebContentProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.open_url.models import ( + WebContent, + WebContentProvider, +) +from onyx.tools.tool_implementations.web_search.models import ( + WebSearchProvider, + WebSearchResult, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/tools/tool_implementations/web_search/clients/google_pse_client.py b/backend/onyx/tools/tool_implementations/web_search/clients/google_pse_client.py index 01e9b6d690..87cffe6013 100644 --- a/backend/onyx/tools/tool_implementations/web_search/clients/google_pse_client.py +++ b/backend/onyx/tools/tool_implementations/web_search/clients/google_pse_client.py @@ -6,8 +6,10 @@ from typing import Any import requests from fastapi import HTTPException -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.web_search.models import ( + WebSearchProvider, + WebSearchResult, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/tools/tool_implementations/web_search/clients/searxng_client.py b/backend/onyx/tools/tool_implementations/web_search/clients/searxng_client.py index aec0374585..f98d4ddadd 100644 --- a/backend/onyx/tools/tool_implementations/web_search/clients/searxng_client.py +++ b/backend/onyx/tools/tool_implementations/web_search/clients/searxng_client.py @@ -1,8 +1,10 @@ import requests from fastapi import HTTPException -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.web_search.models import ( + WebSearchProvider, + WebSearchResult, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/tools/tool_implementations/web_search/clients/serper_client.py b/backend/onyx/tools/tool_implementations/web_search/clients/serper_client.py index 0a870f64fd..55bfc7eb41 100644 --- a/backend/onyx/tools/tool_implementations/web_search/clients/serper_client.py +++ b/backend/onyx/tools/tool_implementations/web_search/clients/serper_client.py @@ -6,10 +6,14 @@ import requests from fastapi import HTTPException from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.tools.tool_implementations.open_url.models import WebContent -from onyx.tools.tool_implementations.open_url.models import WebContentProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.open_url.models import ( + WebContent, + WebContentProvider, +) +from onyx.tools.tool_implementations.web_search.models import ( + WebSearchProvider, + WebSearchResult, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/tools/tool_implementations/web_search/clients/tavily_client.py b/backend/onyx/tools/tool_implementations/web_search/clients/tavily_client.py index a462de0a07..984d5cacdc 100644 --- a/backend/onyx/tools/tool_implementations/web_search/clients/tavily_client.py +++ b/backend/onyx/tools/tool_implementations/web_search/clients/tavily_client.py @@ -5,8 +5,10 @@ from typing import Any import requests from fastapi import HTTPException -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.web_search.models import ( + WebSearchProvider, + WebSearchResult, +) from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder diff --git a/backend/onyx/tools/tool_implementations/web_search/models.py b/backend/onyx/tools/tool_implementations/web_search/models.py index 51302b2cda..42a844dcfe 100644 --- a/backend/onyx/tools/tool_implementations/web_search/models.py +++ b/backend/onyx/tools/tool_implementations/web_search/models.py @@ -3,8 +3,7 @@ from collections.abc import Sequence from datetime import datetime from enum import Enum -from pydantic import BaseModel -from pydantic import field_validator +from pydantic import BaseModel, field_validator from onyx.utils.url import normalize_url diff --git a/backend/onyx/tools/tool_implementations/web_search/providers.py b/backend/onyx/tools/tool_implementations/web_search/providers.py index d58906f95d..eaea8a14f9 100644 --- a/backend/onyx/tools/tool_implementations/web_search/providers.py +++ b/backend/onyx/tools/tool_implementations/web_search/providers.py @@ -1,16 +1,16 @@ from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import InternetSearchProvider -from onyx.db.web_search import fetch_active_web_content_provider -from onyx.db.web_search import fetch_active_web_search_provider +from onyx.db.web_search import ( + fetch_active_web_content_provider, + fetch_active_web_search_provider, +) from onyx.tools.tool_implementations.open_url.firecrawl import FirecrawlClient from onyx.tools.tool_implementations.open_url.models import WebContentProvider from onyx.tools.tool_implementations.open_url.onyx_web_crawler import ( DEFAULT_MAX_HTML_SIZE_BYTES, -) -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import ( DEFAULT_MAX_PDF_SIZE_BYTES, + OnyxWebCrawler, ) -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import OnyxWebCrawler from onyx.tools.tool_implementations.open_url.tavily import TavilyExtractClient from onyx.tools.tool_implementations.web_search.clients.brave_client import BraveClient from onyx.tools.tool_implementations.web_search.clients.exa_client import ExaClient @@ -26,12 +26,13 @@ from onyx.tools.tool_implementations.web_search.clients.serper_client import ( from onyx.tools.tool_implementations.web_search.clients.tavily_client import ( TavilyClient, ) -from onyx.tools.tool_implementations.web_search.models import DEFAULT_MAX_RESULTS -from onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider +from onyx.tools.tool_implementations.web_search.models import ( + DEFAULT_MAX_RESULTS, + WebContentProviderConfig, + WebSearchProvider, +) from onyx.utils.logger import setup_logger -from shared_configs.enums import WebContentProviderType -from shared_configs.enums import WebSearchProviderType +from shared_configs.enums import WebContentProviderType, WebSearchProviderType logger = setup_logger() diff --git a/backend/onyx/tools/tool_implementations/web_search/utils.py b/backend/onyx/tools/tool_implementations/web_search/utils.py index fd3ad6bce8..39d4f2d496 100644 --- a/backend/onyx/tools/tool_implementations/web_search/utils.py +++ b/backend/onyx/tools/tool_implementations/web_search/utils.py @@ -1,13 +1,13 @@ from onyx.configs.constants import DocumentSource -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection -from onyx.context.search.models import SearchDoc +from onyx.context.search.models import InferenceChunk, InferenceSection, SearchDoc from onyx.tools.tool_implementations.open_url.models import WebContent from onyx.tools.tool_implementations.open_url.snippet_matcher import ( find_snippet_in_content, ) -from onyx.tools.tool_implementations.web_search.models import WEB_SEARCH_PREFIX -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.web_search.models import ( + WEB_SEARCH_PREFIX, + WebSearchResult, +) TRUNCATED_CONTENT_SUFFIX = " [...truncated]" TRUNCATED_CONTENT_PREFIX = "[...truncated] " diff --git a/backend/onyx/tools/tool_implementations/web_search/web_search_tool.py b/backend/onyx/tools/tool_implementations/web_search/web_search_tool.py index c950d2f8c1..121e0d0cf6 100644 --- a/backend/onyx/tools/tool_implementations/web_search/web_search_tool.py +++ b/backend/onyx/tools/tool_implementations/web_search/web_search_tool.py @@ -10,29 +10,31 @@ from onyx.context.search.utils import convert_inference_sections_to_search_docs from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.web_search import fetch_active_web_search_provider from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import SearchToolDocumentsDelta -from onyx.server.query_and_chat.streaming_models import SearchToolQueriesDelta -from onyx.server.query_and_chat.streaming_models import SearchToolStart +from onyx.server.query_and_chat.streaming_models import ( + Packet, + SearchToolDocumentsDelta, + SearchToolQueriesDelta, + SearchToolStart, +) from onyx.tools.interface import Tool -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolResponse -from onyx.tools.models import WebSearchToolOverrideKwargs +from onyx.tools.models import ( + ToolCallException, + ToolResponse, + WebSearchToolOverrideKwargs, +) from onyx.tools.tool_implementations.utils import ( convert_inference_sections_to_llm_string, ) -from onyx.tools.tool_implementations.web_search.models import DEFAULT_MAX_RESULTS -from onyx.tools.tool_implementations.web_search.models import WebSearchResult -from onyx.tools.tool_implementations.web_search.providers import ( - build_search_provider_from_config, +from onyx.tools.tool_implementations.web_search.models import ( + DEFAULT_MAX_RESULTS, + WebSearchResult, ) from onyx.tools.tool_implementations.web_search.providers import ( + build_search_provider_from_config, provider_requires_api_key, ) from onyx.tools.tool_implementations.web_search.utils import ( filter_web_search_results_with_no_title_or_snippet, -) -from onyx.tools.tool_implementations.web_search.utils import ( inference_section_from_internet_search_result, ) from onyx.utils.logger import setup_logger diff --git a/backend/onyx/tools/tool_runner.py b/backend/onyx/tools/tool_runner.py index 0897cf4280..5cc1723449 100644 --- a/backend/onyx/tools/tool_runner.py +++ b/backend/onyx/tools/tool_runner.py @@ -7,29 +7,33 @@ from onyx.chat.models import ChatMessageSimple from onyx.configs.constants import MessageType from onyx.context.search.models import SearchDocsResponse from onyx.db.memory import UserMemoryContext -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import PacketException -from onyx.server.query_and_chat.streaming_models import SectionEnd +from onyx.server.query_and_chat.streaming_models import ( + Packet, + PacketException, + SectionEnd, +) from onyx.tools.interface import Tool -from onyx.tools.models import ChatFile -from onyx.tools.models import ChatMinimalTextMessage -from onyx.tools.models import OpenURLToolOverrideKwargs -from onyx.tools.models import ParallelToolCallResponse -from onyx.tools.models import PythonToolOverrideKwargs -from onyx.tools.models import SearchToolOverrideKwargs -from onyx.tools.models import ToolCallException -from onyx.tools.models import ToolCallKickoff -from onyx.tools.models import ToolExecutionException -from onyx.tools.models import ToolResponse -from onyx.tools.models import WebSearchToolOverrideKwargs +from onyx.tools.models import ( + ChatFile, + ChatMinimalTextMessage, + OpenURLToolOverrideKwargs, + ParallelToolCallResponse, + PythonToolOverrideKwargs, + SearchToolOverrideKwargs, + ToolCallException, + ToolCallKickoff, + ToolExecutionException, + ToolResponse, + WebSearchToolOverrideKwargs, +) from onyx.tools.tool_implementations.coding_agent.coding_agent_tool import ( CodingAgentTool, -) -from onyx.tools.tool_implementations.coding_agent.coding_agent_tool import ( CodingAgentToolOverrideKwargs, ) -from onyx.tools.tool_implementations.memory.memory_tool import MemoryTool -from onyx.tools.tool_implementations.memory.memory_tool import MemoryToolOverrideKwargs +from onyx.tools.tool_implementations.memory.memory_tool import ( + MemoryTool, + MemoryToolOverrideKwargs, +) from onyx.tools.tool_implementations.open_url.open_url_tool import OpenURLTool from onyx.tools.tool_implementations.python.python_tool import PythonTool from onyx.tools.tool_implementations.search.search_tool import SearchTool diff --git a/backend/onyx/tools/utils.py b/backend/onyx/tools/utils.py index 423a00ab6c..212e04c297 100644 --- a/backend/onyx/tools/utils.py +++ b/backend/onyx/tools/utils.py @@ -9,8 +9,7 @@ from onyx.db.connector import check_connectors_exist from onyx.db.document import check_docs_exist from onyx.db.models import LLMProvider from onyx.llm.constants import LlmProviderNames -from onyx.llm.model_capabilities import find_model_obj -from onyx.llm.model_capabilities import get_model_map +from onyx.llm.model_capabilities import find_model_obj, get_model_map from onyx.tools.interface import Tool diff --git a/backend/onyx/tracing/braintrust_tracing_processor.py b/backend/onyx/tracing/braintrust_tracing_processor.py index 7d9478bf84..61f02923d2 100644 --- a/backend/onyx/tracing/braintrust_tracing_processor.py +++ b/backend/onyx/tracing/braintrust_tracing_processor.py @@ -1,7 +1,5 @@ import datetime -from typing import Any -from typing import Dict -from typing import Optional +from typing import Any, Dict, Optional import braintrust from braintrust import NOOP_SPAN @@ -9,10 +7,12 @@ from braintrust import NOOP_SPAN from onyx.llm.cost import calculate_llm_cost_cents from .framework.processor_interface import TracingProcessor -from .framework.span_data import AgentSpanData -from .framework.span_data import FunctionSpanData -from .framework.span_data import GenerationSpanData -from .framework.span_data import SpanData +from .framework.span_data import ( + AgentSpanData, + FunctionSpanData, + GenerationSpanData, + SpanData, +) from .framework.spans import Span from .framework.traces import Trace diff --git a/backend/onyx/tracing/dynamic_processor.py b/backend/onyx/tracing/dynamic_processor.py index 40b45691ba..fe37b20f68 100644 --- a/backend/onyx/tracing/dynamic_processor.py +++ b/backend/onyx/tracing/dynamic_processor.py @@ -16,10 +16,12 @@ from onyx.configs.app_configs import TRACING_CONFIG_CACHE_TTL_SECONDS from onyx.tracing.framework.processor_interface import TracingProcessor from onyx.tracing.framework.spans import Span from onyx.tracing.framework.traces import Trace -from onyx.tracing.provider_config import BraintrustConfig -from onyx.tracing.provider_config import EffectiveTracingConfig -from onyx.tracing.provider_config import LangfuseConfig -from onyx.tracing.provider_config import resolve_effective_tracing_config +from onyx.tracing.provider_config import ( + BraintrustConfig, + EffectiveTracingConfig, + LangfuseConfig, + resolve_effective_tracing_config, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/tracing/framework/__init__.py b/backend/onyx/tracing/framework/__init__.py index 238d0e9597..735506a858 100644 --- a/backend/onyx/tracing/framework/__init__.py +++ b/backend/onyx/tracing/framework/__init__.py @@ -1,7 +1,6 @@ from .processor_interface import TracingProcessor from .provider import DefaultTraceProvider -from .setup import get_trace_provider -from .setup import set_trace_provider +from .setup import get_trace_provider, set_trace_provider def add_trace_processor(span_processor: TracingProcessor) -> None: diff --git a/backend/onyx/tracing/framework/_error_tracing.py b/backend/onyx/tracing/framework/_error_tracing.py index 5bd90b0f2e..a14f21ec93 100644 --- a/backend/onyx/tracing/framework/_error_tracing.py +++ b/backend/onyx/tracing/framework/_error_tracing.py @@ -3,8 +3,7 @@ from typing import Any from onyx.utils.logger import setup_logger from .create import get_current_span -from .spans import Span -from .spans import SpanError +from .spans import Span, SpanError logger = setup_logger(__name__) diff --git a/backend/onyx/tracing/framework/create.py b/backend/onyx/tracing/framework/create.py index 90512e7be8..ca71106f58 100644 --- a/backend/onyx/tracing/framework/create.py +++ b/backend/onyx/tracing/framework/create.py @@ -1,18 +1,13 @@ from __future__ import annotations -from collections.abc import Iterator -from collections.abc import Mapping -from collections.abc import Sequence +from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING from onyx.utils.logger import setup_logger from .setup import get_trace_provider -from .span_data import AgentSpanData -from .span_data import FunctionSpanData -from .span_data import GenerationSpanData +from .span_data import AgentSpanData, FunctionSpanData, GenerationSpanData from .spans import Span from .traces import Trace diff --git a/backend/onyx/tracing/framework/processor_interface.py b/backend/onyx/tracing/framework/processor_interface.py index 1780c09875..f97edf8a80 100644 --- a/backend/onyx/tracing/framework/processor_interface.py +++ b/backend/onyx/tracing/framework/processor_interface.py @@ -1,6 +1,5 @@ import abc -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING if TYPE_CHECKING: from .spans import Span diff --git a/backend/onyx/tracing/framework/provider.py b/backend/onyx/tracing/framework/provider.py index aeef785e3b..06da22b906 100644 --- a/backend/onyx/tracing/framework/provider.py +++ b/backend/onyx/tracing/framework/provider.py @@ -2,23 +2,16 @@ from __future__ import annotations import threading import uuid -from abc import ABC -from abc import abstractmethod -from datetime import datetime -from datetime import timezone +from abc import ABC, abstractmethod +from datetime import datetime, timezone from typing import Any from onyx.utils.logger import setup_logger from .processor_interface import TracingProcessor from .scope import Scope -from .spans import NoOpSpan -from .spans import Span -from .spans import SpanImpl -from .spans import TSpanData -from .traces import NoOpTrace -from .traces import Trace -from .traces import TraceImpl +from .spans import NoOpSpan, Span, SpanImpl, TSpanData +from .traces import NoOpTrace, Trace, TraceImpl logger = setup_logger(__name__) diff --git a/backend/onyx/tracing/framework/scope.py b/backend/onyx/tracing/framework/scope.py index 83a0384711..cdde06ad85 100644 --- a/backend/onyx/tracing/framework/scope.py +++ b/backend/onyx/tracing/framework/scope.py @@ -1,6 +1,5 @@ import contextvars -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING if TYPE_CHECKING: from .spans import Span diff --git a/backend/onyx/tracing/framework/span_data.py b/backend/onyx/tracing/framework/span_data.py index e8834f88d2..6298d412cd 100644 --- a/backend/onyx/tracing/framework/span_data.py +++ b/backend/onyx/tracing/framework/span_data.py @@ -1,6 +1,5 @@ import abc -from collections.abc import Mapping -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any diff --git a/backend/onyx/tracing/framework/spans.py b/backend/onyx/tracing/framework/spans.py index c1d517206d..2c00da212a 100644 --- a/backend/onyx/tracing/framework/spans.py +++ b/backend/onyx/tracing/framework/spans.py @@ -3,9 +3,7 @@ from __future__ import annotations import abc import contextvars from types import TracebackType -from typing import Any -from typing import Generic -from typing import TypeVar +from typing import Any, Generic, TypeVar from typing_extensions import TypedDict diff --git a/backend/onyx/tracing/framework/traces.py b/backend/onyx/tracing/framework/traces.py index 78a7675c90..38c0d8e7a7 100644 --- a/backend/onyx/tracing/framework/traces.py +++ b/backend/onyx/tracing/framework/traces.py @@ -3,8 +3,7 @@ from __future__ import annotations import abc import contextvars from types import TracebackType -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING from . import util from .scope import Scope diff --git a/backend/onyx/tracing/framework/util.py b/backend/onyx/tracing/framework/util.py index 6f4a69bb3c..5971917df4 100644 --- a/backend/onyx/tracing/framework/util.py +++ b/backend/onyx/tracing/framework/util.py @@ -1,6 +1,5 @@ import uuid -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone def time_iso() -> str: diff --git a/backend/onyx/tracing/langfuse_tracing_processor.py b/backend/onyx/tracing/langfuse_tracing_processor.py index 1085436f25..5455986ac4 100644 --- a/backend/onyx/tracing/langfuse_tracing_processor.py +++ b/backend/onyx/tracing/langfuse_tracing_processor.py @@ -5,18 +5,18 @@ from __future__ import annotations import logging import threading from datetime import datetime -from typing import Any -from typing import Optional -from typing import Union +from typing import Any, Optional, Union from langfuse import Langfuse from langfuse._client.span import LangfuseObservationWrapper from onyx.tracing.framework.processor_interface import TracingProcessor -from onyx.tracing.framework.span_data import AgentSpanData -from onyx.tracing.framework.span_data import FunctionSpanData -from onyx.tracing.framework.span_data import GenerationSpanData -from onyx.tracing.framework.span_data import SpanData +from onyx.tracing.framework.span_data import ( + AgentSpanData, + FunctionSpanData, + GenerationSpanData, + SpanData, +) from onyx.tracing.framework.spans import Span from onyx.tracing.framework.traces import Trace diff --git a/backend/onyx/tracing/llm_utils.py b/backend/onyx/tracing/llm_utils.py index 4c34bfc85e..6013580797 100644 --- a/backend/onyx/tracing/llm_utils.py +++ b/backend/onyx/tracing/llm_utils.py @@ -1,11 +1,8 @@ from __future__ import annotations -from collections.abc import Iterator -from collections.abc import Mapping -from collections.abc import Sequence +from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager -from typing import Any -from typing import cast +from typing import Any, cast from onyx.llm.interfaces import LLM from onyx.llm.model_response import ModelResponse diff --git a/backend/onyx/tracing/provider_config.py b/backend/onyx/tracing/provider_config.py index f4c17da10e..fe56d6a860 100644 --- a/backend/onyx/tracing/provider_config.py +++ b/backend/onyx/tracing/provider_config.py @@ -6,12 +6,14 @@ from __future__ import annotations from dataclasses import dataclass -from onyx.configs.app_configs import BRAINTRUST_API_KEY -from onyx.configs.app_configs import BRAINTRUST_API_URL -from onyx.configs.app_configs import BRAINTRUST_PROJECT -from onyx.configs.app_configs import LANGFUSE_HOST -from onyx.configs.app_configs import LANGFUSE_PUBLIC_KEY -from onyx.configs.app_configs import LANGFUSE_SECRET_KEY +from onyx.configs.app_configs import ( + BRAINTRUST_API_KEY, + BRAINTRUST_API_URL, + BRAINTRUST_PROJECT, + LANGFUSE_HOST, + LANGFUSE_PUBLIC_KEY, + LANGFUSE_SECRET_KEY, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import TracingProviderConfig from onyx.db.tracing import fetch_all_tracing_providers diff --git a/backend/onyx/utils/batching.py b/backend/onyx/utils/batching.py index 6f49a63a98..a674e4d35f 100644 --- a/backend/onyx/utils/batching.py +++ b/backend/onyx/utils/batching.py @@ -1,6 +1,4 @@ -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterable +from collections.abc import Callable, Generator, Iterable from itertools import islice from typing import TypeVar diff --git a/backend/onyx/utils/callbacks.py b/backend/onyx/utils/callbacks.py index ffb46a2cf6..77b442a033 100644 --- a/backend/onyx/utils/callbacks.py +++ b/backend/onyx/utils/callbacks.py @@ -1,5 +1,4 @@ -from typing import Generic -from typing import TypeVar +from typing import Generic, TypeVar T = TypeVar("T") diff --git a/backend/onyx/utils/client_ip.py b/backend/onyx/utils/client_ip.py index 645a93aea3..c6f76f853b 100644 --- a/backend/onyx/utils/client_ip.py +++ b/backend/onyx/utils/client_ip.py @@ -19,8 +19,7 @@ import ipaddress from contextvars import ContextVar from fastapi import Request -from starlette.middleware.base import BaseHTTPMiddleware -from starlette.middleware.base import RequestResponseEndpoint +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.requests import Request as StarletteRequest from starlette.responses import Response diff --git a/backend/onyx/utils/csv_utils.py b/backend/onyx/utils/csv_utils.py index ae2b6b7c9f..39173e6b18 100644 --- a/backend/onyx/utils/csv_utils.py +++ b/backend/onyx/utils/csv_utils.py @@ -1,9 +1,6 @@ import csv import io -from collections.abc import Generator -from collections.abc import Iterable -from collections.abc import Iterator -from collections.abc import Mapping +from collections.abc import Generator, Iterable, Iterator, Mapping from pydantic import BaseModel diff --git a/backend/onyx/utils/datetime.py b/backend/onyx/utils/datetime.py index bfa3659f3e..a39c7616c3 100644 --- a/backend/onyx/utils/datetime.py +++ b/backend/onyx/utils/datetime.py @@ -1,5 +1,4 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone def datetime_to_utc(dt: datetime) -> datetime: diff --git a/backend/onyx/utils/encryption.py b/backend/onyx/utils/encryption.py index 18b550722b..7f912b7e8a 100644 --- a/backend/onyx/utils/encryption.py +++ b/backend/onyx/utils/encryption.py @@ -1,10 +1,8 @@ from typing import Any -from urllib.parse import urlparse -from urllib.parse import urlunparse +from urllib.parse import urlparse, urlunparse from onyx.configs.app_configs import ENCRYPTION_KEY_SECRET -from onyx.configs.constants import MASK_CREDENTIAL_CHAR -from onyx.configs.constants import MASK_CREDENTIAL_LONG_RE +from onyx.configs.constants import MASK_CREDENTIAL_CHAR, MASK_CREDENTIAL_LONG_RE from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, ) diff --git a/backend/onyx/utils/external_endpoint.py b/backend/onyx/utils/external_endpoint.py index fd01e62a0e..0adbdbde2d 100644 --- a/backend/onyx/utils/external_endpoint.py +++ b/backend/onyx/utils/external_endpoint.py @@ -28,12 +28,10 @@ endpoint; ``None`` means "no signal": import json import time -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar import httpx -from pydantic import BaseModel -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from onyx.utils.logger import setup_logger diff --git a/backend/onyx/utils/gpu_utils.py b/backend/onyx/utils/gpu_utils.py index 8b24f7fa94..71a998b976 100644 --- a/backend/onyx/utils/gpu_utils.py +++ b/backend/onyx/utils/gpu_utils.py @@ -5,10 +5,12 @@ import requests from onyx.utils.logger import setup_logger from onyx.utils.retry_wrapper import retry_builder -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 +from shared_configs.configs import ( + INDEXING_MODEL_SERVER_HOST, + INDEXING_MODEL_SERVER_PORT, + MODEL_SERVER_HOST, + MODEL_SERVER_PORT, +) logger = setup_logger() diff --git a/backend/onyx/utils/headers.py b/backend/onyx/utils/headers.py index 3da83a12ff..d58d1a063d 100644 --- a/backend/onyx/utils/headers.py +++ b/backend/onyx/utils/headers.py @@ -2,8 +2,10 @@ from typing import TypedDict from fastapi.datastructures import Headers -from onyx.configs.model_configs import LITELLM_EXTRA_HEADERS -from onyx.configs.model_configs import LITELLM_PASS_THROUGH_HEADERS +from onyx.configs.model_configs import ( + LITELLM_EXTRA_HEADERS, + LITELLM_PASS_THROUGH_HEADERS, +) from onyx.configs.tool_configs import CUSTOM_TOOL_PASS_THROUGH_HEADERS from onyx.utils.logger import setup_logger diff --git a/backend/onyx/utils/isolated_runner.py b/backend/onyx/utils/isolated_runner.py index f05c32151a..4b5efc1170 100644 --- a/backend/onyx/utils/isolated_runner.py +++ b/backend/onyx/utils/isolated_runner.py @@ -21,9 +21,7 @@ def main() -> None: import pickle - from onyx.utils.process_isolation import STATUS_EXC - from onyx.utils.process_isolation import STATUS_OK - from onyx.utils.process_isolation import STATUS_UNRELAYABLE + from onyx.utils.process_isolation import STATUS_EXC, STATUS_OK, STATUS_UNRELAYABLE fn, args, kwargs = pickle.loads(sys.stdin.buffer.read()) # noqa: S301 try: diff --git a/backend/onyx/utils/jsonriver/__init__.py b/backend/onyx/utils/jsonriver/__init__.py index 02acb52816..933403bd2a 100644 --- a/backend/onyx/utils/jsonriver/__init__.py +++ b/backend/onyx/utils/jsonriver/__init__.py @@ -10,8 +10,7 @@ SPDX-License-Identifier: BSD-3-Clause """ from .parse import _Parser as Parser -from .parse import JsonObject -from .parse import JsonValue +from .parse import JsonObject, JsonValue __all__ = ["Parser", "JsonValue", "JsonObject"] __version__ = "0.0.1" diff --git a/backend/onyx/utils/jsonriver/parse.py b/backend/onyx/utils/jsonriver/parse.py index e396c563cb..648890a8a8 100644 --- a/backend/onyx/utils/jsonriver/parse.py +++ b/backend/onyx/utils/jsonriver/parse.py @@ -10,13 +10,9 @@ from __future__ import annotations import copy from enum import IntEnum -from typing import cast -from typing import Union +from typing import cast, Union -from .tokenize import _Input -from .tokenize import json_token_type_to_string -from .tokenize import JsonTokenType -from .tokenize import Tokenizer +from .tokenize import _Input, json_token_type_to_string, JsonTokenType, Tokenizer # Type definitions for JSON values JsonValue = Union[None, bool, float, str, list["JsonValue"], dict[str, "JsonValue"]] diff --git a/backend/onyx/utils/logger.py b/backend/onyx/utils/logger.py index d793246325..452ceb9035 100644 --- a/backend/onyx/utils/logger.py +++ b/backend/onyx/utils/logger.py @@ -7,17 +7,21 @@ from typing import Any from onyx.utils.platform_utils import is_running_in_container from onyx.utils.tenant import get_tenant_id_short_string -from shared_configs.configs import DEV_LOGGING_ENABLED -from shared_configs.configs import JSON_LOGGING -from shared_configs.configs import LOG_FILE_NAME -from shared_configs.configs import LOG_LEVEL -from shared_configs.configs import LOG_TO_FILE -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.configs import SLACK_CHANNEL_ID -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR -from shared_configs.contextvars import ONYX_REQUEST_ID_CONTEXTVAR +from shared_configs.configs import ( + DEV_LOGGING_ENABLED, + JSON_LOGGING, + LOG_FILE_NAME, + LOG_LEVEL, + LOG_TO_FILE, + MULTI_TENANT, + POSTGRES_DEFAULT_SCHEMA, + SLACK_CHANNEL_ID, +) +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + INDEX_ATTEMPT_INFO_CONTEXTVAR, + ONYX_REQUEST_ID_CONTEXTVAR, +) logging.addLevelName(logging.INFO + 5, "NOTICE") diff --git a/backend/onyx/utils/middleware.py b/backend/onyx/utils/middleware.py index 5804176cbc..c18dba2253 100644 --- a/backend/onyx/utils/middleware.py +++ b/backend/onyx/utils/middleware.py @@ -3,19 +3,17 @@ import hashlib import logging import re import uuid -from collections.abc import Awaitable -from collections.abc import Callable -from datetime import datetime -from datetime import timezone +from collections.abc import Awaitable, Callable +from datetime import datetime, timezone -from fastapi import FastAPI -from fastapi import Request -from fastapi import Response +from fastapi import FastAPI, Request, Response from fastapi.routing import APIRoute -from shared_configs.contextvars import CURRENT_ENDPOINT_CONTEXTVAR -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import ONYX_REQUEST_ID_CONTEXTVAR +from shared_configs.contextvars import ( + CURRENT_ENDPOINT_CONTEXTVAR, + CURRENT_TENANT_ID_CONTEXTVAR, + ONYX_REQUEST_ID_CONTEXTVAR, +) def add_onyx_tenant_id_middleware( diff --git a/backend/onyx/utils/playwright_fetch.py b/backend/onyx/utils/playwright_fetch.py index 451fbf29b8..1f211019cb 100644 --- a/backend/onyx/utils/playwright_fetch.py +++ b/backend/onyx/utils/playwright_fetch.py @@ -14,18 +14,17 @@ Two consumers: from collections.abc import Iterator from contextlib import contextmanager -from playwright.sync_api import BrowserContext -from playwright.sync_api import Playwright -from playwright.sync_api import sync_playwright +from playwright.sync_api import BrowserContext, Playwright, sync_playwright from playwright.sync_api import TimeoutError as PlaywrightTimeoutError from pydantic import BaseModel -from onyx.configs.app_configs import WEB_CONNECTOR_OAUTH_CLIENT_ID -from onyx.configs.app_configs import WEB_CONNECTOR_OAUTH_CLIENT_SECRET -from onyx.configs.app_configs import WEB_CONNECTOR_OAUTH_TOKEN_URL +from onyx.configs.app_configs import ( + WEB_CONNECTOR_OAUTH_CLIENT_ID, + WEB_CONNECTOR_OAUTH_CLIENT_SECRET, + WEB_CONNECTOR_OAUTH_TOKEN_URL, +) from onyx.utils.logger import setup_logger -from onyx.utils.url import SSRFException -from onyx.utils.url import validate_outbound_http_url +from onyx.utils.url import SSRFException, validate_outbound_http_url logger = setup_logger() diff --git a/backend/onyx/utils/postgres_sanitization.py b/backend/onyx/utils/postgres_sanitization.py index 9ce6de11ca..e8087b07da 100644 --- a/backend/onyx/utils/postgres_sanitization.py +++ b/backend/onyx/utils/postgres_sanitization.py @@ -2,9 +2,7 @@ import re from typing import Any from onyx.access.models import ExternalAccess -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import BasicExpertInfo, Document, HierarchyNode from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/utils/process_isolation.py b/backend/onyx/utils/process_isolation.py index cf77bcd22e..8195769244 100644 --- a/backend/onyx/utils/process_isolation.py +++ b/backend/onyx/utils/process_isolation.py @@ -11,8 +11,7 @@ import signal import subprocess import sys from collections.abc import Callable -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar T = TypeVar("T") diff --git a/backend/onyx/utils/pydantic_util.py b/backend/onyx/utils/pydantic_util.py index f4bbdf6568..e4c1105205 100644 --- a/backend/onyx/utils/pydantic_util.py +++ b/backend/onyx/utils/pydantic_util.py @@ -1,9 +1,6 @@ -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar -from pydantic import BaseModel -from pydantic import TypeAdapter -from pydantic import ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/onyx/utils/retry_after.py b/backend/onyx/utils/retry_after.py index e8732739ec..7cc5df6765 100644 --- a/backend/onyx/utils/retry_after.py +++ b/backend/onyx/utils/retry_after.py @@ -1,6 +1,5 @@ import math -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from email.utils import parsedate_to_datetime from onyx.utils.datetime import datetime_to_utc diff --git a/backend/onyx/utils/retry_wrapper.py b/backend/onyx/utils/retry_wrapper.py index 7b5dfc0e53..0d0c3b9c0e 100644 --- a/backend/onyx/utils/retry_wrapper.py +++ b/backend/onyx/utils/retry_wrapper.py @@ -2,18 +2,18 @@ import functools import logging from collections.abc import Callable from logging import Logger -from typing import Any -from typing import cast -from typing import TypeVar +from typing import Any, cast, TypeVar import requests -from tenacity import before_sleep_log +from tenacity import ( + before_sleep_log, + retry_if_exception_type, + stop_after_attempt, + stop_never, + wait_exponential, + wait_random, +) from tenacity import retry as tenacity_retry -from tenacity import retry_if_exception_type -from tenacity import stop_after_attempt -from tenacity import stop_never -from tenacity import wait_exponential -from tenacity import wait_random from tenacity.stop import stop_base from tenacity.wait import wait_base diff --git a/backend/onyx/utils/sensitive.py b/backend/onyx/utils/sensitive.py index 7e792be7a6..876126c0cd 100644 --- a/backend/onyx/utils/sensitive.py +++ b/backend/onyx/utils/sensitive.py @@ -10,14 +10,10 @@ from __future__ import annotations import json from collections.abc import Callable -from typing import Any -from typing import Generic -from typing import NoReturn -from typing import TypeVar +from typing import Any, Generic, NoReturn, TypeVar from unittest.mock import MagicMock -from onyx.utils.encryption import mask_credential_dict -from onyx.utils.encryption import mask_string +from onyx.utils.encryption import mask_credential_dict, mask_string T = TypeVar("T", str, dict[str, Any]) diff --git a/backend/onyx/utils/special_types.py b/backend/onyx/utils/special_types.py index ea9ccf21e5..c6c221a826 100644 --- a/backend/onyx/utils/special_types.py +++ b/backend/onyx/utils/special_types.py @@ -1,5 +1,4 @@ -from collections.abc import Mapping -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TypeAlias JSON_ro: TypeAlias = ( diff --git a/backend/onyx/utils/subclasses.py b/backend/onyx/utils/subclasses.py index 5aeb80369d..886637a27d 100644 --- a/backend/onyx/utils/subclasses.py +++ b/backend/onyx/utils/subclasses.py @@ -5,9 +5,7 @@ import os import pkgutil import sys from types import ModuleType -from typing import List -from typing import Type -from typing import TypeVar +from typing import List, Type, TypeVar T = TypeVar("T") diff --git a/backend/onyx/utils/telemetry.py b/backend/onyx/utils/telemetry.py index 21e6af08c2..9847237ec2 100644 --- a/backend/onyx/utils/telemetry.py +++ b/backend/onyx/utils/telemetry.py @@ -6,24 +6,22 @@ from typing import Any import requests -from onyx.configs.app_configs import DISABLE_TELEMETRY -from onyx.configs.app_configs import ENTERPRISE_EDITION_ENABLED -from onyx.configs.constants import KV_CUSTOMER_UUID_KEY -from onyx.configs.constants import KV_INSTANCE_DOMAIN_KEY -from onyx.configs.constants import MilestoneRecordType -from onyx.db.encrypted_kv_store import load_encrypted_kv -from onyx.db.encrypted_kv_store import upsert_encrypted_kv +from onyx.configs.app_configs import DISABLE_TELEMETRY, ENTERPRISE_EDITION_ENABLED +from onyx.configs.constants import ( + KV_CUSTOMER_UUID_KEY, + KV_INSTANCE_DOMAIN_KEY, + MilestoneRecordType, +) +from onyx.db.encrypted_kv_store import load_encrypted_kv, upsert_encrypted_kv from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import User -from onyx.key_value_store.interface import KvKeyNotFoundError -from onyx.key_value_store.interface import unwrap_str +from onyx.key_value_store.interface import KvKeyNotFoundError, unwrap_str from onyx.utils.logger import setup_logger from onyx.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, + noop_fallback, ) -from onyx.utils.variable_functionality import noop_fallback -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA from shared_configs.contextvars import get_current_tenant_id logger = setup_logger() diff --git a/backend/onyx/utils/threadpool_concurrency.py b/backend/onyx/utils/threadpool_concurrency.py index b71b3702b2..2c1eb98b79 100644 --- a/backend/onyx/utils/threadpool_concurrency.py +++ b/backend/onyx/utils/threadpool_concurrency.py @@ -5,22 +5,15 @@ import contextvars import copy import threading import uuid -from collections.abc import Callable -from collections.abc import Coroutine -from collections.abc import Iterator -from collections.abc import MutableMapping -from collections.abc import Sequence -from concurrent.futures import as_completed -from concurrent.futures import FIRST_COMPLETED -from concurrent.futures import Future -from concurrent.futures import ThreadPoolExecutor -from concurrent.futures import wait -from typing import Any -from typing import cast -from typing import Generic -from typing import overload -from typing import Protocol -from typing import TypeVar +from collections.abc import Callable, Coroutine, Iterator, MutableMapping, Sequence +from concurrent.futures import ( + as_completed, + FIRST_COMPLETED, + Future, + ThreadPoolExecutor, + wait, +) +from typing import Any, cast, Generic, overload, Protocol, TypeVar from pydantic import GetCoreSchemaHandler from pydantic.types import T diff --git a/backend/onyx/utils/timing.py b/backend/onyx/utils/timing.py index 77128e4b07..be4c2b01d1 100644 --- a/backend/onyx/utils/timing.py +++ b/backend/onyx/utils/timing.py @@ -1,16 +1,11 @@ import time -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Callable, Generator, Iterator from functools import wraps from inspect import signature -from typing import Any -from typing import cast -from typing import TypeVar +from typing import Any, cast, TypeVar from onyx.utils.logger import setup_logger -from onyx.utils.telemetry import optional_telemetry -from onyx.utils.telemetry import RecordType +from onyx.utils.telemetry import optional_telemetry, RecordType logger = setup_logger() diff --git a/backend/onyx/utils/url.py b/backend/onyx/utils/url.py index 4ace1d8f48..4e823fd923 100644 --- a/backend/onyx/utils/url.py +++ b/backend/onyx/utils/url.py @@ -2,10 +2,7 @@ import ipaddress import socket import unicodedata from typing import Any -from urllib.parse import parse_qs -from urllib.parse import urlencode -from urllib.parse import urlparse -from urllib.parse import urlunparse +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import requests diff --git a/backend/onyx/utils/variable_functionality.py b/backend/onyx/utils/variable_functionality.py index 2bb4a14977..4005a34fed 100644 --- a/backend/onyx/utils/variable_functionality.py +++ b/backend/onyx/utils/variable_functionality.py @@ -2,16 +2,17 @@ import functools import importlib import inspect import os -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar -from onyx.configs.app_configs import API_SERVER_HOST -from onyx.configs.app_configs import API_SERVER_PROTOCOL -from onyx.configs.app_configs import API_SERVER_URL_OVERRIDE_FOR_HTTP_REQUESTS -from onyx.configs.app_configs import APP_API_PREFIX -from onyx.configs.app_configs import APP_PORT -from onyx.configs.app_configs import DEV_MODE -from onyx.configs.app_configs import ENTERPRISE_EDITION_ENABLED +from onyx.configs.app_configs import ( + API_SERVER_HOST, + API_SERVER_PROTOCOL, + API_SERVER_URL_OVERRIDE_FOR_HTTP_REQUESTS, + APP_API_PREFIX, + APP_PORT, + DEV_MODE, + ENTERPRISE_EDITION_ENABLED, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/onyx/utils/web_content.py b/backend/onyx/utils/web_content.py index 2b441eb4ee..a75dfa6e7c 100644 --- a/backend/onyx/utils/web_content.py +++ b/backend/onyx/utils/web_content.py @@ -1,8 +1,7 @@ from __future__ import annotations import io -from urllib.parse import unquote -from urllib.parse import urlparse +from urllib.parse import unquote, urlparse from bs4.dammit import UnicodeDammit diff --git a/backend/onyx/voice/interface.py b/backend/onyx/voice/interface.py index 86a6b51f8b..32caef4e9f 100644 --- a/backend/onyx/voice/interface.py +++ b/backend/onyx/voice/interface.py @@ -1,5 +1,4 @@ -from abc import ABC -from abc import abstractmethod +from abc import ABC, abstractmethod from collections.abc import AsyncIterator from typing import Protocol diff --git a/backend/onyx/voice/providers/azure.py b/backend/onyx/voice/providers/azure.py index 4b90a36aa9..4df7b69dc6 100644 --- a/backend/onyx/voice/providers/azure.py +++ b/backend/onyx/voice/providers/azure.py @@ -29,18 +29,19 @@ import wave from collections.abc import AsyncIterator from typing import Any from urllib.parse import urlparse -from xml.sax.saxutils import escape -from xml.sax.saxutils import quoteattr +from xml.sax.saxutils import escape, quoteattr import aiohttp from onyx.tracing.flows import LLMFlow from onyx.tracing.llm_utils import traced_llm_call from onyx.utils.logger import setup_logger -from onyx.voice.interface import StreamingSynthesizerProtocol -from onyx.voice.interface import StreamingTranscriberProtocol -from onyx.voice.interface import TranscriptResult -from onyx.voice.interface import VoiceProviderInterface +from onyx.voice.interface import ( + StreamingSynthesizerProtocol, + StreamingTranscriberProtocol, + TranscriptResult, + VoiceProviderInterface, +) # SSML namespace — W3C standard for Speech Synthesis Markup Language. # This is a fixed W3C specification and will not change. diff --git a/backend/onyx/voice/providers/elevenlabs.py b/backend/onyx/voice/providers/elevenlabs.py index 962b40afc4..d2b11ab72d 100644 --- a/backend/onyx/voice/providers/elevenlabs.py +++ b/backend/onyx/voice/providers/elevenlabs.py @@ -22,10 +22,12 @@ import aiohttp from onyx.tracing.flows import LLMFlow from onyx.tracing.llm_utils import traced_llm_call -from onyx.voice.interface import StreamingSynthesizerProtocol -from onyx.voice.interface import StreamingTranscriberProtocol -from onyx.voice.interface import TranscriptResult -from onyx.voice.interface import VoiceProviderInterface +from onyx.voice.interface import ( + StreamingSynthesizerProtocol, + StreamingTranscriberProtocol, + TranscriptResult, + VoiceProviderInterface, +) # Default ElevenLabs API base URL DEFAULT_ELEVENLABS_API_BASE = "https://api.elevenlabs.io" diff --git a/backend/onyx/voice/providers/openai.py b/backend/onyx/voice/providers/openai.py index dd17a2b8cf..ad2a4fc020 100644 --- a/backend/onyx/voice/providers/openai.py +++ b/backend/onyx/voice/providers/openai.py @@ -12,17 +12,18 @@ import io import json from collections.abc import AsyncIterator from enum import StrEnum -from typing import Any -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING import aiohttp from onyx.tracing.flows import LLMFlow from onyx.tracing.llm_utils import traced_llm_call -from onyx.voice.interface import StreamingSynthesizerProtocol -from onyx.voice.interface import StreamingTranscriberProtocol -from onyx.voice.interface import TranscriptResult -from onyx.voice.interface import VoiceProviderInterface +from onyx.voice.interface import ( + StreamingSynthesizerProtocol, + StreamingTranscriberProtocol, + TranscriptResult, + VoiceProviderInterface, +) if TYPE_CHECKING: from openai import AsyncOpenAI @@ -594,8 +595,7 @@ class OpenAIVoiceProvider(VoiceProviderInterface): async def validate_credentials(self) -> None: """Validate OpenAI API key by listing models.""" - from openai import AuthenticationError - from openai import PermissionDeniedError + from openai import AuthenticationError, PermissionDeniedError client = self._get_client() try: diff --git a/backend/scripts/add_connector_creation_script.py b/backend/scripts/add_connector_creation_script.py index 96ea6f9d96..63e0c8c22f 100644 --- a/backend/scripts/add_connector_creation_script.py +++ b/backend/scripts/add_connector_creation_script.py @@ -1,5 +1,4 @@ -from typing import Any -from typing import Dict +from typing import Any, Dict import requests diff --git a/backend/scripts/coding_agent_test.py b/backend/scripts/coding_agent_test.py index d5dc07f393..55917fd7f7 100644 --- a/backend/scripts/coding_agent_test.py +++ b/backend/scripts/coding_agent_test.py @@ -25,14 +25,14 @@ import queue from uuid import uuid4 from onyx.chat.emitter import Emitter -from onyx.coding_agent.mock_tools import CODING_AGENT_QUERY_KEY -from onyx.coding_agent.mock_tools import CODING_AGENT_REPO_KEY -from onyx.coding_agent.mock_tools import CODING_AGENT_TOOL_NAME -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.coding_agent.mock_tools import ( + CODING_AGENT_QUERY_KEY, + CODING_AGENT_REPO_KEY, + CODING_AGENT_TOOL_NAME, +) +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from onyx.db.models import User -from onyx.llm.factory import get_default_llm -from onyx.llm.factory import get_llm_token_counter +from onyx.llm.factory import get_default_llm, get_llm_token_counter from onyx.server.query_and_chat.placement import Placement from onyx.tools.fake_tools.coding_agent import run_coding_agent_call from onyx.tools.models import ToolCallKickoff diff --git a/backend/scripts/debugging/debug_usage_limits.py b/backend/scripts/debugging/debug_usage_limits.py index 823f417358..f6cccba3a7 100644 --- a/backend/scripts/debugging/debug_usage_limits.py +++ b/backend/scripts/debugging/debug_usage_limits.py @@ -14,9 +14,7 @@ Environment variables required: import json import os import sys -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import jwt import requests diff --git a/backend/scripts/debugging/onyx_db.py b/backend/scripts/debugging/onyx_db.py index d34b32bccb..e6f8592596 100644 --- a/backend/scripts/debugging/onyx_db.py +++ b/backend/scripts/debugging/onyx_db.py @@ -13,14 +13,15 @@ if True: # noqa: E402 from pydantic import BaseModel from sqlalchemy import func - from onyx.db.engine.sql_engine import build_connection_string - from onyx.db.engine.sql_engine import get_session_with_tenant - from onyx.db.engine.sql_engine import SqlEngine - from onyx.db.engine.sql_engine import SYNC_DB_API - from onyx.db.engine.sql_engine import USE_IAM_AUTH + from onyx.db.engine.sql_engine import ( + build_connection_string, + get_session_with_tenant, + SqlEngine, + SYNC_DB_API, + USE_IAM_AUTH, + ) from onyx.db.engine.tenant_utils import get_all_tenant_ids - from onyx.db.models import Document - from onyx.db.models import User + from onyx.db.models import Document, User from onyx.utils.logger import setup_logger from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/scripts/debugging/onyx_redis.py b/backend/scripts/debugging/onyx_redis.py index 21824c39a8..75a30ae9ca 100644 --- a/backend/scripts/debugging/onyx_redis.py +++ b/backend/scripts/debugging/onyx_redis.py @@ -11,22 +11,24 @@ from uuid import UUID from redis import Redis from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email -from onyx.auth.invited_users import get_invited_users -from onyx.auth.invited_users import write_invited_users -from onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX -from onyx.configs.app_configs import REDIS_DB_NUMBER -from onyx.configs.app_configs import REDIS_HOST -from onyx.configs.app_configs import REDIS_PASSWORD -from onyx.configs.app_configs import REDIS_PORT -from onyx.configs.app_configs import REDIS_SSL +from onyx.auth.invited_users import get_invited_users, write_invited_users +from onyx.configs.app_configs import ( + REDIS_AUTH_KEY_PREFIX, + REDIS_DB_NUMBER, + REDIS_HOST, + REDIS_PASSWORD, + REDIS_PORT, + REDIS_SSL, +) from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.users import get_user_by_email from onyx.redis.redis_connector import RedisConnector from onyx.redis.redis_pool import RedisPool -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) # Tool to run helpful operations on Redis in production # This is targeted for internal usage and may not have all the necessary parameters diff --git a/backend/scripts/debugging/opensearch/benchmark_retrieval.py b/backend/scripts/debugging/opensearch/benchmark_retrieval.py index 493d6ca555..4ccbfb960a 100644 --- a/backend/scripts/debugging/opensearch/benchmark_retrieval.py +++ b/backend/scripts/debugging/opensearch/benchmark_retrieval.py @@ -18,8 +18,7 @@ from scripts.debugging.opensearch.embedding_io import load_query_embedding_from_ from onyx.configs.chat_configs import NUM_RETURNED_HITS from onyx.context.search.enums import QueryType from onyx.context.search.models import IndexFilters -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from onyx.db.search_settings import get_current_search_settings from onyx.document_index.interfaces_new import TenantState from onyx.document_index.opensearch.opensearch_document_index import ( @@ -27,8 +26,10 @@ from onyx.document_index.opensearch.opensearch_document_index import ( ) from onyx.indexing.models import IndexingSetting from shared_configs.configs import MULTI_TENANT -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) DEFAULT_N = 50 diff --git a/backend/scripts/debugging/opensearch/embed_and_save.py b/backend/scripts/debugging/opensearch/embed_and_save.py index 859a0fadf4..60ca060ec3 100644 --- a/backend/scripts/debugging/opensearch/embed_and_save.py +++ b/backend/scripts/debugging/opensearch/embed_and_save.py @@ -15,8 +15,7 @@ from scripts.debugging.opensearch.constants import DEV_TENANT_ID from scripts.debugging.opensearch.embedding_io import save_query_embedding_to_file from onyx.context.search.utils import get_query_embedding -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/scripts/debugging/opensearch/opensearch_debug.py b/backend/scripts/debugging/opensearch/opensearch_debug.py index 7728a45afa..6e8bc895aa 100644 --- a/backend/scripts/debugging/opensearch/opensearch_debug.py +++ b/backend/scripts/debugging/opensearch/opensearch_debug.py @@ -26,8 +26,10 @@ import os import sys from typing import Any -from onyx.document_index.opensearch.client import OpenSearchClient -from onyx.document_index.opensearch.client import OpenSearchIndexClient +from onyx.document_index.opensearch.client import ( + OpenSearchClient, + OpenSearchIndexClient, +) from onyx.document_index.opensearch.constants import OpenSearchAuthMethod from shared_configs.configs import MULTI_TENANT diff --git a/backend/scripts/dev_celery_reload.py b/backend/scripts/dev_celery_reload.py index 147657615b..f9b7ced536 100644 --- a/backend/scripts/dev_celery_reload.py +++ b/backend/scripts/dev_celery_reload.py @@ -23,8 +23,7 @@ import os import sys from pathlib import Path -from watchfiles import DefaultFilter -from watchfiles import run_process +from watchfiles import DefaultFilter, run_process # Built-in ext-app + craft skill bundles are content shipped to the sandbox, # not imported by the worker, so editing them shouldn't trigger a restart. diff --git a/backend/scripts/env_inventory.py b/backend/scripts/env_inventory.py index c2ce250882..1c0f31f58d 100644 --- a/backend/scripts/env_inventory.py +++ b/backend/scripts/env_inventory.py @@ -40,9 +40,7 @@ import json import re import sys from collections import defaultdict -from dataclasses import asdict -from dataclasses import dataclass -from dataclasses import field +from dataclasses import asdict, dataclass, field from pathlib import Path # --- repo layout -------------------------------------------------------------- diff --git a/backend/scripts/force_delete_connector_by_id.py b/backend/scripts/force_delete_connector_by_id.py index 312c17d0b1..5c3ae89f1f 100755 --- a/backend/scripts/force_delete_connector_by_id.py +++ b/backend/scripts/force_delete_connector_by_id.py @@ -21,19 +21,23 @@ sys.path.append(parent_dir) # Now import Onyx modules from onyx.configs.constants import DocumentSource from onyx.db.connector import fetch_connector_by_id -from onyx.db.connector_credential_pair import get_connector_credential_pair -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id +from onyx.db.connector_credential_pair import ( + get_connector_credential_pair, + get_connector_credential_pair_from_id, +) from onyx.db.document import get_documents_for_connector_credential_pair from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.index_attempt import cancel_indexing_attempts_for_ccpair -from onyx.db.index_attempt import delete_index_attempts -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import DocumentSet__ConnectorCredentialPair -from onyx.db.models import UserGroup__ConnectorCredentialPair -from onyx.db.permission_sync_attempt import ( - delete_doc_permission_sync_attempts__no_commit, +from onyx.db.index_attempt import ( + cancel_indexing_attempts_for_ccpair, + delete_index_attempts, +) +from onyx.db.models import ( + ConnectorCredentialPair, + DocumentSet__ConnectorCredentialPair, + UserGroup__ConnectorCredentialPair, ) from onyx.db.permission_sync_attempt import ( + delete_doc_permission_sync_attempts__no_commit, delete_external_group_permission_sync_attempts__no_commit, ) from onyx.document_index.factory import get_all_document_indices diff --git a/backend/scripts/hard_delete_chats.py b/backend/scripts/hard_delete_chats.py index 6ddcfbc99d..ba147dd408 100644 --- a/backend/scripts/hard_delete_chats.py +++ b/backend/scripts/hard_delete_chats.py @@ -7,8 +7,10 @@ print(parent_dir) sys.path.append(parent_dir) from onyx.db.chat import delete_chat_session # noqa: E402 -from onyx.db.engine.sql_engine import get_session_with_current_tenant # noqa: E402 -from onyx.db.engine.sql_engine import SqlEngine # noqa: E402 +from onyx.db.engine.sql_engine import ( # noqa: E402 + get_session_with_current_tenant, + SqlEngine, +) from onyx.db.models import ChatSession # noqa: E402 diff --git a/backend/scripts/orphan_doc_cleanup_script.py b/backend/scripts/orphan_doc_cleanup_script.py index f93f5dc641..7f923f68bd 100644 --- a/backend/scripts/orphan_doc_cleanup_script.py +++ b/backend/scripts/orphan_doc_cleanup_script.py @@ -12,13 +12,17 @@ parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(parent_dir) from onyx.context.search.models import IndexFilters # noqa: E402 -from onyx.db.document import delete_documents_complete__no_commit # noqa: E402 -from onyx.db.document import get_document # noqa: E402 +from onyx.db.document import ( # noqa: E402 + delete_documents_complete__no_commit, + get_document, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant # noqa: E402 from onyx.db.search_settings import get_current_search_settings # noqa: E402 from onyx.db.tag import delete_orphan_tags_batched # noqa: E402 -from onyx.document_index.interfaces_new import DocumentSectionRequest # noqa: E402 -from onyx.document_index.interfaces_new import TenantState # noqa: E402 +from onyx.document_index.interfaces_new import ( # noqa: E402 + DocumentSectionRequest, + TenantState, +) from onyx.document_index.vespa.vespa_document_index import ( # noqa: E402 VespaDocumentIndex, ) diff --git a/backend/scripts/query_time_check/seed_dummy_docs.py b/backend/scripts/query_time_check/seed_dummy_docs.py index 17e68c7e3a..42d4216cce 100644 --- a/backend/scripts/query_time_check/seed_dummy_docs.py +++ b/backend/scripts/query_time_check/seed_dummy_docs.py @@ -18,15 +18,11 @@ from onyx.connectors.models import Document from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.search_settings import get_current_search_settings from onyx.document_index.document_index_utils import get_multipass_config -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces_new import IndexingMetadata, TenantState from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import IndexChunk +from onyx.indexing.models import ChunkEmbedding, DocMetadataAwareIndexChunk, IndexChunk from onyx.utils.timing import log_function_time -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA from shared_configs.model_server_models import Embedding TOTAL_DOC_SETS = 8 diff --git a/backend/scripts/reencrypt_secrets.py b/backend/scripts/reencrypt_secrets.py index be6c318f11..04385f5b8a 100644 --- a/backend/scripts/reencrypt_secrets.py +++ b/backend/scripts/reencrypt_secrets.py @@ -24,8 +24,10 @@ import sys parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(parent_dir) -from onyx.db.engine.sql_engine import get_session_with_tenant # noqa: E402 -from onyx.db.engine.sql_engine import SqlEngine # noqa: E402 +from onyx.db.engine.sql_engine import ( # noqa: E402 + get_session_with_tenant, + SqlEngine, +) from onyx.db.engine.tenant_utils import get_all_tenant_ids # noqa: E402 from onyx.db.rotate_encryption_key import rotate_encryption_key # noqa: E402 from onyx.utils.variable_functionality import global_version # noqa: E402 diff --git a/backend/scripts/reset_postgres.py b/backend/scripts/reset_postgres.py index 5e9dae123c..bc890faf70 100644 --- a/backend/scripts/reset_postgres.py +++ b/backend/scripts/reset_postgres.py @@ -10,11 +10,13 @@ from onyx.db.engine.sql_engine import get_sqlalchemy_engine parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(parent_dir) -from onyx.configs.app_configs import POSTGRES_DB # noqa: E402 -from onyx.configs.app_configs import POSTGRES_HOST # noqa: E402 -from onyx.configs.app_configs import POSTGRES_PASSWORD # noqa: E402 -from onyx.configs.app_configs import POSTGRES_PORT # noqa: E402 -from onyx.configs.app_configs import POSTGRES_USER # noqa: E402 +from onyx.configs.app_configs import ( # noqa: E402 + POSTGRES_DB, + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, +) from onyx.db.credentials import create_initial_public_credential # noqa: E402 diff --git a/backend/scripts/rotate_llm_provider_keys.py b/backend/scripts/rotate_llm_provider_keys.py index 21784b9dec..5d2e5ac83c 100644 --- a/backend/scripts/rotate_llm_provider_keys.py +++ b/backend/scripts/rotate_llm_provider_keys.py @@ -40,12 +40,16 @@ sys.path.append(parent_dir) from sqlalchemy import select # noqa: E402 -from onyx.db.engine.sql_engine import get_session_with_tenant # noqa: E402 -from onyx.db.engine.sql_engine import SqlEngine # noqa: E402 +from onyx.db.engine.sql_engine import ( # noqa: E402 + get_session_with_tenant, + SqlEngine, +) from onyx.db.engine.tenant_utils import get_all_tenant_ids # noqa: E402 -from onyx.db.models import CloudEmbeddingProvider # noqa: E402 -from onyx.db.models import LLMProvider # noqa: E402 -from onyx.db.models import VoiceProvider # noqa: E402 +from onyx.db.models import ( # noqa: E402 + CloudEmbeddingProvider, + LLMProvider, + VoiceProvider, +) from onyx.utils.variable_functionality import ( # noqa: E402 set_is_ee_based_on_env_variable, ) diff --git a/backend/scripts/run_enterpriserag_bench_questions.py b/backend/scripts/run_enterpriserag_bench_questions.py index ab4ed5ea1e..a7071783ed 100644 --- a/backend/scripts/run_enterpriserag_bench_questions.py +++ b/backend/scripts/run_enterpriserag_bench_questions.py @@ -6,12 +6,9 @@ import json import logging import sys import time -from dataclasses import asdict -from dataclasses import dataclass +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any -from typing import TypedDict -from typing import TypeGuard +from typing import Any, TypedDict, TypeGuard import aiohttp diff --git a/backend/scripts/save_load_state.py b/backend/scripts/save_load_state.py index 611fb0c0dd..1ddfe41da4 100644 --- a/backend/scripts/save_load_state.py +++ b/backend/scripts/save_load_state.py @@ -9,12 +9,14 @@ import requests from alembic import command from alembic.config import Config -from onyx.configs.app_configs import DOCUMENT_INDEX_NAME -from onyx.configs.app_configs import POSTGRES_DB -from onyx.configs.app_configs import POSTGRES_HOST -from onyx.configs.app_configs import POSTGRES_PASSWORD -from onyx.configs.app_configs import POSTGRES_PORT -from onyx.configs.app_configs import POSTGRES_USER +from onyx.configs.app_configs import ( + DOCUMENT_INDEX_NAME, + POSTGRES_DB, + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, +) from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT from onyx.utils.logger import setup_logger diff --git a/backend/scripts/seed_dev_license.py b/backend/scripts/seed_dev_license.py index 92b8831d17..715840d2bb 100644 --- a/backend/scripts/seed_dev_license.py +++ b/backend/scripts/seed_dev_license.py @@ -18,8 +18,10 @@ sys.path.append(parent_dir) from ee.onyx.db.license import upsert_license # noqa: E402 from ee.onyx.utils.license import verify_license_signature # noqa: E402 -from onyx.db.engine.sql_engine import get_session_with_current_tenant # noqa: E402 -from onyx.db.engine.sql_engine import SqlEngine # noqa: E402 +from onyx.db.engine.sql_engine import ( # noqa: E402 + get_session_with_current_tenant, + SqlEngine, +) _PEM_BEGIN = "-----BEGIN ONYX LICENSE-----" _PEM_END = "-----END ONYX LICENSE-----" diff --git a/backend/scripts/sources_selection_analysis.py b/backend/scripts/sources_selection_analysis.py index b7b4e80e4e..d3615dc13b 100644 --- a/backend/scripts/sources_selection_analysis.py +++ b/backend/scripts/sources_selection_analysis.py @@ -5,8 +5,7 @@ import sys import time from datetime import datetime from os import listdir -from os.path import isfile -from os.path import join +from os.path import isfile, join from typing import Optional import requests diff --git a/backend/scripts/tenant_cleanup/analyze_current_tenants.py b/backend/scripts/tenant_cleanup/analyze_current_tenants.py index 6f54acb26a..e961bbbb62 100644 --- a/backend/scripts/tenant_cleanup/analyze_current_tenants.py +++ b/backend/scripts/tenant_cleanup/analyze_current_tenants.py @@ -12,9 +12,7 @@ import json import os import subprocess import sys -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any diff --git a/backend/scripts/tenant_cleanup/cleanup_tenants.py b/backend/scripts/tenant_cleanup/cleanup_tenants.py index 4c4c276acf..42079eb377 100644 --- a/backend/scripts/tenant_cleanup/cleanup_tenants.py +++ b/backend/scripts/tenant_cleanup/cleanup_tenants.py @@ -26,16 +26,17 @@ import json import signal import subprocess import sys -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from pathlib import Path -from scripts.tenant_cleanup.cleanup_utils import confirm_step -from scripts.tenant_cleanup.cleanup_utils import execute_control_plane_query -from scripts.tenant_cleanup.cleanup_utils import find_worker_pod -from scripts.tenant_cleanup.cleanup_utils import get_tenant_status -from scripts.tenant_cleanup.cleanup_utils import read_tenant_ids_from_csv -from scripts.tenant_cleanup.cleanup_utils import TenantNotFoundInControlPlaneError +from scripts.tenant_cleanup.cleanup_utils import ( + confirm_step, + execute_control_plane_query, + find_worker_pod, + get_tenant_status, + read_tenant_ids_from_csv, + TenantNotFoundInControlPlaneError, +) def signal_handler(signum: int, frame: object) -> None: # noqa: ARG001 diff --git a/backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py b/backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py index 81aff3d011..98d85d21d1 100644 --- a/backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py +++ b/backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py @@ -27,16 +27,17 @@ Examples: import subprocess import sys -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, ThreadPoolExecutor from pathlib import Path from threading import Lock from typing import Any -from scripts.tenant_cleanup.cleanup_utils import confirm_step -from scripts.tenant_cleanup.cleanup_utils import find_worker_pod -from scripts.tenant_cleanup.cleanup_utils import get_tenant_status -from scripts.tenant_cleanup.cleanup_utils import read_tenant_ids_from_csv +from scripts.tenant_cleanup.cleanup_utils import ( + confirm_step, + find_worker_pod, + get_tenant_status, + read_tenant_ids_from_csv, +) # Global lock for thread-safe printing _print_lock: Lock = Lock() diff --git a/backend/scripts/tenant_cleanup/no_bastion_analyze_tenants.py b/backend/scripts/tenant_cleanup/no_bastion_analyze_tenants.py index 629c46f14d..80a29b857e 100644 --- a/backend/scripts/tenant_cleanup/no_bastion_analyze_tenants.py +++ b/backend/scripts/tenant_cleanup/no_bastion_analyze_tenants.py @@ -15,14 +15,14 @@ import csv import json import subprocess import sys -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any -from scripts.tenant_cleanup.no_bastion_cleanup_utils import find_background_pod -from scripts.tenant_cleanup.no_bastion_cleanup_utils import find_worker_pod +from scripts.tenant_cleanup.no_bastion_cleanup_utils import ( + find_background_pod, + find_worker_pod, +) def collect_tenant_data( diff --git a/backend/scripts/tenant_cleanup/no_bastion_cleanup_tenants.py b/backend/scripts/tenant_cleanup/no_bastion_cleanup_tenants.py index ed1d9e09b5..6aba3af76e 100644 --- a/backend/scripts/tenant_cleanup/no_bastion_cleanup_tenants.py +++ b/backend/scripts/tenant_cleanup/no_bastion_cleanup_tenants.py @@ -17,20 +17,18 @@ import json import signal import subprocess import sys -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor -from datetime import datetime -from datetime import timezone +from concurrent.futures import as_completed, ThreadPoolExecutor +from datetime import datetime, timezone from pathlib import Path from threading import Lock -from scripts.tenant_cleanup.no_bastion_cleanup_utils import confirm_step -from scripts.tenant_cleanup.no_bastion_cleanup_utils import execute_control_plane_delete -from scripts.tenant_cleanup.no_bastion_cleanup_utils import find_background_pod -from scripts.tenant_cleanup.no_bastion_cleanup_utils import find_worker_pod -from scripts.tenant_cleanup.no_bastion_cleanup_utils import get_tenant_status -from scripts.tenant_cleanup.no_bastion_cleanup_utils import read_tenant_ids_from_csv from scripts.tenant_cleanup.no_bastion_cleanup_utils import ( + confirm_step, + execute_control_plane_delete, + find_background_pod, + find_worker_pod, + get_tenant_status, + read_tenant_ids_from_csv, TenantNotFoundInControlPlaneError, ) diff --git a/backend/scripts/tenant_cleanup/no_bastion_mark_connectors.py b/backend/scripts/tenant_cleanup/no_bastion_mark_connectors.py index 4d9fcaccb5..9beea1e9d5 100644 --- a/backend/scripts/tenant_cleanup/no_bastion_mark_connectors.py +++ b/backend/scripts/tenant_cleanup/no_bastion_mark_connectors.py @@ -14,18 +14,17 @@ Usage: import subprocess import sys -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, ThreadPoolExecutor from pathlib import Path from threading import Lock from typing import Any -from scripts.tenant_cleanup.no_bastion_cleanup_utils import confirm_step -from scripts.tenant_cleanup.no_bastion_cleanup_utils import find_background_pod -from scripts.tenant_cleanup.no_bastion_cleanup_utils import find_worker_pod -from scripts.tenant_cleanup.no_bastion_cleanup_utils import get_tenant_status -from scripts.tenant_cleanup.no_bastion_cleanup_utils import read_tenant_ids_from_csv from scripts.tenant_cleanup.no_bastion_cleanup_utils import ( + confirm_step, + find_background_pod, + find_worker_pod, + get_tenant_status, + read_tenant_ids_from_csv, TenantNotFoundInControlPlaneError, ) diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py b/backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py index 3ef56fd7bc..8860c40d58 100644 --- a/backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py @@ -13,13 +13,10 @@ Output: import json import sys -from sqlalchemy import func -from sqlalchemy import select +from sqlalchemy import func, select -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Document +from onyx.db.engine.sql_engine import get_session_with_tenant, SqlEngine +from onyx.db.models import ConnectorCredentialPair, Document def check_documents_deleted(tenant_id: str) -> dict: diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py b/backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py index ebc39e6ee9..4567f0cb66 100644 --- a/backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py @@ -12,8 +12,7 @@ import sys from sqlalchemy import text -from onyx.db.engine.sql_engine import get_session_with_shared_schema -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_shared_schema, SqlEngine def drop_data_plane_schema(tenant_id: str) -> dict[str, str]: diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/execute_connector_deletion.py b/backend/scripts/tenant_cleanup/on_pod_scripts/execute_connector_deletion.py index 4fae39b536..4fe4fd31fc 100644 --- a/backend/scripts/tenant_cleanup/on_pod_scripts/execute_connector_deletion.py +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/execute_connector_deletion.py @@ -31,13 +31,13 @@ from typing import Any from sqlalchemy.orm import Session from onyx.background.celery.versioned_apps.client import app as client_app -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryTask -from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.connector_credential_pair import get_connector_credential_pairs -from onyx.db.connector_credential_pair import update_connector_credential_pair_from_id -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.configs.constants import OnyxCeleryPriority, OnyxCeleryTask +from onyx.db.connector_credential_pair import ( + get_connector_credential_pair_from_id, + get_connector_credential_pairs, + update_connector_credential_pair_from_id, +) +from onyx.db.engine.sql_engine import get_session_with_tenant, SqlEngine from onyx.db.enums import ConnectorCredentialPairStatus from onyx.db.index_attempt import cancel_indexing_attempts_for_ccpair diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py index c0e6968838..8d702ff1cd 100644 --- a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py @@ -29,8 +29,7 @@ import sys from sqlalchemy import select -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_tenant, SqlEngine from onyx.db.models import ConnectorCredentialPair diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py index f60adc0f07..2fb1453446 100644 --- a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py @@ -10,8 +10,7 @@ Usage: import json import sys -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_tenant, SqlEngine from onyx.db.search_settings import get_current_search_settings diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py index 0384cf86f6..d467a21bff 100644 --- a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py @@ -15,8 +15,7 @@ import sys from sqlalchemy import select -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_tenant, SqlEngine from onyx.db.models import User diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py b/backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py index 52cfbb55c7..b29bfa16c4 100644 --- a/backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py @@ -6,8 +6,7 @@ from sqlalchemy import text from sqlalchemy.exc import ProgrammingError from sqlalchemy.orm import Session -from onyx.db.engine.sql_engine import get_session_with_shared_schema -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_shared_schema, SqlEngine def get_tenant_activity_summary(session: Session) -> list[dict[str, Any]]: diff --git a/backend/scripts/update_connector_prune_freq.py b/backend/scripts/update_connector_prune_freq.py index a3894155ac..fcbefc0b0e 100644 --- a/backend/scripts/update_connector_prune_freq.py +++ b/backend/scripts/update_connector_prune_freq.py @@ -13,8 +13,7 @@ Example: import argparse import os import sys -from typing import Any -from typing import cast +from typing import Any, cast parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(parent_dir) @@ -23,8 +22,10 @@ from sqlalchemy import update # noqa: E402 from sqlalchemy.engine import CursorResult # noqa: E402 from sqlalchemy.exc import ProgrammingError # noqa: E402 -from onyx.db.engine.sql_engine import get_session_with_tenant # noqa: E402 -from onyx.db.engine.sql_engine import SqlEngine # noqa: E402 +from onyx.db.engine.sql_engine import ( # noqa: E402 + get_session_with_tenant, + SqlEngine, +) from onyx.db.engine.tenant_utils import get_all_tenant_ids # noqa: E402 from onyx.db.models import Connector # noqa: E402 from onyx.utils.variable_functionality import global_version # noqa: E402 diff --git a/backend/scripts/update_recommended_models.py b/backend/scripts/update_recommended_models.py index 7272ed489d..1aa3dc56c2 100644 --- a/backend/scripts/update_recommended_models.py +++ b/backend/scripts/update_recommended_models.py @@ -31,12 +31,9 @@ import re import sys import urllib.request from dataclasses import dataclass -from datetime import date -from datetime import datetime -from datetime import timezone +from datetime import date, datetime, timezone from pathlib import Path -from typing import Any -from typing import Literal +from typing import Any, Literal SCRIPT_DIR = Path(__file__).resolve().parent BACKEND_DIR = SCRIPT_DIR.parent diff --git a/backend/shared_configs/configs.py b/backend/shared_configs/configs.py index 54e66df19f..5a4a4c0342 100644 --- a/backend/shared_configs/configs.py +++ b/backend/shared_configs/configs.py @@ -1,6 +1,5 @@ import os -from typing import Any -from typing import List +from typing import Any, List from urllib.parse import urlparse # Used for logging diff --git a/backend/shared_configs/contextvars.py b/backend/shared_configs/contextvars.py index 81251bfac8..013f558f06 100644 --- a/backend/shared_configs/contextvars.py +++ b/backend/shared_configs/contextvars.py @@ -1,7 +1,6 @@ import contextvars -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA # Context variable for the current tenant id CURRENT_TENANT_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = ( diff --git a/backend/shared_configs/model_server_models.py b/backend/shared_configs/model_server_models.py index a38a2fd629..eda43fa102 100644 --- a/backend/shared_configs/model_server_models.py +++ b/backend/shared_configs/model_server_models.py @@ -1,8 +1,6 @@ from pydantic import BaseModel -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType -from shared_configs.enums import RerankerProvider +from shared_configs.enums import EmbeddingProvider, EmbedTextType, RerankerProvider Embedding = list[float] diff --git a/backend/tests/api/test_api.py b/backend/tests/api/test_api.py index 2b4fbe2afc..29d620d026 100644 --- a/backend/tests/api/test_api.py +++ b/backend/tests/api/test_api.py @@ -6,8 +6,7 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from onyx.configs.constants import DEV_VERSION_PATTERN -from onyx.configs.constants import STABLE_VERSION_PATTERN +from onyx.configs.constants import DEV_VERSION_PATTERN, STABLE_VERSION_PATTERN from onyx.main import fetch_versioned_implementation from onyx.utils.logger import setup_logger diff --git a/backend/tests/common/craft/skill_table_isolation.py b/backend/tests/common/craft/skill_table_isolation.py index 5cb078a248..d1913ee6d6 100644 --- a/backend/tests/common/craft/skill_table_isolation.py +++ b/backend/tests/common/craft/skill_table_isolation.py @@ -5,15 +5,16 @@ from __future__ import annotations from typing import Any from sqlalchemy import select -from sqlalchemy.orm import class_mapper -from sqlalchemy.orm import Session +from sqlalchemy.orm import class_mapper, Session -from onyx.db.models import ExternalApp -from onyx.db.models import ExternalAppPolicy -from onyx.db.models import ExternalAppUserCredential -from onyx.db.models import Skill -from onyx.db.models import Skill__User -from onyx.db.models import Skill__UserGroup +from onyx.db.models import ( + ExternalApp, + ExternalAppPolicy, + ExternalAppUserCredential, + Skill, + Skill__User, + Skill__UserGroup, +) # Parent -> child order (FKs all point child -> parent). Restore/insert in this # order; delete in reverse so FK constraints stay satisfied. diff --git a/backend/tests/common/craft/stubs.py b/backend/tests/common/craft/stubs.py index a95df4a636..af8f29cf81 100644 --- a/backend/tests/common/craft/stubs.py +++ b/backend/tests/common/craft/stubs.py @@ -50,22 +50,22 @@ Usage from __future__ import annotations import contextlib -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterable -from typing import Any -from typing import cast +from collections.abc import Callable, Generator, Iterable +from typing import Any, cast from uuid import UUID -from onyx.server.features.build.sandbox.base import SandboxEvent -from onyx.server.features.build.sandbox.base import SandboxManager -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import FilesystemEntry -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import SandboxInfo -from onyx.server.features.build.sandbox.models import SnapshotResult -from onyx.server.features.build.sandbox.serve_transport import PromptSlot -from onyx.server.features.build.sandbox.serve_transport import ServeConnectionInfo +from onyx.server.features.build.sandbox.base import SandboxEvent, SandboxManager +from onyx.server.features.build.sandbox.models import ( + FileSet, + FilesystemEntry, + LLMProviderConfig, + SandboxInfo, + SnapshotResult, +) +from onyx.server.features.build.sandbox.serve_transport import ( + PromptSlot, + ServeConnectionInfo, +) _UNSET = object() diff --git a/backend/tests/common/craft/users.py b/backend/tests/common/craft/users.py index 350d0712dd..2d732d756f 100644 --- a/backend/tests/common/craft/users.py +++ b/backend/tests/common/craft/users.py @@ -6,9 +6,11 @@ import httpx from onyx.auth.schemas import UserRole from tests.integration.common_utils.constants import GENERAL_HEADERS -from tests.integration.common_utils.managers.user import build_email -from tests.integration.common_utils.managers.user import DEFAULT_PASSWORD -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import ( + build_email, + DEFAULT_PASSWORD, + UserManager, +) from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index 1bcb468eb4..806639dd9f 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -4,8 +4,10 @@ from collections.abc import Generator import pytest -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) @pytest.fixture() diff --git a/backend/tests/daily/conftest.py b/backend/tests/daily/conftest.py index 8563943dad..7def2409ef 100644 --- a/backend/tests/daily/conftest.py +++ b/backend/tests/daily/conftest.py @@ -5,11 +5,15 @@ import os # module-level imports below pull in EE versioned implementations. os.environ["LICENSE_ENFORCEMENT_ENABLED"] = "false" -from collections.abc import AsyncGenerator # noqa: E402 -from collections.abc import Generator # noqa: E402 +from collections.abc import ( + AsyncGenerator, # noqa: E402 + Generator, # noqa: E402 +) from contextlib import asynccontextmanager # noqa: E402 -from unittest.mock import MagicMock # noqa: E402 -from unittest.mock import patch # noqa: E402 +from unittest.mock import ( + MagicMock, # noqa: E402 + patch, # noqa: E402 +) import pytest # noqa: E402 from dotenv import load_dotenv # noqa: E402 diff --git a/backend/tests/daily/connectors/airtable/test_airtable_basic.py b/backend/tests/daily/connectors/airtable/test_airtable_basic.py index 6f4af53b59..684b92f30a 100644 --- a/backend/tests/daily/connectors/airtable/test_airtable_basic.py +++ b/backend/tests/daily/connectors/airtable/test_airtable_basic.py @@ -7,10 +7,7 @@ from pydantic import BaseModel from onyx.configs.constants import DocumentSource from onyx.connectors.airtable.airtable_connector import AirtableConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, HierarchyNode, ImageSection, TextSection from tests.utils.secret_names import TestSecret pytestmark = pytest.mark.secrets(TestSecret.AIRTABLE_ACCESS_TOKEN) diff --git a/backend/tests/daily/connectors/blob/test_blob_connector.py b/backend/tests/daily/connectors/blob/test_blob_connector.py index a97f5dcdb8..d5901802c4 100644 --- a/backend/tests/daily/connectors/blob/test_blob_connector.py +++ b/backend/tests/daily/connectors/blob/test_blob_connector.py @@ -1,19 +1,13 @@ import os -from unittest.mock import MagicMock -from unittest.mock import patch -from urllib.parse import parse_qs -from urllib.parse import unquote -from urllib.parse import urlparse +from unittest.mock import MagicMock, patch +from urllib.parse import parse_qs, unquote, urlparse import pytest from onyx.configs.constants import BlobType from onyx.connectors.blob.connector import BlobStorageConnector from onyx.connectors.cross_connector_utils.tabular_section_utils import is_tabular_file -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, HierarchyNode, TabularSection, TextSection from onyx.file_processing.extract_file_text import get_file_ext from onyx.file_processing.file_types import OnyxFileExtensions from tests.daily.connectors.utils import set_test_staging_callback diff --git a/backend/tests/daily/connectors/box/test_box_basic.py b/backend/tests/daily/connectors/box/test_box_basic.py index d270b4290b..50722c89b4 100644 --- a/backend/tests/daily/connectors/box/test_box_basic.py +++ b/backend/tests/daily/connectors/box/test_box_basic.py @@ -13,13 +13,14 @@ from unittest.mock import MagicMock import pytest from ee.onyx.external_permissions.box.group_sync import box_group_sync -from onyx.connectors.box.connector import box_all_enterprise_users_group_id -from onyx.connectors.box.connector import box_group_id -from onyx.connectors.box.connector import BOX_ROOT_FOLDER_ID -from onyx.connectors.box.connector import BoxConnector +from onyx.connectors.box.connector import ( + box_all_enterprise_users_group_id, + box_group_id, + BOX_ROOT_FOLDER_ID, + BoxConnector, +) from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, TextSection from onyx.db.models import ConnectorCredentialPair from tests.daily.connectors.utils import load_all_from_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/canvas/test_canvas_perm_sync_daily.py b/backend/tests/daily/connectors/canvas/test_canvas_perm_sync_daily.py index 0d61bee78e..5112d4f57f 100644 --- a/backend/tests/daily/connectors/canvas/test_canvas_perm_sync_daily.py +++ b/backend/tests/daily/connectors/canvas/test_canvas_perm_sync_daily.py @@ -5,25 +5,26 @@ from unittest.mock import MagicMock import pytest from ee.onyx.connectors.perm_sync_valid import validate_canvas_perm_sync -from ee.onyx.external_permissions.canvas.access import build_course_permission_context -from ee.onyx.external_permissions.canvas.access import CanvasCoursePermissionContext +from ee.onyx.external_permissions.canvas.access import ( + build_course_permission_context, + CanvasCoursePermissionContext, +) 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 onyx.access.models import DocExternalAccess -from onyx.access.models import ExternalAccess -from onyx.connectors.canvas.connector import canvas_all_users_group_id -from onyx.connectors.canvas.connector import canvas_course_group_id -from onyx.connectors.canvas.connector import canvas_group_group_id -from onyx.connectors.canvas.connector import canvas_section_group_id -from onyx.connectors.canvas.connector import CanvasAnnouncement -from onyx.connectors.canvas.connector import CanvasAssignment -from onyx.connectors.canvas.connector import CanvasConnector +from onyx.access.models import DocExternalAccess, ExternalAccess +from onyx.connectors.canvas.connector import ( + canvas_all_users_group_id, + canvas_course_group_id, + canvas_group_group_id, + canvas_section_group_id, + CanvasAnnouncement, + CanvasAssignment, + CanvasConnector, +) from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument +from onyx.connectors.models import HierarchyNode, SlimDocument from onyx.db.models import ConnectorCredentialPair -from onyx.db.utils import DocumentRow -from onyx.db.utils import SortOrder +from onyx.db.utils import DocumentRow, SortOrder from tests.daily.connectors.utils import load_all_from_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/coda/test_coda_connector.py b/backend/tests/daily/connectors/coda/test_coda_connector.py index 3cf9fc6667..e0bcebecfd 100644 --- a/backend/tests/daily/connectors/coda/test_coda_connector.py +++ b/backend/tests/daily/connectors/coda/test_coda_connector.py @@ -8,8 +8,7 @@ import pytest from onyx.configs.constants import DocumentSource from onyx.connectors.coda.connector import CodaConnector from onyx.connectors.exceptions import CredentialInvalidError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode def connector_doc_generator( diff --git a/backend/tests/daily/connectors/confluence/test_confluence_basic.py b/backend/tests/daily/connectors/confluence/test_confluence_basic.py index ef19669d3c..da3ddd92bd 100644 --- a/backend/tests/daily/connectors/confluence/test_confluence_basic.py +++ b/backend/tests/daily/connectors/confluence/test_confluence_basic.py @@ -1,8 +1,7 @@ import os import time from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/daily/connectors/confluence/test_confluence_permissions_basic.py b/backend/tests/daily/connectors/confluence/test_confluence_permissions_basic.py index 639ce2326e..c5d8a1bd37 100644 --- a/backend/tests/daily/connectors/confluence/test_confluence_permissions_basic.py +++ b/backend/tests/daily/connectors/confluence/test_confluence_permissions_basic.py @@ -1,7 +1,6 @@ import os import time -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -12,8 +11,7 @@ from onyx.connectors.confluence.connector import ConfluenceConnector from onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider from onyx.connectors.models import HierarchyNode from onyx.db.models import ConnectorCredentialPair -from onyx.db.utils import DocumentRow -from onyx.db.utils import SortOrder +from onyx.db.utils import DocumentRow, SortOrder from tests.daily.connectors.utils import load_all_from_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/confluence/test_confluence_resolver.py b/backend/tests/daily/connectors/confluence/test_confluence_resolver.py index d459df9901..99472ecdb2 100644 --- a/backend/tests/daily/connectors/confluence/test_confluence_resolver.py +++ b/backend/tests/daily/connectors/confluence/test_confluence_resolver.py @@ -3,19 +3,20 @@ import os import re import time -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.configs.constants import DocumentSource from onyx.connectors.confluence.connector import ConfluenceConnector from onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + EntityFailure, + HierarchyNode, +) from tests.daily.connectors.utils import load_all_from_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/confluence/test_confluence_user_email_overrides.py b/backend/tests/daily/connectors/confluence/test_confluence_user_email_overrides.py index 64ca69bbfa..377d2046f2 100644 --- a/backend/tests/daily/connectors/confluence/test_confluence_user_email_overrides.py +++ b/backend/tests/daily/connectors/confluence/test_confluence_user_email_overrides.py @@ -1,8 +1,7 @@ import types from unittest.mock import patch -from onyx.connectors.confluence.onyx_confluence import ConfluenceUser -from onyx.connectors.confluence.onyx_confluence import OnyxConfluence +from onyx.connectors.confluence.onyx_confluence import ConfluenceUser, OnyxConfluence from onyx.connectors.interfaces import CredentialsProviderInterface diff --git a/backend/tests/daily/connectors/conftest.py b/backend/tests/daily/connectors/conftest.py index 88a00b57af..9f74aa1593 100644 --- a/backend/tests/daily/connectors/conftest.py +++ b/backend/tests/daily/connectors/conftest.py @@ -1,6 +1,5 @@ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/daily/connectors/discord/test_discord_connector.py b/backend/tests/daily/connectors/discord/test_discord_connector.py index af33462a7b..5a73f55361 100644 --- a/backend/tests/daily/connectors/discord/test_discord_connector.py +++ b/backend/tests/daily/connectors/discord/test_discord_connector.py @@ -3,9 +3,7 @@ import time import pytest from onyx.connectors.discord.connector import DiscordConnector -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, DocumentSource, HierarchyNode from tests.utils.secret_names import TestSecret pytestmark = pytest.mark.secrets(TestSecret.DISCORD_CONNECTOR_BOT_TOKEN) diff --git a/backend/tests/daily/connectors/file/test_file_connector.py b/backend/tests/daily/connectors/file/test_file_connector.py index abf60ed6e4..82a5ce8937 100644 --- a/backend/tests/daily/connectors/file/test_file_connector.py +++ b/backend/tests/daily/connectors/file/test_file_connector.py @@ -1,8 +1,6 @@ import io -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest diff --git a/backend/tests/daily/connectors/fireflies/test_fireflies_connector.py b/backend/tests/daily/connectors/fireflies/test_fireflies_connector.py index bd8e807929..4ea16899a7 100644 --- a/backend/tests/daily/connectors/fireflies/test_fireflies_connector.py +++ b/backend/tests/daily/connectors/fireflies/test_fireflies_connector.py @@ -7,8 +7,7 @@ import pytest from onyx.configs.constants import DocumentSource from onyx.connectors.fireflies.connector import FirefliesConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode from tests.utils.secret_names import TestSecret pytestmark = pytest.mark.secrets(TestSecret.FIREFLIES_API_KEY) diff --git a/backend/tests/daily/connectors/gmail/conftest.py b/backend/tests/daily/connectors/gmail/conftest.py index e6f9115f9c..92e60b79fb 100644 --- a/backend/tests/daily/connectors/gmail/conftest.py +++ b/backend/tests/daily/connectors/gmail/conftest.py @@ -7,15 +7,9 @@ import pytest from onyx.connectors.gmail.connector import GmailConnector from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, -) -from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY -from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_DICT_TOKEN_KEY, DB_CREDENTIALS_PRIMARY_ADMIN_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( GoogleOAuthAuthenticationMethod, ) from tests.load_env_vars import load_env_vars diff --git a/backend/tests/daily/connectors/gmail/test_gmail_connector.py b/backend/tests/daily/connectors/gmail/test_gmail_connector.py index 9535aaf058..ce651e4311 100644 --- a/backend/tests/daily/connectors/gmail/test_gmail_connector.py +++ b/backend/tests/daily/connectors/gmail/test_gmail_connector.py @@ -1,15 +1,11 @@ from collections.abc import Callable -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch import pytest from onyx.connectors.gmail.connector import GmailConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument +from onyx.connectors.models import Document, HierarchyNode, SlimDocument from tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/gong/test_gong.py b/backend/tests/daily/connectors/gong/test_gong.py index d72d1d7831..206ae6a1b0 100644 --- a/backend/tests/daily/connectors/gong/test_gong.py +++ b/backend/tests/daily/connectors/gong/test_gong.py @@ -1,6 +1,5 @@ import time -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/daily/connectors/google_drive/conftest.py b/backend/tests/daily/connectors/google_drive/conftest.py index fcb703d0ee..d6af9ad6f3 100644 --- a/backend/tests/daily/connectors/google_drive/conftest.py +++ b/backend/tests/daily/connectors/google_drive/conftest.py @@ -8,15 +8,9 @@ import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, -) -from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY -from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_DICT_TOKEN_KEY, DB_CREDENTIALS_PRIMARY_ADMIN_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( GoogleOAuthAuthenticationMethod, ) from tests.load_env_vars import load_env_vars diff --git a/backend/tests/daily/connectors/google_drive/consts_and_utils.py b/backend/tests/daily/connectors/google_drive/consts_and_utils.py index 0389767ca6..15e0c082e3 100644 --- a/backend/tests/daily/connectors/google_drive/consts_and_utils.py +++ b/backend/tests/daily/connectors/google_drive/consts_and_utils.py @@ -1,17 +1,12 @@ import time from collections.abc import Sequence -from dataclasses import dataclass -from dataclasses import field -from dataclasses import replace +from dataclasses import dataclass, field, replace from urllib.parse import urlparse from onyx.connectors.google_drive.connector import GoogleDriveConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, HierarchyNode, TextSection from onyx.db.enums import HierarchyNodeType -from tests.daily.connectors.utils import ConnectorOutput -from tests.daily.connectors.utils import load_all_from_connector +from tests.daily.connectors.utils import ConnectorOutput, load_all_from_connector ALL_FILES = list(range(0, 60)) SHARED_DRIVE_FILES = list(range(20, 25)) diff --git a/backend/tests/daily/connectors/google_drive/test_admin_oauth.py b/backend/tests/daily/connectors/google_drive/test_admin_oauth.py index 5c060d6120..e5a20cbde6 100644 --- a/backend/tests/daily/connectors/google_drive/test_admin_oauth.py +++ b/backend/tests/daily/connectors/google_drive/test_admin_oauth.py @@ -1,69 +1,46 @@ from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector -from tests.daily.connectors.google_drive.consts_and_utils import _pick -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FOLDER_3_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_MY_DRIVE_ID from tests.daily.connectors.google_drive.consts_and_utils import ( + _pick, + ADMIN_EMAIL, + ADMIN_FILE_IDS, + ADMIN_FOLDER_3_FILE_IDS, + ADMIN_MY_DRIVE_ID, ADMIN_SHORTCUT_FIXTURE_FOLDER_IDS, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_expected_docs_in_retrieved_docs, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_hierarchy_nodes_match_expected, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_resource_key_shortcut_target_in_retrieved_docs, -) -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_ID -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_URL -from tests.daily.connectors.google_drive.consts_and_utils import ( + FOLDER_1_1_FILE_IDS, + FOLDER_1_1_URL, + FOLDER_1_2_FILE_IDS, + FOLDER_1_2_URL, + FOLDER_1_FILE_IDS, + FOLDER_2_1_FILE_IDS, + FOLDER_2_1_URL, + FOLDER_2_2_FILE_IDS, + FOLDER_2_2_URL, + FOLDER_2_FILE_IDS, + FOLDER_2_URL, + FOLDER_3_ID, + FOLDER_3_URL, get_expected_hierarchy_for_shared_drives, -) -from tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs -from tests.daily.connectors.google_drive.consts_and_utils import ( + load_connector_outputs, PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_ONLY_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import PILL_FOLDER_ID -from tests.daily.connectors.google_drive.consts_and_utils import ( + PILL_FOLDER_ID, RESTRICTED_ACCESS_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FOLDER_ID -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_URL -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import ( + SECTIONS_FILE_IDS, + SECTIONS_FOLDER_ID, + SHARED_DRIVE_1_FILE_IDS, + SHARED_DRIVE_1_URL, + SHARED_DRIVE_2_FILE_IDS, TEST_USER_1_EXTRA_DRIVE_1_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_EXTRA_DRIVE_2_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_EXTRA_FOLDER_ID, ) from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/google_drive/test_drive_perm_sync.py b/backend/tests/daily/connectors/google_drive/test_drive_perm_sync.py index a8fb982621..25c71bbb53 100644 --- a/backend/tests/daily/connectors/google_drive/test_drive_perm_sync.py +++ b/backend/tests/daily/connectors/google_drive/test_drive_perm_sync.py @@ -3,8 +3,7 @@ import json import os from collections import defaultdict from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -13,66 +12,37 @@ from ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sy from onyx.access.models import DocExternalAccess from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.db.models import ConnectorCredentialPair -from onyx.db.utils import DocumentRow -from onyx.db.utils import SortOrder +from onyx.db.utils import DocumentRow, SortOrder from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface -from tests.daily.connectors.google_drive.consts_and_utils import _pick -from tests.daily.connectors.google_drive.consts_and_utils import ACCESS_MAPPING -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_MY_DRIVE_ID from tests.daily.connectors.google_drive.consts_and_utils import ( + _pick, + ACCESS_MAPPING, + ADMIN_EMAIL, + ADMIN_MY_DRIVE_ID, ADMIN_SHORTCUT_FIXTURE_FOLDER_IDS, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_hierarchy_nodes_match_expected, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_resource_key_shortcut_target_in_retrieved_docs, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( EXTERNAL_SHARED_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_ID -from tests.daily.connectors.google_drive.consts_and_utils import ( + FOLDER_3_ID, get_expected_hierarchy_for_shared_drives, -) -from tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs -from tests.daily.connectors.google_drive.consts_and_utils import ( + load_connector_outputs, PERM_SYNC_DRIVE_ACCESS_MAPPING, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_ONLY_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import PILL_FOLDER_ID -from tests.daily.connectors.google_drive.consts_and_utils import PUBLIC_RANGE -from tests.daily.connectors.google_drive.consts_and_utils import ( + PILL_FOLDER_ID, + PUBLIC_RANGE, RESTRICTED_ACCESS_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_DRIVE_B_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_DRIVE_B_ID -from tests.daily.connectors.google_drive.consts_and_utils import ( + TEST_USER_1_DRIVE_B_ID, TEST_USER_1_EXTRA_DRIVE_1_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_EXTRA_DRIVE_2_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_EXTRA_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_MY_DRIVE_FOLDER_ID, + TEST_USER_1_MY_DRIVE_ID, + TEST_USER_2_MY_DRIVE, + TEST_USER_3_MY_DRIVE_ID, ) -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_MY_DRIVE_ID -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_MY_DRIVE -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_MY_DRIVE_ID from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/google_drive/test_gdrive_shortcuts_daily.py b/backend/tests/daily/connectors/google_drive/test_gdrive_shortcuts_daily.py index 8a323d00db..613d48ed32 100644 --- a/backend/tests/daily/connectors/google_drive/test_gdrive_shortcuts_daily.py +++ b/backend/tests/daily/connectors/google_drive/test_gdrive_shortcuts_daily.py @@ -1,24 +1,16 @@ from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL +from onyx.connectors.models import Document, TextSection from tests.daily.connectors.google_drive.consts_and_utils import ( + ADMIN_EMAIL, assert_resource_key_shortcut_target_in_retrieved_docs, -) -from tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs -from tests.daily.connectors.google_drive.consts_and_utils import ( + load_connector_outputs, RESOURCE_KEY_SHORTCUT_TARGET_DOC_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( RESOURCE_KEY_SHORTCUT_TARGET_NAME, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( SHORTCUTS_GALORE_FOLDER_ID, ) from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/google_drive/test_link_visibility_filter.py b/backend/tests/daily/connectors/google_drive/test_link_visibility_filter.py index 072ef9905b..6ab4bde1b2 100644 --- a/backend/tests/daily/connectors/google_drive/test_link_visibility_filter.py +++ b/backend/tests/daily/connectors/google_drive/test_link_visibility_filter.py @@ -4,8 +4,7 @@ from unittest.mock import patch from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.connectors.google_drive.file_retrieval import has_link_only_permission -from onyx.connectors.google_drive.models import DriveRetrievalStage -from onyx.connectors.google_drive.models import RetrievedDriveFile +from onyx.connectors.google_drive.models import DriveRetrievalStage, RetrievedDriveFile def _stub_run_functions( diff --git a/backend/tests/daily/connectors/google_drive/test_map_test_ids.py b/backend/tests/daily/connectors/google_drive/test_map_test_ids.py index c13d442257..45433749d4 100644 --- a/backend/tests/daily/connectors/google_drive/test_map_test_ids.py +++ b/backend/tests/daily/connectors/google_drive/test_map_test_ids.py @@ -7,22 +7,24 @@ import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector from tests.daily.connectors.google_drive.conftest import build_credentials -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import file_name_template -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_FILE_IDS +from tests.daily.connectors.google_drive.consts_and_utils import ( + ADMIN_EMAIL, + ADMIN_FILE_IDS, + file_name_template, + FOLDER_1_1_FILE_IDS, + FOLDER_1_2_FILE_IDS, + FOLDER_1_FILE_IDS, + FOLDER_2_1_FILE_IDS, + FOLDER_2_2_FILE_IDS, + FOLDER_2_FILE_IDS, + FOLDER_3_FILE_IDS, + load_connector_outputs, + SHARED_DRIVE_1_FILE_IDS, + SHARED_DRIVE_2_FILE_IDS, + TEST_USER_1_FILE_IDS, + TEST_USER_2_FILE_IDS, + TEST_USER_3_FILE_IDS, +) from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/google_drive/test_resolver.py b/backend/tests/daily/connectors/google_drive/test_resolver.py index 8d1094e43f..e9dc71cb42 100644 --- a/backend/tests/daily/connectors/google_drive/test_resolver.py +++ b/backend/tests/daily/connectors/google_drive/test_resolver.py @@ -8,16 +8,18 @@ from unittest.mock import patch import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + HierarchyNode, +) from tests.daily.connectors.google_drive.consts_and_utils import ( + ADMIN_EMAIL, ALL_EXPECTED_HIERARCHY_NODES, + FOLDER_1_ID, + SHARED_DRIVE_1_ID, ) -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_ID -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_ID from tests.utils.secret_names import TestSecret _DRIVE_ID_MAPPING_PATH = os.path.join( diff --git a/backend/tests/daily/connectors/google_drive/test_sections.py b/backend/tests/daily/connectors/google_drive/test_sections.py index b00240a7a1..73b2350447 100644 --- a/backend/tests/daily/connectors/google_drive/test_sections.py +++ b/backend/tests/daily/connectors/google_drive/test_sections.py @@ -1,13 +1,14 @@ from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs -from tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FOLDER_URL +from tests.daily.connectors.google_drive.consts_and_utils import ( + ADMIN_EMAIL, + load_connector_outputs, + SECTIONS_FOLDER_URL, +) from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/google_drive/test_service_acct.py b/backend/tests/daily/connectors/google_drive/test_service_acct.py index d68796e5f2..05c09e5bb9 100644 --- a/backend/tests/daily/connectors/google_drive/test_service_acct.py +++ b/backend/tests/daily/connectors/google_drive/test_service_acct.py @@ -1,108 +1,69 @@ from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from urllib.parse import urlparse import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.connectors.google_utils.google_utils import execute_paginated_retrieval -from tests.daily.connectors.google_drive.consts_and_utils import _pick -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FOLDER_3_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_MY_DRIVE_ID from tests.daily.connectors.google_drive.consts_and_utils import ( + _pick, + ADMIN_EMAIL, + ADMIN_FILE_IDS, + ADMIN_FOLDER_3_FILE_IDS, + ADMIN_MY_DRIVE_ID, ADMIN_SHORTCUT_FIXTURE_FOLDER_IDS, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_expected_docs_in_retrieved_docs, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_hierarchy_nodes_match_expected, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_resource_key_shortcut_target_in_retrieved_docs, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( EXTERNAL_SHARED_DOC_SINGLETON, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( EXTERNAL_SHARED_DOCS_IN_FOLDER, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( EXTERNAL_SHARED_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( EXTERNAL_SHARED_FOLDER_URL, -) -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_ID -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_URL -from tests.daily.connectors.google_drive.consts_and_utils import ( + FOLDER_1_1_FILE_IDS, + FOLDER_1_1_URL, + FOLDER_1_2_FILE_IDS, + FOLDER_1_2_URL, + FOLDER_1_FILE_IDS, + FOLDER_2_1_FILE_IDS, + FOLDER_2_1_URL, + FOLDER_2_2_FILE_IDS, + FOLDER_2_2_URL, + FOLDER_2_FILE_IDS, + FOLDER_2_URL, + FOLDER_3_ID, + FOLDER_3_URL, get_expected_hierarchy_for_shared_drives, -) -from tests.daily.connectors.google_drive.consts_and_utils import id_to_name -from tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs -from tests.daily.connectors.google_drive.consts_and_utils import ( + id_to_name, + load_connector_outputs, MISC_SHARED_DRIVE_FNAMES, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( PERM_SYNC_DRIVE_ADMIN_ONLY_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import PILL_FOLDER_ID -from tests.daily.connectors.google_drive.consts_and_utils import ( + PILL_FOLDER_ID, RESTRICTED_ACCESS_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( RESTRICTED_ACCESS_FOLDER_URL, -) -from tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FOLDER_ID -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_URL -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import ( + SECTIONS_FILE_IDS, + SECTIONS_FOLDER_ID, + SHARED_DRIVE_1_FILE_IDS, + SHARED_DRIVE_1_URL, + SHARED_DRIVE_2_FILE_IDS, TEST_USER_1_DRIVE_B_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_DRIVE_B_ID -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import ( + TEST_USER_1_DRIVE_B_ID, + TEST_USER_1_EMAIL, TEST_USER_1_EXTRA_DRIVE_1_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_EXTRA_DRIVE_2_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( TEST_USER_1_EXTRA_FOLDER_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import ( + TEST_USER_1_FILE_IDS, TEST_USER_1_MY_DRIVE_FOLDER_ID, + TEST_USER_1_MY_DRIVE_ID, + TEST_USER_2_EMAIL, + TEST_USER_2_FILE_IDS, + TEST_USER_2_MY_DRIVE, + TEST_USER_3_EMAIL, + TEST_USER_3_FILE_IDS, + TEST_USER_3_MY_DRIVE_ID, ) -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_MY_DRIVE_ID -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_MY_DRIVE -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_MY_DRIVE_ID from tests.utils.secret_names import TestSecret pytestmark = pytest.mark.secrets(TestSecret.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR) diff --git a/backend/tests/daily/connectors/google_drive/test_user_1_oauth.py b/backend/tests/daily/connectors/google_drive/test_user_1_oauth.py index c8cbfd855e..23f4a0cd41 100644 --- a/backend/tests/daily/connectors/google_drive/test_user_1_oauth.py +++ b/backend/tests/daily/connectors/google_drive/test_user_1_oauth.py @@ -1,49 +1,36 @@ from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.connectors.models import Document -from tests.daily.connectors.google_drive.consts_and_utils import _clear_parents -from tests.daily.connectors.google_drive.consts_and_utils import _pick -from tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FOLDER_3_FILE_IDS from tests.daily.connectors.google_drive.consts_and_utils import ( + _clear_parents, + _pick, + ADMIN_FOLDER_3_FILE_IDS, assert_expected_docs_in_retrieved_docs, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( assert_hierarchy_nodes_match_expected, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( DONWLOAD_REVOKED_FILE_ID, -) -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_ID -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_ID -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_ID -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_URL -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_ID -from tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_URL -from tests.daily.connectors.google_drive.consts_and_utils import ( + FOLDER_1_1_FILE_IDS, + FOLDER_1_1_ID, + FOLDER_1_2_FILE_IDS, + FOLDER_1_2_ID, + FOLDER_1_FILE_IDS, + FOLDER_1_ID, + FOLDER_1_URL, + FOLDER_3_ID, + FOLDER_3_URL, get_expected_hierarchy_for_test_user_1, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( get_expected_hierarchy_for_test_user_1_my_drive_only, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( get_expected_hierarchy_for_test_user_1_shared_drives_only, -) -from tests.daily.connectors.google_drive.consts_and_utils import ( get_expected_hierarchy_for_test_user_1_shared_with_me_only, + load_connector_outputs, + SHARED_DRIVE_1_FILE_IDS, + SHARED_DRIVE_1_ID, + TEST_USER_1_EMAIL, + TEST_USER_1_FILE_IDS, ) -from tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS -from tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_ID -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_EMAIL -from tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_FILE_IDS from tests.daily.connectors.utils import ConnectorOutput from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/highspot/test_highspot_connector.py b/backend/tests/daily/connectors/highspot/test_highspot_connector.py index e3d57fe297..bc1d4b8bfa 100644 --- a/backend/tests/daily/connectors/highspot/test_highspot_connector.py +++ b/backend/tests/daily/connectors/highspot/test_highspot_connector.py @@ -3,15 +3,13 @@ import os import time from datetime import datetime from pathlib import Path -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.configs.constants import DocumentSource from onyx.connectors.highspot.connector import HighspotConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode from tests.utils.secret_names import TestSecret pytestmark = pytest.mark.secrets( diff --git a/backend/tests/daily/connectors/hubspot/test_hubspot_connector.py b/backend/tests/daily/connectors/hubspot/test_hubspot_connector.py index 5d61b2aaaf..c78fd5126d 100644 --- a/backend/tests/daily/connectors/hubspot/test_hubspot_connector.py +++ b/backend/tests/daily/connectors/hubspot/test_hubspot_connector.py @@ -1,16 +1,15 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.connectors.hubspot.connector import AVAILABLE_OBJECT_TYPES -from onyx.connectors.hubspot.connector import HubSpotConnector -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.hubspot.connector import AVAILABLE_OBJECT_TYPES, HubSpotConnector +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, +) from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/notion/test_notion_connector.py b/backend/tests/daily/connectors/notion/test_notion_connector.py index 60cf104041..55eb8752eb 100644 --- a/backend/tests/daily/connectors/notion/test_notion_connector.py +++ b/backend/tests/daily/connectors/notion/test_notion_connector.py @@ -3,8 +3,7 @@ import time import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode from onyx.connectors.notion.connector import NotionConnector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/outline/test_outline_connector.py b/backend/tests/daily/connectors/outline/test_outline_connector.py index 1070c77ba2..a399305682 100644 --- a/backend/tests/daily/connectors/outline/test_outline_connector.py +++ b/backend/tests/daily/connectors/outline/test_outline_connector.py @@ -5,11 +5,12 @@ from typing import Any import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.exceptions import ConnectorValidationError, CredentialExpiredError +from onyx.connectors.models import ( + ConnectorMissingCredentialError, + Document, + HierarchyNode, +) from onyx.connectors.outline.connector import OutlineConnector diff --git a/backend/tests/daily/connectors/salesforce/test_salesforce_connector.py b/backend/tests/daily/connectors/salesforce/test_salesforce_connector.py index 2610e1db07..28bad22e01 100644 --- a/backend/tests/daily/connectors/salesforce/test_salesforce_connector.py +++ b/backend/tests/daily/connectors/salesforce/test_salesforce_connector.py @@ -1,16 +1,13 @@ import json import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from pathlib import Path -from typing import Any -from typing import cast +from typing import Any, cast import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode from onyx.connectors.salesforce.connector import SalesforceConnector from onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/sharepoint/test_sharepoint_connector.py b/backend/tests/daily/connectors/sharepoint/test_sharepoint_connector.py index 2aafaca1e7..d6cee8c658 100644 --- a/backend/tests/daily/connectors/sharepoint/test_sharepoint_connector.py +++ b/backend/tests/daily/connectors/sharepoint/test_sharepoint_connector.py @@ -1,21 +1,23 @@ import os import time from dataclasses import dataclass -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.sharepoint.connector import SharepointAuthMethod -from onyx.connectors.sharepoint.connector import SharepointConnector +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + HierarchyNode, + ImageSection, +) +from onyx.connectors.sharepoint.connector import ( + SharepointAuthMethod, + SharepointConnector, +) from onyx.db.enums import HierarchyNodeType from tests.daily.connectors.utils import load_all_from_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/slab/test_slab_connector.py b/backend/tests/daily/connectors/slab/test_slab_connector.py index a7f0dd24c0..61465300b4 100644 --- a/backend/tests/daily/connectors/slab/test_slab_connector.py +++ b/backend/tests/daily/connectors/slab/test_slab_connector.py @@ -5,8 +5,7 @@ from pathlib import Path import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode from onyx.connectors.slab.connector import SlabConnector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/slack/test_slack_connector.py b/backend/tests/daily/connectors/slack/test_slack_connector.py index ac9f375ccf..1df1a29cba 100644 --- a/backend/tests/daily/connectors/slack/test_slack_connector.py +++ b/backend/tests/daily/connectors/slack/test_slack_connector.py @@ -4,9 +4,11 @@ import pytest from onyx.connectors.slack.connector import SlackConnector from onyx.db.enums import HierarchyNodeType -from tests.daily.connectors.utils import load_all_from_connector -from tests.daily.connectors.utils import to_sections -from tests.daily.connectors.utils import to_text_sections +from tests.daily.connectors.utils import ( + load_all_from_connector, + to_sections, + to_text_sections, +) from tests.utils.secret_names import TestSecret pytestmark = pytest.mark.secrets(TestSecret.SLACK_BOT_TOKEN) diff --git a/backend/tests/daily/connectors/slack/test_slack_perm_sync.py b/backend/tests/daily/connectors/slack/test_slack_perm_sync.py index d16c34c590..db8d651f95 100644 --- a/backend/tests/daily/connectors/slack/test_slack_perm_sync.py +++ b/backend/tests/daily/connectors/slack/test_slack_perm_sync.py @@ -1,11 +1,9 @@ import time -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument +from onyx.connectors.models import HierarchyNode, SlimDocument from onyx.connectors.slack.connector import SlackConnector from tests.daily.connectors.utils import load_all_from_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/connectors/utils.py b/backend/tests/daily/connectors/utils.py index 2e6d4444cb..253eb0e6c7 100644 --- a/backend/tests/daily/connectors/utils.py +++ b/backend/tests/daily/connectors/utils.py @@ -1,21 +1,24 @@ from collections.abc import Iterator -from typing import IO -from typing import TypeVar +from typing import IO, TypeVar from pydantic import BaseModel from onyx.connectors.connector_runner import CheckpointOutputWrapper -from onyx.connectors.interfaces import BaseConnector -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointedConnectorWithPermSync -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.interfaces import ( + BaseConnector, + CheckpointedConnector, + CheckpointedConnectorWithPermSync, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + Document, + HierarchyNode, + ImageSection, + TabularSection, + TextSection, +) _ITERATION_LIMIT = 100_000 diff --git a/backend/tests/daily/connectors/web/test_web_connector.py b/backend/tests/daily/connectors/web/test_web_connector.py index 35f20cbe17..5c72ea92b3 100644 --- a/backend/tests/daily/connectors/web/test_web_connector.py +++ b/backend/tests/daily/connectors/web/test_web_connector.py @@ -2,10 +2,8 @@ from concurrent.futures import ThreadPoolExecutor import pytest -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS -from onyx.connectors.web.connector import WebConnector +from onyx.connectors.models import Document, HierarchyNode +from onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS, WebConnector EXPECTED_QUOTE = ( "If you can't explain it to a six year old, you don't understand it yourself." diff --git a/backend/tests/daily/connectors/zendesk/test_zendesk_connector.py b/backend/tests/daily/connectors/zendesk/test_zendesk_connector.py index 623b8e2ca5..1172049e87 100644 --- a/backend/tests/daily/connectors/zendesk/test_zendesk_connector.py +++ b/backend/tests/daily/connectors/zendesk/test_zendesk_connector.py @@ -6,8 +6,7 @@ from typing import cast import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.models import Document, HierarchyNode from onyx.connectors.zendesk.connector import ZendeskConnector from tests.daily.connectors.utils import load_all_from_connector from tests.utils.secret_names import TestSecret diff --git a/backend/tests/daily/embedding/test_embeddings.py b/backend/tests/daily/embedding/test_embeddings.py index 4579678c18..a9f13e4a9e 100644 --- a/backend/tests/daily/embedding/test_embeddings.py +++ b/backend/tests/daily/embedding/test_embeddings.py @@ -1,8 +1,10 @@ import pytest -from tenacity import retry -from tenacity import retry_if_exception_type -from tenacity import stop_after_attempt -from tenacity import wait_exponential +from tenacity import ( + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) from onyx.natural_language_processing.search_nlp_models import EmbeddingModel from shared_configs.enums import EmbedTextType diff --git a/backend/tests/daily/targeted_reindex/conftest.py b/backend/tests/daily/targeted_reindex/conftest.py index 0a529e0feb..7f81876a17 100644 --- a/backend/tests/daily/targeted_reindex/conftest.py +++ b/backend/tests/daily/targeted_reindex/conftest.py @@ -16,8 +16,7 @@ from collections.abc import Generator import pytest from sqlalchemy.orm import Session -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/daily/targeted_reindex/helpers.py b/backend/tests/daily/targeted_reindex/helpers.py index 3129f901c4..34ba7bcf1d 100644 --- a/backend/tests/daily/targeted_reindex/helpers.py +++ b/backend/tests/daily/targeted_reindex/helpers.py @@ -15,27 +15,21 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_PRIMARY_ADMIN_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( GoogleOAuthAuthenticationMethod, ) from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError -from onyx.db.models import TargetedReindexJob -from onyx.db.models import TargetedReindexJobTarget +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, IndexingStatus +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + IndexAttempt, + IndexAttemptError, + TargetedReindexJob, + TargetedReindexJobTarget, +) from onyx.db.search_settings import get_current_search_settings _ADMIN_EMAIL = "admin@onyx-test.com" diff --git a/backend/tests/daily/targeted_reindex/test_drive_targeted_reindex_integration.py b/backend/tests/daily/targeted_reindex/test_drive_targeted_reindex_integration.py index 8a524be5c7..52c5a75d93 100644 --- a/backend/tests/daily/targeted_reindex/test_drive_targeted_reindex_integration.py +++ b/backend/tests/daily/targeted_reindex/test_drive_targeted_reindex_integration.py @@ -33,13 +33,17 @@ from onyx.background.celery.tasks.docprocessing.targeted_reindex_task import ( from onyx.db.enums import IndexingStatus from onyx.db.models import Document as DBDocument from onyx.db.models import IndexAttemptError -from onyx.db.targeted_reindex import create_targeted_reindex_job -from onyx.db.targeted_reindex import get_targeted_reindex_job -from onyx.db.targeted_reindex import resolve_error_ids_to_targets -from tests.daily.targeted_reindex.helpers import cleanup_targeted_reindex_state -from tests.daily.targeted_reindex.helpers import make_drive_cc_pair -from tests.daily.targeted_reindex.helpers import make_failed_index_attempt -from tests.daily.targeted_reindex.helpers import make_index_attempt_error +from onyx.db.targeted_reindex import ( + create_targeted_reindex_job, + get_targeted_reindex_job, + resolve_error_ids_to_targets, +) +from tests.daily.targeted_reindex.helpers import ( + cleanup_targeted_reindex_state, + make_drive_cc_pair, + make_failed_index_attempt, + make_index_attempt_error, +) from tests.utils.secret_names import TestSecret _DRIVE_ID_MAPPING_PATH = os.path.join( diff --git a/backend/tests/evals/connector_filter_eval/conftest.py b/backend/tests/evals/connector_filter_eval/conftest.py index 7581bab650..52d4f82835 100644 --- a/backend/tests/evals/connector_filter_eval/conftest.py +++ b/backend/tests/evals/connector_filter_eval/conftest.py @@ -5,33 +5,34 @@ import pytest from sqlalchemy.orm import Session from onyx.db.enums import LLMModelFlowType -from onyx.db.llm import fetch_existing_llm_providers -from onyx.db.llm import fetch_llm_provider_view -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import ( + fetch_existing_llm_providers, + fetch_llm_provider_view, + update_default_provider, + upsert_llm_provider, +) from onyx.llm.constants import LlmProviderNames -from onyx.llm.factory import get_default_llm -from onyx.llm.factory import llm_from_provider +from onyx.llm.factory import get_default_llm, llm_from_provider from onyx.llm.interfaces import LLM -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) # This suite lives outside tests/external_dependency_unit (so the EDU CI # matrix doesn't discover it) but runs against the same live dependencies, # so it reuses that tree's fixtures via explicit imports. -from tests.external_dependency_unit.answer.conftest import ( # noqa: F401 +from tests.external_dependency_unit.answer.conftest import ( # noqa: F401 # noqa: F401 # noqa: F401 mock_external_deps, -) -from tests.external_dependency_unit.answer.conftest import mock_file_store # noqa: F401 -from tests.external_dependency_unit.answer.conftest import mock_gpu_status # noqa: F401 -from tests.external_dependency_unit.answer.conftest import ( # noqa: F401 + mock_file_store, # noqa: F401 + mock_gpu_status, # noqa: F401 mock_nlp_embeddings_post, -) -from tests.external_dependency_unit.answer.conftest import ( # noqa: F401 mock_vespa_query, ) -from tests.external_dependency_unit.conftest import db_session # noqa: F401 -from tests.external_dependency_unit.conftest import full_deployment_setup # noqa: F401 +from tests.external_dependency_unit.conftest import ( + db_session, # noqa: F401 + full_deployment_setup, # noqa: F401 +) # This suite makes real LLM calls, so it is disabled by default and only runs # when explicitly opted in. diff --git a/backend/tests/evals/connector_filter_eval/scope_eval_cases.py b/backend/tests/evals/connector_filter_eval/scope_eval_cases.py index 2282b1656b..dadebbe01b 100644 --- a/backend/tests/evals/connector_filter_eval/scope_eval_cases.py +++ b/backend/tests/evals/connector_filter_eval/scope_eval_cases.py @@ -15,8 +15,7 @@ carrying) across turns. from __future__ import annotations -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.configs.constants import DocumentSource from onyx.secondary_llm_flows.source_filter import SearchCycle diff --git a/backend/tests/evals/connector_filter_eval/test_filter_extraction_regression.py b/backend/tests/evals/connector_filter_eval/test_filter_extraction_regression.py index 3070d1eac8..8c757fc449 100644 --- a/backend/tests/evals/connector_filter_eval/test_filter_extraction_regression.py +++ b/backend/tests/evals/connector_filter_eval/test_filter_extraction_regression.py @@ -26,14 +26,15 @@ from collections import defaultdict from pydantic import BaseModel -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MessageType +from onyx.configs.constants import DocumentSource, MessageType from onyx.llm.interfaces import LLM from onyx.secondary_llm_flows.source_filter import decide_search_scope from onyx.tools.models import ChatMinimalTextMessage from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel -from tests.evals.connector_filter_eval.scope_eval_cases import SCOPE_EVAL_CASES -from tests.evals.connector_filter_eval.scope_eval_cases import ScopeEvalCase +from tests.evals.connector_filter_eval.scope_eval_cases import ( + SCOPE_EVAL_CASES, + ScopeEvalCase, +) _ATTEMPTS_PER_CASE = int(os.environ.get("CONNECTOR_FILTER_EVAL_ATTEMPTS", "2")) _MAX_WORKERS = 4 diff --git a/backend/tests/external_dependency_unit/answer/conftest.py b/backend/tests/external_dependency_unit/answer/conftest.py index 0623adadc4..3653477464 100644 --- a/backend/tests/external_dependency_unit/answer/conftest.py +++ b/backend/tests/external_dependency_unit/answer/conftest.py @@ -1,18 +1,17 @@ import os -from collections.abc import Iterator -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from sqlalchemy.orm import Session -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import update_default_provider, upsert_llm_provider from onyx.llm.constants import LlmProviderNames -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) # Counter for generating unique file IDs in mock file store _mock_file_id_counter = 0 diff --git a/backend/tests/external_dependency_unit/answer/stream_test_assertions.py b/backend/tests/external_dependency_unit/answer/stream_test_assertions.py index ea5f70b9f4..a789a36cc6 100644 --- a/backend/tests/external_dependency_unit/answer/stream_test_assertions.py +++ b/backend/tests/external_dependency_unit/answer/stream_test_assertions.py @@ -2,15 +2,16 @@ from __future__ import annotations from typing import cast -from onyx.chat.models import AnswerStreamPart -from onyx.chat.models import CreateChatSessionID +from onyx.chat.models import AnswerStreamPart, CreateChatSessionID from onyx.context.search.models import SearchDoc from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import ImageGenerationFinal -from onyx.server.query_and_chat.streaming_models import OpenUrlDocuments -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import SearchToolDocumentsDelta +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseStart, + ImageGenerationFinal, + OpenUrlDocuments, + Packet, + SearchToolDocumentsDelta, +) def assert_answer_stream_part_correct( diff --git a/backend/tests/external_dependency_unit/answer/stream_test_builder.py b/backend/tests/external_dependency_unit/answer/stream_test_builder.py index cf8d9b63ac..8f40cb10f0 100644 --- a/backend/tests/external_dependency_unit/answer/stream_test_builder.py +++ b/backend/tests/external_dependency_unit/answer/stream_test_builder.py @@ -4,23 +4,22 @@ from collections.abc import Iterator from onyx.chat.models import AnswerStreamPart from onyx.context.search.models import SearchDoc -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ReasoningDone -from onyx.server.query_and_chat.streaming_models import ReasoningStart +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseStart, + OverallStop, + Packet, + ReasoningDone, + ReasoningStart, +) from tests.external_dependency_unit.answer.stream_test_assertions import ( assert_answer_stream_part_correct, ) from tests.external_dependency_unit.answer.stream_test_utils import ( create_packet_with_agent_response_delta, -) -from tests.external_dependency_unit.answer.stream_test_utils import ( create_packet_with_reasoning_delta, + create_placement, ) -from tests.external_dependency_unit.answer.stream_test_utils import create_placement -from tests.external_dependency_unit.mock_llm import LLMResponse -from tests.external_dependency_unit.mock_llm import MockLLMController +from tests.external_dependency_unit.mock_llm import LLMResponse, MockLLMController class StreamTestBuilder: diff --git a/backend/tests/external_dependency_unit/answer/stream_test_utils.py b/backend/tests/external_dependency_unit/answer/stream_test_utils.py index 75d1377287..267daec442 100644 --- a/backend/tests/external_dependency_unit/answer/stream_test_utils.py +++ b/backend/tests/external_dependency_unit/answer/stream_test_utils.py @@ -10,15 +10,18 @@ from onyx.chat.models import AnswerStreamPart from onyx.chat.process_message import handle_stream_message_objects from onyx.configs.constants import DocumentSource from onyx.context.search.models import SearchDoc -from onyx.db.models import ChatSession -from onyx.db.models import User +from onyx.db.models import ChatSession, User from onyx.llm.override_models import LLMOverride -from onyx.server.query_and_chat.models import ChatSessionCreationRequest -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.server.query_and_chat.models import ( + ChatSessionCreationRequest, + SendMessageRequest, +) from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ReasoningDelta +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + Packet, + ReasoningDelta, +) from tests.external_dependency_unit.mock_content_provider import MockWebContent from tests.external_dependency_unit.mock_search_provider import MockWebSearchResult diff --git a/backend/tests/external_dependency_unit/answer/test_answer_without_openai.py b/backend/tests/external_dependency_unit/answer/test_answer_without_openai.py index af9b3fb205..3251e4cc29 100644 --- a/backend/tests/external_dependency_unit/answer/test_answer_without_openai.py +++ b/backend/tests/external_dependency_unit/answer/test_answer_without_openai.py @@ -5,23 +5,27 @@ from uuid import uuid4 from sqlalchemy.orm import Session -from onyx.chat.models import AnswerStreamPart -from onyx.chat.models import StreamingError +from onyx.chat.models import AnswerStreamPart, StreamingError from onyx.chat.process_message import handle_stream_message_objects from onyx.db.chat import create_chat_session from onyx.db.enums import LLMModelFlowType -from onyx.db.llm import fetch_existing_llm_providers -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import ( + fetch_existing_llm_providers, + remove_llm_provider, + update_default_provider, + upsert_llm_provider, +) from onyx.llm.constants import LlmProviderNames -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest -from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.models import SendMessageRequest -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) +from onyx.server.query_and_chat.models import MessageResponseIDInfo, SendMessageRequest +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, + Packet, +) from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/answer/test_current_datetime_replacement.py b/backend/tests/external_dependency_unit/answer/test_current_datetime_replacement.py index 19f530aaf9..a55123ac8e 100644 --- a/backend/tests/external_dependency_unit/answer/test_current_datetime_replacement.py +++ b/backend/tests/external_dependency_unit/answer/test_current_datetime_replacement.py @@ -3,14 +3,12 @@ from datetime import datetime from sqlalchemy.orm import Session -from onyx.chat.models import AnswerStreamPart -from onyx.chat.models import StreamingError +from onyx.chat.models import AnswerStreamPart, StreamingError from onyx.chat.process_message import handle_stream_message_objects from onyx.db.chat import create_chat_session from onyx.db.models import User from onyx.db.persona import get_persona_by_id -from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.server.query_and_chat.models import MessageResponseIDInfo, SendMessageRequest from onyx.server.query_and_chat.streaming_models import AgentResponseDelta from tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/answer/test_stream_chat_message.py b/backend/tests/external_dependency_unit/answer/test_stream_chat_message.py index eac523c07c..179e834b79 100644 --- a/backend/tests/external_dependency_unit/answer/test_stream_chat_message.py +++ b/backend/tests/external_dependency_unit/answer/test_stream_chat_message.py @@ -9,60 +9,62 @@ from sqlalchemy.orm import Session from onyx.chat.models import CreateChatSessionID from onyx.configs.constants import DocumentSource from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.streaming_models import AgentResponseStart -from onyx.server.query_and_chat.streaming_models import GeneratedImage -from onyx.server.query_and_chat.streaming_models import ImageGenerationFinal -from onyx.server.query_and_chat.streaming_models import ImageGenerationToolHeartbeat -from onyx.server.query_and_chat.streaming_models import ImageGenerationToolStart -from onyx.server.query_and_chat.streaming_models import OpenUrlDocuments -from onyx.server.query_and_chat.streaming_models import OpenUrlStart -from onyx.server.query_and_chat.streaming_models import OpenUrlUrls -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ReasoningDone -from onyx.server.query_and_chat.streaming_models import ReasoningStart -from onyx.server.query_and_chat.streaming_models import SearchToolDocumentsDelta -from onyx.server.query_and_chat.streaming_models import SearchToolQueriesDelta -from onyx.server.query_and_chat.streaming_models import SearchToolStart -from onyx.server.query_and_chat.streaming_models import SectionEnd -from onyx.server.query_and_chat.streaming_models import TopLevelBranching +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseStart, + GeneratedImage, + ImageGenerationFinal, + ImageGenerationToolHeartbeat, + ImageGenerationToolStart, + OpenUrlDocuments, + OpenUrlStart, + OpenUrlUrls, + OverallStop, + Packet, + ReasoningDone, + ReasoningStart, + SearchToolDocumentsDelta, + SearchToolQueriesDelta, + SearchToolStart, + SectionEnd, + TopLevelBranching, +) from tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider from tests.external_dependency_unit.answer.stream_test_assertions import ( assert_answer_stream_part_correct, ) from tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder -from tests.external_dependency_unit.answer.stream_test_utils import create_chat_session from tests.external_dependency_unit.answer.stream_test_utils import ( + create_chat_session, create_packet_with_agent_response_delta, -) -from tests.external_dependency_unit.answer.stream_test_utils import ( create_packet_with_reasoning_delta, -) -from tests.external_dependency_unit.answer.stream_test_utils import create_placement -from tests.external_dependency_unit.answer.stream_test_utils import ( + create_placement, mock_web_content_to_search_doc, -) -from tests.external_dependency_unit.answer.stream_test_utils import ( mock_web_search_result_to_search_doc, + submit_query, + tokenise, ) -from tests.external_dependency_unit.answer.stream_test_utils import submit_query -from tests.external_dependency_unit.answer.stream_test_utils import tokenise from tests.external_dependency_unit.conftest import create_test_user -from tests.external_dependency_unit.mock_content_provider import MockWebContent from tests.external_dependency_unit.mock_content_provider import ( + MockWebContent, use_mock_content_provider, ) from tests.external_dependency_unit.mock_image_provider import ( use_mock_image_generation_provider, ) -from tests.external_dependency_unit.mock_llm import LLMAnswerResponse -from tests.external_dependency_unit.mock_llm import LLMReasoningResponse -from tests.external_dependency_unit.mock_llm import LLMToolCallResponse -from tests.external_dependency_unit.mock_llm import use_mock_llm -from tests.external_dependency_unit.mock_search_pipeline import MockInternalSearchResult -from tests.external_dependency_unit.mock_search_pipeline import use_mock_search_pipeline -from tests.external_dependency_unit.mock_search_provider import MockWebSearchResult -from tests.external_dependency_unit.mock_search_provider import use_mock_web_provider +from tests.external_dependency_unit.mock_llm import ( + LLMAnswerResponse, + LLMReasoningResponse, + LLMToolCallResponse, + use_mock_llm, +) +from tests.external_dependency_unit.mock_search_pipeline import ( + MockInternalSearchResult, + use_mock_search_pipeline, +) +from tests.external_dependency_unit.mock_search_provider import ( + MockWebSearchResult, + use_mock_web_provider, +) def test_stream_chat_with_answer( diff --git a/backend/tests/external_dependency_unit/answer/test_stream_chat_message_objects.py b/backend/tests/external_dependency_unit/answer/test_stream_chat_message_objects.py index 9eb568586c..5535ff5bf1 100644 --- a/backend/tests/external_dependency_unit/answer/test_stream_chat_message_objects.py +++ b/backend/tests/external_dependency_unit/answer/test_stream_chat_message_objects.py @@ -5,16 +5,13 @@ from unittest.mock import MagicMock import pytest from sqlalchemy.orm import Session -from onyx.chat.models import AnswerStreamPart -from onyx.chat.models import StreamingError +from onyx.chat.models import AnswerStreamPart, StreamingError from onyx.chat.process_message import handle_stream_message_objects from onyx.db.chat import create_chat_session from onyx.db.models import User from onyx.db.persona import upsert_persona -from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.models import SendMessageRequest -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import Packet +from onyx.server.query_and_chat.models import MessageResponseIDInfo, SendMessageRequest +from onyx.server.query_and_chat.streaming_models import AgentResponseDelta, Packet from tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/auth/test_mobile_sso_bridge.py b/backend/tests/external_dependency_unit/auth/test_mobile_sso_bridge.py index 47effbe8ba..5f46ceab86 100644 --- a/backend/tests/external_dependency_unit/auth/test_mobile_sso_bridge.py +++ b/backend/tests/external_dependency_unit/auth/test_mobile_sso_bridge.py @@ -17,27 +17,19 @@ Coverage: """ import asyncio -from typing import Any -from typing import cast -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch -from urllib.parse import parse_qs -from urllib.parse import urlparse +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch +from urllib.parse import parse_qs, urlparse import httpx -from fastapi import FastAPI -from fastapi import Response +from fastapi import FastAPI, Response from fastapi.testclient import TestClient -from fastapi_users.authentication import AuthenticationBackend -from fastapi_users.authentication import CookieTransport +from fastapi_users.authentication import AuthenticationBackend, CookieTransport from fastapi_users.jwt import decode_jwt from httpx_oauth.oauth2 import BaseOAuth2 from onyx.auth.mobile_sso.code_store import consume_sso_code -from onyx.auth.users import generate_pkce_pair -from onyx.auth.users import get_oauth_router -from onyx.auth.users import STATE_TOKEN_AUDIENCE +from onyx.auth.users import generate_pkce_pair, get_oauth_router, STATE_TOKEN_AUDIENCE from onyx.error_handling.exceptions import register_onyx_exception_handlers _STATE_SECRET = "test-secret" diff --git a/backend/tests/external_dependency_unit/auth/test_mobile_sso_code_store.py b/backend/tests/external_dependency_unit/auth/test_mobile_sso_code_store.py index 869cf233a9..3d2032536f 100644 --- a/backend/tests/external_dependency_unit/auth/test_mobile_sso_code_store.py +++ b/backend/tests/external_dependency_unit/auth/test_mobile_sso_code_store.py @@ -10,12 +10,10 @@ import json import pytest -from onyx.auth.mobile_sso.code_store import consume_sso_code -from onyx.auth.mobile_sso.code_store import store_sso_code +from onyx.auth.mobile_sso.code_store import consume_sso_code, store_sso_code from onyx.auth.pkce import compute_s256_challenge from onyx.auth.users import generate_pkce_pair -from onyx.configs.app_configs import MOBILE_SSO_CODE_PREFIX -from onyx.configs.app_configs import MOBILE_SSO_CODE_TTL_SECONDS +from onyx.configs.app_configs import MOBILE_SSO_CODE_PREFIX, MOBILE_SSO_CODE_TTL_SECONDS from onyx.redis.redis_pool import get_async_redis_connection diff --git a/backend/tests/external_dependency_unit/auth/test_oidc_multi_live_idp.py b/backend/tests/external_dependency_unit/auth/test_oidc_multi_live_idp.py index 00c6bac87f..0052fe42f9 100644 --- a/backend/tests/external_dependency_unit/auth/test_oidc_multi_live_idp.py +++ b/backend/tests/external_dependency_unit/auth/test_oidc_multi_live_idp.py @@ -23,15 +23,12 @@ import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from httpx import ASGITransport -from httpx import AsyncClient +from httpx import ASGITransport, AsyncClient from sqlalchemy.orm import Session from onyx.auth.users import cookie_transport from onyx.db.enums import SSOProviderType -from onyx.db.models import SSOProvider -from onyx.db.models import User -from onyx.db.models import User__UserGroup +from onyx.db.models import SSOProvider, User, User__UserGroup from onyx.db.sso_provider import create_sso_provider from onyx.error_handling.exceptions import register_onyx_exception_handlers from onyx.server import oidc_multi diff --git a/backend/tests/external_dependency_unit/auth/test_session_rejection_contract.py b/backend/tests/external_dependency_unit/auth/test_session_rejection_contract.py index ba9edace24..0108bf809c 100644 --- a/backend/tests/external_dependency_unit/auth/test_session_rejection_contract.py +++ b/backend/tests/external_dependency_unit/auth/test_session_rejection_contract.py @@ -6,15 +6,12 @@ The classification matrix lives in test_session_token_strategy.py. import json import secrets -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import cast import pytest from fastapi import FastAPI -from httpx import ASGITransport -from httpx import AsyncClient +from httpx import ASGITransport, AsyncClient from sqlalchemy.orm import Session from onyx.auth.users import TenantAwareRedisStrategy diff --git a/backend/tests/external_dependency_unit/auth/test_session_token_strategy.py b/backend/tests/external_dependency_unit/auth/test_session_token_strategy.py index f187fa953a..839edf3ec5 100644 --- a/backend/tests/external_dependency_unit/auth/test_session_token_strategy.py +++ b/backend/tests/external_dependency_unit/auth/test_session_token_strategy.py @@ -7,30 +7,27 @@ request-level 403 contract lives in test_session_rejection_contract.py. import contextvars import json import secrets -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import cast import pytest from fastapi import Request from sqlalchemy.orm import Session -from onyx.auth.session_tokens import get_session_rejection -from onyx.auth.session_tokens import record_session_rejection -from onyx.auth.session_tokens import SESSION_TOKEN_GRACE_PERIOD_SECONDS -from onyx.auth.session_tokens import SessionRejection -from onyx.auth.session_tokens import SessionRejectionReason -from onyx.auth.session_tokens import SessionTokenValue -from onyx.auth.users import TenantAwareRedisStrategy -from onyx.auth.users import UserManager -from onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX -from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS +from onyx.auth.session_tokens import ( + get_session_rejection, + record_session_rejection, + SESSION_TOKEN_GRACE_PERIOD_SECONDS, + SessionRejection, + SessionRejectionReason, + SessionTokenValue, +) +from onyx.auth.users import TenantAwareRedisStrategy, UserManager +from onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX, SESSION_EXPIRE_TIME_SECONDS from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.db.auth import SQLAlchemyUserAdminDB from onyx.db.engine.async_sql_engine import get_async_session_context_manager -from onyx.db.models import OAuthAccount -from onyx.db.models import User +from onyx.db.models import OAuthAccount, User from onyx.redis.redis_pool import get_raw_redis_client from onyx.server.manage.users import get_current_auth_token_expiry_redis from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/auth/test_sso_admin_api.py b/backend/tests/external_dependency_unit/auth/test_sso_admin_api.py index 2374f22ee6..73b12fa307 100644 --- a/backend/tests/external_dependency_unit/auth/test_sso_admin_api.py +++ b/backend/tests/external_dependency_unit/auth/test_sso_admin_api.py @@ -13,8 +13,7 @@ from uuid import uuid4 import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from onyx.auth.users import current_user @@ -23,8 +22,7 @@ from onyx.db.engine.sql_engine import get_session from onyx.db.enums import SSOProviderType from onyx.db.models import SSOProvider from onyx.error_handling.error_codes import OnyxErrorCode -from onyx.error_handling.exceptions import OnyxError -from onyx.error_handling.exceptions import register_onyx_exception_handlers +from onyx.error_handling.exceptions import OnyxError, register_onyx_exception_handlers from onyx.server.manage.sso.api import admin_router from onyx.utils.encryption import is_masked_credential diff --git a/backend/tests/external_dependency_unit/background/test_periodic_task_claim.py b/backend/tests/external_dependency_unit/background/test_periodic_task_claim.py index d3ee5e211d..1e9ace2534 100644 --- a/backend/tests/external_dependency_unit/background/test_periodic_task_claim.py +++ b/backend/tests/external_dependency_unit/background/test_periodic_task_claim.py @@ -10,22 +10,20 @@ the task runs, so long-running tasks don't hold connections. import time from collections.abc import Generator -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from concurrent.futures import as_completed, ThreadPoolExecutor +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock from uuid import uuid4 import pytest -from onyx.background.periodic_poller import _PeriodicTaskDef -from onyx.background.periodic_poller import _try_claim_task -from onyx.background.periodic_poller import _try_run_periodic_task -from onyx.background.periodic_poller import PERIODIC_TASK_KV_PREFIX -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.background.periodic_poller import ( + _PeriodicTaskDef, + _try_claim_task, + _try_run_periodic_task, + PERIODIC_TASK_KV_PREFIX, +) +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from onyx.db.models import KVStore from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/external_dependency_unit/background/test_startup_recovery.py b/backend/tests/external_dependency_unit/background/test_startup_recovery.py index 3195f997a3..31de91bd9d 100644 --- a/backend/tests/external_dependency_unit/background/test_startup_recovery.py +++ b/backend/tests/external_dependency_unit/background/test_startup_recovery.py @@ -11,10 +11,8 @@ the correct files. """ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch -from uuid import UUID -from uuid import uuid4 +from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 import pytest import sqlalchemy as sa diff --git a/backend/tests/external_dependency_unit/cache/test_cache_backend_parity.py b/backend/tests/external_dependency_unit/cache/test_cache_backend_parity.py index 7f1f0164d2..b5d4b56f98 100644 --- a/backend/tests/external_dependency_unit/cache/test_cache_backend_parity.py +++ b/backend/tests/external_dependency_unit/cache/test_cache_backend_parity.py @@ -8,9 +8,7 @@ in conftest.py. import time from uuid import uuid4 -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import TTL_KEY_NOT_FOUND -from onyx.cache.interface import TTL_NO_EXPIRY +from onyx.cache.interface import CacheBackend, TTL_KEY_NOT_FOUND, TTL_NO_EXPIRY def _key() -> str: diff --git a/backend/tests/external_dependency_unit/cache/test_encrypted_kv_store_accessor.py b/backend/tests/external_dependency_unit/cache/test_encrypted_kv_store_accessor.py index 498d075aa5..e6b2a0fa9e 100644 --- a/backend/tests/external_dependency_unit/cache/test_encrypted_kv_store_accessor.py +++ b/backend/tests/external_dependency_unit/cache/test_encrypted_kv_store_accessor.py @@ -5,9 +5,11 @@ from collections.abc import Generator import pytest from sqlalchemy import delete -from onyx.db.encrypted_kv_store import delete_encrypted_kv -from onyx.db.encrypted_kv_store import load_encrypted_kv -from onyx.db.encrypted_kv_store import upsert_encrypted_kv +from onyx.db.encrypted_kv_store import ( + delete_encrypted_kv, + load_encrypted_kv, + upsert_encrypted_kv, +) from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.models import EncryptedKeyValueStore from onyx.key_value_store.interface import KvKeyNotFoundError diff --git a/backend/tests/external_dependency_unit/cache/test_kv_store_cache_layer.py b/backend/tests/external_dependency_unit/cache/test_kv_store_cache_layer.py index f50820e52e..95ff8e9ca8 100644 --- a/backend/tests/external_dependency_unit/cache/test_kv_store_cache_layer.py +++ b/backend/tests/external_dependency_unit/cache/test_kv_store_cache_layer.py @@ -18,11 +18,9 @@ from sqlalchemy import delete from onyx.cache.interface import CacheBackend from onyx.cache.postgres_backend import PostgresCacheBackend from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.models import CacheStore -from onyx.db.models import KVStore +from onyx.db.models import CacheStore, KVStore from onyx.key_value_store.interface import KvKeyNotFoundError -from onyx.key_value_store.store import PgRedisKVStore -from onyx.key_value_store.store import REDIS_KEY_PREFIX +from onyx.key_value_store.store import PgRedisKVStore, REDIS_KEY_PREFIX from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE diff --git a/backend/tests/external_dependency_unit/cache/test_postgres_cache_backend.py b/backend/tests/external_dependency_unit/cache/test_postgres_cache_backend.py index 54975c54a5..d58c937d50 100644 --- a/backend/tests/external_dependency_unit/cache/test_postgres_cache_backend.py +++ b/backend/tests/external_dependency_unit/cache/test_postgres_cache_backend.py @@ -10,10 +10,11 @@ from uuid import uuid4 from sqlalchemy import select -from onyx.cache.interface import TTL_KEY_NOT_FOUND -from onyx.cache.interface import TTL_NO_EXPIRY -from onyx.cache.postgres_backend import cleanup_expired_cache_entries -from onyx.cache.postgres_backend import PostgresCacheBackend +from onyx.cache.interface import TTL_KEY_NOT_FOUND, TTL_NO_EXPIRY +from onyx.cache.postgres_backend import ( + cleanup_expired_cache_entries, + PostgresCacheBackend, +) from onyx.db.models import CacheStore diff --git a/backend/tests/external_dependency_unit/cache/test_stream_buffer_backends.py b/backend/tests/external_dependency_unit/cache/test_stream_buffer_backends.py index 76f0bb28c4..ea332de650 100644 --- a/backend/tests/external_dependency_unit/cache/test_stream_buffer_backends.py +++ b/backend/tests/external_dependency_unit/cache/test_stream_buffer_backends.py @@ -5,9 +5,7 @@ the Postgres (lite) backend.""" from uuid import uuid4 from onyx.cache.interface import CacheBackend -from onyx.chat.stream_buffer import _chunk_key -from onyx.chat.stream_buffer import read_stream_chunks -from onyx.chat.stream_buffer import StreamBufferWriter +from onyx.chat.stream_buffer import _chunk_key, read_stream_chunks, StreamBufferWriter def test_roundtrip_and_done(cache: CacheBackend) -> None: diff --git a/backend/tests/external_dependency_unit/celery/test_created_at_backfill.py b/backend/tests/external_dependency_unit/celery/test_created_at_backfill.py index bcb653012f..1ce32a7ac7 100644 --- a/backend/tests/external_dependency_unit/celery/test_created_at_backfill.py +++ b/backend/tests/external_dependency_unit/celery/test_created_at_backfill.py @@ -9,19 +9,19 @@ Validates that the pruning-side collection of a source creation time works: Uses a mock SlimConnector plus a real PostgreSQL database. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from uuid import uuid4 from sqlalchemy.orm import Session from onyx.background.celery.celery_utils import extract_ids_from_runnable_connector -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument +from onyx.connectors.interfaces import ( + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnector, +) +from onyx.connectors.models import HierarchyNode, SlimDocument from onyx.db.document import backfill_docs_created_at__no_commit from onyx.db.models import Document as DbDocument from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface diff --git a/backend/tests/external_dependency_unit/celery/test_docfetching_priority.py b/backend/tests/external_dependency_unit/celery/test_docfetching_priority.py index 8694c96460..dbd0dd7c3e 100644 --- a/backend/tests/external_dependency_unit/celery/test_docfetching_priority.py +++ b/backend/tests/external_dependency_unit/celery/test_docfetching_priority.py @@ -8,10 +8,8 @@ previously completed indexing. Uses real Redis for locking and real database objects for CC pairs and search settings. """ -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -20,17 +18,20 @@ from sqlalchemy.orm import Session from onyx.background.celery.tasks.docfetching.task_creation_utils import ( try_creating_docfetching_task, ) -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCeleryPriority +from onyx.configs.constants import DocumentSource, OnyxCeleryPriority from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import IndexModelStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import SearchSettings +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + EmbeddingPrecision, + IndexModelStatus, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + SearchSettings, +) from onyx.redis.redis_pool import get_redis_client from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE diff --git a/backend/tests/external_dependency_unit/celery/test_docprocessing_priority.py b/backend/tests/external_dependency_unit/celery/test_docprocessing_priority.py index 6081b11b7e..0d6a27a6e0 100644 --- a/backend/tests/external_dependency_unit/celery/test_docprocessing_priority.py +++ b/backend/tests/external_dependency_unit/celery/test_docprocessing_priority.py @@ -7,29 +7,30 @@ get the correct priority based on last_successful_index_time. Uses real database objects for CC pairs, search settings, and index attempts. """ -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session from onyx.background.indexing.run_docfetching import connector_document_extraction -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import OnyxCeleryPriority +from onyx.configs.constants import DocumentSource, OnyxCeleryPriority from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import IndexAttempt -from onyx.db.models import SearchSettings +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + EmbeddingPrecision, + IndexingStatus, + IndexModelStatus, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + IndexAttempt, + SearchSettings, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE diff --git a/backend/tests/external_dependency_unit/celery/test_document_sync_revoke_cleanup.py b/backend/tests/external_dependency_unit/celery/test_document_sync_revoke_cleanup.py index b295a83b46..b0b49bfe36 100644 --- a/backend/tests/external_dependency_unit/celery/test_document_sync_revoke_cleanup.py +++ b/backend/tests/external_dependency_unit/celery/test_document_sync_revoke_cleanup.py @@ -13,12 +13,14 @@ from uuid import uuid4 import pytest from onyx.background.celery.apps.app_base import on_task_revoked -from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_PREFIX -from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_TASKSET_KEY -from onyx.background.celery.tasks.vespa.document_sync import get_document_sync_remaining -from onyx.background.celery.tasks.vespa.document_sync import is_document_sync_fenced -from onyx.background.celery.tasks.vespa.document_sync import reset_document_sync -from onyx.background.celery.tasks.vespa.document_sync import set_document_sync_fence +from onyx.background.celery.tasks.vespa.document_sync import ( + DOCUMENT_SYNC_PREFIX, + DOCUMENT_SYNC_TASKSET_KEY, + get_document_sync_remaining, + is_document_sync_fenced, + reset_document_sync, + set_document_sync_fence, +) from onyx.background.celery.tasks.vespa.tasks import monitor_document_sync_taskset from onyx.redis.redis_pool import get_redis_client from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/tests/external_dependency_unit/celery/test_persona_file_sync.py b/backend/tests/external_dependency_unit/celery/test_persona_file_sync.py index a16b84b35a..542a2e9e1d 100644 --- a/backend/tests/external_dependency_unit/celery/test_persona_file_sync.py +++ b/backend/tests/external_dependency_unit/celery/test_persona_file_sync.py @@ -20,27 +20,18 @@ since we only need to verify the arguments passed to update_single. from collections.abc import Generator from contextlib import contextmanager from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch -from unittest.mock import PropertyMock +from unittest.mock import MagicMock, patch, PropertyMock from uuid import uuid4 from sqlalchemy.orm import Session from onyx.background.celery.tasks.user_file_processing.tasks import ( check_for_user_file_project_sync, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( process_single_user_file_project_sync, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( user_file_project_sync_lock_key, ) from onyx.db.enums import UserFileStatus -from onyx.db.models import Persona -from onyx.db.models import Persona__UserFile -from onyx.db.models import User -from onyx.db.models import UserFile +from onyx.db.models import Persona, Persona__UserFile, User, UserFile from onyx.db.persona import upsert_persona from onyx.document_index.interfaces_new import MetadataUpdateRequest from onyx.redis.redis_pool import get_redis_client @@ -308,8 +299,7 @@ class TestSyncTaskWritesPersonaIds: tenant_context: None, # noqa: ARG002 ) -> None: """A file linked to both a project and a persona gets both IDs.""" - from onyx.db.models import Project__UserFile - from onyx.db.models import UserProject + from onyx.db.models import Project__UserFile, UserProject user = create_test_user(db_session, "sync_both") uf = _create_completed_user_file( diff --git a/backend/tests/external_dependency_unit/celery/test_pruning_hierarchy_nodes.py b/backend/tests/external_dependency_unit/celery/test_pruning_hierarchy_nodes.py index d30d55952d..49b5398102 100644 --- a/backend/tests/external_dependency_unit/celery/test_pruning_hierarchy_nodes.py +++ b/backend/tests/external_dependency_unit/celery/test_pruning_hierarchy_nodes.py @@ -22,31 +22,34 @@ from sqlalchemy.orm import Session from onyx.access.models import ExternalAccess from onyx.background.celery.celery_utils import extract_ids_from_runnable_connector from onyx.configs.constants import DocumentSource -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.interfaces import SlimConnectorWithPermSync +from onyx.connectors.interfaces import ( + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, + SlimConnectorWithPermSync, +) from onyx.connectors.models import HierarchyNode as PydanticHierarchyNode -from onyx.connectors.models import InputType -from onyx.connectors.models import SlimDocument -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import HierarchyNodeType -from onyx.db.hierarchy import delete_orphaned_hierarchy_nodes -from onyx.db.hierarchy import ensure_source_node_exists -from onyx.db.hierarchy import get_all_hierarchy_nodes_for_source -from onyx.db.hierarchy import get_hierarchy_node_by_raw_id -from onyx.db.hierarchy import link_hierarchy_nodes_to_documents -from onyx.db.hierarchy import remove_stale_hierarchy_node_cc_pair_entries -from onyx.db.hierarchy import reparent_orphaned_hierarchy_nodes -from onyx.db.hierarchy import update_document_parent_hierarchy_nodes -from onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries -from onyx.db.hierarchy import upsert_hierarchy_nodes_batch -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential +from onyx.connectors.models import InputType, SlimDocument +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, HierarchyNodeType +from onyx.db.hierarchy import ( + delete_orphaned_hierarchy_nodes, + ensure_source_node_exists, + get_all_hierarchy_nodes_for_source, + get_hierarchy_node_by_raw_id, + link_hierarchy_nodes_to_documents, + remove_stale_hierarchy_node_cc_pair_entries, + reparent_orphaned_hierarchy_nodes, + update_document_parent_hierarchy_nodes, + upsert_hierarchy_node_cc_pair_entries, + upsert_hierarchy_nodes_batch, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + HierarchyNodeByConnectorCredentialPair, +) from onyx.db.models import Document as DbDocument from onyx.db.models import HierarchyNode as DBHierarchyNode -from onyx.db.models import HierarchyNodeByConnectorCredentialPair from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface from onyx.kg.models import KGStage diff --git a/backend/tests/external_dependency_unit/celery/test_user_file_delete_queue.py b/backend/tests/external_dependency_unit/celery/test_user_file_delete_queue.py index 6560d27989..e40eddb6f9 100644 --- a/backend/tests/external_dependency_unit/celery/test_user_file_delete_queue.py +++ b/backend/tests/external_dependency_unit/celery/test_user_file_delete_queue.py @@ -25,29 +25,23 @@ on the task class so no real broker is needed. from collections.abc import Generator from contextlib import contextmanager from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch -from unittest.mock import PropertyMock +from unittest.mock import MagicMock, patch, PropertyMock from uuid import uuid4 from sqlalchemy.orm import Session from onyx.background.celery.tasks.user_file_processing.tasks import ( _user_file_delete_lock_key, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( _user_file_delete_queued_key, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( check_for_user_file_delete, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( process_single_user_file_delete, ) -from onyx.configs.constants import CELERY_USER_FILE_DELETE_TASK_EXPIRES -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import USER_FILE_DELETE_MAX_QUEUE_DEPTH +from onyx.configs.constants import ( + CELERY_USER_FILE_DELETE_TASK_EXPIRES, + OnyxCeleryQueues, + OnyxCeleryTask, + USER_FILE_DELETE_MAX_QUEUE_DEPTH, +) from onyx.db.enums import UserFileStatus from onyx.db.models import UserFile from onyx.redis.redis_pool import get_redis_client diff --git a/backend/tests/external_dependency_unit/celery/test_user_file_indexing_adapter.py b/backend/tests/external_dependency_unit/celery/test_user_file_indexing_adapter.py index 45362bc878..383c92cf69 100644 --- a/backend/tests/external_dependency_unit/celery/test_user_file_indexing_adapter.py +++ b/backend/tests/external_dependency_unit/celery/test_user_file_indexing_adapter.py @@ -9,26 +9,25 @@ Uses real PostgreSQL for UserFile/Persona/UserProject rows. Mocks the LLM tokenizer and file store since they are not relevant here. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, TextSection from onyx.db.enums import UserFileStatus -from onyx.db.models import Persona -from onyx.db.models import Persona__UserFile -from onyx.db.models import Project__UserFile -from onyx.db.models import User -from onyx.db.models import UserFile -from onyx.db.models import UserProject +from onyx.db.models import ( + Persona, + Persona__UserFile, + Project__UserFile, + User, + UserFile, + UserProject, +) from onyx.indexing.adapters.user_file_indexing_adapter import UserFileIndexingAdapter from onyx.indexing.indexing_pipeline import DocumentBatchPrepareContext -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import IndexChunk +from onyx.indexing.models import ChunkEmbedding, IndexChunk from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/celery/test_user_file_processing_queue.py b/backend/tests/external_dependency_unit/celery/test_user_file_processing_queue.py index 9e50ae67c2..2c47987062 100644 --- a/backend/tests/external_dependency_unit/celery/test_user_file_processing_queue.py +++ b/backend/tests/external_dependency_unit/celery/test_user_file_processing_queue.py @@ -25,27 +25,23 @@ on the task class so no real broker is needed. from collections.abc import Generator from contextlib import contextmanager from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch -from unittest.mock import PropertyMock +from unittest.mock import MagicMock, patch, PropertyMock from uuid import uuid4 from sqlalchemy.orm import Session -from onyx.background.celery.tasks.user_file_processing.tasks import _user_file_lock_key from onyx.background.celery.tasks.user_file_processing.tasks import ( + _user_file_lock_key, _user_file_queued_key, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( check_user_file_processing, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( process_single_user_file, ) -from onyx.configs.constants import CELERY_USER_FILE_PROCESSING_TASK_EXPIRES -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import USER_FILE_PROCESSING_MAX_QUEUE_DEPTH +from onyx.configs.constants import ( + CELERY_USER_FILE_PROCESSING_TASK_EXPIRES, + OnyxCeleryQueues, + OnyxCeleryTask, + USER_FILE_PROCESSING_MAX_QUEUE_DEPTH, +) from onyx.db.enums import UserFileStatus from onyx.db.models import UserFile from onyx.redis.redis_pool import get_redis_client diff --git a/backend/tests/external_dependency_unit/chat/test_build_python_chat_files.py b/backend/tests/external_dependency_unit/chat/test_build_python_chat_files.py index a5e130431b..5fbea08ee4 100644 --- a/backend/tests/external_dependency_unit/chat/test_build_python_chat_files.py +++ b/backend/tests/external_dependency_unit/chat/test_build_python_chat_files.py @@ -15,8 +15,7 @@ import pytest from sqlalchemy.orm import Session from onyx.chat.chat_utils import build_python_chat_files_from_search_docs -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.constants import DocumentSource, FileOrigin from onyx.context.search.models import SearchDoc from onyx.file_store.file_store import get_default_file_store diff --git a/backend/tests/external_dependency_unit/chat/test_lazy_chat_file_loading.py b/backend/tests/external_dependency_unit/chat/test_lazy_chat_file_loading.py index 225f4fea0e..1845a92b39 100644 --- a/backend/tests/external_dependency_unit/chat/test_lazy_chat_file_loading.py +++ b/backend/tests/external_dependency_unit/chat/test_lazy_chat_file_loading.py @@ -18,18 +18,17 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.chat.chat_utils import load_all_chat_files -from onyx.chat.chat_utils import load_chat_file +from onyx.chat.chat_utils import load_all_chat_files, load_chat_file from onyx.chat.models import ChatLoadedFile -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import MessageType -from onyx.db.chat import create_chat_session -from onyx.db.chat import create_new_chat_message -from onyx.db.chat import get_or_create_root_message +from onyx.configs.constants import FileOrigin, MessageType +from onyx.db.chat import ( + create_chat_session, + create_new_chat_message, + get_or_create_root_message, +) from onyx.file_store import file_store as file_store_module from onyx.file_store.file_store import get_default_file_store -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import FileDescriptor +from onyx.file_store.models import ChatFileType, FileDescriptor from onyx.tools.models import ChatFile from tests.external_dependency_unit.conftest import create_test_user @@ -404,8 +403,7 @@ class TestLoadAllChatFilesLazy: def test_max_workers_capped_at_16(self) -> None: """Defense-in-depth: even with 200 files passed in, the thread pool is capped at 16 workers.""" - from typing import Any - from typing import cast + from typing import Any, cast captured: dict[str, int] = {} diff --git a/backend/tests/external_dependency_unit/chat/test_spreadsheet_preview.py b/backend/tests/external_dependency_unit/chat/test_spreadsheet_preview.py index 234f958cfb..2c68c53be1 100644 --- a/backend/tests/external_dependency_unit/chat/test_spreadsheet_preview.py +++ b/backend/tests/external_dependency_unit/chat/test_spreadsheet_preview.py @@ -15,9 +15,11 @@ from onyx.configs.constants import FileOrigin from onyx.file_processing.file_types import SPREADSHEET_MIME_TYPE from onyx.file_store.file_store import get_default_file_store from onyx.file_store.models import ChatFileType -from onyx.server.query_and_chat.chat_utils import is_spreadsheet_mime_type -from onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type -from onyx.server.query_and_chat.chat_utils import parse_spreadsheet_for_preview +from onyx.server.query_and_chat.chat_utils import ( + is_spreadsheet_mime_type, + mime_type_to_chat_file_type, + parse_spreadsheet_for_preview, +) def _build_xlsx() -> bytes: diff --git a/backend/tests/external_dependency_unit/chat/test_user_reminder_message_type.py b/backend/tests/external_dependency_unit/chat/test_user_reminder_message_type.py index 342b887254..f9d5d975bf 100644 --- a/backend/tests/external_dependency_unit/chat/test_user_reminder_message_type.py +++ b/backend/tests/external_dependency_unit/chat/test_user_reminder_message_type.py @@ -14,12 +14,9 @@ from onyx.chat.llm_step import translate_history_to_llm_format from onyx.chat.models import ChatMessageSimple from onyx.configs.constants import MessageType from onyx.llm.interfaces import LLMConfig -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import SystemMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ChatCompletionMessage, SystemMessage, UserMessage from onyx.prompts.chat_prompts import CODE_BLOCK_MARKDOWN -from onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE -from onyx.prompts.constants import SYSTEM_REMINDER_TAG_OPEN +from onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE, SYSTEM_REMINDER_TAG_OPEN def _ensure_list( diff --git a/backend/tests/external_dependency_unit/conftest.py b/backend/tests/external_dependency_unit/conftest.py index 4edfb38a71..79294908c5 100644 --- a/backend/tests/external_dependency_unit/conftest.py +++ b/backend/tests/external_dependency_unit/conftest.py @@ -5,11 +5,9 @@ import pytest from fastapi_users.password import PasswordHelper from sqlalchemy.orm import Session -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from onyx.db.enums import AccountType -from onyx.db.models import User -from onyx.db.models import UserRole +from onyx.db.models import User, UserRole from onyx.file_store.file_store import get_default_file_store from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/external_dependency_unit/connectors/confluence/test_confluence_group_sync.py b/backend/tests/external_dependency_unit/connectors/confluence/test_confluence_group_sync.py index 49e3a49764..8f28ac449a 100644 --- a/backend/tests/external_dependency_unit/connectors/confluence/test_confluence_group_sync.py +++ b/backend/tests/external_dependency_unit/connectors/confluence/test_confluence_group_sync.py @@ -5,11 +5,8 @@ from sqlalchemy.orm import Session from ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus +from onyx.db.models import Connector, ConnectorCredentialPair, Credential from shared_configs.contextvars import get_current_tenant_id from tests.daily.connectors.confluence.models import ExternalUserGroupSet diff --git a/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_group_sync.py b/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_group_sync.py index fcc2f9a9e1..e11aee60ab 100644 --- a/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_group_sync.py +++ b/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_group_sync.py @@ -1,6 +1,5 @@ from collections.abc import Generator -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import Mock, patch from sqlalchemy import select from sqlalchemy.orm import Session @@ -12,16 +11,16 @@ from ee.onyx.db.external_perm import ExternalUserGroup from onyx.access.utils import build_ext_group_name_for_onyx from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import AccountType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import PublicExternalUserGroup -from onyx.db.models import User -from onyx.db.models import User__ExternalUserGroupId -from onyx.db.models import UserRole +from onyx.db.enums import AccessType, AccountType, ConnectorCredentialPairStatus +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + PublicExternalUserGroup, + User, + User__ExternalUserGroupId, + UserRole, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_hierarchy_walker.py b/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_hierarchy_walker.py index 856f0d6f72..c36468a22a 100644 --- a/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_hierarchy_walker.py +++ b/backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_hierarchy_walker.py @@ -29,11 +29,9 @@ After the planned fixes both tests should pass. """ from collections.abc import Callable -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import Mock, patch from uuid import uuid4 from sqlalchemy import delete @@ -44,27 +42,27 @@ from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.constants import DocumentSource from onyx.connectors.connector_runner import ConnectorRunner from onyx.connectors.google_drive.connector import GoogleDriveConnector -from onyx.connectors.google_drive.models import DriveRetrievalStage -from onyx.connectors.google_drive.models import GoogleDriveCheckpoint -from onyx.connectors.google_drive.models import GoogleDriveFileType -from onyx.connectors.google_drive.models import RetrievedDriveFile -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch +from onyx.connectors.google_drive.models import ( + DriveRetrievalStage, + GoogleDriveCheckpoint, + GoogleDriveFileType, + RetrievedDriveFile, +) +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + SecondsSinceUnixEpoch, +) from onyx.connectors.models import HierarchyNode as PydanticHierarchyNode from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import HierarchyNodeType -from onyx.db.hierarchy import ensure_source_node_exists -from onyx.db.hierarchy import get_hierarchy_node_by_raw_id -from onyx.db.hierarchy import get_source_hierarchy_node -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import HierarchyNode -from onyx.utils.threadpool_concurrency import ThreadSafeDict -from onyx.utils.threadpool_concurrency import ThreadSafeSet +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, HierarchyNodeType +from onyx.db.hierarchy import ( + ensure_source_node_exists, + get_hierarchy_node_by_raw_id, + get_source_hierarchy_node, +) +from onyx.db.models import Connector, ConnectorCredentialPair, Credential, HierarchyNode +from onyx.utils.threadpool_concurrency import ThreadSafeDict, ThreadSafeSet SOURCE = DocumentSource.GOOGLE_DRIVE ADMIN_EMAIL = "admin@example.com" diff --git a/backend/tests/external_dependency_unit/connectors/jira/test_jira_doc_sync.py b/backend/tests/external_dependency_unit/connectors/jira/test_jira_doc_sync.py index 90f1b4f409..bad0570c2e 100644 --- a/backend/tests/external_dependency_unit/connectors/jira/test_jira_doc_sync.py +++ b/backend/tests/external_dependency_unit/connectors/jira/test_jira_doc_sync.py @@ -8,13 +8,9 @@ from ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync from onyx.access.models import DocExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.utils import DocumentRow -from onyx.db.utils import SortOrder +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus +from onyx.db.models import Connector, ConnectorCredentialPair, Credential +from onyx.db.utils import DocumentRow, SortOrder # In order to get these tests to run, use the credentials from Bitwarden. # Search up "ENV vars for local and Github tests", and find the Jira relevant key-value pairs. diff --git a/backend/tests/external_dependency_unit/connectors/jira/test_jira_group_sync.py b/backend/tests/external_dependency_unit/connectors/jira/test_jira_group_sync.py index 3093d7d3b0..27c1442dc7 100644 --- a/backend/tests/external_dependency_unit/connectors/jira/test_jira_group_sync.py +++ b/backend/tests/external_dependency_unit/connectors/jira/test_jira_group_sync.py @@ -6,11 +6,8 @@ from sqlalchemy.orm import Session from ee.onyx.external_permissions.jira.group_sync import jira_group_sync from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus +from onyx.db.models import Connector, ConnectorCredentialPair, Credential from shared_configs.contextvars import get_current_tenant_id from tests.daily.connectors.confluence.models import ExternalUserGroupSet diff --git a/backend/tests/external_dependency_unit/connectors/test_gmail_multi_credential.py b/backend/tests/external_dependency_unit/connectors/test_gmail_multi_credential.py index 11aaf1c8fe..05de024268 100644 --- a/backend/tests/external_dependency_unit/connectors/test_gmail_multi_credential.py +++ b/backend/tests/external_dependency_unit/connectors/test_gmail_multi_credential.py @@ -1,10 +1,8 @@ """Creating a Gmail credential preserves existing Gmail credential rows.""" from collections.abc import Generator -from dataclasses import dataclass -from dataclasses import field -from uuid import UUID -from uuid import uuid4 +from dataclasses import dataclass, field +from uuid import UUID, uuid4 import pytest from sqlalchemy import select @@ -12,8 +10,7 @@ from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole from onyx.configs.constants import DocumentSource -from onyx.db.models import Credential -from onyx.db.models import User +from onyx.db.models import Credential, User from onyx.server.documents.credential import create_credential_from_model from onyx.server.documents.models import CredentialBase from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/craft/conftest.py b/backend/tests/external_dependency_unit/craft/conftest.py index 33233dee90..a0b4705e88 100644 --- a/backend/tests/external_dependency_unit/craft/conftest.py +++ b/backend/tests/external_dependency_unit/craft/conftest.py @@ -5,12 +5,9 @@ from __future__ import annotations import hashlib import io import zipfile -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterable +from collections.abc import Callable, Generator, Iterable from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from fastapi_users.password import PasswordHelper @@ -18,35 +15,46 @@ from sqlalchemy import text from sqlalchemy.orm import Session from onyx.configs.constants import FileOrigin -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine -from onyx.db.enums import AccountType -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SkillSharePermission -from onyx.db.llm import fetch_default_llm_model -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import Skill -from onyx.db.models import Skill__UserGroup -from onyx.db.models import User -from onyx.db.models import UserGroup -from onyx.db.models import UserRole +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine +from onyx.db.enums import ( + AccountType, + BuildSessionStatus, + SandboxStatus, + SkillSharePermission, +) +from onyx.db.llm import ( + fetch_default_llm_model, + fetch_existing_llm_provider, + remove_llm_provider, + update_default_provider, + upsert_llm_provider, +) +from onyx.db.models import ( + BuildSession, + Sandbox, + Skill, + Skill__UserGroup, + User, + UserGroup, + UserRole, +) from onyx.file_store.file_store import get_default_file_store from onyx.llm.constants import LlmProviderNames -from onyx.server.features.build.db.sandbox import create_sandbox__no_commit -from onyx.server.features.build.db.sandbox import update_sandbox_status__no_commit +from onyx.server.features.build.db.sandbox import ( + create_sandbox__no_commit, + update_sandbox_status__no_commit, +) from onyx.server.features.build.session.manager import SessionManager -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from tests.common.craft.skill_table_isolation import restore_skill_tables -from tests.common.craft.skill_table_isolation import snapshot_skill_tables +from tests.common.craft.skill_table_isolation import ( + restore_skill_tables, + snapshot_skill_tables, +) from tests.common.craft.stubs import StubSandboxManager diff --git a/backend/tests/external_dependency_unit/craft/db_helpers.py b/backend/tests/external_dependency_unit/craft/db_helpers.py index 8bcf9c1396..401377d299 100644 --- a/backend/tests/external_dependency_unit/craft/db_helpers.py +++ b/backend/tests/external_dependency_unit/craft/db_helpers.py @@ -18,38 +18,40 @@ from __future__ import annotations from datetime import datetime from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 from fastapi_users.password import PasswordHelper -from sqlalchemy import delete -from sqlalchemy import update +from sqlalchemy import delete, update from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource -from onyx.db.enums import AccessType -from onyx.db.enums import AccountType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SkillSharePermission -from onyx.db.models import ActionApproval -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import ExternalApp -from onyx.db.models import ExternalAppPolicy -from onyx.db.models import ExternalAppUserCredential -from onyx.db.models import Sandbox -from onyx.db.models import Skill -from onyx.db.models import Skill__User -from onyx.db.models import Skill__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup -from onyx.db.models import UserGroup__ConnectorCredentialPair -from onyx.db.models import UserRole +from onyx.db.enums import ( + AccessType, + AccountType, + ConnectorCredentialPairStatus, + EndpointPolicy, + ExternalAppType, + SandboxStatus, + SkillSharePermission, +) +from onyx.db.models import ( + ActionApproval, + Connector, + ConnectorCredentialPair, + Credential, + ExternalApp, + ExternalAppPolicy, + ExternalAppUserCredential, + Sandbox, + Skill, + Skill__User, + Skill__UserGroup, + User, + User__UserGroup, + UserGroup, + UserGroup__ConnectorCredentialPair, + UserRole, +) def force_approval_created_at( diff --git a/backend/tests/external_dependency_unit/craft/test_action_approval_db.py b/backend/tests/external_dependency_unit/craft/test_action_approval_db.py index e386fbfd41..e5f63e97f6 100644 --- a/backend/tests/external_dependency_unit/craft/test_action_approval_db.py +++ b/backend/tests/external_dependency_unit/craft/test_action_approval_db.py @@ -9,27 +9,25 @@ from __future__ import annotations import datetime as dt from collections.abc import Callable from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import EndpointPolicy -from onyx.db.models import ActionApproval -from onyx.db.models import BuildSession -from onyx.server.features.build.db.action_approval import get_action_approval -from onyx.server.features.build.db.action_approval import get_action_approval_for_user -from onyx.server.features.build.db.action_approval import insert_action_approval +from onyx.db.enums import ApprovalDecision, EndpointPolicy +from onyx.db.models import ActionApproval, BuildSession from onyx.server.features.build.db.action_approval import ( + get_action_approval, + get_action_approval_for_user, + insert_action_approval, list_session_pending_action_approvals, + try_record_decision, +) +from tests.common.craft.payloads import action_entry, default_action_entries +from tests.external_dependency_unit.craft.db_helpers import ( + force_approval_created_at, + make_user, ) -from onyx.server.features.build.db.action_approval import try_record_decision -from tests.common.craft.payloads import action_entry -from tests.common.craft.payloads import default_action_entries -from tests.external_dependency_unit.craft.db_helpers import force_approval_created_at -from tests.external_dependency_unit.craft.db_helpers import make_user def _seed_pending( diff --git a/backend/tests/external_dependency_unit/craft/test_action_matching.py b/backend/tests/external_dependency_unit/craft/test_action_matching.py index 56c61e3f43..4e7d8cad89 100644 --- a/backend/tests/external_dependency_unit/craft/test_action_matching.py +++ b/backend/tests/external_dependency_unit/craft/test_action_matching.py @@ -19,30 +19,30 @@ from __future__ import annotations import json from collections.abc import Iterator -from contextlib import AbstractContextManager -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager from typing import Callable import pytest from mitmproxy import http from sqlalchemy.orm import Session -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.models import ExternalApp -from onyx.db.models import ExternalAppPolicy -from onyx.db.models import User +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.db.models import ExternalApp, ExternalAppPolicy, User from onyx.external_apps.credentials import app_is_available from onyx.external_apps.matching import engine as matching_engine -from onyx.external_apps.matching.engine import AllMatchedActions -from onyx.external_apps.matching.engine import MatchedAction -from onyx.external_apps.matching.engine import recognize_actions -from onyx.external_apps.matching.engine import WHOLE_DOMAIN_ACTION_TYPE +from onyx.external_apps.matching.engine import ( + AllMatchedActions, + MatchedAction, + recognize_actions, + WHOLE_DOMAIN_ACTION_TYPE, +) from onyx.external_apps.matching.request import ProxiedRequest from onyx.sandbox_proxy import request_evaluator as request_evaluator_mod from onyx.sandbox_proxy.request_evaluator import ExternalAppRequestEvaluator -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_skill, +) def _connect_app( diff --git a/backend/tests/external_dependency_unit/craft/test_approvals_api.py b/backend/tests/external_dependency_unit/craft/test_approvals_api.py index bacb23724c..91cd1e1e12 100644 --- a/backend/tests/external_dependency_unit/craft/test_approvals_api.py +++ b/backend/tests/external_dependency_unit/craft/test_approvals_api.py @@ -4,41 +4,42 @@ constructed ``User`` and the test ``db_session`` (no ``TestClient``).""" from __future__ import annotations from collections.abc import Callable -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from uuid import UUID -from uuid import uuid4 +from datetime import datetime, timedelta, timezone +from uuid import UUID, uuid4 import pytest import redis from sqlalchemy.orm import Session from onyx.cache.factory import get_cache_backend -from onyx.db.enums import ApprovalDecidedVia -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import EndpointPolicy +from onyx.db.enums import ApprovalDecidedVia, ApprovalDecision, EndpointPolicy from onyx.db.models import BuildSession from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.external_apps.matching.engine import MatchedAction from onyx.sandbox_proxy import approval_cache -from onyx.server.features.build.approvals.api import DecisionBody -from onyx.server.features.build.approvals.api import list_live_approvals -from onyx.server.features.build.approvals.api import submit_decision -from onyx.server.features.build.approvals.api import submit_session_grant +from onyx.server.features.build.approvals.api import ( + DecisionBody, + list_live_approvals, + submit_decision, + submit_session_grant, +) from onyx.server.features.build.configs import SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS -from onyx.server.features.build.db.action_approval import get_action_approval -from onyx.server.features.build.db.action_approval import get_action_approval_for_user -from onyx.server.features.build.db.action_approval import insert_action_approval -from onyx.server.features.build.db.action_approval import try_record_decision +from onyx.server.features.build.db.action_approval import ( + get_action_approval, + get_action_approval_for_user, + insert_action_approval, + try_record_decision, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.common.craft.payloads import action_entry from tests.common.craft.payloads import default_action_entries as _default_actions -from tests.external_dependency_unit.craft.db_helpers import force_approval_created_at -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import ( + force_approval_created_at, + make_external_app, + make_skill, + make_user, +) def _stub_send_wake_noop(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/backend/tests/external_dependency_unit/craft/test_background_snapshots.py b/backend/tests/external_dependency_unit/craft/test_background_snapshots.py index 0fa39e2a9f..bd6ef5c270 100644 --- a/backend/tests/external_dependency_unit/craft/test_background_snapshots.py +++ b/backend/tests/external_dependency_unit/craft/test_background_snapshots.py @@ -17,8 +17,7 @@ from __future__ import annotations import datetime import logging from collections.abc import Generator -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy import update @@ -27,12 +26,8 @@ from sqlalchemy.orm import Session from onyx.background.celery.tasks.build import tasks as tasks_module from onyx.background.celery.tasks.build.tasks import cleanup_idle_sandboxes_task from onyx.configs.constants import OnyxRedisLocks -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SandboxStatus -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import Snapshot -from onyx.db.models import User +from onyx.db.enums import BuildSessionStatus, SandboxStatus +from onyx.db.models import BuildSession, Sandbox, Snapshot, User from onyx.redis.redis_pool import get_redis_client from onyx.server.features.build.sandbox.models import SnapshotResult from onyx.server.features.build.session import ( @@ -40,8 +35,7 @@ from onyx.server.features.build.session import ( ) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.common.craft.stubs import StubSandboxManager -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import make_sandbox, make_user # --------------------------------------------------------------------------- # Fixtures diff --git a/backend/tests/external_dependency_unit/craft/test_build_llm_provider_access.py b/backend/tests/external_dependency_unit/craft/test_build_llm_provider_access.py index 9c67836aec..71def571b3 100644 --- a/backend/tests/external_dependency_unit/craft/test_build_llm_provider_access.py +++ b/backend/tests/external_dependency_unit/craft/test_build_llm_provider_access.py @@ -12,22 +12,30 @@ from uuid import uuid4 from sqlalchemy import delete from sqlalchemy.orm import Session -from onyx.db.llm import fetch_first_accessible_llm_provider_by_type -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import upsert_llm_provider -from onyx.db.models import LLMProvider -from onyx.db.models import LLMProvider__Persona -from onyx.db.models import LLMProvider__UserGroup -from onyx.db.models import User -from onyx.db.models import UserRole +from onyx.db.llm import ( + fetch_first_accessible_llm_provider_by_type, + remove_llm_provider, + upsert_llm_provider, +) +from onyx.db.models import ( + LLMProvider, + LLMProvider__Persona, + LLMProvider__UserGroup, + User, + UserRole, +) from onyx.server.features.build.db.build_session import ( fetch_all_supported_build_llm_providers, ) -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest -from tests.external_dependency_unit.craft.db_helpers import add_user_to_group -from tests.external_dependency_unit.craft.db_helpers import make_group -from tests.external_dependency_unit.craft.db_helpers import make_user +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) +from tests.external_dependency_unit.craft.db_helpers import ( + add_user_to_group, + make_group, + make_user, +) from tests.external_dependency_unit.db.agent_sharing_helpers import create_test_persona diff --git a/backend/tests/external_dependency_unit/craft/test_built_in_skill_seeding.py b/backend/tests/external_dependency_unit/craft/test_built_in_skill_seeding.py index 69de9b1d5b..9b6d585bb9 100644 --- a/backend/tests/external_dependency_unit/craft/test_built_in_skill_seeding.py +++ b/backend/tests/external_dependency_unit/craft/test_built_in_skill_seeding.py @@ -16,20 +16,17 @@ from pathlib import Path from uuid import uuid4 import pytest -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session -from onyx.db.models import Skill -from onyx.db.models import User -from onyx.db.skill import fetch_skill -from onyx.db.skill import list_skills -from onyx.db.skill import SkillAccessPolicy -from onyx.skills.built_in import BUILT_IN_SKILLS -from onyx.skills.built_in import BuiltInSkillDefinition -from tests.external_dependency_unit.craft.db_helpers import make_built_in_skill_row -from tests.external_dependency_unit.craft.db_helpers import make_skill +from onyx.db.models import Skill, User +from onyx.db.skill import fetch_skill, list_skills, SkillAccessPolicy +from onyx.skills.built_in import BUILT_IN_SKILLS, BuiltInSkillDefinition +from tests.external_dependency_unit.craft.db_helpers import ( + make_built_in_skill_row, + make_skill, +) @pytest.fixture(autouse=True) diff --git a/backend/tests/external_dependency_unit/craft/test_connectable_apps.py b/backend/tests/external_dependency_unit/craft/test_connectable_apps.py index 6008f9f3e7..d4dd02f349 100644 --- a/backend/tests/external_dependency_unit/craft/test_connectable_apps.py +++ b/backend/tests/external_dependency_unit/craft/test_connectable_apps.py @@ -6,10 +6,12 @@ from onyx.db.external_app import get_connectable_apps_for_user from onyx.server.features.build.sandbox.util.agent_instructions import ( build_connectable_apps_list, ) -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import make_user_credential +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_skill, + make_user, + make_user_credential, +) # auth_template value carrying a user-supplied placeholder. _USER_TOKEN_TEMPLATE = {"Authorization": "Bearer {token}"} diff --git a/backend/tests/external_dependency_unit/craft/test_credential_encryption.py b/backend/tests/external_dependency_unit/craft/test_credential_encryption.py index df185e431d..8377ec5862 100644 --- a/backend/tests/external_dependency_unit/craft/test_credential_encryption.py +++ b/backend/tests/external_dependency_unit/craft/test_credential_encryption.py @@ -15,10 +15,12 @@ from sqlalchemy.orm import Session from onyx.db.enums import ExternalAppType from onyx.db.models import User from onyx.utils.sensitive import SensitiveValue -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import make_user_credential +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_skill, + make_user, + make_user_credential, +) def _raw_column_bytes(db_session: Session, query: str, **params: object) -> bytes: diff --git a/backend/tests/external_dependency_unit/craft/test_credential_injection.py b/backend/tests/external_dependency_unit/craft/test_credential_injection.py index dd682f7977..69b3ee63d8 100644 --- a/backend/tests/external_dependency_unit/craft/test_credential_injection.py +++ b/backend/tests/external_dependency_unit/craft/test_credential_injection.py @@ -4,10 +4,12 @@ from sqlalchemy.orm import Session from onyx.db.enums import ExternalAppType from onyx.external_apps.credentials import resolve_injection_headers -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import make_user_credential +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_skill, + make_user, + make_user_credential, +) _BEARER = {"Authorization": "Bearer {access_token}"} diff --git a/backend/tests/external_dependency_unit/craft/test_custom_external_app_create.py b/backend/tests/external_dependency_unit/craft/test_custom_external_app_create.py index dd539bfdf5..491306a0f0 100644 --- a/backend/tests/external_dependency_unit/craft/test_custom_external_app_create.py +++ b/backend/tests/external_dependency_unit/craft/test_custom_external_app_create.py @@ -7,25 +7,24 @@ from uuid import uuid4 import pytest from fastapi import UploadFile -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from onyx.db.enums import ExternalAppType -from onyx.db.models import ExternalApp -from onyx.db.models import Skill -from onyx.db.models import User +from onyx.db.models import ExternalApp, Skill, User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.build.external_apps.api import create_built_in_external_app -from onyx.server.features.build.external_apps.api import create_custom_external_app -from onyx.server.features.build.external_apps.api import replace_custom_app_bundle -from onyx.server.features.build.external_apps.api import update_external_app_admin +from onyx.server.features.build.external_apps.api import ( + create_built_in_external_app, + create_custom_external_app, + replace_custom_app_bundle, + update_external_app_admin, +) from onyx.server.features.build.external_apps.models import ( CreateBuiltInExternalAppRequest, + ExternalAppAdminResponse, + UpdateExternalAppRequest, ) -from onyx.server.features.build.external_apps.models import ExternalAppAdminResponse -from onyx.server.features.build.external_apps.models import UpdateExternalAppRequest from onyx.utils.encryption import is_masked_credential _AUTH_TEMPLATE = {"Authorization": "Bearer {api_key}"} diff --git a/backend/tests/external_dependency_unit/craft/test_ensure_sandbox_running.py b/backend/tests/external_dependency_unit/craft/test_ensure_sandbox_running.py index 5fc66340a4..f974487ab1 100644 --- a/backend/tests/external_dependency_unit/craft/test_ensure_sandbox_running.py +++ b/backend/tests/external_dependency_unit/craft/test_ensure_sandbox_running.py @@ -20,15 +20,13 @@ manager method, so each test must declare the slice of the manager it uses. from __future__ import annotations from collections.abc import Callable -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session from onyx.db.enums import SandboxStatus -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import Sandbox, User from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id from onyx.server.features.build.sandbox.models import SandboxInfo from onyx.server.features.build.session.errors import SandboxProvisioningError diff --git a/backend/tests/external_dependency_unit/craft/test_external_app_admin_update.py b/backend/tests/external_dependency_unit/craft/test_external_app_admin_update.py index 312fbabd49..fb29a9b6fb 100644 --- a/backend/tests/external_dependency_unit/craft/test_external_app_admin_update.py +++ b/backend/tests/external_dependency_unit/craft/test_external_app_admin_update.py @@ -14,12 +14,12 @@ from sqlalchemy.orm import Session import onyx.server.features.build.external_apps.api as api from onyx.db.enums import ExternalAppType from onyx.db.external_app import get_external_app_by_id -from onyx.db.models import ExternalApp -from onyx.db.models import Skill -from onyx.db.models import User +from onyx.db.models import ExternalApp, Skill, User from onyx.server.features.build.external_apps.models import UpdateExternalAppRequest -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import reset_built_in_skill_row +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + reset_built_in_skill_row, +) _AUTH_TEMPLATE = {"Authorization": "Bearer {token}"} _PATTERNS = ["https://slack.com/api/.*"] diff --git a/backend/tests/external_dependency_unit/craft/test_external_app_fileset.py b/backend/tests/external_dependency_unit/craft/test_external_app_fileset.py index 5cf36be446..28941258e2 100644 --- a/backend/tests/external_dependency_unit/craft/test_external_app_fileset.py +++ b/backend/tests/external_dependency_unit/craft/test_external_app_fileset.py @@ -17,17 +17,17 @@ from __future__ import annotations from sqlalchemy.orm import Session -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.models import Skill -from onyx.db.models import User +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.db.models import Skill, User from onyx.external_apps.providers.slack import SlackAction from onyx.skills.built_in import SLACK from onyx.skills.push import build_skills_fileset_for_user -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import make_user_credential -from tests.external_dependency_unit.craft.db_helpers import reset_built_in_skill_row +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_user, + make_user_credential, + reset_built_in_skill_row, +) # Slack ships on-disk skill content; require a single user-supplied key. _AUTH_TEMPLATE = {"Authorization": "Bearer {token}"} diff --git a/backend/tests/external_dependency_unit/craft/test_external_app_sandbox_push.py b/backend/tests/external_dependency_unit/craft/test_external_app_sandbox_push.py index ad2eaf6756..9c9995d9ec 100644 --- a/backend/tests/external_dependency_unit/craft/test_external_app_sandbox_push.py +++ b/backend/tests/external_dependency_unit/craft/test_external_app_sandbox_push.py @@ -1,7 +1,6 @@ from __future__ import annotations -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session @@ -11,12 +10,14 @@ from onyx.db.enums import ExternalAppType from onyx.db.models import User from onyx.server.features.build.external_apps.models import ( CreateBuiltInExternalAppRequest, + UpsertUserCredentialsRequest, +) +from tests.external_dependency_unit.craft.db_helpers import ( + make_built_in_skill_row, + make_external_app, + make_sandbox, + make_user, ) -from onyx.server.features.build.external_apps.models import UpsertUserCredentialsRequest -from tests.external_dependency_unit.craft.db_helpers import make_built_in_skill_row -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_user def test_credential_upsert_refreshes_only_the_calling_user( diff --git a/backend/tests/external_dependency_unit/craft/test_external_app_skill_visibility.py b/backend/tests/external_dependency_unit/craft/test_external_app_skill_visibility.py index 466b25a7e8..c40fa490c2 100644 --- a/backend/tests/external_dependency_unit/craft/test_external_app_skill_visibility.py +++ b/backend/tests/external_dependency_unit/craft/test_external_app_skill_visibility.py @@ -7,18 +7,21 @@ from uuid import UUID import pytest from sqlalchemy.orm import Session -from onyx.db.models import User -from onyx.db.models import UserRole -from onyx.db.skill import fetch_skill -from onyx.db.skill import list_skills -from onyx.db.skill import set_skill_enabled_for_user -from onyx.db.skill import SkillAccessPolicy +from onyx.db.models import User, UserRole +from onyx.db.skill import ( + fetch_skill, + list_skills, + set_skill_enabled_for_user, + SkillAccessPolicy, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import make_user_credential +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_skill, + make_user, + make_user_credential, +) _AUTH_TEMPLATE = {"token": "{token}", "account": "{account}"} _FULL_CREDS = {"token": "t", "account": "a"} diff --git a/backend/tests/external_dependency_unit/craft/test_granted_scopes_persistence.py b/backend/tests/external_dependency_unit/craft/test_granted_scopes_persistence.py index d5fa136bd8..9115b2eaae 100644 --- a/backend/tests/external_dependency_unit/craft/test_granted_scopes_persistence.py +++ b/backend/tests/external_dependency_unit/craft/test_granted_scopes_persistence.py @@ -12,12 +12,16 @@ import time from sqlalchemy.orm import Session from onyx.db.enums import ExternalAppType -from onyx.db.external_app import get_external_app_user_credential -from onyx.db.external_app import upsert_external_app_user_credential +from onyx.db.external_app import ( + get_external_app_user_credential, + upsert_external_app_user_credential, +) from onyx.db.models import ExternalApp -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_skill, + make_user, +) _BEARER = {"Authorization": "Bearer {access_token}"} diff --git a/backend/tests/external_dependency_unit/craft/test_idle_cleanup.py b/backend/tests/external_dependency_unit/craft/test_idle_cleanup.py index ee0dae719e..ad41dd2479 100644 --- a/backend/tests/external_dependency_unit/craft/test_idle_cleanup.py +++ b/backend/tests/external_dependency_unit/craft/test_idle_cleanup.py @@ -23,12 +23,8 @@ from sqlalchemy.orm import Session from onyx.background.celery.tasks.build import tasks as tasks_module from onyx.background.celery.tasks.build.tasks import cleanup_idle_sandboxes_task from onyx.configs.constants import OnyxRedisLocks -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SandboxStatus -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import Snapshot -from onyx.db.models import User +from onyx.db.enums import BuildSessionStatus, SandboxStatus +from onyx.db.models import BuildSession, Sandbox, Snapshot, User from onyx.redis.redis_pool import get_redis_client from onyx.server.features.build.db.build_session import skills_are_stale from onyx.server.features.build.sandbox.models import SnapshotResult @@ -38,8 +34,7 @@ from onyx.server.features.build.session import ( from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR from tests.common.craft.stubs import StubSandboxManager -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import make_sandbox, make_user # --------------------------------------------------------------------------- # Fixtures diff --git a/backend/tests/external_dependency_unit/craft/test_interrupt_message.py b/backend/tests/external_dependency_unit/craft/test_interrupt_message.py index bdf011227a..4442524c53 100644 --- a/backend/tests/external_dependency_unit/craft/test_interrupt_message.py +++ b/backend/tests/external_dependency_unit/craft/test_interrupt_message.py @@ -15,9 +15,7 @@ from typing import Callable from sqlalchemy.orm import Session from onyx.db.enums import SandboxStatus -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import BuildSession, Sandbox, User from onyx.server.features.build.session.manager import SessionManager from tests.common.craft.stubs import StubSandboxManager diff --git a/backend/tests/external_dependency_unit/craft/test_managed_external_apps.py b/backend/tests/external_dependency_unit/craft/test_managed_external_apps.py index 09308f3b12..58ceb1c1ed 100644 --- a/backend/tests/external_dependency_unit/craft/test_managed_external_apps.py +++ b/backend/tests/external_dependency_unit/craft/test_managed_external_apps.py @@ -20,18 +20,18 @@ from sqlalchemy.orm import Session import onyx.server.features.build.external_apps.api as api from onyx.db.enums import ExternalAppType -from onyx.db.external_app import create_external_app -from onyx.db.external_app import get_built_in_external_app -from onyx.db.models import Skill -from onyx.db.models import User +from onyx.db.external_app import create_external_app, get_built_in_external_app +from onyx.db.models import Skill, User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.external_apps.providers.registry import fetch_onyx_managed_built_in_apps -from onyx.external_apps.providers.registry import PROVIDERS +from onyx.external_apps.providers.registry import ( + fetch_onyx_managed_built_in_apps, + PROVIDERS, +) from onyx.server.features.build.external_apps.models import ( CreateBuiltInExternalAppRequest, + UpdateExternalAppRequest, ) -from onyx.server.features.build.external_apps.models import UpdateExternalAppRequest from onyx.skills.built_in import EXTERNAL_APP_BUILT_IN_SKILL_IDS _BUILT_IN_SLUGS = list(EXTERNAL_APP_BUILT_IN_SKILL_IDS.values()) diff --git a/backend/tests/external_dependency_unit/craft/test_provisioning_lock.py b/backend/tests/external_dependency_unit/craft/test_provisioning_lock.py index 953eb00948..148f7ba5a2 100644 --- a/backend/tests/external_dependency_unit/craft/test_provisioning_lock.py +++ b/backend/tests/external_dependency_unit/craft/test_provisioning_lock.py @@ -12,8 +12,7 @@ from __future__ import annotations import threading import time from collections.abc import Callable -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from kubernetes.client.rest import ApiException diff --git a/backend/tests/external_dependency_unit/craft/test_sandbox_lifecycle.py b/backend/tests/external_dependency_unit/craft/test_sandbox_lifecycle.py index 02d28208c9..7735f72b00 100644 --- a/backend/tests/external_dependency_unit/craft/test_sandbox_lifecycle.py +++ b/backend/tests/external_dependency_unit/craft/test_sandbox_lifecycle.py @@ -14,36 +14,37 @@ from __future__ import annotations import datetime from typing import Callable -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from fastapi import HTTPException from sqlalchemy.orm import Session -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SandboxStatus -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import User -from onyx.server.features.build.db.sandbox import create_sandbox__no_commit -from onyx.server.features.build.db.sandbox import create_snapshot__no_commit -from onyx.server.features.build.db.sandbox import get_running_sandboxes -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import FilesystemEntry -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import SandboxInfo +from onyx.db.enums import BuildSessionStatus, SandboxStatus +from onyx.db.models import BuildSession, Sandbox, User +from onyx.server.features.build.db.sandbox import ( + create_sandbox__no_commit, + create_snapshot__no_commit, + get_running_sandboxes, +) +from onyx.server.features.build.sandbox.models import ( + FileSet, + FilesystemEntry, + LLMProviderConfig, + SandboxInfo, +) from onyx.server.features.build.sandbox.user_library import USER_LIBRARY_MOUNT_PATH from onyx.server.features.build.session.api import restore_session from onyx.server.features.build.session.manager import SessionManager -from onyx.server.features.build.session.sandbox_lifecycle import is_sandbox_idle -from onyx.server.features.build.session.sandbox_lifecycle import provision_sandbox +from onyx.server.features.build.session.sandbox_lifecycle import ( + is_sandbox_idle, + provision_sandbox, +) from onyx.skills.push import SKILLS_MOUNT_PATH from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.common.craft.payloads import default_llm_config from tests.common.craft.stubs import StubSandboxManager -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import make_sandbox, make_user class TestProvisionTransitions: diff --git a/backend/tests/external_dependency_unit/craft/test_sandbox_pat.py b/backend/tests/external_dependency_unit/craft/test_sandbox_pat.py index 7e50e3c695..88c3507739 100644 --- a/backend/tests/external_dependency_unit/craft/test_sandbox_pat.py +++ b/backend/tests/external_dependency_unit/craft/test_sandbox_pat.py @@ -2,9 +2,7 @@ from __future__ import annotations -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from uuid import uuid4 @@ -12,14 +10,9 @@ import pytest from sqlalchemy.orm import Session from onyx.auth.pat import hash_pat -from onyx.db.enums import PatType -from onyx.db.enums import Permission -from onyx.db.enums import SandboxStatus -from onyx.db.models import PersonalAccessToken -from onyx.db.models import Sandbox -from onyx.db.models import User -from onyx.db.pat import create_pat -from onyx.db.pat import list_user_pats +from onyx.db.enums import PatType, Permission, SandboxStatus +from onyx.db.models import PersonalAccessToken, Sandbox, User +from onyx.db.pat import create_pat, list_user_pats from onyx.server.features.build.db.sandbox import ensure_sandbox_pat from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, diff --git a/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py b/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py index 1afc32ecaa..d93f2cea72 100644 --- a/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py +++ b/backend/tests/external_dependency_unit/craft/test_scheduled_task_executor.py @@ -21,8 +21,7 @@ import datetime import threading from collections.abc import Callable from typing import Any -from unittest.mock import patch -from unittest.mock import PropertyMock +from unittest.mock import patch, PropertyMock from uuid import UUID import pytest @@ -30,23 +29,22 @@ from sqlalchemy.orm import Session from onyx.background.celery.tasks.scheduled_tasks.tasks import ( cleanup_stuck_scheduled_runs, -) -from onyx.background.celery.tasks.scheduled_tasks.tasks import ( dispatch_due_scheduled_tasks, ) -from onyx.db.enums import SandboxStatus -from onyx.db.enums import ScheduledTaskErrorClass -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import ScheduledTaskStatus -from onyx.db.enums import ScheduledTaskTriggerSource -from onyx.db.models import Sandbox -from onyx.db.models import ScheduledTask -from onyx.db.models import ScheduledTaskRun -from onyx.db.models import User -from onyx.server.features.build.sandbox.event_schema import Error -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TIMEOUT -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TRANSPORT +from onyx.db.enums import ( + SandboxStatus, + ScheduledTaskErrorClass, + ScheduledTaskRunStatus, + ScheduledTaskStatus, + ScheduledTaskTriggerSource, +) +from onyx.db.models import Sandbox, ScheduledTask, ScheduledTaskRun, User +from onyx.server.features.build.sandbox.event_schema import ( + Error, + PromptResponse, + TURN_ERROR_CODE_TIMEOUT, + TURN_ERROR_CODE_TRANSPORT, +) from onyx.server.features.build.scheduled_tasks.executor import run_scheduled_task_logic from onyx.server.features.build.session.manager import SessionManager from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE diff --git a/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py b/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py index 97dcfb7c2c..c077ab070e 100644 --- a/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py +++ b/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py @@ -11,33 +11,40 @@ from __future__ import annotations from typing import Callable import pytest -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session -from onyx.db.enums import ApprovalDecidedVia -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import ScheduledTaskStatus -from onyx.db.enums import ScheduledTaskTriggerSource -from onyx.db.models import BuildSession -from onyx.db.models import ExternalApp -from onyx.db.models import ScheduledTask -from onyx.db.models import ScheduledTaskPreApprovedApp -from onyx.db.models import User -from onyx.db.scheduled_task import create_scheduled_task -from onyx.db.scheduled_task import get_live_scheduled_run_grants -from onyx.db.scheduled_task import insert_run -from onyx.db.scheduled_task import mark_run_status -from onyx.db.scheduled_task import update_scheduled_task +from onyx.db.enums import ( + ApprovalDecidedVia, + ApprovalDecision, + ScheduledTaskRunStatus, + ScheduledTaskStatus, + ScheduledTaskTriggerSource, +) +from onyx.db.models import ( + BuildSession, + ExternalApp, + ScheduledTask, + ScheduledTaskPreApprovedApp, + User, +) +from onyx.db.scheduled_task import ( + create_scheduled_task, + get_live_scheduled_run_grants, + insert_run, + mark_run_status, + update_scheduled_task, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.server.features.build.db.action_approval import insert_action_approval from onyx.server.features.build.scheduled_tasks import api as scheduled_tasks_api from tests.common.craft.payloads import default_action_entries -from tests.external_dependency_unit.craft.db_helpers import make_external_app -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import ( + make_external_app, + make_skill, + make_user, +) def _make_app(db_session: Session) -> int: diff --git a/backend/tests/external_dependency_unit/craft/test_session_lifecycle.py b/backend/tests/external_dependency_unit/craft/test_session_lifecycle.py index ccc4b4141c..9c9af071ec 100644 --- a/backend/tests/external_dependency_unit/craft/test_session_lifecycle.py +++ b/backend/tests/external_dependency_unit/craft/test_session_lifecycle.py @@ -16,30 +16,25 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.configs.constants import FileOrigin -from onyx.configs.constants import MessageType -from onyx.db.enums import ArtifactType -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SessionOrigin -from onyx.db.models import Artifact -from onyx.db.models import BuildMessage -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import Snapshot -from onyx.db.models import User +from onyx.configs.constants import FileOrigin, MessageType +from onyx.db.enums import ArtifactType, BuildSessionStatus, SandboxStatus, SessionOrigin +from onyx.db.models import Artifact, BuildMessage, BuildSession, Sandbox, Snapshot, User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import get_default_file_store from onyx.redis.redis_pool import get_redis_client -from onyx.server.features.build.db.build_session import allocate_nextjs_port -from onyx.server.features.build.db.build_session import get_user_build_sessions -from onyx.server.features.build.db.build_session import skills_are_stale +from onyx.server.features.build.db.build_session import ( + allocate_nextjs_port, + get_user_build_sessions, + skills_are_stale, +) from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id from onyx.server.features.build.sandbox.models import SandboxInfo from onyx.server.features.build.sandbox.user_library import USER_LIBRARY_MOUNT_PATH -from onyx.server.features.build.session.api import reload_session_skills -from onyx.server.features.build.session.api import restore_session +from onyx.server.features.build.session.api import ( + reload_session_skills, + restore_session, +) from onyx.server.features.build.session.manager import SessionManager from onyx.server.features.build.session.sandbox_lifecycle import hydrate_managed_content from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE diff --git a/backend/tests/external_dependency_unit/craft/test_skill_affected_users.py b/backend/tests/external_dependency_unit/craft/test_skill_affected_users.py index a7b4bf795f..faa1ed51b8 100644 --- a/backend/tests/external_dependency_unit/craft/test_skill_affected_users.py +++ b/backend/tests/external_dependency_unit/craft/test_skill_affected_users.py @@ -5,17 +5,18 @@ from uuid import uuid4 from sqlalchemy.orm import Session from onyx.db.enums import SandboxStatus -from onyx.db.models import OAuthAccount -from onyx.db.models import User +from onyx.db.models import OAuthAccount, User from onyx.db.skill import affected_user_ids_for_skill from onyx.server.features.build.db.sandbox import get_sandbox_user_map -from tests.external_dependency_unit.craft.db_helpers import add_user_to_group -from tests.external_dependency_unit.craft.db_helpers import make_group -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_group -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_user +from tests.external_dependency_unit.craft.db_helpers import ( + add_user_to_group, + make_group, + make_sandbox, + make_skill, + make_user, + share_skill_with_group, + share_skill_with_user, +) class TestAffectedUserIdsForSkill: diff --git a/backend/tests/external_dependency_unit/craft/test_skill_api_access.py b/backend/tests/external_dependency_unit/craft/test_skill_api_access.py index 08f3c663e4..70ded3bee8 100644 --- a/backend/tests/external_dependency_unit/craft/test_skill_api_access.py +++ b/backend/tests/external_dependency_unit/craft/test_skill_api_access.py @@ -11,32 +11,30 @@ import pytest from fastapi import UploadFile from sqlalchemy.orm import Session -from onyx.db.enums import SkillAccessLevel -from onyx.db.enums import SkillSharePermission -from onyx.db.models import User -from onyx.db.models import UserRole -from onyx.db.models import UserSkillPreference +from onyx.db.enums import SkillAccessLevel, SkillSharePermission +from onyx.db.models import User, UserRole, UserSkillPreference from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.skill.api import create_custom_skill -from onyx.server.features.skill.api import create_custom_skill_from_editor -from onyx.server.features.skill.api import fetch_skill_for_current_user -from onyx.server.features.skill.api import patch_current_user_skill -from onyx.server.features.skill.api import remove_current_user_skill_file -from onyx.server.features.skill.api import replace_current_user_skill_bundle -from onyx.server.features.skill.api import set_skill_enabled_for_current_user -from onyx.server.features.skill.api import upload_current_user_skill_files -from onyx.server.features.skill.models import SkillEnableRequest -from onyx.server.features.skill.models import SkillPatchRequest -from onyx.skills.bundle import build_single_file_bundle -from onyx.skills.bundle import build_skill_md -from onyx.skills.bundle import SKILL_MD_NAME -from tests.external_dependency_unit.craft.db_helpers import add_user_to_group -from tests.external_dependency_unit.craft.db_helpers import make_group -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_group -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_user +from onyx.server.features.skill.api import ( + create_custom_skill, + create_custom_skill_from_editor, + fetch_skill_for_current_user, + patch_current_user_skill, + remove_current_user_skill_file, + replace_current_user_skill_bundle, + set_skill_enabled_for_current_user, + upload_current_user_skill_files, +) +from onyx.server.features.skill.models import SkillEnableRequest, SkillPatchRequest +from onyx.skills.bundle import build_single_file_bundle, build_skill_md, SKILL_MD_NAME +from tests.external_dependency_unit.craft.db_helpers import ( + add_user_to_group, + make_group, + make_skill, + make_user, + share_skill_with_group, + share_skill_with_user, +) def _upload(filename: str, content: bytes = b"bundle") -> UploadFile: diff --git a/backend/tests/external_dependency_unit/craft/test_skill_db_mutations.py b/backend/tests/external_dependency_unit/craft/test_skill_db_mutations.py index 282398e205..2e5f4f9cad 100644 --- a/backend/tests/external_dependency_unit/craft/test_skill_db_mutations.py +++ b/backend/tests/external_dependency_unit/craft/test_skill_db_mutations.py @@ -1,7 +1,6 @@ from __future__ import annotations -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy import select @@ -9,21 +8,23 @@ from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole from onyx.db.enums import SkillSharePermission -from onyx.db.models import Skill__User -from onyx.db.models import Skill__UserGroup -from onyx.db.models import UserSkillPreference -from onyx.db.skill import replace_skill_shares -from onyx.db.skill import skill_user_states -from onyx.db.skill import transfer_skill_ownership +from onyx.db.models import Skill__User, Skill__UserGroup, UserSkillPreference +from onyx.db.skill import ( + replace_skill_shares, + skill_user_states, + transfer_skill_ownership, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.server.features.skill.response_helpers import skill_response_for_user -from tests.external_dependency_unit.craft.db_helpers import make_built_in_skill_row -from tests.external_dependency_unit.craft.db_helpers import make_group -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_group -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_user +from tests.external_dependency_unit.craft.db_helpers import ( + make_built_in_skill_row, + make_group, + make_skill, + make_user, + share_skill_with_group, + share_skill_with_user, +) def _direct_share_permissions( diff --git a/backend/tests/external_dependency_unit/craft/test_skill_ownership_transfer.py b/backend/tests/external_dependency_unit/craft/test_skill_ownership_transfer.py index 82929ace5f..0a2af48822 100644 --- a/backend/tests/external_dependency_unit/craft/test_skill_ownership_transfer.py +++ b/backend/tests/external_dependency_unit/craft/test_skill_ownership_transfer.py @@ -4,20 +4,18 @@ import pytest from sqlalchemy import select from sqlalchemy.orm import Session -from onyx.db.enums import AccountType -from onyx.db.enums import SkillSharePermission -from onyx.db.models import Skill -from onyx.db.models import Skill__User -from onyx.db.models import User -from onyx.db.models import UserRole +from onyx.db.enums import AccountType, SkillSharePermission +from onyx.db.models import Skill, Skill__User, User, UserRole from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.server.features.skill.api import transfer_current_user_skill_ownership from onyx.server.features.skill.models import TransferSkillOwnershipRequest -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_user +from tests.external_dependency_unit.craft.db_helpers import ( + make_sandbox, + make_skill, + make_user, + share_skill_with_user, +) def _owned_skill(db_session: Session, owner: User) -> Skill: diff --git a/backend/tests/external_dependency_unit/craft/test_skill_push_fault.py b/backend/tests/external_dependency_unit/craft/test_skill_push_fault.py index 19bfba3cee..c2ea1db2b4 100644 --- a/backend/tests/external_dependency_unit/craft/test_skill_push_fault.py +++ b/backend/tests/external_dependency_unit/craft/test_skill_push_fault.py @@ -9,19 +9,13 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SessionOrigin -from onyx.db.models import BuildSession -from onyx.db.models import Skill -from onyx.db.models import User +from onyx.db.enums import BuildSessionStatus, SandboxStatus, SessionOrigin +from onyx.db.models import BuildSession, Skill, User from onyx.server.features.build.db.build_session import skills_are_stale from onyx.server.features.build.sandbox.models import FatalWriteError -from onyx.skills.push import compute_skills_hash -from onyx.skills.push import push_skills_for_users +from onyx.skills.push import compute_skills_hash, push_skills_for_users from tests.common.craft.stubs import StubSandboxManager -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import make_sandbox, make_user def test_one_failing_sandbox_does_not_abort_push_to_others( diff --git a/backend/tests/external_dependency_unit/craft/test_skill_visibility.py b/backend/tests/external_dependency_unit/craft/test_skill_visibility.py index fe33c09177..977ebcc63b 100644 --- a/backend/tests/external_dependency_unit/craft/test_skill_visibility.py +++ b/backend/tests/external_dependency_unit/craft/test_skill_visibility.py @@ -6,17 +6,20 @@ import pytest from sqlalchemy.orm import Session from onyx.db.enums import SkillSharePermission -from onyx.db.models import User -from onyx.db.models import UserRole -from onyx.db.skill import fetch_skill -from onyx.db.skill import list_skills -from onyx.db.skill import set_skill_public_permission -from onyx.db.skill import SkillAccessPolicy -from tests.external_dependency_unit.craft.db_helpers import add_user_to_group -from tests.external_dependency_unit.craft.db_helpers import make_group -from tests.external_dependency_unit.craft.db_helpers import make_skill -from tests.external_dependency_unit.craft.db_helpers import make_user -from tests.external_dependency_unit.craft.db_helpers import share_skill_with_group +from onyx.db.models import User, UserRole +from onyx.db.skill import ( + fetch_skill, + list_skills, + set_skill_public_permission, + SkillAccessPolicy, +) +from tests.external_dependency_unit.craft.db_helpers import ( + add_user_to_group, + make_group, + make_skill, + make_user, + share_skill_with_group, +) def _user_skills(user: User, db_session: Session): diff --git a/backend/tests/external_dependency_unit/craft/test_skills_fileset.py b/backend/tests/external_dependency_unit/craft/test_skills_fileset.py index 64dd1b9472..0c323a4e21 100644 --- a/backend/tests/external_dependency_unit/craft/test_skills_fileset.py +++ b/backend/tests/external_dependency_unit/craft/test_skills_fileset.py @@ -10,17 +10,17 @@ import pytest from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource -from onyx.db.models import Skill -from onyx.db.models import User -from onyx.db.models import UserGroup +from onyx.db.models import Skill, User, UserGroup from onyx.db.skill import set_skill_enabled_for_user from onyx.skills import built_in as built_in_module from onyx.skills.built_in import BuiltInSkillDefinition from onyx.skills.push import build_skills_fileset_for_user -from tests.external_dependency_unit.craft.db_helpers import add_user_to_group -from tests.external_dependency_unit.craft.db_helpers import make_built_in_skill_row -from tests.external_dependency_unit.craft.db_helpers import make_group -from tests.external_dependency_unit.craft.db_helpers import reset_built_in_skill_row +from tests.external_dependency_unit.craft.db_helpers import ( + add_user_to_group, + make_built_in_skill_row, + make_group, + reset_built_in_skill_row, +) from tests.external_dependency_unit.indexing_helpers import make_cc_pair _FRONTMATTER = "---\nname: {slug}\ndescription: {slug}\n---\n" diff --git a/backend/tests/external_dependency_unit/craft/test_streaming_persistence.py b/backend/tests/external_dependency_unit/craft/test_streaming_persistence.py index 2f2359dcc2..c09cbb3cda 100644 --- a/backend/tests/external_dependency_unit/craft/test_streaming_persistence.py +++ b/backend/tests/external_dependency_unit/craft/test_streaming_persistence.py @@ -15,19 +15,20 @@ from typing import Any from sqlalchemy.orm import Session from onyx.configs.constants import MessageType -from onyx.db.models import BuildMessage -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.db.models import User -from onyx.server.features.build.db.build_session import create_message -from onyx.server.features.build.db.build_session import get_session_messages -from onyx.server.features.build.db.build_session import upsert_agent_plan +from onyx.db.models import BuildMessage, BuildSession, Sandbox, User +from onyx.server.features.build.db.build_session import ( + create_message, + get_session_messages, + upsert_agent_plan, +) from onyx.server.features.build.db.sandbox import get_sandbox_by_user_id -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import AgentThoughtChunk -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import ToolCallProgress -from onyx.server.features.build.sandbox.event_schema import ToolCallStart +from onyx.server.features.build.sandbox.event_schema import ( + AgentMessageChunk, + AgentThoughtChunk, + PromptResponse, + ToolCallProgress, + ToolCallStart, +) from onyx.server.features.build.sandbox.sse import SSEKeepalive from onyx.server.features.build.session.manager import SessionManager from onyx.server.features.build.session.streaming import BuildStreamingState diff --git a/backend/tests/external_dependency_unit/craft/test_terminated_sandbox_fast_fail.py b/backend/tests/external_dependency_unit/craft/test_terminated_sandbox_fast_fail.py index 3e15ba8ad0..1ecbf90212 100644 --- a/backend/tests/external_dependency_unit/craft/test_terminated_sandbox_fast_fail.py +++ b/backend/tests/external_dependency_unit/craft/test_terminated_sandbox_fast_fail.py @@ -13,8 +13,7 @@ import pytest from sqlalchemy.orm import Session from onyx.db.enums import SandboxStatus -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import Sandbox, User from onyx.server.features.build.sandbox.opencode.event_bus import PodEventBus from tests.common.craft.stubs import StubSandboxManager diff --git a/backend/tests/external_dependency_unit/craft/test_turn_heartbeat_refresh.py b/backend/tests/external_dependency_unit/craft/test_turn_heartbeat_refresh.py index 71b4e7e22a..ef4c826c5c 100644 --- a/backend/tests/external_dependency_unit/craft/test_turn_heartbeat_refresh.py +++ b/backend/tests/external_dependency_unit/craft/test_turn_heartbeat_refresh.py @@ -4,22 +4,17 @@ reaper's only activity signal — at turn start and periodically while events fl from __future__ import annotations import datetime -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest -from acp.schema import AgentMessageChunk -from acp.schema import PromptResponse -from acp.schema import TextContentBlock +from acp.schema import AgentMessageChunk, PromptResponse, TextContentBlock from sqlalchemy import update from sqlalchemy.orm import Session -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import Sandbox, User from onyx.server.features.build.session import streaming as streaming_module from tests.common.craft.stubs import StubSandboxManager -from tests.external_dependency_unit.craft.db_helpers import make_sandbox -from tests.external_dependency_unit.craft.db_helpers import make_user +from tests.external_dependency_unit.craft.db_helpers import make_sandbox, make_user def _agent_message_chunk(text: str) -> AgentMessageChunk: diff --git a/backend/tests/external_dependency_unit/craft/test_user_library_fileset.py b/backend/tests/external_dependency_unit/craft/test_user_library_fileset.py index 02f533a111..8596568a4a 100644 --- a/backend/tests/external_dependency_unit/craft/test_user_library_fileset.py +++ b/backend/tests/external_dependency_unit/craft/test_user_library_fileset.py @@ -8,19 +8,20 @@ import pytest from sqlalchemy.orm import Session from onyx.db.enums import SandboxStatus -from onyx.db.models import Sandbox -from onyx.db.models import User -from onyx.server.features.build.db.user_library import create_directory_record -from onyx.server.features.build.db.user_library import fetch_user_file_for_user -from onyx.server.features.build.db.user_library import get_or_create_craft_connector -from onyx.server.features.build.db.user_library import set_sync_disabled -from onyx.server.features.build.db.user_library import store_user_file -from onyx.server.features.build.sandbox.user_library import build_user_library_fileset -from onyx.server.features.build.sandbox.user_library import hydrate_user_library +from onyx.db.models import Sandbox, User +from onyx.server.features.build.db.user_library import ( + create_directory_record, + fetch_user_file_for_user, + get_or_create_craft_connector, + set_sync_disabled, + store_user_file, +) from onyx.server.features.build.sandbox.user_library import ( + build_user_library_fileset, + hydrate_user_library, sync_user_library_to_active_sandboxes, + USER_LIBRARY_MOUNT_PATH, ) -from onyx.server.features.build.sandbox.user_library import USER_LIBRARY_MOUNT_PATH from tests.common.craft.stubs import StubSandboxManager diff --git a/backend/tests/external_dependency_unit/db/agent_sharing_helpers.py b/backend/tests/external_dependency_unit/db/agent_sharing_helpers.py index ba889afbd8..471920bf73 100644 --- a/backend/tests/external_dependency_unit/db/agent_sharing_helpers.py +++ b/backend/tests/external_dependency_unit/db/agent_sharing_helpers.py @@ -5,12 +5,14 @@ from uuid import uuid4 from sqlalchemy.orm import Session from onyx.db.enums import PersonaSharePermission -from onyx.db.models import Persona -from onyx.db.models import Persona__User -from onyx.db.models import Persona__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import ( + Persona, + Persona__User, + Persona__UserGroup, + User, + User__UserGroup, + UserGroup, +) def create_test_persona( diff --git a/backend/tests/external_dependency_unit/db/conftest.py b/backend/tests/external_dependency_unit/db/conftest.py index 4d37532a90..ba300a73cd 100644 --- a/backend/tests/external_dependency_unit/db/conftest.py +++ b/backend/tests/external_dependency_unit/db/conftest.py @@ -8,17 +8,14 @@ Requires a running Postgres instance. Run with:: python -m dotenv -f .vscode/.env run -- pytest tests/external_dependency_unit/db/ """ -from collections.abc import Callable -from collections.abc import Generator -from uuid import UUID -from uuid import uuid4 +from collections.abc import Callable, Generator +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session from ee.onyx.db.scim import ScimDAL -from onyx.db.models import ScimToken -from onyx.db.models import UserGroup +from onyx.db.models import ScimToken, UserGroup @pytest.fixture diff --git a/backend/tests/external_dependency_unit/db/test_admin_banner_notification.py b/backend/tests/external_dependency_unit/db/test_admin_banner_notification.py index 3f274cc771..90d5ec65cb 100644 --- a/backend/tests/external_dependency_unit/db/test_admin_banner_notification.py +++ b/backend/tests/external_dependency_unit/db/test_admin_banner_notification.py @@ -10,9 +10,7 @@ import pytest from sqlalchemy.orm import Session from onyx.configs.constants import NotificationType -from onyx.db.admin_banner import clear_admin_banner -from onyx.db.admin_banner import get_admin_banner -from onyx.db.admin_banner import set_admin_banner +from onyx.db.admin_banner import clear_admin_banner, get_admin_banner, set_admin_banner from onyx.db.models import UserRole from onyx.db.notification import get_notifications from onyx.server.features.admin_banner import api as admin_banner_api diff --git a/backend/tests/external_dependency_unit/db/test_agent_sharing_groups.py b/backend/tests/external_dependency_unit/db/test_agent_sharing_groups.py index ba4e9c6a82..a1fcbe2821 100644 --- a/backend/tests/external_dependency_unit/db/test_agent_sharing_groups.py +++ b/backend/tests/external_dependency_unit/db/test_agent_sharing_groups.py @@ -9,21 +9,21 @@ from sqlalchemy.orm import Session from ee.onyx.db.persona import transfer_persona_ownership as ee_transfer from ee.onyx.db.persona import update_persona_access as ee_update_persona_access from ee.onyx.db.user_group import _handle_owned_personas_for_group_deletion__no_commit -from onyx.db.enums import PersonaAccessLevel -from onyx.db.enums import PersonaSharePermission -from onyx.db.enums import PersonaSharingStatus -from onyx.db.models import Persona -from onyx.db.models import Persona__UserGroup -from onyx.db.models import User -from onyx.db.persona import fetch_persona_by_id_for_user -from onyx.db.persona_sharing import derive_persona_sharing_status -from onyx.db.persona_sharing import get_persona_access_level -from tests.external_dependency_unit.conftest import create_test_user -from tests.external_dependency_unit.db.agent_sharing_helpers import create_test_persona -from tests.external_dependency_unit.db.agent_sharing_helpers import ( - create_test_user_group, +from onyx.db.enums import ( + PersonaAccessLevel, + PersonaSharePermission, + PersonaSharingStatus, ) +from onyx.db.models import Persona, Persona__UserGroup, User +from onyx.db.persona import fetch_persona_by_id_for_user +from onyx.db.persona_sharing import ( + derive_persona_sharing_status, + get_persona_access_level, +) +from tests.external_dependency_unit.conftest import create_test_user from tests.external_dependency_unit.db.agent_sharing_helpers import ( + create_test_persona, + create_test_user_group, share_persona_with_group, ) diff --git a/backend/tests/external_dependency_unit/db/test_agent_sharing_knowledge_guard.py b/backend/tests/external_dependency_unit/db/test_agent_sharing_knowledge_guard.py index dc58cbd9ed..b4ba67d84c 100644 --- a/backend/tests/external_dependency_unit/db/test_agent_sharing_knowledge_guard.py +++ b/backend/tests/external_dependency_unit/db/test_agent_sharing_knowledge_guard.py @@ -10,13 +10,11 @@ from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole from onyx.db.enums import PersonaSharePermission -from onyx.db.models import DocumentSet -from onyx.db.models import Persona -from onyx.db.models import User +from onyx.db.models import DocumentSet, Persona, User from onyx.db.persona import upsert_persona from tests.external_dependency_unit.conftest import create_test_user -from tests.external_dependency_unit.db.agent_sharing_helpers import create_test_persona from tests.external_dependency_unit.db.agent_sharing_helpers import ( + create_test_persona, share_persona_with_user, ) diff --git a/backend/tests/external_dependency_unit/db/test_agent_sharing_lifecycle.py b/backend/tests/external_dependency_unit/db/test_agent_sharing_lifecycle.py index df49af98a9..abc1085b52 100644 --- a/backend/tests/external_dependency_unit/db/test_agent_sharing_lifecycle.py +++ b/backend/tests/external_dependency_unit/db/test_agent_sharing_lifecycle.py @@ -8,13 +8,12 @@ from sqlalchemy.orm import Session from onyx.db.enums import PersonaSharePermission from onyx.db.models import Persona -from onyx.db.persona import remove_user_from_persona_shares -from onyx.db.persona import update_persona_shared +from onyx.db.persona import remove_user_from_persona_shares, update_persona_shared from onyx.db.persona_sharing import persona_ownership_is_vacant from onyx.db.users import delete_user_from_db from tests.external_dependency_unit.conftest import create_test_user -from tests.external_dependency_unit.db.agent_sharing_helpers import create_test_persona from tests.external_dependency_unit.db.agent_sharing_helpers import ( + create_test_persona, share_persona_with_user, ) diff --git a/backend/tests/external_dependency_unit/db/test_agent_sharing_permissions.py b/backend/tests/external_dependency_unit/db/test_agent_sharing_permissions.py index 19af3c93fa..a0a69eb62e 100644 --- a/backend/tests/external_dependency_unit/db/test_agent_sharing_permissions.py +++ b/backend/tests/external_dependency_unit/db/test_agent_sharing_permissions.py @@ -8,18 +8,20 @@ from fastapi import HTTPException from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.enums import PersonaAccessLevel -from onyx.db.enums import PersonaSharePermission -from onyx.db.enums import PersonaSharingStatus -from onyx.db.models import Persona__User -from onyx.db.models import User -from onyx.db.persona import fetch_persona_by_id_for_user -from onyx.db.persona import update_persona_access -from onyx.db.persona_sharing import derive_persona_sharing_status -from onyx.db.persona_sharing import get_persona_access_level +from onyx.db.enums import ( + PersonaAccessLevel, + PersonaSharePermission, + PersonaSharingStatus, +) +from onyx.db.models import Persona__User, User +from onyx.db.persona import fetch_persona_by_id_for_user, update_persona_access +from onyx.db.persona_sharing import ( + derive_persona_sharing_status, + get_persona_access_level, +) from tests.external_dependency_unit.conftest import create_test_user -from tests.external_dependency_unit.db.agent_sharing_helpers import create_test_persona from tests.external_dependency_unit.db.agent_sharing_helpers import ( + create_test_persona, share_persona_with_user, ) diff --git a/backend/tests/external_dependency_unit/db/test_agent_sharing_transfer.py b/backend/tests/external_dependency_unit/db/test_agent_sharing_transfer.py index be21ea060a..f6bb88dba3 100644 --- a/backend/tests/external_dependency_unit/db/test_agent_sharing_transfer.py +++ b/backend/tests/external_dependency_unit/db/test_agent_sharing_transfer.py @@ -9,14 +9,13 @@ import pytest from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.enums import AccountType -from onyx.db.enums import PersonaSharePermission +from onyx.db.enums import AccountType, PersonaSharePermission from onyx.db.models import Persona__User from onyx.db.persona import transfer_persona_ownership from onyx.db.persona_sharing import persona_ownership_is_vacant from tests.external_dependency_unit.conftest import create_test_user -from tests.external_dependency_unit.db.agent_sharing_helpers import create_test_persona from tests.external_dependency_unit.db.agent_sharing_helpers import ( + create_test_persona, share_persona_with_user, ) diff --git a/backend/tests/external_dependency_unit/db/test_api_key_permissions.py b/backend/tests/external_dependency_unit/db/test_api_key_permissions.py index b29023b446..fb0e4a8b12 100644 --- a/backend/tests/external_dependency_unit/db/test_api_key_permissions.py +++ b/backend/tests/external_dependency_unit/db/test_api_key_permissions.py @@ -8,10 +8,12 @@ from sqlalchemy import select from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.api_key import insert_api_key -from onyx.db.api_key import regenerate_api_key -from onyx.db.api_key import remove_api_key -from onyx.db.api_key import update_api_key +from onyx.db.api_key import ( + insert_api_key, + regenerate_api_key, + remove_api_key, + update_api_key, +) from onyx.db.models import User from onyx.server.api_key.models import APIKeyArgs diff --git a/backend/tests/external_dependency_unit/db/test_chat_session_eager_load.py b/backend/tests/external_dependency_unit/db/test_chat_session_eager_load.py index 9f5bf80187..ee7b05c443 100644 --- a/backend/tests/external_dependency_unit/db/test_chat_session_eager_load.py +++ b/backend/tests/external_dependency_unit/db/test_chat_session_eager_load.py @@ -1,10 +1,8 @@ from sqlalchemy import inspect from sqlalchemy.orm import Session -from onyx.db.chat import create_chat_session -from onyx.db.chat import get_chat_session_by_id -from onyx.db.models import Persona -from onyx.db.models import UserProject +from onyx.db.chat import create_chat_session, get_chat_session_by_id +from onyx.db.models import Persona, UserProject from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/db/test_connector_state_snapshot_query.py b/backend/tests/external_dependency_unit/db/test_connector_state_snapshot_query.py index 68980dbdb7..e3cd5d404d 100644 --- a/backend/tests/external_dependency_unit/db/test_connector_state_snapshot_query.py +++ b/backend/tests/external_dependency_unit/db/test_connector_state_snapshot_query.py @@ -1,17 +1,17 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from sqlalchemy.orm import Session -from onyx.configs.constants import DEFAULT_CC_PAIR_ID -from onyx.configs.constants import DocumentSource -from onyx.db.connector_credential_pair import ConnectorStateSnapshot -from onyx.db.connector_credential_pair import get_connector_state_snapshots -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingMode -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from onyx.configs.constants import DEFAULT_CC_PAIR_ID, DocumentSource +from onyx.db.connector_credential_pair import ( + ConnectorStateSnapshot, + get_connector_state_snapshots, +) +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, IndexingMode +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) def test_connector_state_snapshot_query( diff --git a/backend/tests/external_dependency_unit/db/test_document_file_id_lookup.py b/backend/tests/external_dependency_unit/db/test_document_file_id_lookup.py index 5e4ce93106..c90cfc8380 100644 --- a/backend/tests/external_dependency_unit/db/test_document_file_id_lookup.py +++ b/backend/tests/external_dependency_unit/db/test_document_file_id_lookup.py @@ -19,8 +19,7 @@ import pytest from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection +from onyx.context.search.models import InferenceChunk, InferenceSection from onyx.context.search.utils import populate_file_ids_on_sections from onyx.db.document import get_document_id_to_file_id_map from onyx.db.models import Document as DBDocument diff --git a/backend/tests/external_dependency_unit/db/test_index_attempt_stage_metrics.py b/backend/tests/external_dependency_unit/db/test_index_attempt_stage_metrics.py index dadf92867a..465a5f936c 100644 --- a/backend/tests/external_dependency_unit/db/test_index_attempt_stage_metrics.py +++ b/backend/tests/external_dependency_unit/db/test_index_attempt_stage_metrics.py @@ -22,8 +22,7 @@ import math import statistics import time from collections.abc import Generator -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, ThreadPoolExecutor from uuid import uuid4 import pytest @@ -33,19 +32,27 @@ from sqlalchemy.orm import Session from onyx.db.document import prepare_to_modify_documents from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import IndexingStatus -from onyx.db.index_attempt_metrics import get_stage_metrics_for_attempt -from onyx.db.index_attempt_metrics import record_single_event -from onyx.db.index_attempt_metrics import record_stage_aggregate -from onyx.db.index_attempt_metrics import StageEventBuffer +from onyx.db.index_attempt_metrics import ( + get_stage_metrics_for_attempt, + record_single_event, + record_stage_aggregate, + StageEventBuffer, +) from onyx.db.index_attempt_metrics_models import IndexAttemptStage -from onyx.db.models import ConnectorCredentialPair +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + IndexAttemptStageMetric, +) from onyx.db.models import Document as DbDocument -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptStageMetric -from onyx.server.documents.models import IndexAttemptStageMetricSnapshot -from onyx.server.documents.models import synthesize_unaccounted -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from onyx.server.documents.models import ( + IndexAttemptStageMetricSnapshot, + synthesize_unaccounted, +) +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) @pytest.fixture diff --git a/backend/tests/external_dependency_unit/db/test_notification.py b/backend/tests/external_dependency_unit/db/test_notification.py index d2ce273433..a81e68c689 100644 --- a/backend/tests/external_dependency_unit/db/test_notification.py +++ b/backend/tests/external_dependency_unit/db/test_notification.py @@ -1,21 +1,20 @@ from collections.abc import Callable -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import pytest from sqlalchemy import select from sqlalchemy.orm import Session from onyx.configs.constants import NotificationType -from onyx.db.models import Notification -from onyx.db.models import User -from onyx.db.notification import batch_create_notifications -from onyx.db.notification import count_notifications -from onyx.db.notification import create_notification -from onyx.db.notification import delete_notifications_by_additional_data -from onyx.db.notification import dismiss_user_notifications -from onyx.db.notification import get_notifications +from onyx.db.models import Notification, User +from onyx.db.notification import ( + batch_create_notifications, + count_notifications, + create_notification, + delete_notifications_by_additional_data, + dismiss_user_notifications, + get_notifications, +) from onyx.server.features.notifications import api as notifications_api from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/db/test_orphan_tag_cleanup.py b/backend/tests/external_dependency_unit/db/test_orphan_tag_cleanup.py index 561f7c510b..9a4a841e73 100644 --- a/backend/tests/external_dependency_unit/db/test_orphan_tag_cleanup.py +++ b/backend/tests/external_dependency_unit/db/test_orphan_tag_cleanup.py @@ -13,12 +13,12 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import Document -from onyx.db.models import Document__Tag -from onyx.db.models import Tag -from onyx.db.tag import _delete_orphan_tags_batch -from onyx.db.tag import delete_orphan_tags__no_commit -from onyx.db.tag import delete_orphan_tags_batched +from onyx.db.models import Document, Document__Tag, Tag +from onyx.db.tag import ( + _delete_orphan_tags_batch, + delete_orphan_tags__no_commit, + delete_orphan_tags_batched, +) def _seed_orphan_tags(db_session: Session, count: int) -> list[int]: diff --git a/backend/tests/external_dependency_unit/db/test_port_flow_e2e.py b/backend/tests/external_dependency_unit/db/test_port_flow_e2e.py index bf2a5b2c57..6e1373d08a 100644 --- a/backend/tests/external_dependency_unit/db/test_port_flow_e2e.py +++ b/backend/tests/external_dependency_unit/db/test_port_flow_e2e.py @@ -13,11 +13,8 @@ SearchSettings row in teardown (this dev DB has a leftover PRESENT row + a stale FUTURE row, so we never rely on the live current/secondary settings being ours). """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -25,40 +22,44 @@ from sqlalchemy.orm import Session from onyx.access.models import DocumentAccess from onyx.background.celery.tasks.port import tasks as port_task -from onyx.background.celery.tasks.port.tasks import run_check_for_port -from onyx.background.celery.tasks.port.tasks import run_port_attempt +from onyx.background.celery.tasks.port.tasks import run_check_for_port, run_port_attempt from onyx.configs.constants import DocumentSource -from onyx.configs.model_configs import ASYM_PASSAGE_PREFIX -from onyx.configs.model_configs import ASYM_QUERY_PREFIX +from onyx.configs.model_configs import ASYM_PASSAGE_PREFIX, ASYM_QUERY_PREFIX from onyx.context.search.models import SavedSearchSettings from onyx.db import swap_index -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.enums import PortAttemptStatus -from onyx.db.enums import SwitchoverType -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import PortAttempt -from onyx.db.models import SearchSettings +from onyx.db.enums import ( + EmbeddingPrecision, + IndexingStatus, + IndexModelStatus, + PortAttemptStatus, + SwitchoverType, +) +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + PortAttempt, + SearchSettings, +) from onyx.db.port_attempt import get_port_attempt -from onyx.db.search_settings import create_search_settings -from onyx.db.search_settings import get_current_search_settings +from onyx.db.search_settings import create_search_settings, get_current_search_settings from onyx.document_index.interfaces_new import TenantState from onyx.document_index.opensearch.client import OpenSearchIndexClient from onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE from onyx.document_index.opensearch.opensearch_document_index import ( generate_opensearch_filtered_access_control_list, ) -from onyx.document_index.opensearch.schema import DocumentChunk -from onyx.document_index.opensearch.schema import DocumentSchema -from onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from onyx.document_index.opensearch.schema import ( + DocumentChunk, + DocumentSchema, + get_opensearch_doc_chunk_id, +) +from shared_configs.configs import MODEL_SERVER_HOST, POSTGRES_DEFAULT_SCHEMA from shared_configs.contextvars import get_current_tenant_id -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import seed_cc_pair_documents +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, + seed_cc_pair_documents, +) # The port re-embeds PRESENT -> FUTURE against the local embedding model server; # ext-dep shards run with it disabled, so this composition test only runs where diff --git a/backend/tests/external_dependency_unit/db/test_port_metadata_defer_e2e.py b/backend/tests/external_dependency_unit/db/test_port_metadata_defer_e2e.py index 976c07b109..2c6dc01f56 100644 --- a/backend/tests/external_dependency_unit/db/test_port_metadata_defer_e2e.py +++ b/backend/tests/external_dependency_unit/db/test_port_metadata_defer_e2e.py @@ -18,8 +18,7 @@ Uses real Postgres + OpenSearch. Mirrors test_port_flow_e2e's index lifecycle. """ from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import uuid4 import pytest @@ -27,37 +26,46 @@ from sqlalchemy.orm import Session from onyx.access.models import DocumentAccess from onyx.configs.constants import DocumentSource -from onyx.db.document import count_secondary_only_sync_pending_documents -from onyx.db.document import document_has_indexable_cc_pair -from onyx.db.document import mark_document_as_synced -from onyx.db.document import mark_document_synced_secondary_pending -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import EmbeddingPrecision -from onyx.db.models import ConnectorCredentialPair +from onyx.db.document import ( + count_secondary_only_sync_pending_documents, + document_has_indexable_cc_pair, + mark_document_as_synced, + mark_document_synced_secondary_pending, +) +from onyx.db.enums import ConnectorCredentialPairStatus, EmbeddingPrecision +from onyx.db.models import ( + ConnectorCredentialPair, + DocumentByConnectorCredentialPair, + SearchSettings, +) from onyx.db.models import Document as DbDocument -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import SearchSettings -from onyx.db.port_attempt import any_future_port_in_progress -from onyx.db.port_attempt import create_port_attempt -from onyx.db.port_attempt import mark_port_in_progress -from onyx.document_index.interfaces_new import MetadataUpdateRequest -from onyx.document_index.interfaces_new import SecondaryIndexDocumentMissingError -from onyx.document_index.interfaces_new import TenantState +from onyx.db.port_attempt import ( + any_future_port_in_progress, + create_port_attempt, + mark_port_in_progress, +) +from onyx.document_index.interfaces_new import ( + MetadataUpdateRequest, + SecondaryIndexDocumentMissingError, + TenantState, +) from onyx.document_index.opensearch.client import OpenSearchIndexClient from onyx.document_index.opensearch.opensearch_document_index import ( generate_opensearch_filtered_access_control_list, -) -from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, + OpenSearchIndexPair, +) +from onyx.document_index.opensearch.schema import ( + DocumentChunk, + DocumentSchema, + get_opensearch_doc_chunk_id, ) -from onyx.document_index.opensearch.opensearch_document_index import OpenSearchIndexPair -from onyx.document_index.opensearch.schema import DocumentChunk -from onyx.document_index.opensearch.schema import DocumentSchema -from onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_future_search_settings +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, + make_future_search_settings, +) _VECTOR_DIM = 8 # tiny: this test never re-embeds, it only updates metadata _CHUNKS_PER_DOC = 2 diff --git a/backend/tests/external_dependency_unit/db/test_port_swap_criterion.py b/backend/tests/external_dependency_unit/db/test_port_swap_criterion.py index 1b6026a979..b8d603197b 100644 --- a/backend/tests/external_dependency_unit/db/test_port_swap_criterion.py +++ b/backend/tests/external_dependency_unit/db/test_port_swap_criterion.py @@ -28,29 +28,39 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.db import swap_index from onyx.db.document import mark_document_synced_secondary_pending -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import PortAttemptStatus -from onyx.db.enums import SwitchoverType -from onyx.db.models import ConnectorCredentialPair +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + PortAttemptStatus, + SwitchoverType, +) +from onyx.db.models import ( + ConnectorCredentialPair, + DocumentByConnectorCredentialPair, + PortAttempt, + SearchSettings, +) from onyx.db.models import Document as DbDocument -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import PortAttempt -from onyx.db.models import SearchSettings -from onyx.db.port_attempt import cancel_active_port_attempts -from onyx.db.port_attempt import create_port_attempt -from onyx.db.port_attempt import get_active_port_attempt -from onyx.db.port_attempt import mark_port_canceled -from onyx.db.port_attempt import mark_port_in_progress -from onyx.db.port_attempt import mark_port_succeeded -from onyx.db.port_attempt import request_port_cancel -from onyx.db.swap_index import _port_swap_ready -from onyx.db.swap_index import _required_cc_pairs_for_switchover -from onyx.db.swap_index import check_and_perform_index_swap +from onyx.db.port_attempt import ( + cancel_active_port_attempts, + create_port_attempt, + get_active_port_attempt, + mark_port_canceled, + mark_port_in_progress, + mark_port_succeeded, + request_port_cancel, +) +from onyx.db.swap_index import ( + _port_swap_ready, + _required_cc_pairs_for_switchover, + check_and_perform_index_swap, +) from onyx.kg.models import KGStage -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair_and_future -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_future_search_settings +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + cleanup_cc_pair_and_future, + make_cc_pair, + make_future_search_settings, +) _PENDING_DOC_PREFIX = "swapdoc-" diff --git a/backend/tests/external_dependency_unit/db/test_reindex_port.py b/backend/tests/external_dependency_unit/db/test_reindex_port.py index a49f503c3a..d703694394 100644 --- a/backend/tests/external_dependency_unit/db/test_reindex_port.py +++ b/backend/tests/external_dependency_unit/db/test_reindex_port.py @@ -17,12 +17,9 @@ covered by applying the migration, which this repo does not test programmaticall """ from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -32,64 +29,70 @@ from sqlalchemy.orm import Session from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT from onyx.background.celery.tasks.docprocessing.utils import should_index from onyx.background.celery.tasks.port import tasks as port_task -from onyx.background.celery.tasks.port.tasks import run_check_for_port -from onyx.background.celery.tasks.port.tasks import run_port_attempt -from onyx.configs.app_configs import INDEX_BATCH_SIZE -from onyx.configs.app_configs import MAX_CONCURRENT_PORT_ATTEMPTS -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.background.celery.tasks.port.tasks import run_check_for_port, run_port_attempt +from onyx.configs.app_configs import INDEX_BATCH_SIZE, MAX_CONCURRENT_PORT_ATTEMPTS +from onyx.configs.constants import OnyxCeleryQueues, OnyxCeleryTask from onyx.context.search.models import SavedSearchSettings from onyx.db import port_attempt as port_attempt_db from onyx.db.connector_credential_pair import get_last_successful_attempt_poll_range_end -from onyx.db.document import document_has_indexable_cc_pair -from onyx.db.document import filter_existing_cc_pair_document_ids -from onyx.db.document import get_document_ids_for_cc_pair_batch -from onyx.db.document import get_max_document_id_for_cc_pair -from onyx.db.document import mark_document_as_modified -from onyx.db.document import mark_document_as_synced -from onyx.db.document import mark_document_synced_secondary_pending -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.enums import PortAttemptStatus +from onyx.db.document import ( + document_has_indexable_cc_pair, + filter_existing_cc_pair_document_ids, + get_document_ids_for_cc_pair_batch, + get_max_document_id_for_cc_pair, + mark_document_as_modified, + mark_document_as_synced, + mark_document_synced_secondary_pending, +) +from onyx.db.enums import ( + ConnectorCredentialPairStatus, + EmbeddingPrecision, + IndexingStatus, + IndexModelStatus, + PortAttemptStatus, +) from onyx.db.index_attempt import ( count_unique_active_cc_pairs_with_successful_index_attempts, + count_unique_cc_pairs_with_successful_index_attempts, + create_synthetic_seed_attempt, + get_latest_successful_index_attempt_for_cc_pair_id, + mock_successful_index_attempt, +) +from onyx.db.models import ( + ConnectorCredentialPair, + DocumentByConnectorCredentialPair, + IndexAttempt, + PortAttempt, + SearchSettings, ) -from onyx.db.index_attempt import count_unique_cc_pairs_with_successful_index_attempts -from onyx.db.index_attempt import create_synthetic_seed_attempt -from onyx.db.index_attempt import get_latest_successful_index_attempt_for_cc_pair_id -from onyx.db.index_attempt import mock_successful_index_attempt -from onyx.db.models import ConnectorCredentialPair from onyx.db.models import Document as DbDocument -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import PortAttempt -from onyx.db.models import SearchSettings -from onyx.db.port_attempt import cancel_active_port_attempts -from onyx.db.port_attempt import commit_port_cursor -from onyx.db.port_attempt import count_consecutive_failed_port_attempts_no_progress -from onyx.db.port_attempt import create_port_attempt -from onyx.db.port_attempt import get_active_port_attempt -from onyx.db.port_attempt import get_latest_port_attempt -from onyx.db.port_attempt import mark_port_canceled -from onyx.db.port_attempt import mark_port_failed -from onyx.db.port_attempt import mark_port_in_progress -from onyx.db.port_attempt import mark_port_succeeded -from onyx.db.port_attempt import port_backfill_has_pending_work -from onyx.db.port_attempt import request_port_cancel -from onyx.db.search_settings import create_search_settings -from onyx.db.search_settings import get_current_search_settings +from onyx.db.port_attempt import ( + cancel_active_port_attempts, + commit_port_cursor, + count_consecutive_failed_port_attempts_no_progress, + create_port_attempt, + get_active_port_attempt, + get_latest_port_attempt, + mark_port_canceled, + mark_port_failed, + mark_port_in_progress, + mark_port_succeeded, + port_backfill_has_pending_work, + request_port_cancel, +) +from onyx.db.search_settings import create_search_settings, get_current_search_settings from onyx.db.swap_index import _port_swap_ready from onyx.document_index.opensearch import port_copy from onyx.document_index.opensearch.port_copy import copy_present_chunks_to_future from onyx.indexing.port_reembed import ReembedStrategy from onyx.kg.models import KGStage from shared_configs.contextvars import get_current_tenant_id -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair_and_future -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_future_search_settings +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + cleanup_cc_pair_and_future, + make_cc_pair, + make_future_search_settings, +) from tests.external_dependency_unit.indexing_helpers import ( seed_cc_pair_documents as _seed_cc_pair_documents, ) diff --git a/backend/tests/external_dependency_unit/db/test_rotate_encryption_key.py b/backend/tests/external_dependency_unit/db/test_rotate_encryption_key.py index b37a37cf84..741aa708c5 100644 --- a/backend/tests/external_dependency_unit/db/test_rotate_encryption_key.py +++ b/backend/tests/external_dependency_unit/db/test_rotate_encryption_key.py @@ -15,23 +15,25 @@ from collections.abc import Generator from unittest.mock import patch import pytest -from sqlalchemy import LargeBinary -from sqlalchemy import select -from sqlalchemy import text +from sqlalchemy import LargeBinary, select, text from sqlalchemy.orm import Session -from ee.onyx.utils.encryption import _decrypt_bytes -from ee.onyx.utils.encryption import _encrypt_string -from ee.onyx.utils.encryption import _get_trimmed_key +from ee.onyx.utils.encryption import _decrypt_bytes, _encrypt_string, _get_trimmed_key from onyx.configs.constants import DocumentSource -from onyx.db.models import Credential -from onyx.db.models import EncryptedJson -from onyx.db.models import EncryptedString -from onyx.db.models import InternetSearchProvider -from onyx.db.rotate_encryption_key import _discover_encrypted_columns -from onyx.db.rotate_encryption_key import rotate_encryption_key -from onyx.utils.variable_functionality import fetch_versioned_implementation -from onyx.utils.variable_functionality import global_version +from onyx.db.models import ( + Credential, + EncryptedJson, + EncryptedString, + InternetSearchProvider, +) +from onyx.db.rotate_encryption_key import ( + _discover_encrypted_columns, + rotate_encryption_key, +) +from onyx.utils.variable_functionality import ( + fetch_versioned_implementation, + global_version, +) EE_MODULE = "ee.onyx.utils.encryption" ROTATE_MODULE = "onyx.db.rotate_encryption_key" diff --git a/backend/tests/external_dependency_unit/db/test_run_targeted_reindex.py b/backend/tests/external_dependency_unit/db/test_run_targeted_reindex.py index 3739f9bc5f..84595d1a81 100644 --- a/backend/tests/external_dependency_unit/db/test_run_targeted_reindex.py +++ b/backend/tests/external_dependency_unit/db/test_run_targeted_reindex.py @@ -15,34 +15,41 @@ Drive connector reindex run. from collections.abc import Generator from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from sqlalchemy.orm import Session from onyx.background.indexing.run_targeted_reindex import process_targets_for_cc_pair from onyx.connectors.interfaces import Resolver -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + DocumentSource, + HierarchyNode, + TextSection, +) from onyx.db.enums import IndexingStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError -from onyx.db.models import TargetedReindexJob -from onyx.db.models import TargetedReindexJobTarget +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + IndexAttemptError, + TargetedReindexJob, + TargetedReindexJobTarget, +) from onyx.db.search_settings import get_current_search_settings -from onyx.db.targeted_reindex import create_targeted_reindex_job -from onyx.db.targeted_reindex import resolve_error_ids_to_targets -from onyx.db.targeted_reindex import targets_to_connector_failures -from onyx.db.targeted_reindex import TargetSpec +from onyx.db.targeted_reindex import ( + create_targeted_reindex_job, + resolve_error_ids_to_targets, + targets_to_connector_failures, + TargetSpec, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) @pytest.fixture diff --git a/backend/tests/external_dependency_unit/db/test_sso_provider_store.py b/backend/tests/external_dependency_unit/db/test_sso_provider_store.py index 416aa5f213..718e701efc 100644 --- a/backend/tests/external_dependency_unit/db/test_sso_provider_store.py +++ b/backend/tests/external_dependency_unit/db/test_sso_provider_store.py @@ -11,12 +11,14 @@ from sqlalchemy.orm import Session from onyx.db.enums import SSOProviderType from onyx.db.models import SSOProvider -from onyx.db.sso_provider import create_sso_provider -from onyx.db.sso_provider import fetch_sso_provider_by_name -from onyx.db.sso_provider import fetch_sso_providers -from onyx.db.sso_provider import set_sso_provider_enabled -from onyx.db.sso_provider import update_sso_provider -from onyx.db.sso_provider import validate_sso_provider_name +from onyx.db.sso_provider import ( + create_sso_provider, + fetch_sso_provider_by_name, + fetch_sso_providers, + set_sso_provider_enabled, + update_sso_provider, + validate_sso_provider_name, +) _NAME_PREFIX = "testsso" _GOOGLE_CONFIG = {"client_id": "client-id", "client_secret": "super-secret"} diff --git a/backend/tests/external_dependency_unit/db/test_swap_credentials_hierarchy.py b/backend/tests/external_dependency_unit/db/test_swap_credentials_hierarchy.py index 349060043d..574ef525eb 100644 --- a/backend/tests/external_dependency_unit/db/test_swap_credentials_hierarchy.py +++ b/backend/tests/external_dependency_unit/db/test_swap_credentials_hierarchy.py @@ -8,20 +8,19 @@ connector with hierarchy nodes has its credential swapped (e.g. Confluence). from uuid import uuid4 -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import HierarchyNodeType -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import HierarchyNode -from onyx.db.models import HierarchyNodeByConnectorCredentialPair +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, HierarchyNodeType +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + HierarchyNode, + HierarchyNodeByConnectorCredentialPair, +) def test_cc_pair_credential_swap_cascades_to_hierarchy_join( diff --git a/backend/tests/external_dependency_unit/db/test_sync_priority.py b/backend/tests/external_dependency_unit/db/test_sync_priority.py index e9b77845f4..7b51f54e2e 100644 --- a/backend/tests/external_dependency_unit/db/test_sync_priority.py +++ b/backend/tests/external_dependency_unit/db/test_sync_priority.py @@ -20,21 +20,25 @@ from onyx.background.celery.tasks.vespa.document_sync import ( from onyx.configs.constants import OnyxCeleryPriority from onyx.db.document import ( construct_document_id_select_by_needs_sync_or_secondary_pending, + count_documents_by_needs_sync, + count_documents_by_needs_sync_or_secondary_pending, + mark_document_as_modified, + mark_document_synced_secondary_pending, ) -from onyx.db.document import count_documents_by_needs_sync -from onyx.db.document import count_documents_by_needs_sync_or_secondary_pending -from onyx.db.document import mark_document_as_modified -from onyx.db.document import mark_document_synced_secondary_pending from onyx.db.models import ConnectorCredentialPair from onyx.db.models import Document as DbDocument -from onyx.db.port_attempt import any_future_port_in_progress -from onyx.db.port_attempt import create_port_attempt -from onyx.db.port_attempt import mark_port_in_progress -from onyx.db.port_attempt import mark_port_succeeded +from onyx.db.port_attempt import ( + any_future_port_in_progress, + create_port_attempt, + mark_port_in_progress, + mark_port_succeeded, +) from onyx.kg.models import KGStage -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair_and_future -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_future_search_settings +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair_and_future, + make_cc_pair, + make_future_search_settings, +) _DOC_PREFIX = "syncdoc-" diff --git a/backend/tests/external_dependency_unit/db/test_tag_race_condition.py b/backend/tests/external_dependency_unit/db/test_tag_race_condition.py index 02e6a3b135..43d5acf762 100644 --- a/backend/tests/external_dependency_unit/db/test_tag_race_condition.py +++ b/backend/tests/external_dependency_unit/db/test_tag_race_condition.py @@ -6,9 +6,7 @@ UniqueViolation errors, which would occur if the upsert logic isn't properly implemented. """ -from concurrent.futures import as_completed -from concurrent.futures import Future -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, Future, ThreadPoolExecutor from typing import Union from uuid import uuid4 @@ -17,10 +15,8 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import Document -from onyx.db.models import Tag -from onyx.db.tag import create_or_add_document_tag -from onyx.db.tag import create_or_add_document_tag_list +from onyx.db.models import Document, Tag +from onyx.db.tag import create_or_add_document_tag, create_or_add_document_tag_list def _create_test_document(db_session: Session, doc_id: str) -> Document: diff --git a/backend/tests/external_dependency_unit/db/test_targeted_reindex_create.py b/backend/tests/external_dependency_unit/db/test_targeted_reindex_create.py index e2fd4c58db..c77355951e 100644 --- a/backend/tests/external_dependency_unit/db/test_targeted_reindex_create.py +++ b/backend/tests/external_dependency_unit/db/test_targeted_reindex_create.py @@ -16,19 +16,25 @@ import pytest from sqlalchemy.orm import Session from onyx.db.enums import IndexingStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError -from onyx.db.models import TargetedReindexJob -from onyx.db.models import TargetedReindexJobTarget +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + IndexAttemptError, + TargetedReindexJob, + TargetedReindexJobTarget, +) from onyx.db.search_settings import get_current_search_settings -from onyx.db.targeted_reindex import create_targeted_reindex_job -from onyx.db.targeted_reindex import get_targeted_reindex_job -from onyx.db.targeted_reindex import MAX_TARGETS_PER_REQUEST -from onyx.db.targeted_reindex import resolve_error_ids_to_targets -from onyx.db.targeted_reindex import TargetSpec -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from onyx.db.targeted_reindex import ( + create_targeted_reindex_job, + get_targeted_reindex_job, + MAX_TARGETS_PER_REQUEST, + resolve_error_ids_to_targets, + TargetSpec, +) +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) @pytest.fixture diff --git a/backend/tests/external_dependency_unit/db/test_targeted_reindex_filter.py b/backend/tests/external_dependency_unit/db/test_targeted_reindex_filter.py index 37e6d68a35..ab188b5805 100644 --- a/backend/tests/external_dependency_unit/db/test_targeted_reindex_filter.py +++ b/backend/tests/external_dependency_unit/db/test_targeted_reindex_filter.py @@ -13,32 +13,32 @@ from sqlalchemy.orm import Session from onyx.db.connector_credential_pair import get_last_successful_attempt_poll_range_end from onyx.db.enums import IndexingStatus -from onyx.db.index_attempt import cancel_indexing_attempts_for_ccpair -from onyx.db.index_attempt import count_index_attempts_for_cc_pair from onyx.db.index_attempt import ( + cancel_indexing_attempts_for_ccpair, + count_index_attempts_for_cc_pair, count_unique_active_cc_pairs_with_successful_index_attempts, + count_unique_cc_pairs_with_successful_index_attempts, + get_in_progress_index_attempts, + get_index_attempts_for_cc_pair, + get_last_attempt, + get_last_attempt_for_cc_pair, + get_latest_index_attempt_for_cc_pair_id, + get_latest_index_attempts, + get_latest_index_attempts_by_status, + get_latest_successful_index_attempt_for_cc_pair_id, + get_latest_successful_index_attempts_parallel, + get_paginated_index_attempts_for_cc_pair_id, + get_recent_attempts_for_cc_pair, + get_recent_completed_attempts_for_cc_pair, ) -from onyx.db.index_attempt import count_unique_cc_pairs_with_successful_index_attempts -from onyx.db.index_attempt import get_in_progress_index_attempts -from onyx.db.index_attempt import get_index_attempts_for_cc_pair -from onyx.db.index_attempt import get_last_attempt -from onyx.db.index_attempt import get_last_attempt_for_cc_pair -from onyx.db.index_attempt import get_latest_index_attempt_for_cc_pair_id -from onyx.db.index_attempt import get_latest_index_attempts -from onyx.db.index_attempt import get_latest_index_attempts_by_status -from onyx.db.index_attempt import get_latest_successful_index_attempt_for_cc_pair_id -from onyx.db.index_attempt import get_latest_successful_index_attempts_parallel -from onyx.db.index_attempt import get_paginated_index_attempts_for_cc_pair_id -from onyx.db.index_attempt import get_recent_attempts_for_cc_pair -from onyx.db.index_attempt import get_recent_completed_attempts_for_cc_pair from onyx.db.indexing_coordination import IndexingCoordination -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import TargetedReindexJob +from onyx.db.models import ConnectorCredentialPair, IndexAttempt, TargetedReindexJob from onyx.db.search_settings import get_current_search_settings from onyx.server.documents.models import ConnectorCredentialPairIdentifier -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) @pytest.fixture @@ -553,8 +553,7 @@ def test_get_last_successful_poll_range_end_skips_targeted( ) -> None: """The freshness scheduler reads `poll_range_end` from the latest successful full-run attempt, ignoring targeted reindexes.""" - from datetime import datetime - from datetime import timezone + from datetime import datetime, timezone settings = get_current_search_settings(db_session) full_run_end = datetime(2026, 1, 1, tzinfo=timezone.utc) 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 1c61377f94..8425dff9f8 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 @@ -25,17 +25,23 @@ from onyx.background.celery.tasks.docprocessing.targeted_reindex_task import ( ) from onyx.background.indexing.run_targeted_reindex import CCPairReindexResult from onyx.db.enums import IndexingStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError -from onyx.db.models import TargetedReindexJob -from onyx.db.models import TargetedReindexJobTarget +from onyx.db.models import ( + ConnectorCredentialPair, + IndexAttempt, + IndexAttemptError, + TargetedReindexJob, + TargetedReindexJobTarget, +) from onyx.db.search_settings import get_current_search_settings -from onyx.db.targeted_reindex import create_targeted_reindex_job -from onyx.db.targeted_reindex import resolve_error_ids_to_targets -from onyx.db.targeted_reindex import TargetSpec -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from onyx.db.targeted_reindex import ( + create_targeted_reindex_job, + resolve_error_ids_to_targets, + TargetSpec, +) +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) _PROCESSOR_PATH = ( "onyx.background.celery.tasks.docprocessing." diff --git a/backend/tests/external_dependency_unit/db/test_usage_export_pairing.py b/backend/tests/external_dependency_unit/db/test_usage_export_pairing.py index d7e82b634a..10a1c7b461 100644 --- a/backend/tests/external_dependency_unit/db/test_usage_export_pairing.py +++ b/backend/tests/external_dependency_unit/db/test_usage_export_pairing.py @@ -1,14 +1,15 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import UUID from sqlalchemy.orm import Session from ee.onyx.db.usage_export import get_empty_chat_messages_entries__paginated from onyx.configs.constants import MessageType -from onyx.db.chat import create_chat_session -from onyx.db.chat import create_new_chat_message -from onyx.db.chat import get_or_create_root_message +from onyx.db.chat import ( + create_chat_session, + create_new_chat_message, + get_or_create_root_message, +) from onyx.db.models import ChatMessage from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/db/test_user_account_type.py b/backend/tests/external_dependency_unit/db/test_user_account_type.py index 8836531d7f..62b16da63f 100644 --- a/backend/tests/external_dependency_unit/db/test_user_account_type.py +++ b/backend/tests/external_dependency_unit/db/test_user_account_type.py @@ -11,8 +11,10 @@ from sqlalchemy.orm import Session from onyx.db.enums import AccountType from onyx.db.models import UserRole -from onyx.db.users import add_slack_user_if_not_exists -from onyx.db.users import batch_add_ext_perm_user_if_not_exists +from onyx.db.users import ( + add_slack_user_if_not_exists, + batch_add_ext_perm_user_if_not_exists, +) def test_slack_user_creation_sets_account_type_bot(db_session: Session) -> None: diff --git a/backend/tests/external_dependency_unit/discord_bot/conftest.py b/backend/tests/external_dependency_unit/discord_bot/conftest.py index 8d66c5c840..ac03e61131 100644 --- a/backend/tests/external_dependency_unit/discord_bot/conftest.py +++ b/backend/tests/external_dependency_unit/discord_bot/conftest.py @@ -1,15 +1,13 @@ """Fixtures for Discord bot external dependency tests.""" from collections.abc import Generator -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import discord import pytest from sqlalchemy.orm import Session -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR TEST_TENANT_ID: str = "public" diff --git a/backend/tests/external_dependency_unit/discord_bot/test_discord_events.py b/backend/tests/external_dependency_unit/discord_bot/test_discord_events.py index 62338f51c8..35da708688 100644 --- a/backend/tests/external_dependency_unit/discord_bot/test_discord_events.py +++ b/backend/tests/external_dependency_unit/discord_bot/test_discord_events.py @@ -3,20 +3,22 @@ These tests mock the Discord API to test event handling logic. """ -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import discord import pytest -from onyx.onyxbot.discord.handle_commands import get_text_channels -from onyx.onyxbot.discord.handle_commands import handle_dm -from onyx.onyxbot.discord.handle_commands import handle_registration_command -from onyx.onyxbot.discord.handle_commands import handle_sync_channels_command -from onyx.onyxbot.discord.handle_message import process_chat_message -from onyx.onyxbot.discord.handle_message import send_error_response -from onyx.onyxbot.discord.handle_message import send_response +from onyx.onyxbot.discord.handle_commands import ( + get_text_channels, + handle_dm, + handle_registration_command, + handle_sync_channels_command, +) +from onyx.onyxbot.discord.handle_message import ( + process_chat_message, + send_error_response, + send_response, +) class TestGuildRegistrationCommand: diff --git a/backend/tests/external_dependency_unit/document_index/conftest.py b/backend/tests/external_dependency_unit/document_index/conftest.py index fdc7916069..9244779701 100644 --- a/backend/tests/external_dependency_unit/document_index/conftest.py +++ b/backend/tests/external_dependency_unit/document_index/conftest.py @@ -12,17 +12,17 @@ from onyx.access.models import DocumentAccess from onyx.configs.constants import DocumentSource from onyx.connectors.models import Document from onyx.db.enums import EmbeddingPrecision -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces_new import IndexingMetadata, TenantState from onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, ) -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocMetadataAwareIndexChunk +from onyx.indexing.models import ChunkEmbedding, DocMetadataAwareIndexChunk from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) EMBEDDING_DIM = 128 diff --git a/backend/tests/external_dependency_unit/document_index/test_document_index.py b/backend/tests/external_dependency_unit/document_index/test_document_index.py index 3e4114139b..a263c0b897 100644 --- a/backend/tests/external_dependency_unit/document_index/test_document_index.py +++ b/backend/tests/external_dependency_unit/document_index/test_document_index.py @@ -5,28 +5,28 @@ These tests assume OpenSearch is running. import time import uuid -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Generator, Iterator from unittest.mock import patch import pytest from onyx.configs.constants import PUBLIC_DOC_PAT -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import IndexFilters, InferenceChunk from onyx.db.enums import EmbeddingPrecision from onyx.document_index.interfaces_new import DocumentIndex as DocumentIndexNew -from onyx.document_index.interfaces_new import DocumentSectionRequest -from onyx.document_index.interfaces_new import MetadataUpdateRequest -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces_new import ( + DocumentSectionRequest, + MetadataUpdateRequest, + TenantState, +) from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, ) from onyx.indexing.models import DocMetadataAwareIndexChunk from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from tests.external_dependency_unit.document_index.conftest import EMBEDDING_DIM -from tests.external_dependency_unit.document_index.conftest import make_chunk from tests.external_dependency_unit.document_index.conftest import ( + EMBEDDING_DIM, + make_chunk, make_indexing_metadata, ) diff --git a/backend/tests/external_dependency_unit/ee/onyx/background/celery/test_license_notifications_task.py b/backend/tests/external_dependency_unit/ee/onyx/background/celery/test_license_notifications_task.py index 9c43479305..16cc7cd02c 100644 --- a/backend/tests/external_dependency_unit/ee/onyx/background/celery/test_license_notifications_task.py +++ b/backend/tests/external_dependency_unit/ee/onyx/background/celery/test_license_notifications_task.py @@ -10,11 +10,8 @@ And dispatches to ``notify_admins_for_stage`` only when there is a real, verifiable, in-window license. """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/external_dependency_unit/ee/onyx/db/test_license_cache.py b/backend/tests/external_dependency_unit/ee/onyx/db/test_license_cache.py index 0bc2e7d2bd..d541ef4f3b 100644 --- a/backend/tests/external_dependency_unit/ee/onyx/db/test_license_cache.py +++ b/backend/tests/external_dependency_unit/ee/onyx/db/test_license_cache.py @@ -7,21 +7,23 @@ manually during BUSINESS → ENTERPRISE → BUSINESS lifecycle testing. """ from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import pytest from sqlalchemy.orm import Session -from ee.onyx.db.license import get_cached_license_metadata -from ee.onyx.db.license import LICENSE_CACHE_TTL_SECONDS -from ee.onyx.db.license import LICENSE_METADATA_KEY -from ee.onyx.db.license import update_license_cache -from ee.onyx.server.license.models import CustomerTier -from ee.onyx.server.license.models import LicensePayload -from ee.onyx.server.license.models import LicenseSource -from ee.onyx.server.license.models import PlanType +from ee.onyx.db.license import ( + get_cached_license_metadata, + LICENSE_CACHE_TTL_SECONDS, + LICENSE_METADATA_KEY, + update_license_cache, +) +from ee.onyx.server.license.models import ( + CustomerTier, + LicensePayload, + LicenseSource, + PlanType, +) from onyx.redis.redis_pool import get_redis_client from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE diff --git a/backend/tests/external_dependency_unit/ee/onyx/utils/test_license_notifications.py b/backend/tests/external_dependency_unit/ee/onyx/utils/test_license_notifications.py index 2f4ea02344..d755212cff 100644 --- a/backend/tests/external_dependency_unit/ee/onyx/utils/test_license_notifications.py +++ b/backend/tests/external_dependency_unit/ee/onyx/utils/test_license_notifications.py @@ -5,9 +5,7 @@ against a real PostgreSQL database. The email send path is patched so we can assert "fresh insert only" behavior without configuring SendGrid. """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from unittest.mock import patch from uuid import UUID @@ -18,9 +16,7 @@ from sqlalchemy.orm import Session from ee.onyx.utils.license_expiry import ExpiryWarningStage from ee.onyx.utils.license_notifications import notify_admins_for_stage from onyx.configs.constants import NotificationType -from onyx.db.models import Notification -from onyx.db.models import User -from onyx.db.models import UserRole +from onyx.db.models import Notification, User, UserRole from tests.external_dependency_unit.conftest import create_test_user EXPIRES_AT = datetime(2026, 6, 1, 12, 0, 0, tzinfo=timezone.utc) diff --git a/backend/tests/external_dependency_unit/ee/onyx/utils/test_trial_promotion.py b/backend/tests/external_dependency_unit/ee/onyx/utils/test_trial_promotion.py index adc71782f2..7e2d10af82 100644 --- a/backend/tests/external_dependency_unit/ee/onyx/utils/test_trial_promotion.py +++ b/backend/tests/external_dependency_unit/ee/onyx/utils/test_trial_promotion.py @@ -8,19 +8,18 @@ real. import json from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from unittest.mock import patch import pytest from ee.onyx.server.license.models import CustomerTier -from ee.onyx.server.tenants.models import BillingInformation -from ee.onyx.server.tenants.models import SubscriptionStatusResponse -from ee.onyx.server.tenants.tier_management import get_cached_tier -from ee.onyx.server.tenants.tier_management import TENANT_TIER_KEY -from ee.onyx.server.tenants.tier_management import update_tenant_tier +from ee.onyx.server.tenants.models import BillingInformation, SubscriptionStatusResponse +from ee.onyx.server.tenants.tier_management import ( + get_cached_tier, + TENANT_TIER_KEY, + update_tenant_tier, +) from ee.onyx.utils import tier as tier_module from onyx.redis.redis_pool import get_redis_client from onyx.server.settings.models import Tier diff --git a/backend/tests/external_dependency_unit/feature_flags/test_feature_flag_provider_factory.py b/backend/tests/external_dependency_unit/feature_flags/test_feature_flag_provider_factory.py index c1e68949ec..025187c76a 100644 --- a/backend/tests/external_dependency_unit/feature_flags/test_feature_flag_provider_factory.py +++ b/backend/tests/external_dependency_unit/feature_flags/test_feature_flag_provider_factory.py @@ -9,8 +9,7 @@ from uuid import UUID from ee.onyx.feature_flags.posthog_provider import PostHogFeatureFlagProvider from onyx.feature_flags.factory import get_default_feature_flag_provider -from onyx.feature_flags.interface import FeatureFlagProvider -from onyx.feature_flags.interface import NoOpFeatureFlagProvider +from onyx.feature_flags.interface import FeatureFlagProvider, NoOpFeatureFlagProvider class TestNoOpFeatureFlagProvider: diff --git a/backend/tests/external_dependency_unit/file_store/test_file_store_non_mocked.py b/backend/tests/external_dependency_unit/file_store/test_file_store_non_mocked.py index 26fdd9267a..9baa28f4dd 100644 --- a/backend/tests/external_dependency_unit/file_store/test_file_store_non_mocked.py +++ b/backend/tests/external_dependency_unit/file_store/test_file_store_non_mocked.py @@ -2,15 +2,9 @@ import os import time import uuid from collections.abc import Generator -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, ThreadPoolExecutor from io import BytesIO -from typing import Any -from typing import cast -from typing import Dict -from typing import List -from typing import Tuple -from typing import TypedDict +from typing import Any, cast, Dict, List, Tuple, TypedDict from unittest.mock import patch import pytest @@ -57,8 +51,7 @@ class WorkerResult(TypedDict): def _get_all_backend_configs() -> List[BackendConfig]: """Get configurations for all available backends""" - from onyx.configs.app_configs import AWS_REGION_NAME - from onyx.configs.app_configs import S3_ENDPOINT_URL + from onyx.configs.app_configs import AWS_REGION_NAME, S3_ENDPOINT_URL s3_aws_access_key_id = os.environ.get("S3_AWS_ACCESS_KEY_ID_FOR_TEST") s3_aws_secret_access_key = os.environ.get("S3_AWS_SECRET_ACCESS_KEY_FOR_TEST") diff --git a/backend/tests/external_dependency_unit/file_store/test_postgres_file_store_non_mocked.py b/backend/tests/external_dependency_unit/file_store/test_postgres_file_store_non_mocked.py index 662bb7cb65..d7687356ae 100644 --- a/backend/tests/external_dependency_unit/file_store/test_postgres_file_store_non_mocked.py +++ b/backend/tests/external_dependency_unit/file_store/test_postgres_file_store_non_mocked.py @@ -8,23 +8,24 @@ and the underlying pg_largeobject storage. import uuid from collections.abc import Generator -from io import BytesIO -from io import StringIO -from typing import Any -from typing import Dict -from typing import List +from io import BytesIO, StringIO +from typing import Any, Dict, List import pytest from sqlalchemy.orm import Session from onyx.configs.constants import FileOrigin from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.file_content import get_file_content_by_file_id -from onyx.db.file_content import get_file_content_by_file_id_optional -from onyx.file_store.postgres_file_store import _get_raw_connection -from onyx.file_store.postgres_file_store import _read_large_object -from onyx.file_store.postgres_file_store import POSTGRES_BUCKET_SENTINEL -from onyx.file_store.postgres_file_store import PostgresBackedFileStore +from onyx.db.file_content import ( + get_file_content_by_file_id, + get_file_content_by_file_id_optional, +) +from onyx.file_store.postgres_file_store import ( + _get_raw_connection, + _read_large_object, + POSTGRES_BUCKET_SENTINEL, + PostgresBackedFileStore, +) from onyx.utils.logger import setup_logger logger = setup_logger() diff --git a/backend/tests/external_dependency_unit/file_store/test_staging.py b/backend/tests/external_dependency_unit/file_store/test_staging.py index cd47ec68fb..1f9c5137c3 100644 --- a/backend/tests/external_dependency_unit/file_store/test_staging.py +++ b/backend/tests/external_dependency_unit/file_store/test_staging.py @@ -16,11 +16,9 @@ from sqlalchemy.orm import Session from onyx.configs.constants import FileOrigin from onyx.connectors.interfaces import BaseConnector -from onyx.db.file_record import delete_filerecord_by_file_id -from onyx.db.file_record import get_filerecord_by_file_id +from onyx.db.file_record import delete_filerecord_by_file_id, get_filerecord_by_file_id from onyx.file_store.file_store import get_default_file_store -from onyx.file_store.staging import build_raw_file_callback -from onyx.file_store.staging import stage_raw_file +from onyx.file_store.staging import build_raw_file_callback, stage_raw_file @pytest.fixture(scope="function") diff --git a/backend/tests/external_dependency_unit/file_store/test_staging_concurrent_attempt_skip.py b/backend/tests/external_dependency_unit/file_store/test_staging_concurrent_attempt_skip.py index 47adef770c..3f298f03e6 100644 --- a/backend/tests/external_dependency_unit/file_store/test_staging_concurrent_attempt_skip.py +++ b/backend/tests/external_dependency_unit/file_store/test_staging_concurrent_attempt_skip.py @@ -17,13 +17,13 @@ import pytest from sqlalchemy.orm import Session from onyx.db.enums import IndexingStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from onyx.file_store.staging import reap_prior_attempt_staged_files -from onyx.file_store.staging import stage_raw_file +from onyx.db.models import ConnectorCredentialPair, IndexAttempt +from onyx.file_store.staging import reap_prior_attempt_staged_files, stage_raw_file from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) @pytest.fixture diff --git a/backend/tests/external_dependency_unit/full_setup.py b/backend/tests/external_dependency_unit/full_setup.py index 4e0a30089c..e69ae7727d 100644 --- a/backend/tests/external_dependency_unit/full_setup.py +++ b/backend/tests/external_dependency_unit/full_setup.py @@ -4,15 +4,15 @@ import os from pathlib import Path from typing import Optional -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from onyx.db.search_settings import get_active_search_settings -from onyx.document_index.factory import get_all_document_indices -from onyx.document_index.factory import get_default_document_index +from onyx.document_index.factory import ( + get_all_document_indices, + get_default_document_index, +) from onyx.file_store.file_store import get_default_file_store from onyx.indexing.models import IndexingSetting -from onyx.setup import setup_document_indices -from onyx.setup import setup_postgres +from onyx.setup import setup_document_indices, setup_postgres from shared_configs import configs as shared_configs_module from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/external_dependency_unit/hierarchy/test_stub_cross_batch.py b/backend/tests/external_dependency_unit/hierarchy/test_stub_cross_batch.py index 256c34e4a3..64b47c878d 100644 --- a/backend/tests/external_dependency_unit/hierarchy/test_stub_cross_batch.py +++ b/backend/tests/external_dependency_unit/hierarchy/test_stub_cross_batch.py @@ -16,10 +16,12 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.models import HierarchyNode as PydanticHierarchyNode from onyx.db.enums import HierarchyNodeType -from onyx.db.hierarchy import ensure_source_node_exists -from onyx.db.hierarchy import get_hierarchy_node_by_raw_id -from onyx.db.hierarchy import get_source_hierarchy_node -from onyx.db.hierarchy import upsert_hierarchy_nodes_batch +from onyx.db.hierarchy import ( + ensure_source_node_exists, + get_hierarchy_node_by_raw_id, + get_source_hierarchy_node, + upsert_hierarchy_nodes_batch, +) @pytest.fixture() diff --git a/backend/tests/external_dependency_unit/hierarchy/test_upsert_parents_deep_chain.py b/backend/tests/external_dependency_unit/hierarchy/test_upsert_parents_deep_chain.py index 732bc1517e..86afec09de 100644 --- a/backend/tests/external_dependency_unit/hierarchy/test_upsert_parents_deep_chain.py +++ b/backend/tests/external_dependency_unit/hierarchy/test_upsert_parents_deep_chain.py @@ -15,9 +15,11 @@ from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.models import HierarchyNode as PydanticHierarchyNode from onyx.db.enums import HierarchyNodeType -from onyx.db.hierarchy import ensure_source_node_exists -from onyx.db.hierarchy import get_hierarchy_node_by_raw_id -from onyx.db.hierarchy import upsert_hierarchy_nodes_batch +from onyx.db.hierarchy import ( + ensure_source_node_exists, + get_hierarchy_node_by_raw_id, + upsert_hierarchy_nodes_batch, +) from onyx.db.models import HierarchyNode # Notion can produce parent chains exceeding Python's default recursion limit diff --git a/backend/tests/external_dependency_unit/indexing/test_docfetching_orphan_cleanup.py b/backend/tests/external_dependency_unit/indexing/test_docfetching_orphan_cleanup.py index a7494a210c..1ceb334c7a 100644 --- a/backend/tests/external_dependency_unit/indexing/test_docfetching_orphan_cleanup.py +++ b/backend/tests/external_dependency_unit/indexing/test_docfetching_orphan_cleanup.py @@ -23,8 +23,7 @@ depends on JSONB metadata filtering and the file_store client is what actually deletes blob + FileRecord together. """ -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Generator, Iterator from io import BytesIO from typing import Any from unittest.mock import MagicMock @@ -34,30 +33,25 @@ import pytest from sqlalchemy.orm import Session from onyx.background.indexing.run_docfetching import run_docfetching_entrypoint -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin +from onyx.configs.constants import DocumentSource, FileOrigin from onyx.connectors import factory as connector_factory from onyx.connectors.factory import instantiate_connector from onyx.connectors.interfaces import LoadConnector -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import InputType -from onyx.connectors.models import TextSection -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus +from onyx.connectors.models import Document, HierarchyNode, InputType, TextSection +from onyx.db.enums import EmbeddingPrecision, IndexingStatus, IndexModelStatus from onyx.db.file_record import get_filerecord_by_file_id_optional -from onyx.db.models import Credential -from onyx.db.models import FileRecord -from onyx.db.models import IndexAttempt -from onyx.db.models import SearchSettings +from onyx.db.models import Credential, FileRecord, IndexAttempt, SearchSettings from onyx.file_store.file_store import get_default_file_store -from onyx.file_store.staging import build_raw_file_callback -from onyx.file_store.staging import cleanup_staged_files_for_attempt -from onyx.file_store.staging import reap_prior_attempt_staged_files +from onyx.file_store.staging import ( + build_raw_file_callback, + cleanup_staged_files_for_attempt, + reap_prior_attempt_staged_files, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) # --------------------------------------------------------------------------- # Mock connector that stages a file, yields a doc, then crashes diff --git a/backend/tests/external_dependency_unit/indexing/test_document_deletion_file_cleanup.py b/backend/tests/external_dependency_unit/indexing/test_document_deletion_file_cleanup.py index 0755afa403..2fea174ca3 100644 --- a/backend/tests/external_dependency_unit/indexing/test_document_deletion_file_cleanup.py +++ b/backend/tests/external_dependency_unit/indexing/test_document_deletion_file_cleanup.py @@ -10,28 +10,30 @@ best-effort deletes the underlying files after the DB commit. """ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session from onyx.background.celery.tasks.shared.tasks import document_by_cc_pair_cleanup_task -from onyx.connectors.models import Document -from onyx.connectors.models import IndexAttemptMetadata -from onyx.db.document import delete_all_documents_for_connector_credential_pair -from onyx.db.document import upsert_document_by_connector_credential_pair +from onyx.connectors.models import Document, IndexAttemptMetadata +from onyx.db.document import ( + delete_all_documents_for_connector_credential_pair, + upsert_document_by_connector_credential_pair, +) from onyx.db.models import ConnectorCredentialPair from onyx.indexing.indexing_pipeline import index_doc_batch_prepare from onyx.server.onyx_api.ingestion import delete_ingestion_doc from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import get_doc_row -from tests.external_dependency_unit.indexing_helpers import get_filerecord -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_doc -from tests.external_dependency_unit.indexing_helpers import stage_file +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + get_doc_row, + get_filerecord, + make_cc_pair, + make_doc, + stage_file, +) # --------------------------------------------------------------------------- # Helpers (file-local) diff --git a/backend/tests/external_dependency_unit/indexing/test_get_old_index_attempt_ids.py b/backend/tests/external_dependency_unit/indexing/test_get_old_index_attempt_ids.py index aefa16a2c9..722cc0c744 100644 --- a/backend/tests/external_dependency_unit/indexing/test_get_old_index_attempt_ids.py +++ b/backend/tests/external_dependency_unit/indexing/test_get_old_index_attempt_ids.py @@ -18,17 +18,18 @@ import pytest from sqlalchemy import update from sqlalchemy.orm import Session -from onyx.background.indexing.index_attempt_utils import get_old_index_attempt_ids from onyx.background.indexing.index_attempt_utils import ( + get_old_index_attempt_ids, NUM_RECENT_INDEX_ATTEMPTS_TO_KEEP, ) from onyx.configs.constants import NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS from onyx.db.engine.time_utils import get_db_current_time from onyx.db.enums import IndexingStatus -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexAttempt -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from onyx.db.models import ConnectorCredentialPair, IndexAttempt +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) @pytest.fixture diff --git a/backend/tests/external_dependency_unit/indexing/test_index_doc_batch_prepare.py b/backend/tests/external_dependency_unit/indexing/test_index_doc_batch_prepare.py index 12a5579221..e18d571e84 100644 --- a/backend/tests/external_dependency_unit/indexing/test_index_doc_batch_prepare.py +++ b/backend/tests/external_dependency_unit/indexing/test_index_doc_batch_prepare.py @@ -20,12 +20,14 @@ from onyx.configs.constants import FileOrigin from onyx.connectors.models import IndexAttemptMetadata from onyx.db.models import ConnectorCredentialPair from onyx.indexing.indexing_pipeline import index_doc_batch_prepare -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import get_doc_row -from tests.external_dependency_unit.indexing_helpers import get_filerecord -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_doc -from tests.external_dependency_unit.indexing_helpers import stage_file +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + get_doc_row, + get_filerecord, + make_cc_pair, + make_doc, + stage_file, +) # --------------------------------------------------------------------------- # Fixtures diff --git a/backend/tests/external_dependency_unit/indexing/test_persistent_indexing.py b/backend/tests/external_dependency_unit/indexing/test_persistent_indexing.py index ae65146b1c..e0477949d7 100644 --- a/backend/tests/external_dependency_unit/indexing/test_persistent_indexing.py +++ b/backend/tests/external_dependency_unit/indexing/test_persistent_indexing.py @@ -26,27 +26,28 @@ from sqlalchemy.orm import Session from onyx.background.indexing.run_docfetching import run_docfetching_entrypoint from onyx.configs.constants import DocumentSource from onyx.connectors import factory as connector_factory -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import CheckpointOutput -from onyx.connectors.interfaces import GenerateSlimDocumentOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import InputType -from onyx.connectors.models import TextSection -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import IndexingStatus -from onyx.db.enums import IndexModelStatus -from onyx.db.index_attempt import get_index_attempt -from onyx.db.index_attempt import get_index_attempt_errors -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexAttemptError -from onyx.db.models import SearchSettings +from onyx.connectors.interfaces import ( + CheckpointedConnector, + CheckpointOutput, + GenerateSlimDocumentOutput, + SecondsSinceUnixEpoch, +) +from onyx.connectors.models import ( + ConnectorCheckpoint, + ConnectorFailure, + Document, + DocumentFailure, + InputType, + TextSection, +) +from onyx.db.enums import EmbeddingPrecision, IndexingStatus, IndexModelStatus +from onyx.db.index_attempt import get_index_attempt, get_index_attempt_errors +from onyx.db.models import IndexAttempt, IndexAttemptError, SearchSettings from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) # --------------------------------------------------------------------------- # Mock checkpointed connector with configurable failure behavior diff --git a/backend/tests/external_dependency_unit/indexing/test_port_reembed.py b/backend/tests/external_dependency_unit/indexing/test_port_reembed.py index 273789035a..85eb44d46d 100644 --- a/backend/tests/external_dependency_unit/indexing/test_port_reembed.py +++ b/backend/tests/external_dependency_unit/indexing/test_port_reembed.py @@ -15,15 +15,18 @@ from unittest.mock import MagicMock import pytest -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import RETURN_SEPARATOR -from onyx.configs.model_configs import ASYM_PASSAGE_PREFIX -from onyx.configs.model_configs import ASYM_QUERY_PREFIX -from onyx.configs.model_configs import DEFAULT_DOCUMENT_ENCODER_MODEL -from onyx.configs.model_configs import DOC_EMBEDDING_DIM -from onyx.configs.model_configs import NORMALIZE_EMBEDDINGS -from onyx.connectors.models import convert_metadata_dict_to_list_of_strings -from onyx.connectors.models import convert_metadata_list_of_strings_to_dict +from onyx.configs.constants import DocumentSource, RETURN_SEPARATOR +from onyx.configs.model_configs import ( + ASYM_PASSAGE_PREFIX, + ASYM_QUERY_PREFIX, + DEFAULT_DOCUMENT_ENCODER_MODEL, + DOC_EMBEDDING_DIM, + NORMALIZE_EMBEDDINGS, +) +from onyx.connectors.models import ( + convert_metadata_dict_to_list_of_strings, + convert_metadata_list_of_strings_to_dict, +) from onyx.db.models import SearchSettings from onyx.document_index.chunk_content_enrichment import ( generate_enriched_content_for_chunk_embedding, @@ -32,25 +35,26 @@ from onyx.document_index.interfaces_new import TenantState from onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE from onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors from onyx.indexing.chunker import get_metadata_suffix_for_document_index -from onyx.indexing.embedder import DefaultIndexingEmbedder -from onyx.indexing.embedder import IndexingEmbedder -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import IndexChunk -from onyx.indexing.port_reembed import _bare_contents -from onyx.indexing.port_reembed import _reconstruct_source_document -from onyx.indexing.port_reembed import _stored_chunk_to_doc_aware -from onyx.indexing.port_reembed import AugmentationReembedContext -from onyx.indexing.port_reembed import re_embed_chunks -from onyx.indexing.port_reembed import rebuild_semantic_tail -from onyx.indexing.port_reembed import recover_embedding_input -from onyx.indexing.port_reembed import ReembedStrategy -from onyx.indexing.port_reembed import select_reembed_strategy +from onyx.indexing.embedder import DefaultIndexingEmbedder, IndexingEmbedder +from onyx.indexing.models import ChunkEmbedding, DocAwareChunk, IndexChunk +from onyx.indexing.port_reembed import ( + _bare_contents, + _reconstruct_source_document, + _stored_chunk_to_doc_aware, + AugmentationReembedContext, + re_embed_chunks, + rebuild_semantic_tail, + recover_embedding_input, + ReembedStrategy, + select_reembed_strategy, +) from onyx.natural_language_processing.utils import BaseTokenizer from onyx.utils.pydantic_util import shallow_model_dump -from shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA +from shared_configs.configs import ( + DOC_EMBEDDING_CONTEXT_SIZE, + MODEL_SERVER_HOST, + POSTGRES_DEFAULT_SCHEMA, +) def _stored_chunk( diff --git a/backend/tests/external_dependency_unit/indexing_helpers.py b/backend/tests/external_dependency_unit/indexing_helpers.py index 9f35a4b10c..f4631b6471 100644 --- a/backend/tests/external_dependency_unit/indexing_helpers.py +++ b/backend/tests/external_dependency_unit/indexing_helpers.py @@ -14,26 +14,22 @@ from uuid import uuid4 from sqlalchemy.orm import Session -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FileOrigin -from onyx.connectors.models import Document -from onyx.connectors.models import InputType -from onyx.connectors.models import TextSection +from onyx.configs.constants import DocumentSource, FileOrigin +from onyx.connectors.models import Document, InputType, TextSection from onyx.context.search.models import SavedSearchSettings -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexModelStatus +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, IndexModelStatus from onyx.db.file_record import get_filerecord_by_file_id_optional -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + DocumentByConnectorCredentialPair, + FileRecord, + IndexAttempt, + SearchSettings, +) from onyx.db.models import Document as DBDocument -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import FileRecord -from onyx.db.models import IndexAttempt -from onyx.db.models import SearchSettings -from onyx.db.search_settings import create_search_settings -from onyx.db.search_settings import get_current_search_settings +from onyx.db.search_settings import create_search_settings, get_current_search_settings from onyx.file_store.file_store import get_default_file_store from onyx.kg.models import KGStage diff --git a/backend/tests/external_dependency_unit/llm/test_bedrock.py b/backend/tests/external_dependency_unit/llm/test_bedrock.py index 15c039a7b0..8b2a3d8406 100644 --- a/backend/tests/external_dependency_unit/llm/test_bedrock.py +++ b/backend/tests/external_dependency_unit/llm/test_bedrock.py @@ -15,8 +15,7 @@ The auth-error path lives at the API surface in import pytest from onyx.llm.constants import LlmProviderNames -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ChatCompletionMessage, UserMessage from onyx.llm.multi_llm import LitellmLLM from tests.utils.secret_names import TestSecret diff --git a/backend/tests/external_dependency_unit/llm/test_llm_provider.py b/backend/tests/external_dependency_unit/llm/test_llm_provider.py index 3f4cddfaa5..7b89c78722 100644 --- a/backend/tests/external_dependency_unit/llm/test_llm_provider.py +++ b/backend/tests/external_dependency_unit/llm/test_llm_provider.py @@ -6,18 +6,19 @@ that LLM credentials are valid before saving them. """ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session from onyx.db.enums import LLMModelFlowType -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import ( + fetch_existing_llm_provider, + remove_llm_provider, + update_default_provider, + upsert_llm_provider, +) from onyx.db.models import UserRole from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError @@ -29,9 +30,11 @@ from onyx.server.manage.llm.api import ( from onyx.server.manage.llm.api import ( test_llm_configuration as run_test_llm_configuration, ) -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + LLMProviderView, + ModelConfigurationUpsertRequest, +) from onyx.server.manage.llm.models import TestLLMRequest as LLMTestRequest diff --git a/backend/tests/external_dependency_unit/llm/test_llm_provider_api_base.py b/backend/tests/external_dependency_unit/llm/test_llm_provider_api_base.py index 8c7e84f1d6..2130569ba6 100644 --- a/backend/tests/external_dependency_unit/llm/test_llm_provider_api_base.py +++ b/backend/tests/external_dependency_unit/llm/test_llm_provider_api_base.py @@ -11,26 +11,28 @@ also need to control the MULTI_TENANT setting via patching. """ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import ( + fetch_existing_llm_provider, + remove_llm_provider, + upsert_llm_provider, +) from onyx.db.models import UserRole from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.llm.constants import LlmProviderNames -from onyx.server.manage.llm.api import _mask_string -from onyx.server.manage.llm.api import put_llm_provider +from onyx.server.manage.llm.api import _mask_string, put_llm_provider from onyx.server.manage.llm.api import test_llm_configuration as run_llm_config_test -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + LLMProviderView, + ModelConfigurationUpsertRequest, +) from onyx.server.manage.llm.models import TestLLMRequest as LLMTestRequest from tests.external_dependency_unit.mock_llm import LLM diff --git a/backend/tests/external_dependency_unit/llm/test_llm_provider_auto_mode.py b/backend/tests/external_dependency_unit/llm/test_llm_provider_auto_mode.py index 82e54e2d4e..f1ecc2adbf 100644 --- a/backend/tests/external_dependency_unit/llm/test_llm_provider_auto_mode.py +++ b/backend/tests/external_dependency_unit/llm/test_llm_provider_auto_mode.py @@ -7,34 +7,39 @@ is uploaded with is_auto_mode=True. from collections.abc import Generator from datetime import datetime -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session from onyx.db.enums import LLMModelFlowType -from onyx.db.llm import fetch_auto_mode_providers -from onyx.db.llm import fetch_default_llm_model -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import fetch_existing_llm_providers -from onyx.db.llm import fetch_llm_provider_view -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import sync_auto_mode_models -from onyx.db.llm import update_default_provider +from onyx.db.llm import ( + fetch_auto_mode_providers, + fetch_default_llm_model, + fetch_existing_llm_provider, + fetch_existing_llm_providers, + fetch_llm_provider_view, + remove_llm_provider, + sync_auto_mode_models, + update_default_provider, +) from onyx.db.models import UserRole from onyx.llm.constants import LlmProviderNames from onyx.llm.interfaces import LLM -from onyx.llm.well_known_providers.auto_update_models import LLMProviderRecommendation -from onyx.llm.well_known_providers.auto_update_models import LLMRecommendations +from onyx.llm.well_known_providers.auto_update_models import ( + LLMProviderRecommendation, + LLMRecommendations, +) from onyx.llm.well_known_providers.models import SimpleKnownModel from onyx.server.manage.llm.api import put_llm_provider from onyx.server.manage.llm.api import ( test_default_provider as run_test_default_provider, ) -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) def _create_mock_admin() -> MagicMock: diff --git a/backend/tests/external_dependency_unit/llm/test_llm_provider_called.py b/backend/tests/external_dependency_unit/llm/test_llm_provider_called.py index debdd30cf9..9721748b15 100644 --- a/backend/tests/external_dependency_unit/llm/test_llm_provider_called.py +++ b/backend/tests/external_dependency_unit/llm/test_llm_provider_called.py @@ -9,27 +9,33 @@ from fastapi_users.password import PasswordHelper from sqlalchemy.orm import Session from onyx.db.enums import AccountType -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider -from onyx.db.models import User -from onyx.db.models import UserRole +from onyx.db.llm import ( + fetch_existing_llm_provider, + remove_llm_provider, + update_default_provider, + upsert_llm_provider, +) +from onyx.db.models import User, UserRole from onyx.llm.constants import LlmProviderNames from onyx.llm.override_models import LLMOverride -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) from onyx.server.query_and_chat.chat_backend import create_new_chat_session -from onyx.server.query_and_chat.models import ChatSessionCreationRequest -from onyx.server.query_and_chat.models import MessageResponseIDInfo +from onyx.server.query_and_chat.models import ( + ChatSessionCreationRequest, + MessageResponseIDInfo, +) from tests.external_dependency_unit.answer.stream_test_assertions import ( assert_answer_stream_part_correct, ) from tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder -from tests.external_dependency_unit.answer.stream_test_utils import submit_query -from tests.external_dependency_unit.answer.stream_test_utils import tokenise -from tests.external_dependency_unit.mock_llm import LLMAnswerResponse -from tests.external_dependency_unit.mock_llm import MockLLM +from tests.external_dependency_unit.answer.stream_test_utils import ( + submit_query, + tokenise, +) +from tests.external_dependency_unit.mock_llm import LLMAnswerResponse, MockLLM def _create_admin(db_session: Session) -> User: diff --git a/backend/tests/external_dependency_unit/llm/test_llm_provider_default_model_protection.py b/backend/tests/external_dependency_unit/llm/test_llm_provider_default_model_protection.py index f7004ceaeb..25882916cf 100644 --- a/backend/tests/external_dependency_unit/llm/test_llm_provider_default_model_protection.py +++ b/backend/tests/external_dependency_unit/llm/test_llm_provider_default_model_protection.py @@ -11,15 +11,19 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.db.llm import fetch_existing_llm_provider -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import update_default_provider -from onyx.db.llm import update_default_vision_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import ( + fetch_existing_llm_provider, + remove_llm_provider, + update_default_provider, + update_default_vision_provider, + upsert_llm_provider, +) from onyx.llm.constants import LlmProviderNames -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + LLMProviderView, + ModelConfigurationUpsertRequest, +) def _create_test_provider( diff --git a/backend/tests/external_dependency_unit/llm/test_llm_provider_user_default_cleanup.py b/backend/tests/external_dependency_unit/llm/test_llm_provider_user_default_cleanup.py index 8e50701e7e..cd4c1cbc8b 100644 --- a/backend/tests/external_dependency_unit/llm/test_llm_provider_user_default_cleanup.py +++ b/backend/tests/external_dependency_unit/llm/test_llm_provider_user_default_cleanup.py @@ -9,13 +9,14 @@ from uuid import uuid4 from sqlalchemy import delete from sqlalchemy.orm import Session -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import remove_llm_provider, upsert_llm_provider from onyx.db.models import User from onyx.llm.constants import LlmProviderNames -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + LLMProviderView, + ModelConfigurationUpsertRequest, +) from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/llm/test_ollama_streaming.py b/backend/tests/external_dependency_unit/llm/test_ollama_streaming.py index 38204eeb82..532574b661 100644 --- a/backend/tests/external_dependency_unit/llm/test_ollama_streaming.py +++ b/backend/tests/external_dependency_unit/llm/test_ollama_streaming.py @@ -13,8 +13,7 @@ protocol drift. import pytest from onyx.llm.constants import LlmProviderNames -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ChatCompletionMessage, UserMessage from onyx.llm.multi_llm import LitellmLLM from tests.utils.secret_names import TestSecret diff --git a/backend/tests/external_dependency_unit/llm/test_openai_responses_api.py b/backend/tests/external_dependency_unit/llm/test_openai_responses_api.py index b566fa6594..3dc2947201 100644 --- a/backend/tests/external_dependency_unit/llm/test_openai_responses_api.py +++ b/backend/tests/external_dependency_unit/llm/test_openai_responses_api.py @@ -24,8 +24,7 @@ import pytest from onyx.llm.constants import LlmProviderNames from onyx.llm.litellm_singleton import litellm -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ChatCompletionMessage, UserMessage from onyx.llm.multi_llm import LitellmLLM from tests.utils.secret_names import TestSecret diff --git a/backend/tests/external_dependency_unit/llm/test_prompt_caching.py b/backend/tests/external_dependency_unit/llm/test_prompt_caching.py index d1494f3dc4..6555d9a02a 100644 --- a/backend/tests/external_dependency_unit/llm/test_prompt_caching.py +++ b/backend/tests/external_dependency_unit/llm/test_prompt_caching.py @@ -16,10 +16,12 @@ from litellm import completion_cost from sqlalchemy.orm import Session from onyx.llm.model_response import Usage -from onyx.llm.models import AssistantMessage -from onyx.llm.models import ChatCompletionMessage -from onyx.llm.models import SystemMessage -from onyx.llm.models import UserMessage +from onyx.llm.models import ( + AssistantMessage, + ChatCompletionMessage, + SystemMessage, + UserMessage, +) from onyx.llm.multi_llm import LitellmLLM from onyx.llm.prompt_cache.processor import process_with_prompt_cache diff --git a/backend/tests/external_dependency_unit/mock_content_provider.py b/backend/tests/external_dependency_unit/mock_content_provider.py index 3d40ecaf64..8d7147c28b 100644 --- a/backend/tests/external_dependency_unit/mock_content_provider.py +++ b/backend/tests/external_dependency_unit/mock_content_provider.py @@ -1,13 +1,14 @@ import abc -from collections.abc import Generator -from collections.abc import Sequence +from collections.abc import Generator, Sequence from contextlib import contextmanager from unittest.mock import patch from pydantic import BaseModel -from onyx.tools.tool_implementations.open_url.models import WebContent -from onyx.tools.tool_implementations.open_url.models import WebContentProvider +from onyx.tools.tool_implementations.open_url.models import ( + WebContent, + WebContentProvider, +) class MockWebContent(BaseModel): diff --git a/backend/tests/external_dependency_unit/mock_image_provider.py b/backend/tests/external_dependency_unit/mock_image_provider.py index e69ff636d8..3baf42cb10 100644 --- a/backend/tests/external_dependency_unit/mock_image_provider.py +++ b/backend/tests/external_dependency_unit/mock_image_provider.py @@ -8,12 +8,13 @@ from datetime import datetime from typing import Any from unittest.mock import patch -from litellm.types.utils import ImageObject -from litellm.types.utils import ImageResponse +from litellm.types.utils import ImageObject, ImageResponse -from onyx.image_gen.interfaces import ImageGenerationProvider -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.interfaces import ( + ImageGenerationProvider, + ImageGenerationProviderCredentials, + ReferenceImage, +) from onyx.llm.interfaces import LLMConfig diff --git a/backend/tests/external_dependency_unit/mock_llm.py b/backend/tests/external_dependency_unit/mock_llm.py index 5ca45f495a..0dbf841595 100644 --- a/backend/tests/external_dependency_unit/mock_llm.py +++ b/backend/tests/external_dependency_unit/mock_llm.py @@ -3,31 +3,30 @@ from __future__ import annotations import abc import threading import time -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Generator, Iterator from contextlib import contextmanager from enum import Enum -from typing import Any -from typing import cast -from typing import Generic -from typing import Literal -from typing import TypeVar +from typing import Any, cast, Generic, Literal, TypeVar from unittest.mock import patch from pydantic import BaseModel -from onyx.llm.interfaces import LanguageModelInput -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMConfig -from onyx.llm.interfaces import LLMUserIdentity -from onyx.llm.interfaces import ReasoningEffort -from onyx.llm.interfaces import ToolChoiceOptions -from onyx.llm.model_response import ChatCompletionDeltaToolCall -from onyx.llm.model_response import Delta -from onyx.llm.model_response import FunctionCall -from onyx.llm.model_response import ModelResponse -from onyx.llm.model_response import ModelResponseStream -from onyx.llm.model_response import StreamingChoice +from onyx.llm.interfaces import ( + LanguageModelInput, + LLM, + LLMConfig, + LLMUserIdentity, + ReasoningEffort, + ToolChoiceOptions, +) +from onyx.llm.model_response import ( + ChatCompletionDeltaToolCall, + Delta, + FunctionCall, + ModelResponse, + ModelResponseStream, + StreamingChoice, +) T = TypeVar("T") diff --git a/backend/tests/external_dependency_unit/mock_search_pipeline.py b/backend/tests/external_dependency_unit/mock_search_pipeline.py index b8aba11472..55c38aadee 100644 --- a/backend/tests/external_dependency_unit/mock_search_pipeline.py +++ b/backend/tests/external_dependency_unit/mock_search_pipeline.py @@ -1,20 +1,19 @@ -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from pydantic import BaseModel from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource -from onyx.context.search.models import ChunkSearchRequest -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import PersonaSearchInfo -from onyx.context.search.models import SearchDoc -from onyx.db.models import SearchSettings -from onyx.db.models import User +from onyx.context.search.models import ( + ChunkSearchRequest, + InferenceChunk, + PersonaSearchInfo, + SearchDoc, +) +from onyx.db.models import SearchSettings, User from onyx.document_index.interfaces_new import DocumentIndex from onyx.federated_connectors.federated_retrieval import FederatedRetrievalInfo from onyx.llm.interfaces import LLM diff --git a/backend/tests/external_dependency_unit/mock_search_provider.py b/backend/tests/external_dependency_unit/mock_search_provider.py index c98ef88aa3..02412c886d 100644 --- a/backend/tests/external_dependency_unit/mock_search_provider.py +++ b/backend/tests/external_dependency_unit/mock_search_provider.py @@ -1,7 +1,6 @@ import abc from collections import defaultdict -from collections.abc import Generator -from collections.abc import Sequence +from collections.abc import Generator, Sequence from contextlib import contextmanager from unittest.mock import patch @@ -10,8 +9,10 @@ from sqlalchemy.orm import Session from onyx.db.models import InternetSearchProvider from onyx.db.web_search import fetch_web_search_provider_by_name -from onyx.tools.tool_implementations.web_search.models import WebSearchProvider -from onyx.tools.tool_implementations.web_search.models import WebSearchResult +from onyx.tools.tool_implementations.web_search.models import ( + WebSearchProvider, + WebSearchResult, +) from shared_configs.enums import WebSearchProviderType diff --git a/backend/tests/external_dependency_unit/natural_language_processing/test_get_query_embeddings.py b/backend/tests/external_dependency_unit/natural_language_processing/test_get_query_embeddings.py index 4457b37f7c..e07b2a5844 100644 --- a/backend/tests/external_dependency_unit/natural_language_processing/test_get_query_embeddings.py +++ b/backend/tests/external_dependency_unit/natural_language_processing/test_get_query_embeddings.py @@ -1,7 +1,6 @@ """Tests that the cache is wired correctly into ``get_query_embeddings``.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from sqlalchemy.orm import Session diff --git a/backend/tests/external_dependency_unit/natural_language_processing/test_query_embedding_cache.py b/backend/tests/external_dependency_unit/natural_language_processing/test_query_embedding_cache.py index d4066aa749..59d9613707 100644 --- a/backend/tests/external_dependency_unit/natural_language_processing/test_query_embedding_cache.py +++ b/backend/tests/external_dependency_unit/natural_language_processing/test_query_embedding_cache.py @@ -7,11 +7,9 @@ from prometheus_client import REGISTRY from redis.exceptions import RedisError from onyx.cache.factory import get_cache_backend -from onyx.natural_language_processing.query_embedding_cache import _build_key from onyx.natural_language_processing.query_embedding_cache import ( + _build_key, cache_query_embeddings, -) -from onyx.natural_language_processing.query_embedding_cache import ( get_cached_query_embeddings, ) from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/external_dependency_unit/opensearch/test_assistant_knowledge_filter.py b/backend/tests/external_dependency_unit/opensearch/test_assistant_knowledge_filter.py index 0a2bc62330..5b7166da0f 100644 --- a/backend/tests/external_dependency_unit/opensearch/test_assistant_knowledge_filter.py +++ b/backend/tests/external_dependency_unit/opensearch/test_assistant_knowledge_filter.py @@ -10,11 +10,13 @@ from typing import Any from onyx.configs.constants import DocumentSource from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.opensearch.schema import ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME -from onyx.document_index.opensearch.schema import DOCUMENT_ID_FIELD_NAME -from onyx.document_index.opensearch.schema import DOCUMENT_SETS_FIELD_NAME -from onyx.document_index.opensearch.schema import PERSONAS_FIELD_NAME -from onyx.document_index.opensearch.schema import USER_PROJECTS_FIELD_NAME +from onyx.document_index.opensearch.schema import ( + ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME, + DOCUMENT_ID_FIELD_NAME, + DOCUMENT_SETS_FIELD_NAME, + PERSONAS_FIELD_NAME, + USER_PROJECTS_FIELD_NAME, +) from onyx.document_index.opensearch.search import DocumentQuery from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA diff --git a/backend/tests/external_dependency_unit/opensearch/test_opensearch_client.py b/backend/tests/external_dependency_unit/opensearch/test_opensearch_client.py index a946fe87cc..cb03502c0a 100644 --- a/backend/tests/external_dependency_unit/opensearch/test_opensearch_client.py +++ b/backend/tests/external_dependency_unit/opensearch/test_opensearch_client.py @@ -7,51 +7,49 @@ using real schemas, pipelines, and search queries from the codebase. import re import uuid from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import MagicMock import pytest -from opensearchpy import ConflictError -from opensearchpy import NotFoundError +from opensearchpy import ConflictError, NotFoundError from opensearchpy.helpers import BulkIndexError import onyx.document_index.opensearch.client as client_module from onyx.access.models import DocumentAccess from onyx.access.utils import prefix_user_email from onyx.configs.constants import DocumentSource -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import TimeRange +from onyx.context.search.models import IndexFilters, TimeRange from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.opensearch.client import OpenSearchDocumentMissingError -from onyx.document_index.opensearch.client import OpenSearchIndexClient -from onyx.document_index.opensearch.client import OpenSearchIndexError -from onyx.document_index.opensearch.client import OpenSearchServerSideTimeout -from onyx.document_index.opensearch.client import OpenSearchUpdateError -from onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout -from onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE -from onyx.document_index.opensearch.constants import HybridSearchNormalizationPipeline -from onyx.document_index.opensearch.constants import HybridSearchSubqueryConfiguration -from onyx.document_index.opensearch.constants import OpenSearchSearchType +from onyx.document_index.opensearch.client import ( + OpenSearchDocumentMissingError, + OpenSearchIndexClient, + OpenSearchIndexError, + OpenSearchServerSideTimeout, + OpenSearchUpdateError, + wait_for_opensearch_with_timeout, +) +from onyx.document_index.opensearch.constants import ( + DEFAULT_MAX_CHUNK_SIZE, + HybridSearchNormalizationPipeline, + HybridSearchSubqueryConfiguration, + OpenSearchSearchType, +) from onyx.document_index.opensearch.opensearch_document_index import ( generate_opensearch_filtered_access_control_list, ) -from onyx.document_index.opensearch.schema import ACCESS_CONTROL_LIST_FIELD_NAME -from onyx.document_index.opensearch.schema import CONTENT_FIELD_NAME -from onyx.document_index.opensearch.schema import DocumentChunk -from onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors -from onyx.document_index.opensearch.schema import DocumentSchema -from onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id -from onyx.document_index.opensearch.search import DocumentQuery +from onyx.document_index.opensearch.schema import ( + ACCESS_CONTROL_LIST_FIELD_NAME, + CONTENT_FIELD_NAME, + DocumentChunk, + DocumentChunkWithoutVectors, + DocumentSchema, + get_opensearch_doc_chunk_id, +) from onyx.document_index.opensearch.search import ( + DocumentQuery, get_min_max_normalization_pipeline_name_and_config, -) -from onyx.document_index.opensearch.search import ( get_normalization_pipeline_name_and_config, -) -from onyx.document_index.opensearch.search import ( get_zscore_normalization_pipeline_name_and_config, ) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA diff --git a/backend/tests/external_dependency_unit/permission_sync/test_cc_pair_sync_attempts_routes.py b/backend/tests/external_dependency_unit/permission_sync/test_cc_pair_sync_attempts_routes.py index 2da7c41aa3..8c25c4a202 100644 --- a/backend/tests/external_dependency_unit/permission_sync/test_cc_pair_sync_attempts_routes.py +++ b/backend/tests/external_dependency_unit/permission_sync/test_cc_pair_sync_attempts_routes.py @@ -17,31 +17,36 @@ test that needs ``applicable=True`` uses the ``enable_ee`` fixture from the root ``backend/tests/conftest.py``. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import PermissionSyncStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import User -from onyx.db.models import UserRole -from onyx.db.permission_sync_attempt import create_doc_permission_sync_attempt -from onyx.db.permission_sync_attempt import create_external_group_sync_attempt +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + PermissionSyncStatus, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + User, + UserRole, +) from onyx.db.permission_sync_attempt import ( + create_doc_permission_sync_attempt, + create_external_group_sync_attempt, get_relevant_external_group_sync_attempts_for_cc_pair, ) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.documents.cc_pair import get_cc_pair_external_group_sync_attempts -from onyx.server.documents.cc_pair import get_cc_pair_permission_sync_attempts +from onyx.server.documents.cc_pair import ( + get_cc_pair_external_group_sync_attempts, + get_cc_pair_permission_sync_attempts, +) from tests.external_dependency_unit.conftest import create_test_user # Every applicable=True path here depends on the EE-only ``sync_params`` diff --git a/backend/tests/external_dependency_unit/permission_sync/test_doc_permission_sync_attempt.py b/backend/tests/external_dependency_unit/permission_sync/test_doc_permission_sync_attempt.py index c16d87689e..31e91a4e41 100644 --- a/backend/tests/external_dependency_unit/permission_sync/test_doc_permission_sync_attempt.py +++ b/backend/tests/external_dependency_unit/permission_sync/test_doc_permission_sync_attempt.py @@ -5,28 +5,27 @@ Tests the basic CRUD operations for document permission sync attempts, including creation, status updates, progress tracking, and querying. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import PermissionSyncStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.permission_sync_attempt import complete_doc_permission_sync_attempt -from onyx.db.permission_sync_attempt import create_doc_permission_sync_attempt -from onyx.db.permission_sync_attempt import get_doc_permission_sync_attempt +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + PermissionSyncStatus, +) +from onyx.db.models import Connector, ConnectorCredentialPair, Credential from onyx.db.permission_sync_attempt import ( + complete_doc_permission_sync_attempt, + create_doc_permission_sync_attempt, + get_doc_permission_sync_attempt, get_recent_doc_permission_sync_attempts_for_cc_pair, + mark_doc_permission_sync_attempt_failed, + mark_doc_permission_sync_attempt_in_progress, ) -from onyx.db.permission_sync_attempt import mark_doc_permission_sync_attempt_failed -from onyx.db.permission_sync_attempt import mark_doc_permission_sync_attempt_in_progress from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/permission_sync/test_external_group_permission_sync_attempt.py b/backend/tests/external_dependency_unit/permission_sync/test_external_group_permission_sync_attempt.py index 4067d67fab..e7b2b513d7 100644 --- a/backend/tests/external_dependency_unit/permission_sync/test_external_group_permission_sync_attempt.py +++ b/backend/tests/external_dependency_unit/permission_sync/test_external_group_permission_sync_attempt.py @@ -6,29 +6,32 @@ including creation, status updates, progress tracking, and querying. Supports both connector-specific and global group sync attempts. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import PermissionSyncStatus -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import ExternalGroupPermissionSyncAttempt -from onyx.db.permission_sync_attempt import complete_external_group_sync_attempt -from onyx.db.permission_sync_attempt import create_external_group_sync_attempt -from onyx.db.permission_sync_attempt import get_external_group_sync_attempt +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + PermissionSyncStatus, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + ExternalGroupPermissionSyncAttempt, +) from onyx.db.permission_sync_attempt import ( + complete_external_group_sync_attempt, + create_external_group_sync_attempt, + get_external_group_sync_attempt, get_recent_external_group_sync_attempts_for_cc_pair, + mark_external_group_sync_attempt_failed, + mark_external_group_sync_attempt_in_progress, ) -from onyx.db.permission_sync_attempt import mark_external_group_sync_attempt_failed -from onyx.db.permission_sync_attempt import mark_external_group_sync_attempt_in_progress from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/redis/test_lock_context.py b/backend/tests/external_dependency_unit/redis/test_lock_context.py index bb4b336459..9eff784d60 100644 --- a/backend/tests/external_dependency_unit/redis/test_lock_context.py +++ b/backend/tests/external_dependency_unit/redis/test_lock_context.py @@ -9,8 +9,7 @@ from collections.abc import Generator import pytest -from onyx.redis.lock_context import redis_shared_lock -from onyx.redis.lock_context import RedisSharedLockAcquisitionError +from onyx.redis.lock_context import redis_shared_lock, RedisSharedLockAcquisitionError from onyx.redis.redis_pool import get_shared_redis_client from onyx.utils.logger import setup_logger diff --git a/backend/tests/external_dependency_unit/redis/test_tenant_redis.py b/backend/tests/external_dependency_unit/redis/test_tenant_redis.py index e186a4db83..f7bdc8a560 100644 --- a/backend/tests/external_dependency_unit/redis/test_tenant_redis.py +++ b/backend/tests/external_dependency_unit/redis/test_tenant_redis.py @@ -22,8 +22,7 @@ from uuid import uuid4 import pytest from redis import Redis -from onyx.redis.redis_pool import get_raw_redis_client -from onyx.redis.redis_pool import redis_pool +from onyx.redis.redis_pool import get_raw_redis_client, redis_pool from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/tests/external_dependency_unit/sandbox_proxy/test_approval_cache.py b/backend/tests/external_dependency_unit/sandbox_proxy/test_approval_cache.py index 3a0a0ebf92..9702aa7908 100644 --- a/backend/tests/external_dependency_unit/sandbox_proxy/test_approval_cache.py +++ b/backend/tests/external_dependency_unit/sandbox_proxy/test_approval_cache.py @@ -3,21 +3,22 @@ import threading import time -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from onyx.cache.factory import get_cache_backend from onyx.db.enums import ApprovalDecision -from onyx.sandbox_proxy.approval_cache import _wake_key -from onyx.sandbox_proxy.approval_cache import announce_approval -from onyx.sandbox_proxy.approval_cache import announce_key -from onyx.sandbox_proxy.approval_cache import cache_session_grant_actions -from onyx.sandbox_proxy.approval_cache import cached_session_grants_cover -from onyx.sandbox_proxy.approval_cache import pop_announcement -from onyx.sandbox_proxy.approval_cache import send_wake -from onyx.sandbox_proxy.approval_cache import wait_for_wake +from onyx.sandbox_proxy.approval_cache import ( + _wake_key, + announce_approval, + announce_key, + cache_session_grant_actions, + cached_session_grants_cover, + pop_announcement, + send_wake, + wait_for_wake, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE # --------------------------------------------------------------------------- diff --git a/backend/tests/external_dependency_unit/sandbox_proxy/test_gate_claim_arbiter.py b/backend/tests/external_dependency_unit/sandbox_proxy/test_gate_claim_arbiter.py index 58d14c183f..73fdaf6f13 100644 --- a/backend/tests/external_dependency_unit/sandbox_proxy/test_gate_claim_arbiter.py +++ b/backend/tests/external_dependency_unit/sandbox_proxy/test_gate_claim_arbiter.py @@ -9,17 +9,13 @@ from __future__ import annotations import datetime as dt from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 from sqlalchemy.orm import Session -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import BuildSessionStatus -from onyx.db.models import ActionApproval -from onyx.db.models import BuildSession -from onyx.sandbox_proxy.addons.gate import _IdentityResolver -from onyx.sandbox_proxy.addons.gate import GateAddon +from onyx.db.enums import ApprovalDecision, BuildSessionStatus +from onyx.db.models import ActionApproval, BuildSession +from onyx.sandbox_proxy.addons.gate import _IdentityResolver, GateAddon from onyx.sandbox_proxy.credential_injection import CredentialInjectionDispatcher from onyx.sandbox_proxy.identity import ResolvedSandbox from onyx.sandbox_proxy.request_evaluator import RequestEvaluator diff --git a/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_docker_lookup.py b/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_docker_lookup.py index ee100e7bcf..e69028e7d5 100644 --- a/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_docker_lookup.py +++ b/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_docker_lookup.py @@ -14,17 +14,17 @@ from __future__ import annotations import os import time from collections.abc import Generator -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from docker import DockerClient -from docker.errors import APIError -from docker.errors import NotFound +from docker.errors import APIError, NotFound from docker.models.containers import Container -from onyx.sandbox_proxy.identity_docker import _identity_from_container -from onyx.sandbox_proxy.identity_docker import DockerEventsLookup +from onyx.sandbox_proxy.identity_docker import ( + _identity_from_container, + DockerEventsLookup, +) _DOCKER_SOCKET = os.environ.get("SANDBOX_DOCKER_SOCKET", "/var/run/docker.sock") _TEST_NETWORK = "onyx-craft-sandbox-test" diff --git a/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_resolver.py b/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_resolver.py index a44987ee4e..31b28addc0 100644 --- a/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_resolver.py +++ b/backend/tests/external_dependency_unit/sandbox_proxy/test_identity_resolver.py @@ -7,17 +7,14 @@ enum drift, ordering change) actually fails the test. import datetime as dt from collections.abc import Generator -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session from onyx.db.enums import BuildSessionStatus -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox -from onyx.sandbox_proxy.identity import IdentityResolver -from onyx.sandbox_proxy.identity import SandboxIdentity +from onyx.db.models import BuildSession, Sandbox +from onyx.sandbox_proxy.identity import IdentityResolver, SandboxIdentity from shared_configs.contextvars import POSTGRES_DEFAULT_SCHEMA from tests.external_dependency_unit.conftest import create_test_user from tests.unit.sandbox_proxy.conftest import StaticLookup diff --git a/backend/tests/external_dependency_unit/sandbox_proxy/test_llm_provider_key_resolver.py b/backend/tests/external_dependency_unit/sandbox_proxy/test_llm_provider_key_resolver.py index c96fcab217..40c2fa64c7 100644 --- a/backend/tests/external_dependency_unit/sandbox_proxy/test_llm_provider_key_resolver.py +++ b/backend/tests/external_dependency_unit/sandbox_proxy/test_llm_provider_key_resolver.py @@ -13,8 +13,7 @@ from uuid import uuid4 from sqlalchemy.orm import Session -from onyx.db.llm import remove_llm_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import remove_llm_provider, upsert_llm_provider from onyx.sandbox_proxy.credential_injection import InjectionContext from onyx.sandbox_proxy.identity import ResolvedSandbox from onyx.sandbox_proxy.resolvers.llm_provider_key import LLMProviderKeyResolver diff --git a/backend/tests/external_dependency_unit/sandbox_proxy/test_onyx_pat_resolver.py b/backend/tests/external_dependency_unit/sandbox_proxy/test_onyx_pat_resolver.py index a6da136bd5..3e9c670070 100644 --- a/backend/tests/external_dependency_unit/sandbox_proxy/test_onyx_pat_resolver.py +++ b/backend/tests/external_dependency_unit/sandbox_proxy/test_onyx_pat_resolver.py @@ -13,14 +13,15 @@ from unittest.mock import MagicMock import pytest from sqlalchemy.orm import Session -from onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME -from onyx.auth.constants import API_KEY_HEADER_NAME +from onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME, API_KEY_HEADER_NAME from onyx.sandbox_proxy.credential_injection import InjectionContext from onyx.sandbox_proxy.identity import ResolvedSandbox from onyx.sandbox_proxy.resolvers import onyx_pat as onyx_pat_mod from onyx.sandbox_proxy.resolvers.onyx_pat import OnyxPatResolver -from onyx.server.features.build.db.sandbox import create_sandbox__no_commit -from onyx.server.features.build.db.sandbox import ensure_sandbox_pat +from onyx.server.features.build.db.sandbox import ( + create_sandbox__no_commit, + ensure_sandbox_pat, +) from shared_configs.contextvars import POSTGRES_DEFAULT_SCHEMA from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/search_settings/test_index_swap_workflow.py b/backend/tests/external_dependency_unit/search_settings/test_index_swap_workflow.py index 038aef0c5d..19cd25cb3e 100644 --- a/backend/tests/external_dependency_unit/search_settings/test_index_swap_workflow.py +++ b/backend/tests/external_dependency_unit/search_settings/test_index_swap_workflow.py @@ -18,19 +18,19 @@ from sqlalchemy.orm import Session from onyx.connectors.models import IndexAttemptMetadata from onyx.context.search.models import SavedSearchSettings -from onyx.db.enums import EmbeddingPrecision -from onyx.db.enums import SwitchoverType -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import IndexModelStatus +from onyx.db.enums import EmbeddingPrecision, SwitchoverType +from onyx.db.models import ConnectorCredentialPair, IndexModelStatus from onyx.db.search_settings import create_search_settings from onyx.db.swap_index import check_and_perform_index_swap from onyx.indexing.indexing_pipeline import index_doc_batch_prepare -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import get_doc_row -from tests.external_dependency_unit.indexing_helpers import get_filerecord -from tests.external_dependency_unit.indexing_helpers import make_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_doc -from tests.external_dependency_unit.indexing_helpers import stage_file +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + get_doc_row, + get_filerecord, + make_cc_pair, + make_doc, + stage_file, +) # --------------------------------------------------------------------------- # Helpers (file-local) diff --git a/backend/tests/external_dependency_unit/search_settings/test_search_settings.py b/backend/tests/external_dependency_unit/search_settings/test_search_settings.py index 013c3e22f6..1d6b3a9429 100644 --- a/backend/tests/external_dependency_unit/search_settings/test_search_settings.py +++ b/backend/tests/external_dependency_unit/search_settings/test_search_settings.py @@ -1,34 +1,41 @@ """Tests that search settings with contextual RAG are properly propagated to the indexing pipeline's LLM configuration.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from sqlalchemy.orm import Session -from onyx.context.search.models import SavedSearchSettings -from onyx.context.search.models import SearchSettingsCreationRequest -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import EmbeddingPrecision -from onyx.db.llm import fetch_default_contextual_rag_model -from onyx.db.llm import update_default_contextual_model -from onyx.db.llm import upsert_llm_provider -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexModelStatus -from onyx.db.models import SearchSettings -from onyx.db.search_settings import create_search_settings -from onyx.db.search_settings import update_search_settings +from onyx.context.search.models import ( + SavedSearchSettings, + SearchSettingsCreationRequest, +) +from onyx.db.enums import ConnectorCredentialPairStatus, EmbeddingPrecision +from onyx.db.llm import ( + fetch_default_contextual_rag_model, + update_default_contextual_model, + upsert_llm_provider, +) +from onyx.db.models import IndexAttempt, IndexModelStatus, SearchSettings +from onyx.db.search_settings import create_search_settings, update_search_settings from onyx.db.swap_index import check_and_perform_index_swap -from onyx.indexing.indexing_pipeline import IndexingPipelineResult -from onyx.indexing.indexing_pipeline import run_indexing_pipeline -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest -from onyx.server.manage.search_settings import set_new_search_settings -from onyx.server.manage.search_settings import update_saved_search_settings +from onyx.indexing.indexing_pipeline import ( + IndexingPipelineResult, + run_indexing_pipeline, +) +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) +from onyx.server.manage.search_settings import ( + set_new_search_settings, + update_saved_search_settings, +) from shared_configs.configs import PRESERVED_SEARCH_FIELDS -from tests.external_dependency_unit.indexing_helpers import cleanup_cc_pair -from tests.external_dependency_unit.indexing_helpers import make_cc_pair +from tests.external_dependency_unit.indexing_helpers import ( + cleanup_cc_pair, + make_cc_pair, +) TEST_PROVIDER_NAME = "test-contextual-provider" TEST_MODEL_NAME = "test-contextual-model" diff --git a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_control.py b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_control.py index f0f02a5f95..68392ae641 100644 --- a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_control.py +++ b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_control.py @@ -9,14 +9,15 @@ from uuid import uuid4 from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.mcp import get_mcp_servers_accessible_to_user -from onyx.db.mcp import user_can_access_mcp_server -from onyx.db.models import MCPServer -from onyx.db.models import MCPServer__User -from onyx.db.models import MCPServer__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.mcp import get_mcp_servers_accessible_to_user, user_can_access_mcp_server +from onyx.db.models import ( + MCPServer, + MCPServer__User, + MCPServer__UserGroup, + User, + User__UserGroup, + UserGroup, +) from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_edge_cases.py b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_edge_cases.py index d90dad6ca2..58e5c58c17 100644 --- a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_edge_cases.py +++ b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_access_edge_cases.py @@ -8,14 +8,15 @@ from uuid import uuid4 from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.mcp import get_mcp_servers_accessible_to_user -from onyx.db.mcp import user_can_access_mcp_server -from onyx.db.models import MCPServer -from onyx.db.models import MCPServer__User -from onyx.db.models import MCPServer__UserGroup -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.mcp import get_mcp_servers_accessible_to_user, user_can_access_mcp_server +from onyx.db.models import ( + MCPServer, + MCPServer__User, + MCPServer__UserGroup, + User, + User__UserGroup, + UserGroup, +) from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_grant_reconciliation.py b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_grant_reconciliation.py index 1c6860c381..68ac80cdb4 100644 --- a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_grant_reconciliation.py +++ b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_grant_reconciliation.py @@ -6,10 +6,7 @@ from sqlalchemy import select from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.models import MCPServer -from onyx.db.models import MCPServer__User -from onyx.db.models import MCPServer__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import MCPServer, MCPServer__User, MCPServer__UserGroup, UserGroup from onyx.server.features.mcp.api import update_mcp_server_simple from onyx.server.features.mcp.models import MCPServerSimpleUpdateRequest from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_persona_gate.py b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_persona_gate.py index a996a39244..089e9ca1eb 100644 --- a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_persona_gate.py +++ b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_persona_gate.py @@ -8,13 +8,15 @@ import pytest from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.models import MCPServer -from onyx.db.models import MCPServer__UserGroup -from onyx.db.models import Persona -from onyx.db.models import Tool -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import ( + MCPServer, + MCPServer__UserGroup, + Persona, + Tool, + User, + User__UserGroup, + UserGroup, +) from onyx.db.persona import upsert_persona from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_user_group_deletion.py b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_user_group_deletion.py index a35c7ac5a4..de3f4d9c1c 100644 --- a/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_user_group_deletion.py +++ b/backend/tests/external_dependency_unit/server/features/mcp/test_mcp_user_group_deletion.py @@ -8,9 +8,7 @@ from sqlalchemy import select from sqlalchemy.orm import Session from ee.onyx.db.user_group import prepare_user_group_for_deletion -from onyx.db.models import MCPServer -from onyx.db.models import MCPServer__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import MCPServer, MCPServer__UserGroup, UserGroup pytestmark = pytest.mark.usefixtures("enable_ee") diff --git a/backend/tests/external_dependency_unit/server/features/mcp/test_per_user_api_token_credentials.py b/backend/tests/external_dependency_unit/server/features/mcp/test_per_user_api_token_credentials.py index a7d19a1b61..369f4707f7 100644 --- a/backend/tests/external_dependency_unit/server/features/mcp/test_per_user_api_token_credentials.py +++ b/backend/tests/external_dependency_unit/server/features/mcp/test_per_user_api_token_credentials.py @@ -12,18 +12,23 @@ from uuid import uuid4 from sqlalchemy.orm import Session from onyx.auth.schemas import UserRole -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPTransport -from onyx.db.mcp import extract_connection_data -from onyx.db.mcp import get_user_connection_config +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPTransport, +) +from onyx.db.mcp import extract_connection_data, get_user_connection_config from onyx.db.models import User -from onyx.server.features.mcp.api import _upsert_mcp_server -from onyx.server.features.mcp.api import HEADER_SUBSTITUTIONS -from onyx.server.features.mcp.api import save_user_credentials -from onyx.server.features.mcp.models import MCPAuthTemplate -from onyx.server.features.mcp.models import MCPToolCreateRequest -from onyx.server.features.mcp.models import MCPUserCredentialsRequest +from onyx.server.features.mcp.api import ( + _upsert_mcp_server, + HEADER_SUBSTITUTIONS, + save_user_credentials, +) +from onyx.server.features.mcp.models import ( + MCPAuthTemplate, + MCPToolCreateRequest, + MCPUserCredentialsRequest, +) from onyx.utils.encryption import mask_string from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/server/security/test_security_settings_put.py b/backend/tests/external_dependency_unit/server/security/test_security_settings_put.py index 1cbca54208..0bf17a9ffe 100644 --- a/backend/tests/external_dependency_unit/server/security/test_security_settings_put.py +++ b/backend/tests/external_dependency_unit/server/security/test_security_settings_put.py @@ -9,13 +9,11 @@ import asyncio import json import threading from collections.abc import Generator -from typing import Any -from typing import cast +from typing import Any, cast import pytest from fastapi import Request -from sqlalchemy import delete -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from onyx.db.engine.sql_engine import get_session_with_current_tenant @@ -27,9 +25,11 @@ from onyx.server.security import api as security_api from onyx.server.security import store as security_store from onyx.server.security.api import put_security_settings_endpoint from onyx.server.security.models import SecuritySettingsOverrides -from onyx.server.security.store import _build_env_defaults -from onyx.server.security.store import _install_cache_for_test -from onyx.server.security.store import invalidate_security_cache +from onyx.server.security.store import ( + _build_env_defaults, + _install_cache_for_test, + invalidate_security_cache, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/external_dependency_unit/server/security/test_security_settings_store.py b/backend/tests/external_dependency_unit/server/security/test_security_settings_store.py index a28f1c2c88..e458c163e4 100644 --- a/backend/tests/external_dependency_unit/server/security/test_security_settings_store.py +++ b/backend/tests/external_dependency_unit/server/security/test_security_settings_store.py @@ -14,11 +14,13 @@ from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import SecuritySettings as SecuritySettingsRow from onyx.server.security import store as security_store from onyx.server.security.models import SecuritySettingsOverrides -from onyx.server.security.store import _build_env_defaults -from onyx.server.security.store import _install_cache_for_test -from onyx.server.security.store import apply_patch -from onyx.server.security.store import get_security_settings -from onyx.server.security.store import invalidate_security_cache +from onyx.server.security.store import ( + _build_env_defaults, + _install_cache_for_test, + apply_patch, + get_security_settings, + invalidate_security_cache, +) from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_crud.py b/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_crud.py index 89d6215eed..ea54cfa78d 100644 --- a/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_crud.py +++ b/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_crud.py @@ -9,8 +9,7 @@ from uuid import uuid4 from sqlalchemy.orm import Session -from onyx.db.slack_bot import insert_slack_bot -from onyx.db.slack_bot import update_slack_bot +from onyx.db.slack_bot import insert_slack_bot, update_slack_bot from onyx.server.manage.models import SlackBot from onyx.utils.sensitive import SensitiveValue diff --git a/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_federated_search.py b/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_federated_search.py index 1b08c62a61..4f24a7b8d1 100644 --- a/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_federated_search.py +++ b/backend/tests/external_dependency_unit/slack_bot/test_slack_bot_federated_search.py @@ -2,15 +2,14 @@ # ruff: noqa: ARG005 import os from typing import Any -from unittest.mock import MagicMock -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import MagicMock, Mock, patch from uuid import uuid4 -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.db.llm import update_default_provider, upsert_llm_provider +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) # Set environment variables to disable model server for testing os.environ["DISABLE_MODEL_SERVER"] = "true" @@ -23,16 +22,18 @@ from sqlalchemy.orm import Session from onyx.configs.constants import FederatedConnectorSource from onyx.context.search.federated.slack_search import fetch_and_cache_channel_metadata -from onyx.db.models import DocumentSet -from onyx.db.models import FederatedConnector -from onyx.db.models import FederatedConnector__DocumentSet -from onyx.db.models import LLMProvider -from onyx.db.models import Persona -from onyx.db.models import Persona__DocumentSet -from onyx.db.models import Persona__Tool -from onyx.db.models import SlackBot -from onyx.db.models import SlackChannelConfig -from onyx.db.models import User +from onyx.db.models import ( + DocumentSet, + FederatedConnector, + FederatedConnector__DocumentSet, + LLMProvider, + Persona, + Persona__DocumentSet, + Persona__Tool, + SlackBot, + SlackChannelConfig, + User, +) from onyx.db.tools import get_builtin_tool from onyx.llm.constants import LlmProviderNames from onyx.onyxbot.slack.listener import process_message diff --git a/backend/tests/external_dependency_unit/tenant_work_gating/test_gate_generator.py b/backend/tests/external_dependency_unit/tenant_work_gating/test_gate_generator.py index 52e5d63cc1..b83198dda0 100644 --- a/backend/tests/external_dependency_unit/tenant_work_gating/test_gate_generator.py +++ b/backend/tests/external_dependency_unit/tenant_work_gating/test_gate_generator.py @@ -11,8 +11,7 @@ Requires a running Redis instance. Run with:: """ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -20,8 +19,7 @@ from ee.onyx.background.celery.tasks.cloud import tasks as cloud_tasks from onyx.configs.constants import ONYX_CLOUD_TENANT_ID from onyx.redis import redis_tenant_work_gating as twg from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_tenant_work_gating import _SET_KEY -from onyx.redis.redis_tenant_work_gating import mark_tenant_active +from onyx.redis.redis_tenant_work_gating import _SET_KEY, mark_tenant_active _TENANT_A = "tenant_aaaa0000-0000-0000-0000-000000000001" _TENANT_B = "tenant_bbbb0000-0000-0000-0000-000000000002" diff --git a/backend/tests/external_dependency_unit/tenant_work_gating/test_tenant_work_gating.py b/backend/tests/external_dependency_unit/tenant_work_gating/test_tenant_work_gating.py index 1b39b0d1ba..6a914df7e4 100644 --- a/backend/tests/external_dependency_unit/tenant_work_gating/test_tenant_work_gating.py +++ b/backend/tests/external_dependency_unit/tenant_work_gating/test_tenant_work_gating.py @@ -15,11 +15,13 @@ import pytest from onyx.configs.constants import ONYX_CLOUD_TENANT_ID from onyx.redis import redis_tenant_work_gating as twg from onyx.redis.redis_pool import get_redis_client -from onyx.redis.redis_tenant_work_gating import _SET_KEY -from onyx.redis.redis_tenant_work_gating import cleanup_expired -from onyx.redis.redis_tenant_work_gating import get_active_tenants -from onyx.redis.redis_tenant_work_gating import mark_tenant_active -from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active +from onyx.redis.redis_tenant_work_gating import ( + _SET_KEY, + cleanup_expired, + get_active_tenants, + mark_tenant_active, + maybe_mark_tenant_active, +) @pytest.fixture(autouse=True) diff --git a/backend/tests/external_dependency_unit/test_telemetry_secret_storage.py b/backend/tests/external_dependency_unit/test_telemetry_secret_storage.py index ce1107b884..287ec1ea53 100644 --- a/backend/tests/external_dependency_unit/test_telemetry_secret_storage.py +++ b/backend/tests/external_dependency_unit/test_telemetry_secret_storage.py @@ -2,11 +2,9 @@ from sqlalchemy.orm import Session -from onyx.configs.constants import KV_CUSTOMER_UUID_KEY -from onyx.configs.constants import KV_INSTANCE_DOMAIN_KEY +from onyx.configs.constants import KV_CUSTOMER_UUID_KEY, KV_INSTANCE_DOMAIN_KEY from onyx.db.encrypted_kv_store import load_encrypted_kv -from onyx.db.models import EncryptedKeyValueStore -from onyx.db.models import KVStore +from onyx.db.models import EncryptedKeyValueStore, KVStore from onyx.key_value_store.interface import unwrap_str from onyx.utils import telemetry from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/tier/test_tier_order.py b/backend/tests/external_dependency_unit/tier/test_tier_order.py index 6135b2af42..ce323e4590 100644 --- a/backend/tests/external_dependency_unit/tier/test_tier_order.py +++ b/backend/tests/external_dependency_unit/tier/test_tier_order.py @@ -5,8 +5,7 @@ import pytest from ee.onyx.configs.license_enforcement_config import PATH_PREFIX_MIN_TIER from ee.onyx.server.middleware.tier_gate import _required_tier from onyx.server.settings.models import Tier -from onyx.server.settings.tier_order import tier_at_least -from onyx.server.settings.tier_order import TIER_RANK +from onyx.server.settings.tier_order import tier_at_least, TIER_RANK class TestTierOrdering: diff --git a/backend/tests/external_dependency_unit/tools/test_mcp_passthrough_oauth.py b/backend/tests/external_dependency_unit/tools/test_mcp_passthrough_oauth.py index 92ab2ad467..baa05cd4a5 100644 --- a/backend/tests/external_dependency_unit/tools/test_mcp_passthrough_oauth.py +++ b/backend/tests/external_dependency_unit/tools/test_mcp_passthrough_oauth.py @@ -22,19 +22,17 @@ import pytest from sqlalchemy.orm import Session from onyx.chat.emitter import Emitter -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPTransport +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPTransport, +) from onyx.db.mcp import create_mcp_server__no_commit -from onyx.db.models import OAuthAccount -from onyx.db.models import Persona -from onyx.db.models import Tool -from onyx.db.models import User +from onyx.db.models import OAuthAccount, Persona, Tool, User from onyx.llm.factory import get_default_llm from onyx.server.query_and_chat.placement import Placement from onyx.tools.models import CustomToolCallSummary -from onyx.tools.tool_constructor import construct_tools -from onyx.tools.tool_constructor import SearchToolConfig +from onyx.tools.tool_constructor import construct_tools, SearchToolConfig from onyx.tools.tool_implementations.mcp.mcp_tool import MCPTool from tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/tools/test_memory_tool_integration.py b/backend/tests/external_dependency_unit/tools/test_memory_tool_integration.py index b5c8e65c24..fcb20f7443 100644 --- a/backend/tests/external_dependency_unit/tools/test_memory_tool_integration.py +++ b/backend/tests/external_dependency_unit/tools/test_memory_tool_integration.py @@ -4,12 +4,13 @@ import pytest from sqlalchemy import select from sqlalchemy.orm import Session -from onyx.db.memory import add_memory -from onyx.db.memory import get_memories -from onyx.db.memory import MAX_MEMORIES_PER_USER -from onyx.db.memory import update_memory_at_index -from onyx.db.models import Memory -from onyx.db.models import User +from onyx.db.memory import ( + add_memory, + get_memories, + MAX_MEMORIES_PER_USER, + update_memory_at_index, +) +from onyx.db.models import Memory, User from onyx.tools.tool_implementations.memory.models import MemoryToolResponse from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/tools/test_oauth_config_crud.py b/backend/tests/external_dependency_unit/tools/test_oauth_config_crud.py index 4033577e04..23f0c61450 100644 --- a/backend/tests/external_dependency_unit/tools/test_oauth_config_crud.py +++ b/backend/tests/external_dependency_unit/tools/test_oauth_config_crud.py @@ -10,19 +10,19 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.db.models import OAuthConfig -from onyx.db.models import Tool -from onyx.db.oauth_config import create_oauth_config -from onyx.db.oauth_config import delete_oauth_config -from onyx.db.oauth_config import delete_user_oauth_token -from onyx.db.oauth_config import get_oauth_config -from onyx.db.oauth_config import get_oauth_configs -from onyx.db.oauth_config import get_tools_by_oauth_config -from onyx.db.oauth_config import get_user_oauth_token -from onyx.db.oauth_config import update_oauth_config -from onyx.db.oauth_config import upsert_user_oauth_token -from onyx.db.tools import delete_tool__no_commit -from onyx.db.tools import update_tool +from onyx.db.models import OAuthConfig, Tool +from onyx.db.oauth_config import ( + create_oauth_config, + delete_oauth_config, + delete_user_oauth_token, + get_oauth_config, + get_oauth_configs, + get_tools_by_oauth_config, + get_user_oauth_token, + update_oauth_config, + upsert_user_oauth_token, +) +from onyx.db.tools import delete_tool__no_commit, update_tool from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/tools/test_oauth_token_manager.py b/backend/tests/external_dependency_unit/tools/test_oauth_token_manager.py index 7c499fdcbd..02b646cf8e 100644 --- a/backend/tests/external_dependency_unit/tools/test_oauth_token_manager.py +++ b/backend/tests/external_dependency_unit/tools/test_oauth_token_manager.py @@ -7,24 +7,22 @@ All HTTP requests to external OAuth providers are mocked. """ import time -from unittest.mock import Mock -from unittest.mock import patch -from urllib.parse import parse_qs -from urllib.parse import urlparse +from unittest.mock import Mock, patch +from urllib.parse import parse_qs, urlparse from uuid import uuid4 import pytest -from requests import HTTPError -from requests import Response +from requests import HTTPError, Response from sqlalchemy.orm import Session -from onyx.auth.oauth_token_manager import build_oauth_authorization_url -from onyx.auth.oauth_token_manager import exchange_oauth_code_for_token -from onyx.auth.oauth_token_manager import OAuthFlowParams -from onyx.auth.oauth_token_manager import OAuthTokenManager +from onyx.auth.oauth_token_manager import ( + build_oauth_authorization_url, + exchange_oauth_code_for_token, + OAuthFlowParams, + OAuthTokenManager, +) from onyx.db.models import OAuthConfig -from onyx.db.oauth_config import create_oauth_config -from onyx.db.oauth_config import upsert_user_oauth_token +from onyx.db.oauth_config import create_oauth_config, upsert_user_oauth_token from onyx.utils.sensitive import SensitiveValue from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/tools/test_oauth_tool_integration.py b/backend/tests/external_dependency_unit/tools/test_oauth_tool_integration.py index 32792dbb52..bba7f4779d 100644 --- a/backend/tests/external_dependency_unit/tools/test_oauth_tool_integration.py +++ b/backend/tests/external_dependency_unit/tools/test_oauth_tool_integration.py @@ -10,24 +10,17 @@ All external HTTP calls are mocked, but Postgres and Redis are running. import queue from typing import Any -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import Mock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session from onyx.chat.emitter import Emitter -from onyx.db.models import OAuthAccount -from onyx.db.models import OAuthConfig -from onyx.db.models import Persona -from onyx.db.models import Tool -from onyx.db.models import User -from onyx.db.oauth_config import create_oauth_config -from onyx.db.oauth_config import upsert_user_oauth_token +from onyx.db.models import OAuthAccount, OAuthConfig, Persona, Tool, User +from onyx.db.oauth_config import create_oauth_config, upsert_user_oauth_token from onyx.llm.factory import get_default_llm -from onyx.tools.tool_constructor import construct_tools -from onyx.tools.tool_constructor import SearchToolConfig +from onyx.tools.tool_constructor import construct_tools, SearchToolConfig from onyx.tools.tool_implementations.custom.custom_tool import CustomTool from tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider from tests.external_dependency_unit.conftest import create_test_user diff --git a/backend/tests/external_dependency_unit/tools/test_python_tool.py b/backend/tests/external_dependency_unit/tools/test_python_tool.py index f46d44b215..270247c97a 100644 --- a/backend/tests/external_dependency_unit/tools/test_python_tool.py +++ b/backend/tests/external_dependency_unit/tools/test_python_tool.py @@ -926,8 +926,7 @@ import io import json import threading from collections.abc import Generator -from http.server import BaseHTTPRequestHandler -from http.server import HTTPServer +from http.server import BaseHTTPRequestHandler, HTTPServer from typing import Any from unittest.mock import patch @@ -941,24 +940,29 @@ import onyx.tools.tool_implementations.python.code_interpreter_client as ci_mod from onyx.chat.process_message import handle_stream_message_objects from onyx.db.models import Persona from onyx.db.tools import get_builtin_tool -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import FileDescriptor +from onyx.file_store.models import ChatFileType, FileDescriptor from onyx.server.features.projects.api import upload_user_files from onyx.server.query_and_chat.chat_backend import get_chat_session from onyx.server.query_and_chat.models import SendMessageRequest -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import PythonToolDelta -from onyx.server.query_and_chat.streaming_models import PythonToolStart -from onyx.server.query_and_chat.streaming_models import SectionEnd -from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta +from onyx.server.query_and_chat.streaming_models import ( + Packet, + PythonToolDelta, + PythonToolStart, + SectionEnd, + ToolCallArgumentDelta, +) from onyx.tools.tool_implementations.python.python_tool import PythonTool from tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder -from tests.external_dependency_unit.answer.stream_test_utils import create_chat_session -from tests.external_dependency_unit.answer.stream_test_utils import create_placement +from tests.external_dependency_unit.answer.stream_test_utils import ( + create_chat_session, + create_placement, +) from tests.external_dependency_unit.conftest import create_test_user -from tests.external_dependency_unit.mock_llm import LLMAnswerResponse -from tests.external_dependency_unit.mock_llm import LLMToolCallResponse -from tests.external_dependency_unit.mock_llm import use_mock_llm +from tests.external_dependency_unit.mock_llm import ( + LLMAnswerResponse, + LLMToolCallResponse, + use_mock_llm, +) # --------------------------------------------------------------------------- # Mock Code Interpreter Server diff --git a/backend/tests/external_dependency_unit/tools/test_python_tool_server_enabled.py b/backend/tests/external_dependency_unit/tools/test_python_tool_server_enabled.py index 8296067806..48aed209d2 100644 --- a/backend/tests/external_dependency_unit/tools/test_python_tool_server_enabled.py +++ b/backend/tests/external_dependency_unit/tools/test_python_tool_server_enabled.py @@ -8,8 +8,10 @@ from unittest.mock import patch from sqlalchemy.orm import Session -from onyx.db.code_interpreter import fetch_code_interpreter_server -from onyx.db.code_interpreter import update_code_interpreter_server_enabled +from onyx.db.code_interpreter import ( + fetch_code_interpreter_server, + update_code_interpreter_server_enabled, +) from onyx.tools.tool_implementations.python.python_tool import PythonTool diff --git a/backend/tests/external_dependency_unit/tracing/test_llm_span_recording.py b/backend/tests/external_dependency_unit/tracing/test_llm_span_recording.py index 929af2eb1c..9d3a39601c 100644 --- a/backend/tests/external_dependency_unit/tracing/test_llm_span_recording.py +++ b/backend/tests/external_dependency_unit/tracing/test_llm_span_recording.py @@ -5,17 +5,17 @@ from unittest.mock import MagicMock import pytest -from onyx.llm.model_response import ChatCompletionMessageToolCall -from onyx.llm.model_response import Choice +from onyx.llm.model_response import ( + ChatCompletionMessageToolCall, + Choice, + Message, + ModelResponse, + Usage, +) from onyx.llm.model_response import FunctionCall as ModelResponseFunctionCall -from onyx.llm.model_response import Message -from onyx.llm.model_response import ModelResponse -from onyx.llm.model_response import Usage -from onyx.llm.models import FunctionCall -from onyx.llm.models import ToolCall +from onyx.llm.models import FunctionCall, ToolCall from onyx.tracing.framework.span_data import GenerationSpanData -from onyx.tracing.llm_utils import record_llm_response -from onyx.tracing.llm_utils import record_llm_span_output +from onyx.tracing.llm_utils import record_llm_response, record_llm_span_output @pytest.fixture diff --git a/backend/tests/external_dependency_unit/tracing/test_tracing_admin_api.py b/backend/tests/external_dependency_unit/tracing/test_tracing_admin_api.py index 5b6a260e42..737e38c6c8 100644 --- a/backend/tests/external_dependency_unit/tracing/test_tracing_admin_api.py +++ b/backend/tests/external_dependency_unit/tracing/test_tracing_admin_api.py @@ -14,12 +14,16 @@ from onyx.db.models import User from onyx.db.tracing import delete_tracing_provider from onyx.error_handling.exceptions import OnyxError from onyx.server.manage.tracing import api as tracing_api -from onyx.server.manage.tracing.api import adopt_env_tracing_provider -from onyx.server.manage.tracing.api import disconnect_tracing_provider -from onyx.server.manage.tracing.api import list_tracing_providers -from onyx.server.manage.tracing.api import upsert_tracing_provider_endpoint -from onyx.server.manage.tracing.models import TracingProviderTestRequest -from onyx.server.manage.tracing.models import TracingProviderUpsertRequest +from onyx.server.manage.tracing.api import ( + adopt_env_tracing_provider, + disconnect_tracing_provider, + list_tracing_providers, + upsert_tracing_provider_endpoint, +) +from onyx.server.manage.tracing.models import ( + TracingProviderTestRequest, + TracingProviderUpsertRequest, +) from onyx.tracing import provider_config from shared_configs.enums import TracingProviderType diff --git a/backend/tests/external_dependency_unit/tracing/test_tracing_provider_config.py b/backend/tests/external_dependency_unit/tracing/test_tracing_provider_config.py index 1fba88abde..a6d37bf6b2 100644 --- a/backend/tests/external_dependency_unit/tracing/test_tracing_provider_config.py +++ b/backend/tests/external_dependency_unit/tracing/test_tracing_provider_config.py @@ -9,9 +9,11 @@ from collections.abc import Generator import pytest from sqlalchemy.orm import Session -from onyx.db.tracing import delete_tracing_provider -from onyx.db.tracing import fetch_tracing_provider -from onyx.db.tracing import upsert_tracing_provider +from onyx.db.tracing import ( + delete_tracing_provider, + fetch_tracing_provider, + upsert_tracing_provider, +) from onyx.tracing import provider_config from onyx.tracing.provider_config import resolve_effective_tracing_config from shared_configs.enums import TracingProviderType diff --git a/backend/tests/external_dependency_unit/voice/test_openai_streaming.py b/backend/tests/external_dependency_unit/voice/test_openai_streaming.py index 5a9ab6e805..80c199f911 100644 --- a/backend/tests/external_dependency_unit/voice/test_openai_streaming.py +++ b/backend/tests/external_dependency_unit/voice/test_openai_streaming.py @@ -10,9 +10,11 @@ import struct import pytest -from onyx.voice.providers.openai import OPENAI_REALTIME_STT_MODEL -from onyx.voice.providers.openai import OpenAIStreamingTranscriber -from onyx.voice.providers.openai import OpenAIVoiceProvider +from onyx.voice.providers.openai import ( + OPENAI_REALTIME_STT_MODEL, + OpenAIStreamingTranscriber, + OpenAIVoiceProvider, +) from tests.utils.secret_names import TestSecret # 24kHz mono PCM16 — matches the browser's voice WS format. diff --git a/backend/tests/integration/common_utils/document_acl.py b/backend/tests/integration/common_utils/document_acl.py index 6aca1b57ea..a2a0f43057 100644 --- a/backend/tests/integration/common_utils/document_acl.py +++ b/backend/tests/integration/common_utils/document_acl.py @@ -10,15 +10,12 @@ from sqlalchemy.orm import Session from ee.onyx.access.access import _get_access_for_documents from ee.onyx.db.external_perm import fetch_external_groups_for_user -from onyx.access.utils import prefix_external_group -from onyx.access.utils import prefix_user_email +from onyx.access.utils import prefix_external_group, prefix_user_email from onyx.configs.constants import PUBLIC_DOC_PAT -from onyx.db.models import DocumentByConnectorCredentialPair -from onyx.db.models import User +from onyx.db.models import DocumentByConnectorCredentialPair, User from onyx.db.users import fetch_user_by_id from onyx.utils.logger import setup_logger -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestCCPair, DATestUser logger = setup_logger() diff --git a/backend/tests/integration/common_utils/managers/api_key.py b/backend/tests/integration/common_utils/managers/api_key.py index 343f4c3c52..75ca260d15 100644 --- a/backend/tests/integration/common_utils/managers/api_key.py +++ b/backend/tests/integration/common_utils/managers/api_key.py @@ -2,11 +2,9 @@ from uuid import uuid4 from onyx.db.models import UserRole from onyx.server.api_key.models import APIKeyArgs -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser class APIKeyManager: diff --git a/backend/tests/integration/common_utils/managers/build_approvals.py b/backend/tests/integration/common_utils/managers/build_approvals.py index bc2d5b2067..1790312416 100644 --- a/backend/tests/integration/common_utils/managers/build_approvals.py +++ b/backend/tests/integration/common_utils/managers/build_approvals.py @@ -10,8 +10,7 @@ from __future__ import annotations import time from uuid import UUID -from onyx.server.features.build.approvals.api import ApprovalListResponse -from onyx.server.features.build.approvals.api import ApprovalView +from onyx.server.features.build.approvals.api import ApprovalListResponse, ApprovalView from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/build_session.py b/backend/tests/integration/common_utils/managers/build_session.py index 27548eb614..61f7933b5f 100644 --- a/backend/tests/integration/common_utils/managers/build_session.py +++ b/backend/tests/integration/common_utils/managers/build_session.py @@ -8,22 +8,22 @@ Each method calls the API server through the same ``user.headers`` / from __future__ import annotations -from typing import Any -from typing import NamedTuple +from typing import Any, NamedTuple from uuid import UUID -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SharingScope +from onyx.db.enums import SandboxStatus, SharingScope from onyx.server.features.build.interactive_turns.models import InteractiveTurnResponse from onyx.server.features.build.models import UploadResponse from onyx.server.features.build.sandbox.models import DirectoryListing -from onyx.server.features.build.session.models import DetailedSessionResponse -from onyx.server.features.build.session.models import MessageListResponse -from onyx.server.features.build.session.models import MessageResponse -from onyx.server.features.build.session.models import OpencodeHistorySnapshotResponse -from onyx.server.features.build.session.models import SessionListResponse -from onyx.server.features.build.session.models import SessionResponse -from onyx.server.features.build.session.models import SnapshotResponse +from onyx.server.features.build.session.models import ( + DetailedSessionResponse, + MessageListResponse, + MessageResponse, + OpencodeHistorySnapshotResponse, + SessionListResponse, + SessionResponse, + SnapshotResponse, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/cc_pair.py b/backend/tests/integration/common_utils/managers/cc_pair.py index c32780c0fc..75a1dc52dc 100644 --- a/backend/tests/integration/common_utils/managers/cc_pair.py +++ b/backend/tests/integration/common_utils/managers/cc_pair.py @@ -4,21 +4,20 @@ from typing import Any from uuid import uuid4 from onyx.connectors.models import InputType -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.server.documents.models import CCPairFullInfo -from onyx.server.documents.models import ConnectorCredentialPairIdentifier -from onyx.server.documents.models import ConnectorIndexingStatusLite -from onyx.server.documents.models import ConnectorStatus -from onyx.server.documents.models import DocumentSource -from onyx.server.documents.models import DocumentSyncStatus -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import MAX_DELAY +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus +from onyx.server.documents.models import ( + CCPairFullInfo, + ConnectorCredentialPairIdentifier, + ConnectorIndexingStatusLite, + ConnectorStatus, + DocumentSource, + DocumentSyncStatus, +) +from tests.integration.common_utils.constants import API_SERVER_URL, MAX_DELAY from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestCCPair, DATestUser def _cc_pair_creator( diff --git a/backend/tests/integration/common_utils/managers/chat.py b/backend/tests/integration/common_utils/managers/chat.py index 168ef78292..fb4828a50f 100644 --- a/backend/tests/integration/common_utils/managers/chat.py +++ b/backend/tests/integration/common_utils/managers/chat.py @@ -1,30 +1,30 @@ import json -from typing import Any -from typing import cast -from typing import Literal -from typing import TypedDict +from typing import Any, cast, Literal, TypedDict from uuid import UUID import httpx -from onyx.context.search.models import SavedSearchDoc -from onyx.context.search.models import SearchDoc +from onyx.context.search.models import SavedSearchDoc, SearchDoc from onyx.file_store.models import FileDescriptor from onyx.llm.override_models import LLMOverride -from onyx.server.query_and_chat.models import AUTO_PLACE_AFTER_LATEST_MESSAGE -from onyx.server.query_and_chat.models import ChatSessionCreationRequest -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.server.query_and_chat.models import ( + AUTO_PLACE_AFTER_LATEST_MESSAGE, + ChatSessionCreationRequest, + SendMessageRequest, +) from onyx.server.query_and_chat.streaming_models import StreamingType from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestChatMessage -from tests.integration.common_utils.test_models import DATestChatSession -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import ErrorResponse -from tests.integration.common_utils.test_models import StreamedResponse -from tests.integration.common_utils.test_models import ToolCallDebug -from tests.integration.common_utils.test_models import ToolName -from tests.integration.common_utils.test_models import ToolResult +from tests.integration.common_utils.test_models import ( + DATestChatMessage, + DATestChatSession, + DATestUser, + ErrorResponse, + StreamedResponse, + ToolCallDebug, + ToolName, + ToolResult, +) class StreamPacketObj(TypedDict, total=False): diff --git a/backend/tests/integration/common_utils/managers/connector.py b/backend/tests/integration/common_utils/managers/connector.py index 3a34c7fce5..42d3bfa5ad 100644 --- a/backend/tests/integration/common_utils/managers/connector.py +++ b/backend/tests/integration/common_utils/managers/connector.py @@ -3,12 +3,10 @@ from uuid import uuid4 from onyx.connectors.models import InputType from onyx.db.enums import AccessType -from onyx.server.documents.models import ConnectorUpdateRequest -from onyx.server.documents.models import DocumentSource +from onyx.server.documents.models import ConnectorUpdateRequest, DocumentSource from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestConnector -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestConnector, DATestUser class ConnectorManager: diff --git a/backend/tests/integration/common_utils/managers/credential.py b/backend/tests/integration/common_utils/managers/credential.py index 25c7fe84b7..9013ca9da0 100644 --- a/backend/tests/integration/common_utils/managers/credential.py +++ b/backend/tests/integration/common_utils/managers/credential.py @@ -1,12 +1,10 @@ from typing import Any from uuid import uuid4 -from onyx.server.documents.models import CredentialSnapshot -from onyx.server.documents.models import DocumentSource +from onyx.server.documents.models import CredentialSnapshot, DocumentSource from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestCredential, DATestUser class CredentialManager: diff --git a/backend/tests/integration/common_utils/managers/discord_bot.py b/backend/tests/integration/common_utils/managers/discord_bot.py index f448467097..7bdaa9e11e 100644 --- a/backend/tests/integration/common_utils/managers/discord_bot.py +++ b/backend/tests/integration/common_utils/managers/discord_bot.py @@ -1,17 +1,21 @@ """Manager for Discord bot API integration tests.""" -from onyx.db.discord_bot import create_channel_config -from onyx.db.discord_bot import create_guild_config -from onyx.db.discord_bot import register_guild +from onyx.db.discord_bot import ( + create_channel_config, + create_guild_config, + register_guild, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.utils import DiscordChannelView from onyx.server.manage.discord_bot.utils import generate_discord_registration_key from shared_configs.contextvars import get_current_tenant_id from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestDiscordChannelConfig -from tests.integration.common_utils.test_models import DATestDiscordGuildConfig -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestDiscordChannelConfig, + DATestDiscordGuildConfig, + DATestUser, +) DISCORD_BOT_API_URL = f"{API_SERVER_URL}/manage/admin/discord-bot" diff --git a/backend/tests/integration/common_utils/managers/document.py b/backend/tests/integration/common_utils/managers/document.py index 66ec631345..4856335434 100644 --- a/backend/tests/integration/common_utils/managers/document.py +++ b/backend/tests/integration/common_utils/managers/document.py @@ -1,20 +1,19 @@ from uuid import uuid4 -from sqlalchemy import and_ -from sqlalchemy import select +from sqlalchemy import and_, select from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource from onyx.db.enums import AccessType -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import DocumentByConnectorCredentialPair -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import NUM_DOCS +from onyx.db.models import ConnectorCredentialPair, DocumentByConnectorCredentialPair +from tests.integration.common_utils.constants import API_SERVER_URL, NUM_DOCS from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.api_key import DATestAPIKey -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import SimpleTestDocument +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestUser, + SimpleTestDocument, +) from tests.integration.common_utils.vespa import vespa_fixture diff --git a/backend/tests/integration/common_utils/managers/document_search.py b/backend/tests/integration/common_utils/managers/document_search.py index 5f0d28c903..efc9ac289d 100644 --- a/backend/tests/integration/common_utils/managers/document_search.py +++ b/backend/tests/integration/common_utils/managers/document_search.py @@ -1,5 +1,7 @@ -from ee.onyx.server.query_and_chat.models import SearchFullResponse -from ee.onyx.server.query_and_chat.models import SendSearchQueryRequest +from ee.onyx.server.query_and_chat.models import ( + SearchFullResponse, + SendSearchQueryRequest, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/document_set.py b/backend/tests/integration/common_utils/managers/document_set.py index ab65af706d..e45e739c35 100644 --- a/backend/tests/integration/common_utils/managers/document_set.py +++ b/backend/tests/integration/common_utils/managers/document_set.py @@ -1,13 +1,10 @@ import time from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import MAX_DELAY +from tests.integration.common_utils.constants import API_SERVER_URL, MAX_DELAY from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestDocumentSet -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestDocumentSet, DATestUser class DocumentSetManager: diff --git a/backend/tests/integration/common_utils/managers/external_app.py b/backend/tests/integration/common_utils/managers/external_app.py index e0e8f46a82..9d870ab4e7 100644 --- a/backend/tests/integration/common_utils/managers/external_app.py +++ b/backend/tests/integration/common_utils/managers/external_app.py @@ -4,15 +4,14 @@ import zipfile from typing import Any from uuid import uuid4 -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.server.features.build.external_apps.models import ( CreateBuiltInExternalAppRequest, + ExternalAppAdminResponse, + ExternalAppUserResponse, + UpdateExternalAppRequest, + UpsertUserCredentialsRequest, ) -from onyx.server.features.build.external_apps.models import ExternalAppAdminResponse -from onyx.server.features.build.external_apps.models import ExternalAppUserResponse -from onyx.server.features.build.external_apps.models import UpdateExternalAppRequest -from onyx.server.features.build.external_apps.models import UpsertUserCredentialsRequest from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/file.py b/backend/tests/integration/common_utils/managers/file.py index 7df32a8abe..c8c1ac3dfa 100644 --- a/backend/tests/integration/common_utils/managers/file.py +++ b/backend/tests/integration/common_utils/managers/file.py @@ -1,9 +1,6 @@ import io import mimetypes -from typing import cast -from typing import IO -from typing import List -from typing import Tuple +from typing import cast, IO, List, Tuple from onyx.file_store.models import FileDescriptor from onyx.server.documents.models import FileUploadResponse diff --git a/backend/tests/integration/common_utils/managers/image_generation.py b/backend/tests/integration/common_utils/managers/image_generation.py index e473a1edc3..a86a38a757 100644 --- a/backend/tests/integration/common_utils/managers/image_generation.py +++ b/backend/tests/integration/common_utils/managers/image_generation.py @@ -5,8 +5,10 @@ from uuid import uuid4 from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestImageGenerationConfig -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestImageGenerationConfig, + DATestUser, +) def _serialize_custom_config( diff --git a/backend/tests/integration/common_utils/managers/index_attempt.py b/backend/tests/integration/common_utils/managers/index_attempt.py index fc8761c67a..730d250dcd 100644 --- a/backend/tests/integration/common_utils/managers/index_attempt.py +++ b/backend/tests/integration/common_utils/managers/index_attempt.py @@ -1,21 +1,16 @@ import time -from datetime import datetime -from datetime import timedelta +from datetime import datetime, timedelta from urllib.parse import urlencode from onyx.background.indexing.models import IndexAttemptErrorPydantic from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import IndexModelStatus -from onyx.db.models import IndexAttempt -from onyx.db.models import IndexingStatus +from onyx.db.models import IndexAttempt, IndexingStatus from onyx.db.search_settings import get_current_search_settings -from onyx.server.documents.models import IndexAttemptSnapshot -from onyx.server.documents.models import PaginatedReturn -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import MAX_DELAY +from onyx.server.documents.models import IndexAttemptSnapshot, PaginatedReturn +from tests.integration.common_utils.constants import API_SERVER_URL, MAX_DELAY from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestIndexAttempt -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestIndexAttempt, DATestUser class IndexAttemptManager: diff --git a/backend/tests/integration/common_utils/managers/llm_provider.py b/backend/tests/integration/common_utils/managers/llm_provider.py index 17bc898216..1e2e2f1e30 100644 --- a/backend/tests/integration/common_utils/managers/llm_provider.py +++ b/backend/tests/integration/common_utils/managers/llm_provider.py @@ -2,15 +2,15 @@ import os from uuid import uuid4 from onyx.llm.constants import LlmProviderNames -from onyx.server.manage.llm.models import DefaultModel -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from onyx.server.manage.llm.models import ( + DefaultModel, + LLMProviderUpsertRequest, + LLMProviderView, + ModelConfigurationUpsertRequest, +) +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser class LLMProviderManager: diff --git a/backend/tests/integration/common_utils/managers/pat.py b/backend/tests/integration/common_utils/managers/pat.py index 7173a3d4d3..e98ee1ae60 100644 --- a/backend/tests/integration/common_utils/managers/pat.py +++ b/backend/tests/integration/common_utils/managers/pat.py @@ -9,8 +9,7 @@ from onyx.db.enums import Permission from onyx.db.pat import create_pat from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestPAT -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestPAT, DATestUser class PATManager: diff --git a/backend/tests/integration/common_utils/managers/persona.py b/backend/tests/integration/common_utils/managers/persona.py index 17d12f3290..fb16167065 100644 --- a/backend/tests/integration/common_utils/managers/persona.py +++ b/backend/tests/integration/common_utils/managers/persona.py @@ -1,13 +1,16 @@ -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 -from onyx.server.features.persona.models import FullPersonaSnapshot -from onyx.server.features.persona.models import PersonaUpsertRequest +from onyx.server.features.persona.models import ( + FullPersonaSnapshot, + PersonaUpsertRequest, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestPersona -from tests.integration.common_utils.test_models import DATestPersonaLabel -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestPersona, + DATestPersonaLabel, + DATestUser, +) class PersonaManager: diff --git a/backend/tests/integration/common_utils/managers/project.py b/backend/tests/integration/common_utils/managers/project.py index fe574d9e2e..69f7508b3d 100644 --- a/backend/tests/integration/common_utils/managers/project.py +++ b/backend/tests/integration/common_utils/managers/project.py @@ -1,8 +1,10 @@ from typing import List -from onyx.server.features.projects.models import CategorizedFilesSnapshot -from onyx.server.features.projects.models import UserFileSnapshot -from onyx.server.features.projects.models import UserProjectSnapshot +from onyx.server.features.projects.models import ( + CategorizedFilesSnapshot, + UserFileSnapshot, + UserProjectSnapshot, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/query_history.py b/backend/tests/integration/common_utils/managers/query_history.py index 16fe67bfb3..44f44f8a08 100644 --- a/backend/tests/integration/common_utils/managers/query_history.py +++ b/backend/tests/integration/common_utils/managers/query_history.py @@ -5,13 +5,11 @@ from uuid import UUID import httpx -from ee.onyx.server.query_history.models import ChatSessionMinimal -from ee.onyx.server.query_history.models import ChatSessionSnapshot +from ee.onyx.server.query_history.models import ChatSessionMinimal, ChatSessionSnapshot from onyx.configs.constants import QAFeedbackType from onyx.db.enums import TaskStatus from onyx.server.documents.models import PaginatedReturn -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import MAX_DELAY +from tests.integration.common_utils.constants import API_SERVER_URL, MAX_DELAY from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/scim_client.py b/backend/tests/integration/common_utils/managers/scim_client.py index 06c0c938be..e15ec686f8 100644 --- a/backend/tests/integration/common_utils/managers/scim_client.py +++ b/backend/tests/integration/common_utils/managers/scim_client.py @@ -1,7 +1,6 @@ import httpx -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client diff --git a/backend/tests/integration/common_utils/managers/scim_token.py b/backend/tests/integration/common_utils/managers/scim_token.py index 32b73886d7..64498f8e21 100644 --- a/backend/tests/integration/common_utils/managers/scim_token.py +++ b/backend/tests/integration/common_utils/managers/scim_token.py @@ -1,7 +1,6 @@ from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestScimToken -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestScimToken, DATestUser class ScimTokenManager: diff --git a/backend/tests/integration/common_utils/managers/settings.py b/backend/tests/integration/common_utils/managers/settings.py index 6278be946e..20a932235c 100644 --- a/backend/tests/integration/common_utils/managers/settings.py +++ b/backend/tests/integration/common_utils/managers/settings.py @@ -1,11 +1,8 @@ -from typing import Any -from typing import Dict -from typing import Optional +from typing import Any, Dict, Optional from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestSettings -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestSettings, DATestUser class SettingsManager: diff --git a/backend/tests/integration/common_utils/managers/skill.py b/backend/tests/integration/common_utils/managers/skill.py index f6886b3df4..63d8d6776b 100644 --- a/backend/tests/integration/common_utils/managers/skill.py +++ b/backend/tests/integration/common_utils/managers/skill.py @@ -2,23 +2,24 @@ import io import zipfile from collections.abc import Sequence from typing import TypeVar -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx from pydantic import BaseModel from onyx.db.enums import SkillSharePermission -from onyx.server.features.skill.models import SkillCreateRequest -from onyx.server.features.skill.models import SkillEditableDetailResponse -from onyx.server.features.skill.models import SkillGroupShareRequest -from onyx.server.features.skill.models import SkillPatchRequest -from onyx.server.features.skill.models import SkillPreviewResponse -from onyx.server.features.skill.models import SkillResponse -from onyx.server.features.skill.models import SkillShareRequest -from onyx.server.features.skill.models import SkillsList -from onyx.server.features.skill.models import SkillUserShareRequest -from onyx.server.features.skill.models import TransferSkillOwnershipRequest +from onyx.server.features.skill.models import ( + SkillCreateRequest, + SkillEditableDetailResponse, + SkillGroupShareRequest, + SkillPatchRequest, + SkillPreviewResponse, + SkillResponse, + SkillShareRequest, + SkillsList, + SkillUserShareRequest, + TransferSkillOwnershipRequest, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/tenant.py b/backend/tests/integration/common_utils/managers/tenant.py index cd7c95a90a..58fd456885 100644 --- a/backend/tests/integration/common_utils/managers/tenant.py +++ b/backend/tests/integration/common_utils/managers/tenant.py @@ -1,12 +1,9 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import jwt from onyx.server.manage.models import AllUsersResponse -from onyx.server.models import FullUserSnapshot -from onyx.server.models import InvitedUserSnapshot +from onyx.server.models import FullUserSnapshot, InvitedUserSnapshot from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/tool.py b/backend/tests/integration/common_utils/managers/tool.py index 96affaf327..e295c60fc2 100644 --- a/backend/tests/integration/common_utils/managers/tool.py +++ b/backend/tests/integration/common_utils/managers/tool.py @@ -2,8 +2,7 @@ import httpx from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestTool -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestTool, DATestUser class ToolManager: diff --git a/backend/tests/integration/common_utils/managers/user.py b/backend/tests/integration/common_utils/managers/user.py index 300e16cc5c..fa54079993 100644 --- a/backend/tests/integration/common_utils/managers/user.py +++ b/backend/tests/integration/common_utils/managers/user.py @@ -6,15 +6,15 @@ import httpx import pytest from onyx.auth.schemas import UserRole -from onyx.configs.constants import ANONYMOUS_USER_EMAIL -from onyx.configs.constants import ANONYMOUS_USER_UUID -from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME +from onyx.configs.constants import ( + ANONYMOUS_USER_EMAIL, + ANONYMOUS_USER_UUID, + FASTAPI_USERS_AUTH_COOKIE_NAME, +) from onyx.server.documents.models import PaginatedReturn from onyx.server.manage.models import UserInfo -from onyx.server.models import FullUserSnapshot -from onyx.server.models import InvitedUserSnapshot -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from onyx.server.models import FullUserSnapshot, InvitedUserSnapshot +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/common_utils/managers/user_group.py b/backend/tests/integration/common_utils/managers/user_group.py index 361958d625..a8c9462f90 100644 --- a/backend/tests/integration/common_utils/managers/user_group.py +++ b/backend/tests/integration/common_utils/managers/user_group.py @@ -4,11 +4,9 @@ from uuid import uuid4 import httpx from ee.onyx.server.user_group.models import UserGroup -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import MAX_DELAY +from tests.integration.common_utils.constants import API_SERVER_URL, MAX_DELAY from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import DATestUserGroup +from tests.integration.common_utils.test_models import DATestUser, DATestUserGroup class UserGroupManager: diff --git a/backend/tests/integration/common_utils/reset.py b/backend/tests/integration/common_utils/reset.py index b76711d6f5..c3066e853c 100644 --- a/backend/tests/integration/common_utils/reset.py +++ b/backend/tests/integration/common_utils/reset.py @@ -7,13 +7,17 @@ from alembic import command from alembic.config import Config from sqlalchemy.orm import Session -from onyx.configs.app_configs import POSTGRES_HOST -from onyx.configs.app_configs import POSTGRES_PASSWORD -from onyx.configs.app_configs import POSTGRES_PORT -from onyx.configs.app_configs import POSTGRES_USER -from onyx.db.engine.sql_engine import build_connection_string -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SYNC_DB_API +from onyx.configs.app_configs import ( + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, +) +from onyx.db.engine.sql_engine import ( + build_connection_string, + get_session_with_current_tenant, + SYNC_DB_API, +) from onyx.db.swap_index import check_and_perform_index_swap from onyx.file_store.file_store import get_default_file_store from onyx.setup import setup_postgres diff --git a/backend/tests/integration/common_utils/test_document_utils.py b/backend/tests/integration/common_utils/test_document_utils.py index be7ab6128c..063213da9e 100644 --- a/backend/tests/integration/common_utils/test_document_utils.py +++ b/backend/tests/integration/common_utils/test_document_utils.py @@ -1,12 +1,13 @@ import uuid -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from onyx.configs.constants import DocumentSource -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + TextSection, +) def create_test_document( diff --git a/backend/tests/integration/common_utils/test_models.py b/backend/tests/integration/common_utils/test_models.py index 7746787991..87a3f587b5 100644 --- a/backend/tests/integration/common_utils/test_models.py +++ b/backend/tests/integration/common_utils/test_models.py @@ -4,19 +4,18 @@ from enum import Enum from typing import Any from uuid import UUID -from pydantic import BaseModel -from pydantic import Field +from pydantic import BaseModel, Field from onyx.auth.schemas import UserRole -from onyx.configs.constants import MessageType -from onyx.configs.constants import QAFeedbackType -from onyx.context.search.models import SavedSearchDoc -from onyx.context.search.models import SearchDoc +from onyx.configs.constants import MessageType, QAFeedbackType +from onyx.context.search.models import SavedSearchDoc, SearchDoc from onyx.db.enums import AccessType -from onyx.server.documents.models import DocumentSource -from onyx.server.documents.models import IndexAttemptSnapshot -from onyx.server.documents.models import IndexingStatus -from onyx.server.documents.models import InputType +from onyx.server.documents.models import ( + DocumentSource, + IndexAttemptSnapshot, + IndexingStatus, + InputType, +) from onyx.server.query_and_chat.streaming_models import GeneratedImage """ diff --git a/backend/tests/integration/common_utils/timeout.py b/backend/tests/integration/common_utils/timeout.py index a77d5e0696..b0b7e9935a 100644 --- a/backend/tests/integration/common_utils/timeout.py +++ b/backend/tests/integration/common_utils/timeout.py @@ -24,8 +24,7 @@ import multiprocessing import traceback from collections.abc import Callable from multiprocessing import Queue -from typing import Any -from typing import TypeVar +from typing import Any, TypeVar T = TypeVar("T") diff --git a/backend/tests/integration/common_utils/vespa.py b/backend/tests/integration/common_utils/vespa.py index 4c3f7758b2..c2b4e2d895 100644 --- a/backend/tests/integration/common_utils/vespa.py +++ b/backend/tests/integration/common_utils/vespa.py @@ -3,11 +3,13 @@ from typing import Any from opensearchpy import OpenSearch from opensearchpy.exceptions import NotFoundError -from onyx.configs.app_configs import OPENSEARCH_ADMIN_PASSWORD -from onyx.configs.app_configs import OPENSEARCH_ADMIN_USERNAME -from onyx.configs.app_configs import OPENSEARCH_HOST -from onyx.configs.app_configs import OPENSEARCH_REST_API_PORT -from onyx.configs.app_configs import OPENSEARCH_USE_SSL +from onyx.configs.app_configs import ( + OPENSEARCH_ADMIN_PASSWORD, + OPENSEARCH_ADMIN_USERNAME, + OPENSEARCH_HOST, + OPENSEARCH_REST_API_PORT, + OPENSEARCH_USE_SSL, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.search_settings import get_current_search_settings diff --git a/backend/tests/integration/conftest.py b/backend/tests/integration/conftest.py index 13e79340c1..769ec58f75 100644 --- a/backend/tests/integration/conftest.py +++ b/backend/tests/integration/conftest.py @@ -2,8 +2,7 @@ import os import platform import shutil import subprocess -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from typing import Any from unittest.mock import patch @@ -54,16 +53,20 @@ import onyx.main # noqa: E402, F401 from onyx.auth.schemas import UserRole # noqa: E402 from onyx.background.celery.apps.client import celery_app # noqa: E402 from onyx.configs.constants import DocumentSource # noqa: E402 -from onyx.db.engine.sql_engine import get_session_with_current_tenant # noqa: E402 -from onyx.db.engine.sql_engine import SqlEngine # noqa: E402 +from onyx.db.engine.sql_engine import ( # noqa: E402 + get_session_with_current_tenant, + SqlEngine, +) from onyx.db.search_settings import get_current_search_settings # noqa: E402 from onyx.utils.variable_functionality import ( # noqa: E402 fetch_versioned_implementation, ) from shared_configs.configs import MULTI_TENANT # noqa: E402 from tests.integration.common_utils import http_client # noqa: E402 -from tests.integration.common_utils.constants import ADMIN_USER_NAME # noqa: E402 -from tests.integration.common_utils.constants import GENERAL_HEADERS # noqa: E402 +from tests.integration.common_utils.constants import ( # noqa: E402 + ADMIN_USER_NAME, + GENERAL_HEADERS, +) from tests.integration.common_utils.managers.api_key import APIKeyManager # noqa: E402 from tests.integration.common_utils.managers.document import ( # noqa: E402 DocumentManager, @@ -74,19 +77,23 @@ from tests.integration.common_utils.managers.image_generation import ( # noqa: from tests.integration.common_utils.managers.llm_provider import ( # noqa: E402 LLMProviderManager, ) -from tests.integration.common_utils.managers.user import build_email # noqa: E402 -from tests.integration.common_utils.managers.user import DEFAULT_PASSWORD # noqa: E402 -from tests.integration.common_utils.managers.user import UserManager # noqa: E402 -from tests.integration.common_utils.reset import _seed_dev_license_if_set # noqa: E402 -from tests.integration.common_utils.reset import reset_all # noqa: E402 -from tests.integration.common_utils.reset import reset_all_multitenant # noqa: E402 -from tests.integration.common_utils.test_models import DATestAPIKey # noqa: E402 +from tests.integration.common_utils.managers.user import ( # noqa: E402 + build_email, + DEFAULT_PASSWORD, + UserManager, +) +from tests.integration.common_utils.reset import ( # noqa: E402 + _seed_dev_license_if_set, + reset_all, + reset_all_multitenant, +) from tests.integration.common_utils.test_models import ( # noqa: E402 + DATestAPIKey, DATestImageGenerationConfig, + DATestLLMProvider, + DATestUser, + SimpleTestDocument, ) -from tests.integration.common_utils.test_models import DATestLLMProvider # noqa: E402 -from tests.integration.common_utils.test_models import DATestUser # noqa: E402 -from tests.integration.common_utils.test_models import SimpleTestDocument # noqa: E402 from tests.integration.common_utils.vespa import vespa_fixture # noqa: E402 BASIC_USER_NAME = "basic_user" diff --git a/backend/tests/integration/connector_job_tests/github/conftest.py b/backend/tests/integration/connector_job_tests/github/conftest.py index 945f9733b7..f3c8651e1f 100644 --- a/backend/tests/integration/connector_job_tests/github/conftest.py +++ b/backend/tests/integration/connector_job_tests/github/conftest.py @@ -1,7 +1,6 @@ import os from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest @@ -14,10 +13,12 @@ from tests.integration.common_utils.managers.credential import CredentialManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.reset import reset_all -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestConnector -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestConnector, + DATestCredential, + DATestUser, +) GitHubTestEnvSetupTuple = tuple[ DATestUser, # admin_user diff --git a/backend/tests/integration/connector_job_tests/github/test_github_empty_repository.py b/backend/tests/integration/connector_job_tests/github/test_github_empty_repository.py index e5433aa251..173bdfdc7f 100644 --- a/backend/tests/integration/connector_job_tests/github/test_github_empty_repository.py +++ b/backend/tests/integration/connector_job_tests/github/test_github_empty_repository.py @@ -1,8 +1,7 @@ from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import IndexingStatus +from onyx.db.enums import AccessType, IndexingStatus from tests.integration.common_utils.document_acl import get_all_connector_documents from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager diff --git a/backend/tests/integration/connector_job_tests/github/test_github_permission_sync.py b/backend/tests/integration/connector_job_tests/github/test_github_permission_sync.py index 5b60b7d006..8a7ae42869 100644 --- a/backend/tests/integration/connector_job_tests/github/test_github_permission_sync.py +++ b/backend/tests/integration/connector_job_tests/github/test_github_permission_sync.py @@ -1,14 +1,15 @@ import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest from github import Github from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.utils.logger import setup_logger -from tests.integration.common_utils.document_acl import get_all_connector_documents -from tests.integration.common_utils.document_acl import get_user_document_access_via_acl +from tests.integration.common_utils.document_acl import ( + get_all_connector_documents, + get_user_document_access_via_acl, +) from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.connector_job_tests.github.conftest import ( GitHubTestEnvSetupTuple, diff --git a/backend/tests/integration/connector_job_tests/google/google_drive_api_utils.py b/backend/tests/integration/connector_job_tests/google/google_drive_api_utils.py index 0930659c68..2f47ade32c 100644 --- a/backend/tests/integration/connector_job_tests/google/google_drive_api_utils.py +++ b/backend/tests/integration/connector_job_tests/google/google_drive_api_utils.py @@ -3,10 +3,12 @@ from uuid import uuid4 from google.oauth2.service_account import Credentials -from onyx.connectors.google_utils.resources import get_drive_service -from onyx.connectors.google_utils.resources import get_google_docs_service -from onyx.connectors.google_utils.resources import GoogleDocsService -from onyx.connectors.google_utils.resources import GoogleDriveService +from onyx.connectors.google_utils.resources import ( + get_drive_service, + get_google_docs_service, + GoogleDocsService, + GoogleDriveService, +) GOOGLE_SCOPES = { "google_drive": [ diff --git a/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py b/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py index b651260d4c..b097445af6 100644 --- a/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py +++ b/backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py @@ -1,8 +1,7 @@ import json import os from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import uuid4 import pytest @@ -11,8 +10,6 @@ from onyx.configs.constants import DocumentSource from onyx.connectors.google_utils.resources import GoogleDriveService from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_PRIMARY_ADMIN_KEY, ) from onyx.connectors.models import InputType @@ -25,10 +22,12 @@ from tests.integration.common_utils.managers.document_search import ( ) from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestConnector -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestConnector, + DATestCredential, + DATestUser, +) from tests.integration.common_utils.vespa import vespa_fixture from tests.integration.connector_job_tests.google.google_drive_api_utils import ( GoogleDriveManager, diff --git a/backend/tests/integration/connector_job_tests/jira/conftest.py b/backend/tests/integration/connector_job_tests/jira/conftest.py index 017fa81e01..cbe6194f26 100644 --- a/backend/tests/integration/connector_job_tests/jira/conftest.py +++ b/backend/tests/integration/connector_job_tests/jira/conftest.py @@ -1,7 +1,6 @@ import os from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest @@ -12,10 +11,12 @@ from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestConnector -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestConnector, + DATestCredential, + DATestUser, +) JiraTestEnvSetupTuple = tuple[ DATestUser, diff --git a/backend/tests/integration/connector_job_tests/jira/test_jira_permission_sync_full.py b/backend/tests/integration/connector_job_tests/jira/test_jira_permission_sync_full.py index ae98e2a6e9..d36ba9dc53 100644 --- a/backend/tests/integration/connector_job_tests/jira/test_jira_permission_sync_full.py +++ b/backend/tests/integration/connector_job_tests/jira/test_jira_permission_sync_full.py @@ -1,6 +1,5 @@ import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest diff --git a/backend/tests/integration/connector_job_tests/sharepoint/conftest.py b/backend/tests/integration/connector_job_tests/sharepoint/conftest.py index 9d39750f82..b8320d5c4f 100644 --- a/backend/tests/integration/connector_job_tests/sharepoint/conftest.py +++ b/backend/tests/integration/connector_job_tests/sharepoint/conftest.py @@ -1,7 +1,6 @@ import os from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest @@ -15,10 +14,12 @@ from tests.integration.common_utils.managers.credential import CredentialManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.reset import reset_all -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestConnector -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestConnector, + DATestCredential, + DATestUser, +) SharepointTestEnvSetupTuple = tuple[ DATestUser, # admin_user diff --git a/backend/tests/integration/connector_job_tests/sharepoint/test_sharepoint_permissions.py b/backend/tests/integration/connector_job_tests/sharepoint/test_sharepoint_permissions.py index 309fd8389f..606dd9eefe 100644 --- a/backend/tests/integration/connector_job_tests/sharepoint/test_sharepoint_permissions.py +++ b/backend/tests/integration/connector_job_tests/sharepoint/test_sharepoint_permissions.py @@ -4,9 +4,11 @@ import pytest from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.utils.logger import setup_logger -from tests.integration.common_utils.document_acl import get_all_connector_documents -from tests.integration.common_utils.document_acl import get_documents_by_permission_type -from tests.integration.common_utils.document_acl import get_user_document_access_via_acl +from tests.integration.common_utils.document_acl import ( + get_all_connector_documents, + get_documents_by_permission_type, + get_user_document_access_via_acl, +) from tests.integration.connector_job_tests.sharepoint.conftest import ( SharepointTestEnvSetupTuple, ) diff --git a/backend/tests/integration/connector_job_tests/slack/slack_api_utils.py b/backend/tests/integration/connector_job_tests/slack/slack_api_utils.py index a3aa328072..cc6f8ab4d2 100644 --- a/backend/tests/integration/connector_job_tests/slack/slack_api_utils.py +++ b/backend/tests/integration/connector_job_tests/slack/slack_api_utils.py @@ -9,8 +9,7 @@ Assumptions: - chat:write.public """ -from typing import Any -from typing import cast +from typing import Any, cast from uuid import uuid4 from slack_sdk import WebClient diff --git a/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py b/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py index c06068bd39..2037706707 100644 --- a/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py +++ b/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py @@ -1,6 +1,5 @@ import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest @@ -17,14 +16,18 @@ from tests.integration.common_utils.managers.document_search import ( from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestConnector -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestConnector, + DATestCredential, + DATestUser, +) from tests.integration.common_utils.vespa import vespa_fixture -from tests.integration.connector_job_tests.slack.conftest import SLACK_ADMIN_EMAIL -from tests.integration.connector_job_tests.slack.conftest import SLACK_TEST_USER_1_EMAIL -from tests.integration.connector_job_tests.slack.conftest import SLACK_TEST_USER_2_EMAIL +from tests.integration.connector_job_tests.slack.conftest import ( + SLACK_ADMIN_EMAIL, + SLACK_TEST_USER_1_EMAIL, + SLACK_TEST_USER_2_EMAIL, +) from tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager diff --git a/backend/tests/integration/connector_job_tests/slack/test_prune.py b/backend/tests/integration/connector_job_tests/slack/test_prune.py index 77758f7d3f..04bac7739e 100644 --- a/backend/tests/integration/connector_job_tests/slack/test_prune.py +++ b/backend/tests/integration/connector_job_tests/slack/test_prune.py @@ -1,6 +1,5 @@ import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest @@ -16,10 +15,12 @@ from tests.integration.common_utils.managers.document_search import ( ) from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestConnector -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestConnector, + DATestCredential, + DATestUser, +) from tests.integration.common_utils.vespa import vespa_fixture from tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager diff --git a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_google_oauth.py b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_google_oauth.py index 64a898fa07..0e654fae31 100644 --- a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_google_oauth.py +++ b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_google_oauth.py @@ -25,8 +25,7 @@ import uvicorn from fastapi import FastAPI from fastapi.responses import PlainTextResponse from fastmcp import FastMCP -from fastmcp.server.auth import AccessToken -from fastmcp.server.auth import TokenVerifier +from fastmcp.server.auth import AccessToken, TokenVerifier from fastmcp.server.dependencies import get_access_token # Google's tokeninfo endpoint for validating access tokens diff --git a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_oauth.py b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_oauth.py index 8809b8e1f5..3f59ee7ce9 100644 --- a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_oauth.py +++ b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_oauth.py @@ -1,17 +1,11 @@ import os -from collections.abc import Awaitable -from collections.abc import Callable -from collections.abc import Iterable +from collections.abc import Awaitable, Callable, Iterable from typing import Any -from urllib.parse import urlsplit -from urllib.parse import urlunsplit +from urllib.parse import urlsplit, urlunsplit import uvicorn -from fastapi import FastAPI -from fastapi import Request -from fastapi.responses import JSONResponse -from fastapi.responses import PlainTextResponse -from fastapi.responses import Response +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, PlainTextResponse, Response from fastmcp import FastMCP from fastmcp.server.auth.providers.jwt import JWTVerifier from fastmcp.server.dependencies import get_access_token diff --git a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py index 1310a6e94c..315e2be775 100644 --- a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py +++ b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py @@ -1,25 +1,18 @@ import argparse import sys -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import Awaitable -from typing import Callable -from typing import Dict -from typing import Optional +from datetime import datetime, timezone +from typing import Any, Awaitable, Callable, Dict, Optional import bcrypt import uvicorn from fastapi import FastAPI from fastapi.responses import PlainTextResponse from fastmcp import FastMCP -from fastmcp.server.auth.auth import AccessToken -from fastmcp.server.auth.auth import TokenVerifier +from fastmcp.server.auth.auth import AccessToken, TokenVerifier from fastmcp.server.dependencies import get_access_token from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request -from starlette.responses import JSONResponse -from starlette.responses import Response +from starlette.responses import JSONResponse, Response # pip install fastmcp bcrypt fastapi uvicorn diff --git a/backend/tests/integration/mock_services/mcp_test_server/run_mock_oidc_idp.py b/backend/tests/integration/mock_services/mcp_test_server/run_mock_oidc_idp.py index 7556661b8d..4ee389a62b 100644 --- a/backend/tests/integration/mock_services/mcp_test_server/run_mock_oidc_idp.py +++ b/backend/tests/integration/mock_services/mcp_test_server/run_mock_oidc_idp.py @@ -53,12 +53,8 @@ import jwt import uvicorn from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from fastapi import FastAPI -from fastapi import Form -from fastapi import Request -from fastapi.responses import JSONResponse -from fastapi.responses import PlainTextResponse -from fastapi.responses import RedirectResponse +from fastapi import FastAPI, Form, Request +from fastapi.responses import JSONResponse, PlainTextResponse, RedirectResponse KEY_ID = "mock-oidc-key-1" diff --git a/backend/tests/integration/mock_services/mock_connector_server/main.py b/backend/tests/integration/mock_services/mock_connector_server/main.py index a6ffbf57b3..d91967f0fe 100644 --- a/backend/tests/integration/mock_services/mock_connector_server/main.py +++ b/backend/tests/integration/mock_services/mock_connector_server/main.py @@ -1,7 +1,5 @@ -from fastapi import FastAPI -from fastapi import HTTPException -from pydantic import BaseModel -from pydantic import Field +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field # We would like to import these, but it makes building this so much harder/slower # from onyx.connectors.mock_connector.connector import SingleConnectorYield diff --git a/backend/tests/integration/multitenant_tests/discord_bot/test_discord_bot_multitenant.py b/backend/tests/integration/multitenant_tests/discord_bot/test_discord_bot_multitenant.py index 0c25804646..d71823a2bb 100644 --- a/backend/tests/integration/multitenant_tests/discord_bot/test_discord_bot_multitenant.py +++ b/backend/tests/integration/multitenant_tests/discord_bot/test_discord_bot_multitenant.py @@ -9,14 +9,15 @@ from uuid import uuid4 import pytest -from onyx.db.discord_bot import get_guild_config_by_registration_key -from onyx.db.discord_bot import register_guild +from onyx.db.discord_bot import get_guild_config_by_registration_key, register_guild from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.models import UserRole from onyx.onyxbot.discord.cache import DiscordCacheManager -from onyx.server.manage.discord_bot.utils import generate_discord_registration_key -from onyx.server.manage.discord_bot.utils import parse_discord_registration_key -from onyx.server.manage.discord_bot.utils import REGISTRATION_KEY_PREFIX +from onyx.server.manage.discord_bot.utils import ( + generate_discord_registration_key, + parse_discord_registration_key, + REGISTRATION_KEY_PREFIX, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.user import UserManager diff --git a/backend/tests/integration/multitenant_tests/syncing/test_search_permissions.py b/backend/tests/integration/multitenant_tests/syncing/test_search_permissions.py index 735ada2a11..9883b9f041 100644 --- a/backend/tests/integration/multitenant_tests/syncing/test_search_permissions.py +++ b/backend/tests/integration/multitenant_tests/syncing/test_search_permissions.py @@ -8,11 +8,13 @@ from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestChatSession -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import ToolName +from tests.integration.common_utils.test_models import ( + DATestAPIKey, + DATestCCPair, + DATestChatSession, + DATestUser, + ToolName, +) def setup_test_tenants(reset_multitenant: None) -> dict[str, Any]: # noqa: ARG001 diff --git a/backend/tests/integration/multitenant_tests/tenants/test_tenant_provisioning_rollback.py b/backend/tests/integration/multitenant_tests/tenants/test_tenant_provisioning_rollback.py index 5f3bbc7102..84d7dde516 100644 --- a/backend/tests/integration/multitenant_tests/tenants/test_tenant_provisioning_rollback.py +++ b/backend/tests/integration/multitenant_tests/tenants/test_tenant_provisioning_rollback.py @@ -9,13 +9,14 @@ This test verifies the full flow: provisioning failure → rollback → schema c """ import uuid -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from sqlalchemy import text -from ee.onyx.server.tenants.schema_management import create_schema_if_not_exists -from ee.onyx.server.tenants.schema_management import drop_schema +from ee.onyx.server.tenants.schema_management import ( + create_schema_if_not_exists, + drop_schema, +) from onyx.db.engine.sql_engine import get_session_with_shared_schema from shared_configs.configs import TENANT_ID_PREFIX diff --git a/backend/tests/integration/multitenant_tests/test_mobile_bearer_tenant_resolution.py b/backend/tests/integration/multitenant_tests/test_mobile_bearer_tenant_resolution.py index da4b188fe9..8c2f89b47e 100644 --- a/backend/tests/integration/multitenant_tests/test_mobile_bearer_tenant_resolution.py +++ b/backend/tests/integration/multitenant_tests/test_mobile_bearer_tenant_resolution.py @@ -19,8 +19,7 @@ from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.mobile_auth._helpers import bearer -from tests.integration.tests.mobile_auth._helpers import mobile_login +from tests.integration.tests.mobile_auth._helpers import bearer, mobile_login def test_mobile_bearer_resolves_tenant(reset_multitenant: None) -> None: # noqa: ARG001 diff --git a/backend/tests/integration/tests/anonymous_user/test_anonymous_user.py b/backend/tests/integration/tests/anonymous_user/test_anonymous_user.py index 2b95903069..1ed3a70ca1 100644 --- a/backend/tests/integration/tests/anonymous_user/test_anonymous_user.py +++ b/backend/tests/integration/tests/anonymous_user/test_anonymous_user.py @@ -2,8 +2,7 @@ from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.settings import SettingsManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestSettings -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestSettings, DATestUser def test_me_endpoint_returns_anonymous_user_when_enabled( diff --git a/backend/tests/integration/tests/api_key/test_api_key.py b/backend/tests/integration/tests/api_key/test_api_key.py index 2d94c8671c..f1a866afba 100644 --- a/backend/tests/integration/tests/api_key/test_api_key.py +++ b/backend/tests/integration/tests/api_key/test_api_key.py @@ -7,8 +7,7 @@ from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.api_key import APIKeyManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser def test_limited(reset: None) -> None: # noqa: ARG001 diff --git a/backend/tests/integration/tests/chat/test_chat_deletion.py b/backend/tests/integration/tests/chat/test_chat_deletion.py index f9a5d3141d..c5de72a978 100644 --- a/backend/tests/integration/tests/chat/test_chat_deletion.py +++ b/backend/tests/integration/tests/chat/test_chat_deletion.py @@ -3,8 +3,7 @@ import pytest from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.reset import reset_all -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser MESSAGE = "Hi" diff --git a/backend/tests/integration/tests/chat/test_chat_session_access.py b/backend/tests/integration/tests/chat/test_chat_session_access.py index 9abce55dc3..3b5e5928be 100644 --- a/backend/tests/integration/tests/chat/test_chat_session_access.py +++ b/backend/tests/integration/tests/chat/test_chat_session_access.py @@ -4,13 +4,14 @@ import httpx import pytest from onyx.auth.schemas import UserRole -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.chat import ChatSessionManager -from tests.integration.common_utils.managers.user import build_email -from tests.integration.common_utils.managers.user import DEFAULT_PASSWORD -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import ( + build_email, + DEFAULT_PASSWORD, + UserManager, +) from tests.integration.common_utils.reset import reset_all from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/chat_retention/test_chat_retention.py b/backend/tests/integration/tests/chat_retention/test_chat_retention.py index 4d25aed81a..dbf1a9d23c 100644 --- a/backend/tests/integration/tests/chat_retention/test_chat_retention.py +++ b/backend/tests/integration/tests/chat_retention/test_chat_retention.py @@ -1,9 +1,7 @@ import os import time import uuid -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from uuid import UUID import httpx @@ -13,21 +11,20 @@ from sqlalchemy import update from sqlalchemy.orm import Session import ee.onyx.background.celery.tasks.ttl_management.tasks as ttl_tasks -from onyx.configs.constants import CELERY_CHAT_TTL_DELETE_TASK_EXPIRES -from onyx.configs.constants import OnyxRedisLocks -from onyx.db.chat import delete_chat_session -from onyx.db.chat import get_chat_sessions_older_than +from onyx.configs.constants import CELERY_CHAT_TTL_DELETE_TASK_EXPIRES, OnyxRedisLocks +from onyx.db.chat import delete_chat_session, get_chat_sessions_older_than from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.models import ChatMessage -from onyx.db.models import ChatSession +from onyx.db.models import ChatMessage, ChatSession from onyx.redis.redis_pool import get_redis_client from shared_configs.contextvars import get_current_tenant_id from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.settings import SettingsManager -from tests.integration.common_utils.test_models import DATestChatSession -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestSettings -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestChatSession, + DATestLLMProvider, + DATestSettings, + DATestUser, +) RETENTION_SECONDS = 10 diff --git a/backend/tests/integration/tests/cli/conftest.py b/backend/tests/integration/tests/cli/conftest.py index 18a5f070a4..522a1893c9 100644 --- a/backend/tests/integration/tests/cli/conftest.py +++ b/backend/tests/integration/tests/cli/conftest.py @@ -17,8 +17,7 @@ import pytest import uvicorn from fastapi.testclient import TestClient -from tests.integration.common_utils.constants import API_SERVER_HOST -from tests.integration.common_utils.constants import API_SERVER_PORT +from tests.integration.common_utils.constants import API_SERVER_HOST, API_SERVER_PORT # Distinctive token tests assert on; the surrounding text is long enough to # exceed the small --max-output thresholds the truncation tests use. diff --git a/backend/tests/integration/tests/cli/test_cli_commands.py b/backend/tests/integration/tests/cli/test_cli_commands.py index 8012762141..4a839eb977 100644 --- a/backend/tests/integration/tests/cli/test_cli_commands.py +++ b/backend/tests/integration/tests/cli/test_cli_commands.py @@ -58,10 +58,12 @@ from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.document_set import DocumentSetManager from tests.integration.common_utils.managers.pat import PATManager from tests.integration.common_utils.managers.persona import PersonaManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestPersona -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestAPIKey, + DATestLLMProvider, + DATestPersona, + DATestUser, +) from tests.integration.tests.cli.conftest import MOCK_LLM_TOKEN _CLI_BINARY = os.environ.get("ONYX_CLI_BINARY") diff --git a/backend/tests/integration/tests/connector/test_connector_creation.py b/backend/tests/integration/tests/connector/test_connector_creation.py index ff621989e2..9d4dacf444 100644 --- a/backend/tests/integration/tests/connector/test_connector_creation.py +++ b/backend/tests/integration/tests/connector/test_connector_creation.py @@ -1,6 +1,5 @@ import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from onyx.connectors.models import InputType from onyx.server.documents.models import DocumentSource diff --git a/backend/tests/integration/tests/connector/test_connector_deletion.py b/backend/tests/integration/tests/connector/test_connector_deletion.py index a9047e97c6..e563cad981 100644 --- a/backend/tests/integration/tests/connector/test_connector_deletion.py +++ b/backend/tests/integration/tests/connector/test_connector_deletion.py @@ -11,12 +11,10 @@ from uuid import uuid4 from sqlalchemy.orm import Session -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import DocumentFailure +from onyx.connectors.models import ConnectorFailure, DocumentFailure from onyx.db.engine.sql_engine import get_sqlalchemy_engine from onyx.db.enums import IndexingStatus -from onyx.db.index_attempt import create_index_attempt -from onyx.db.index_attempt import create_index_attempt_error +from onyx.db.index_attempt import create_index_attempt, create_index_attempt_error from onyx.db.models import IndexAttempt from onyx.db.search_settings import get_current_search_settings from onyx.server.documents.models import DocumentSource @@ -27,9 +25,11 @@ from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.document_set import DocumentSetManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import DATestUserGroup +from tests.integration.common_utils.test_models import ( + DATestAPIKey, + DATestUser, + DATestUserGroup, +) from tests.integration.common_utils.vespa import vespa_fixture diff --git a/backend/tests/integration/tests/connector/test_last_indexed_time.py b/backend/tests/integration/tests/connector/test_last_indexed_time.py index 4dd3af5067..efc2250fc9 100644 --- a/backend/tests/integration/tests/connector/test_last_indexed_time.py +++ b/backend/tests/integration/tests/connector/test_last_indexed_time.py @@ -15,20 +15,16 @@ Edge cases: 4. COMPLETED_WITH_ERRORS counts as a success for last_indexed purposes. """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from onyx.db.models import IndexingStatus -from onyx.server.documents.models import CCPairFullInfo -from onyx.server.documents.models import ConnectorIndexingStatusLite +from onyx.server.documents.models import CCPairFullInfo, ConnectorIndexingStatusLite from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestCCPair, DATestUser def _wait_for_real_success( diff --git a/backend/tests/integration/tests/craft/conftest.py b/backend/tests/integration/tests/craft/conftest.py index aba4fdbcd1..4b7bcdd42f 100644 --- a/backend/tests/integration/tests/craft/conftest.py +++ b/backend/tests/integration/tests/craft/conftest.py @@ -5,28 +5,26 @@ from __future__ import annotations import contextlib from collections.abc import Generator from typing import NamedTuple -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx import pytest -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.db.engine.sql_engine import get_session_with_tenant, SqlEngine +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from onyx.server.features.build.db.sandbox import get_running_sandboxes from onyx.server.features.build.sandbox.factory import get_sandbox_manager from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.common.craft.users import create_or_login_admin from tests.integration.common_utils.constants import ADMIN_USER_NAME -from tests.integration.common_utils.http_client import RetryingTransport -from tests.integration.common_utils.http_client import set_test_client +from tests.integration.common_utils.http_client import ( + RetryingTransport, + set_test_client, +) from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser class SharedSession(NamedTuple): diff --git a/backend/tests/integration/tests/craft/docker_e2e/conftest.py b/backend/tests/integration/tests/craft/docker_e2e/conftest.py index 129e3110ed..2334e0caf9 100644 --- a/backend/tests/integration/tests/craft/docker_e2e/conftest.py +++ b/backend/tests/integration/tests/craft/docker_e2e/conftest.py @@ -3,18 +3,14 @@ from __future__ import annotations import subprocess -from typing import NamedTuple -from typing import Protocol +from typing import NamedTuple, Protocol from uuid import UUID import pytest from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.enums import SandboxStatus -from onyx.db.external_app import create_external_app -from onyx.db.external_app import get_built_in_external_app +from onyx.db.enums import EndpointPolicy, ExternalAppType, SandboxStatus +from onyx.db.external_app import create_external_app, get_built_in_external_app from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/craft/docker_e2e/test_approval_gate_docker.py b/backend/tests/integration/tests/craft/docker_e2e/test_approval_gate_docker.py index 082a3a0a14..516fe7f495 100644 --- a/backend/tests/integration/tests/craft/docker_e2e/test_approval_gate_docker.py +++ b/backend/tests/integration/tests/craft/docker_e2e/test_approval_gate_docker.py @@ -5,18 +5,15 @@ from __future__ import annotations import json import subprocess from typing import NamedTuple -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from httpx import Response from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import ExternalAppType +from onyx.db.enums import ApprovalDecision, ExternalAppType from onyx.db.external_app import get_built_in_external_app -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.build_approvals import ( @@ -24,9 +21,11 @@ from tests.integration.common_utils.managers.build_approvals import ( ) from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.docker_e2e.conftest import DockerExec -from tests.integration.tests.craft.docker_e2e.conftest import DockerSandbox -from tests.integration.tests.craft.docker_e2e.conftest import ProvisionSandbox +from tests.integration.tests.craft.docker_e2e.conftest import ( + DockerExec, + DockerSandbox, + ProvisionSandbox, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.DOCKER, diff --git a/backend/tests/integration/tests/craft/docker_e2e/test_sandbox_network_posture_docker.py b/backend/tests/integration/tests/craft/docker_e2e/test_sandbox_network_posture_docker.py index 7f7dbc5cbd..b0d18540c5 100644 --- a/backend/tests/integration/tests/craft/docker_e2e/test_sandbox_network_posture_docker.py +++ b/backend/tests/integration/tests/craft/docker_e2e/test_sandbox_network_posture_docker.py @@ -14,17 +14,21 @@ from uuid import uuid4 import pytest -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_PROXY_INJECTED_PLACEHOLDER -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + SANDBOX_BACKEND, + SANDBOX_PROXY_INJECTED_PLACEHOLDER, + SandboxBackend, +) from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( SANDBOX_EXEC_USER, ) from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.docker_e2e.conftest import DockerExec -from tests.integration.tests.craft.docker_e2e.conftest import DockerSandbox -from tests.integration.tests.craft.docker_e2e.conftest import ProvisionSandbox +from tests.integration.tests.craft.docker_e2e.conftest import ( + DockerExec, + DockerSandbox, + ProvisionSandbox, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.DOCKER, diff --git a/backend/tests/integration/tests/craft/docker_e2e/test_serve_streaming_docker.py b/backend/tests/integration/tests/craft/docker_e2e/test_serve_streaming_docker.py index 02c68baebf..d0a93fffde 100644 --- a/backend/tests/integration/tests/craft/docker_e2e/test_serve_streaming_docker.py +++ b/backend/tests/integration/tests/craft/docker_e2e/test_serve_streaming_docker.py @@ -7,28 +7,29 @@ import os import subprocess import time from collections.abc import Generator -from dataclasses import dataclass -from dataclasses import field -from typing import Any -from typing import NamedTuple -from uuid import UUID -from uuid import uuid4 +from dataclasses import dataclass, field +from typing import Any, NamedTuple +from uuid import UUID, uuid4 import pytest from onyx.llm.constants import LlmProviderNames -from onyx.server.features.build.configs import OPENCODE_SERVER_PASSWORD -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + OPENCODE_SERVER_PASSWORD, + SANDBOX_BACKEND, + SandboxBackend, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.docker_e2e.conftest import DockerExec -from tests.integration.tests.craft.docker_e2e.conftest import DockerSandbox -from tests.integration.tests.craft.docker_e2e.conftest import ProvisionSandbox +from tests.integration.tests.craft.docker_e2e.conftest import ( + DockerExec, + DockerSandbox, + ProvisionSandbox, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.DOCKER, diff --git a/backend/tests/integration/tests/craft/docker_e2e/test_snapshot_roundtrip_docker.py b/backend/tests/integration/tests/craft/docker_e2e/test_snapshot_roundtrip_docker.py index 05ac0b0c94..d3ef9813b5 100644 --- a/backend/tests/integration/tests/craft/docker_e2e/test_snapshot_roundtrip_docker.py +++ b/backend/tests/integration/tests/craft/docker_e2e/test_snapshot_roundtrip_docker.py @@ -12,8 +12,7 @@ from __future__ import annotations import datetime import time -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy import select @@ -21,22 +20,19 @@ from sqlalchemy import select from onyx.db.engine.sql_engine import get_session_with_tenant from onyx.db.enums import SandboxStatus from onyx.db.models import Sandbox -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( OPENCODE_DATA_DIR, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( SANDBOX_EXEC_USER, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( SESSIONS_ROOT, ) from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.docker_e2e.conftest import DockerExec -from tests.integration.tests.craft.docker_e2e.conftest import ProvisionSandbox +from tests.integration.tests.craft.docker_e2e.conftest import ( + DockerExec, + ProvisionSandbox, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.DOCKER, diff --git a/backend/tests/integration/tests/craft/docker_e2e/test_workspace_setup_docker.py b/backend/tests/integration/tests/craft/docker_e2e/test_workspace_setup_docker.py index 02efc488cb..45a3a3175e 100644 --- a/backend/tests/integration/tests/craft/docker_e2e/test_workspace_setup_docker.py +++ b/backend/tests/integration/tests/craft/docker_e2e/test_workspace_setup_docker.py @@ -13,17 +13,18 @@ from uuid import uuid4 import pytest -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( SANDBOX_EXEC_USER, ) from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.docker_e2e.conftest import DockerExec -from tests.integration.tests.craft.docker_e2e.conftest import DockerSandbox -from tests.integration.tests.craft.docker_e2e.conftest import ProvisionSandbox +from tests.integration.tests.craft.docker_e2e.conftest import ( + DockerExec, + DockerSandbox, + ProvisionSandbox, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.DOCKER, diff --git a/backend/tests/integration/tests/craft/k8s/conftest.py b/backend/tests/integration/tests/craft/k8s/conftest.py index fe1cab9367..41a8788081 100644 --- a/backend/tests/integration/tests/craft/k8s/conftest.py +++ b/backend/tests/integration/tests/craft/k8s/conftest.py @@ -11,15 +11,16 @@ import httpx import pytest from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from onyx.server.features.build.db.sandbox import get_running_sandboxes from onyx.server.features.build.sandbox.factory import get_sandbox_manager from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.common.craft.users import create_or_login_admin from tests.integration.common_utils.constants import ADMIN_USER_NAME -from tests.integration.common_utils.http_client import RetryingTransport -from tests.integration.common_utils.http_client import set_test_client +from tests.integration.common_utils.http_client import ( + RetryingTransport, + set_test_client, +) from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.test_models import DATestUser from tests.integration.tests.craft.k8s.k8s_fixtures import suite_sandbox_ids diff --git a/backend/tests/integration/tests/craft/k8s/k8s_db_fixtures.py b/backend/tests/integration/tests/craft/k8s/k8s_db_fixtures.py index d8d8a7e9a0..4a422ccad2 100644 --- a/backend/tests/integration/tests/craft/k8s/k8s_db_fixtures.py +++ b/backend/tests/integration/tests/craft/k8s/k8s_db_fixtures.py @@ -7,12 +7,13 @@ from collections.abc import Generator import pytest from sqlalchemy.orm import Session -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR -from tests.common.craft.skill_table_isolation import restore_skill_tables -from tests.common.craft.skill_table_isolation import snapshot_skill_tables +from tests.common.craft.skill_table_isolation import ( + restore_skill_tables, + snapshot_skill_tables, +) # Modules opt into skill-table snapshot/restore via this marker. _SKILL_ISOLATION_MARKER = "craft_skill_isolation" diff --git a/backend/tests/integration/tests/craft/k8s/k8s_fixtures.py b/backend/tests/integration/tests/craft/k8s/k8s_fixtures.py index 412615a5e1..60f1e39da4 100644 --- a/backend/tests/integration/tests/craft/k8s/k8s_fixtures.py +++ b/backend/tests/integration/tests/craft/k8s/k8s_fixtures.py @@ -6,18 +6,12 @@ import json import os import shlex import time -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Sequence -from contextlib import contextmanager -from contextlib import suppress -from dataclasses import dataclass -from dataclasses import field +from collections.abc import Callable, Generator, Sequence +from contextlib import contextmanager, suppress +from dataclasses import dataclass, field from pathlib import PurePosixPath -from typing import NamedTuple -from typing import TYPE_CHECKING -from uuid import UUID -from uuid import uuid4 +from typing import NamedTuple, TYPE_CHECKING +from uuid import UUID, uuid4 import pytest @@ -26,17 +20,16 @@ if TYPE_CHECKING: from tests.integration.common_utils.test_models import DATestUser -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import Sandbox -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SANDBOX_PROXY_PORT -from onyx.server.features.build.db.user_library import delete_user_file -from onyx.server.features.build.db.user_library import list_user_files +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine +from onyx.db.models import ( + ConnectorCredentialPair, + Credential, + Sandbox, + User, + User__UserGroup, +) +from onyx.server.features.build.configs import SANDBOX_NAMESPACE, SANDBOX_PROXY_PORT +from onyx.server.features.build.db.user_library import delete_user_file, list_user_files from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, ) @@ -489,8 +482,7 @@ def _pool_pod( k8s_client: "k8s_client_module.CoreV1Api", ) -> Generator[_PoolPod, None, None]: """Module-scoped sandbox pod shared by all ``running_sandbox()`` calls.""" - from onyx.server.features.build.configs import SANDBOX_BACKEND - from onyx.server.features.build.configs import SandboxBackend + from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend if SANDBOX_BACKEND != SandboxBackend.KUBERNETES: pytest.skip( @@ -530,8 +522,7 @@ def running_sandbox( clean slate. Extra user-owned pods come from ``SandboxHandle.provision_api_user``. """ - from onyx.server.features.build.configs import SANDBOX_BACKEND - from onyx.server.features.build.configs import SandboxBackend + from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend if SANDBOX_BACKEND != SandboxBackend.KUBERNETES: pytest.skip( diff --git a/backend/tests/integration/tests/craft/k8s/test_approval_gate.py b/backend/tests/integration/tests/craft/k8s/test_approval_gate.py index 08dcbaf6a8..7debec7bb6 100644 --- a/backend/tests/integration/tests/craft/k8s/test_approval_gate.py +++ b/backend/tests/integration/tests/craft/k8s/test_approval_gate.py @@ -6,8 +6,7 @@ import json import time from collections.abc import Generator from typing import NamedTuple -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx import pytest @@ -16,21 +15,23 @@ from sqlalchemy.orm import Session from onyx.cache.factory import get_cache_backend from onyx.configs.constants import NotificationType -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import BuildSessionStatus -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.db.models import ActionApproval -from onyx.db.models import BuildSession -from onyx.db.models import Notification +from onyx.db.enums import ( + ApprovalDecision, + BuildSessionStatus, + EndpointPolicy, + ExternalAppType, +) +from onyx.db.models import ActionApproval, BuildSession, Notification from onyx.error_handling.error_codes import OnyxErrorCode from onyx.sandbox_proxy.approval_cache import pop_announcement -from onyx.server.features.build.configs import SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SANDBOX_PROXY_NAMESPACE -from onyx.server.features.build.configs import SANDBOX_PROXY_PORT -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + SANDBOX_APPROVAL_WAIT_TIMEOUT_SECONDS, + SANDBOX_BACKEND, + SANDBOX_NAMESPACE, + SANDBOX_PROXY_NAMESPACE, + SANDBOX_PROXY_PORT, + SandboxBackend, +) from onyx.server.features.build.external_apps.models import ExternalAppAdminResponse from onyx.utils.logger import setup_logger from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE @@ -38,11 +39,13 @@ from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client as http_client from tests.integration.common_utils.managers.external_app import ExternalAppManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.k8s.k8s_fixtures import OwnedLivePod -from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec -from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec_async -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_pod_exec_output -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_proxy_redeploy +from tests.integration.tests.craft.k8s.k8s_fixtures import ( + OwnedLivePod, + pod_exec, + pod_exec_async, + wait_for_pod_exec_output, + wait_for_proxy_redeploy, +) logger = setup_logger() diff --git a/backend/tests/integration/tests/craft/k8s/test_browser.py b/backend/tests/integration/tests/craft/k8s/test_browser.py index 2cc2c2f6ed..cc7cc2aa99 100644 --- a/backend/tests/integration/tests/craft/k8s/test_browser.py +++ b/backend/tests/integration/tests/craft/k8s/test_browser.py @@ -13,11 +13,12 @@ from __future__ import annotations import pytest from kubernetes import client -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SandboxBackend -from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec -from tests.integration.tests.craft.k8s.k8s_fixtures import PoolSession +from onyx.server.features.build.configs import ( + SANDBOX_BACKEND, + SANDBOX_NAMESPACE, + SandboxBackend, +) +from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec, PoolSession pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.KUBERNETES, diff --git a/backend/tests/integration/tests/craft/k8s/test_bun_node_modules_dedup.py b/backend/tests/integration/tests/craft/k8s/test_bun_node_modules_dedup.py index 7eaea9c967..3c60f39552 100644 --- a/backend/tests/integration/tests/craft/k8s/test_bun_node_modules_dedup.py +++ b/backend/tests/integration/tests/craft/k8s/test_bun_node_modules_dedup.py @@ -6,15 +6,16 @@ import os from collections.abc import Generator from contextlib import suppress from typing import NamedTuple -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from kubernetes import client -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + SANDBOX_BACKEND, + SANDBOX_NAMESPACE, + SandboxBackend, +) from onyx.server.features.build.sandbox.base import BUN_CACHE_DIR from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, @@ -23,10 +24,12 @@ from tests.common.craft.payloads import default_llm_config from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.k8s.k8s_fixtures import cleanup_api_user_sandbox_rows -from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec -from tests.integration.tests.craft.k8s.k8s_fixtures import PoolSession -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_pod_deletion +from tests.integration.tests.craft.k8s.k8s_fixtures import ( + cleanup_api_user_sandbox_rows, + pod_exec, + PoolSession, + wait_for_pod_deletion, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.KUBERNETES, diff --git a/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox.py b/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox.py index 9210a685bd..d11eb3334b 100644 --- a/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox.py +++ b/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox.py @@ -8,18 +8,22 @@ import pytest from kubernetes import client from kubernetes.client.rest import ApiException -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + SANDBOX_BACKEND, + SANDBOX_NAMESPACE, + SandboxBackend, +) from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, ) from onyx.utils.logger import setup_logger -from tests.integration.tests.craft.k8s.k8s_fixtures import OwnedLivePod -from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec -from tests.integration.tests.craft.k8s.k8s_fixtures import PoolSession -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_pod_deletion -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_until_healthy +from tests.integration.tests.craft.k8s.k8s_fixtures import ( + OwnedLivePod, + pod_exec, + PoolSession, + wait_for_pod_deletion, + wait_until_healthy, +) logger = setup_logger() diff --git a/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox_file_ops.py b/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox_file_ops.py index 014533746f..5a21857319 100644 --- a/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox_file_ops.py +++ b/backend/tests/integration/tests/craft/k8s/test_kubernetes_sandbox_file_ops.py @@ -9,9 +9,11 @@ import httpx import pytest from kubernetes import client -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + SANDBOX_BACKEND, + SANDBOX_NAMESPACE, + SandboxBackend, +) from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, ) @@ -19,12 +21,14 @@ from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client as http_client from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.k8s.k8s_fixtures import OwnedLivePod -from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec -from tests.integration.tests.craft.k8s.k8s_fixtures import PoolSession -from tests.integration.tests.craft.k8s.k8s_fixtures import SandboxHandle -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_pod_deletion -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_until_healthy +from tests.integration.tests.craft.k8s.k8s_fixtures import ( + OwnedLivePod, + pod_exec, + PoolSession, + SandboxHandle, + wait_for_pod_deletion, + wait_until_healthy, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.KUBERNETES, diff --git a/backend/tests/integration/tests/craft/k8s/test_messages_api_k8s.py b/backend/tests/integration/tests/craft/k8s/test_messages_api_k8s.py index 149fdfe580..27c55fb7d8 100644 --- a/backend/tests/integration/tests/craft/k8s/test_messages_api_k8s.py +++ b/backend/tests/integration/tests/craft/k8s/test_messages_api_k8s.py @@ -6,15 +6,13 @@ import os import time from contextlib import suppress from typing import cast -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from kubernetes import client from onyx.configs.constants import MessageType -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from onyx.server.features.build.interactive_turns.models import InteractiveTurnResponse from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, @@ -22,8 +20,10 @@ from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager im from onyx.server.features.build.session.models import MessageResponse from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.tests.craft.k8s.k8s_fixtures import cleanup_api_user_sandbox_rows -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_pod_deletion +from tests.integration.tests.craft.k8s.k8s_fixtures import ( + cleanup_api_user_sandbox_rows, + wait_for_pod_deletion, +) pytestmark = [ pytest.mark.skipif( diff --git a/backend/tests/integration/tests/craft/k8s/test_session_provisioning_api.py b/backend/tests/integration/tests/craft/k8s/test_session_provisioning_api.py index 69fdf6eeec..a417ceb0af 100644 --- a/backend/tests/integration/tests/craft/k8s/test_session_provisioning_api.py +++ b/backend/tests/integration/tests/craft/k8s/test_session_provisioning_api.py @@ -3,23 +3,26 @@ from __future__ import annotations from contextlib import suppress -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from kubernetes import client from onyx.db.enums import SandboxStatus -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + SANDBOX_BACKEND, + SANDBOX_NAMESPACE, + SandboxBackend, +) from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, ) from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.tests.craft.k8s.k8s_fixtures import cleanup_api_user_sandbox_rows -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_pod_deletion +from tests.integration.tests.craft.k8s.k8s_fixtures import ( + cleanup_api_user_sandbox_rows, + wait_for_pod_deletion, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.KUBERNETES, diff --git a/backend/tests/integration/tests/craft/k8s/test_skill_push.py b/backend/tests/integration/tests/craft/k8s/test_skill_push.py index 213de0fde5..8152bf237f 100644 --- a/backend/tests/integration/tests/craft/k8s/test_skill_push.py +++ b/backend/tests/integration/tests/craft/k8s/test_skill_push.py @@ -4,47 +4,46 @@ from __future__ import annotations import io import zipfile -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from pathlib import Path -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx import pytest from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import ExternalAppType -from onyx.db.enums import SandboxStatus -from onyx.db.enums import SkillSharePermission -from onyx.db.models import Connector -from onyx.db.models import ConnectorCredentialPair -from onyx.db.models import Credential -from onyx.db.models import Sandbox -from onyx.db.models import Skill -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup -from onyx.db.models import UserGroup__ConnectorCredentialPair -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend -from onyx.server.features.skill.models import SkillResponse -from onyx.server.features.skill.models import SkillUserShareRequest -from onyx.skills.built_in import BUILT_IN_SKILLS -from onyx.skills.built_in import BuiltInSkillDefinition -from onyx.skills.push import build_skills_fileset_for_user -from onyx.skills.push import push_skill_to_affected_sandboxes +from onyx.db.enums import ( + AccessType, + ConnectorCredentialPairStatus, + ExternalAppType, + SandboxStatus, + SkillSharePermission, +) +from onyx.db.models import ( + Connector, + ConnectorCredentialPair, + Credential, + Sandbox, + Skill, + User, + User__UserGroup, + UserGroup, + UserGroup__ConnectorCredentialPair, +) +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend +from onyx.server.features.skill.models import SkillResponse, SkillUserShareRequest +from onyx.skills.built_in import BUILT_IN_SKILLS, BuiltInSkillDefinition +from onyx.skills.push import ( + build_skills_fileset_for_user, + push_skill_to_affected_sandboxes, +) from tests.integration.common_utils.managers.external_app import ExternalAppManager from tests.integration.common_utils.managers.skill import SkillManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import DATestUserGroup -from tests.integration.tests.craft.k8s.k8s_fixtures import SandboxHandle -from tests.integration.tests.craft.k8s.k8s_fixtures import WorkspaceProxy +from tests.integration.common_utils.test_models import DATestUser, DATestUserGroup +from tests.integration.tests.craft.k8s.k8s_fixtures import SandboxHandle, WorkspaceProxy pytestmark = [ pytest.mark.skipif( diff --git a/backend/tests/integration/tests/craft/k8s/test_snapshot_restore.py b/backend/tests/integration/tests/craft/k8s/test_snapshot_restore.py index 79985f7bce..6af60710f0 100644 --- a/backend/tests/integration/tests/craft/k8s/test_snapshot_restore.py +++ b/backend/tests/integration/tests/craft/k8s/test_snapshot_restore.py @@ -7,8 +7,7 @@ import shutil import tarfile from collections.abc import Callable from pathlib import Path -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx import pytest @@ -17,9 +16,11 @@ from sqlalchemy.orm import Session from onyx.configs.constants import FileOrigin from onyx.file_store.file_store import get_default_file_store -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SANDBOX_NAMESPACE -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import ( + SANDBOX_BACKEND, + SANDBOX_NAMESPACE, + SandboxBackend, +) from onyx.server.features.build.db.sandbox import create_snapshot__no_commit from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, @@ -30,12 +31,14 @@ from tests.common.craft.payloads import default_llm_config from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.skill import SkillManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.k8s.k8s_fixtures import OwnedLivePod -from tests.integration.tests.craft.k8s.k8s_fixtures import pod_exec -from tests.integration.tests.craft.k8s.k8s_fixtures import PoolSession -from tests.integration.tests.craft.k8s.k8s_fixtures import SandboxHandle -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_for_pod_deletion -from tests.integration.tests.craft.k8s.k8s_fixtures import wait_until_healthy +from tests.integration.tests.craft.k8s.k8s_fixtures import ( + OwnedLivePod, + pod_exec, + PoolSession, + SandboxHandle, + wait_for_pod_deletion, + wait_until_healthy, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.KUBERNETES, diff --git a/backend/tests/integration/tests/craft/k8s/test_user_library_sync.py b/backend/tests/integration/tests/craft/k8s/test_user_library_sync.py index 9f4d267216..fcb007a512 100644 --- a/backend/tests/integration/tests/craft/k8s/test_user_library_sync.py +++ b/backend/tests/integration/tests/craft/k8s/test_user_library_sync.py @@ -7,17 +7,19 @@ from uuid import uuid4 import pytest -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend -from onyx.server.features.build.user_library.api import DeleteFileResponse -from onyx.server.features.build.user_library.api import UploadResponse -from tests.integration.tests.craft.k8s.k8s_fixtures import SandboxHandle -from tests.integration.tests.craft.k8s.k8s_fixtures import WorkspaceProxy -from tests.integration.tests.craft.user_library_http import delete_user_library_file -from tests.integration.tests.craft.user_library_http import list_user_library_tree -from tests.integration.tests.craft.user_library_http import make_zip_bytes -from tests.integration.tests.craft.user_library_http import upload_user_library_files -from tests.integration.tests.craft.user_library_http import upload_user_library_zip +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend +from onyx.server.features.build.user_library.api import ( + DeleteFileResponse, + UploadResponse, +) +from tests.integration.tests.craft.k8s.k8s_fixtures import SandboxHandle, WorkspaceProxy +from tests.integration.tests.craft.user_library_http import ( + delete_user_library_file, + list_user_library_tree, + make_zip_bytes, + upload_user_library_files, + upload_user_library_zip, +) pytestmark = pytest.mark.skipif( SANDBOX_BACKEND != SandboxBackend.KUBERNETES, diff --git a/backend/tests/integration/tests/craft/k8s/test_webapp_preview.py b/backend/tests/integration/tests/craft/k8s/test_webapp_preview.py index f98e6c1e90..2a91f3e455 100644 --- a/backend/tests/integration/tests/craft/k8s/test_webapp_preview.py +++ b/backend/tests/integration/tests/craft/k8s/test_webapp_preview.py @@ -11,8 +11,7 @@ import time import httpx import pytest -from onyx.server.features.build.configs import SANDBOX_BACKEND -from onyx.server.features.build.configs import SandboxBackend +from onyx.server.features.build.configs import SANDBOX_BACKEND, SandboxBackend from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/craft/test_file_ops_api.py b/backend/tests/integration/tests/craft/test_file_ops_api.py index 1f05910858..5c3acfbb33 100644 --- a/backend/tests/integration/tests/craft/test_file_ops_api.py +++ b/backend/tests/integration/tests/craft/test_file_ops_api.py @@ -10,8 +10,10 @@ import pytest from onyx.server.features.build.sandbox.factory import get_sandbox_manager from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.managers.build_session import BuildSessionManager -from tests.integration.common_utils.managers.build_session import SessionWithSandbox +from tests.integration.common_utils.managers.build_session import ( + BuildSessionManager, + SessionWithSandbox, +) from tests.integration.common_utils.test_models import DATestUser from tests.integration.tests.craft.conftest import SharedSession diff --git a/backend/tests/integration/tests/craft/test_messages_api.py b/backend/tests/integration/tests/craft/test_messages_api.py index 0c32698fbf..7fc467b783 100644 --- a/backend/tests/integration/tests/craft/test_messages_api.py +++ b/backend/tests/integration/tests/craft/test_messages_api.py @@ -2,8 +2,7 @@ from __future__ import annotations -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client diff --git a/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py b/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py index e3a53a85ad..1c5df5c08d 100644 --- a/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py +++ b/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py @@ -5,20 +5,19 @@ from __future__ import annotations import time from collections.abc import Generator from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx import pytest from sqlalchemy import select -from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.engine.sql_engine import SqlEngine -from onyx.db.enums import ScheduledTaskRunStatus -from onyx.db.enums import ScheduledTaskStatus -from onyx.db.enums import ScheduledTaskTriggerSource -from onyx.db.models import ScheduledTask -from onyx.db.models import ScheduledTaskRun +from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine +from onyx.db.enums import ( + ScheduledTaskRunStatus, + ScheduledTaskStatus, + ScheduledTaskTriggerSource, +) +from onyx.db.models import ScheduledTask, ScheduledTaskRun from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR from tests.integration.common_utils.constants import API_SERVER_URL diff --git a/backend/tests/integration/tests/craft/test_sessions_api.py b/backend/tests/integration/tests/craft/test_sessions_api.py index 6155707de5..5e0ca8e42d 100644 --- a/backend/tests/integration/tests/craft/test_sessions_api.py +++ b/backend/tests/integration/tests/craft/test_sessions_api.py @@ -2,8 +2,7 @@ from __future__ import annotations -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 from onyx.db.enums import SharingScope from onyx.redis.redis_pool import get_redis_client @@ -14,9 +13,11 @@ from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.build_session import BuildSessionManager from tests.integration.common_utils.managers.settings import SettingsManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestSettings -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestLLMProvider, + DATestSettings, + DATestUser, +) from tests.integration.tests.craft.conftest import SharedSession diff --git a/backend/tests/integration/tests/craft/test_user_library_api.py b/backend/tests/integration/tests/craft/test_user_library_api.py index 975606a54a..3b5e8c2ae3 100644 --- a/backend/tests/integration/tests/craft/test_user_library_api.py +++ b/backend/tests/integration/tests/craft/test_user_library_api.py @@ -4,15 +4,19 @@ from __future__ import annotations from uuid import uuid4 -from onyx.server.features.build.user_library.api import DeleteFileResponse -from onyx.server.features.build.user_library.api import LibraryEntryResponse -from onyx.server.features.build.user_library.api import UploadResponse +from onyx.server.features.build.user_library.api import ( + DeleteFileResponse, + LibraryEntryResponse, + UploadResponse, +) from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.craft.user_library_http import delete_user_library_file -from tests.integration.tests.craft.user_library_http import list_user_library_tree -from tests.integration.tests.craft.user_library_http import make_zip_bytes -from tests.integration.tests.craft.user_library_http import upload_user_library_files -from tests.integration.tests.craft.user_library_http import upload_user_library_zip +from tests.integration.tests.craft.user_library_http import ( + delete_user_library_file, + list_user_library_tree, + make_zip_bytes, + upload_user_library_files, + upload_user_library_zip, +) def _find_doc_by_name( diff --git a/backend/tests/integration/tests/discord_bot/test_discord_bot_api.py b/backend/tests/integration/tests/discord_bot/test_discord_bot_api.py index 906b3f62b2..14d7789f92 100644 --- a/backend/tests/integration/tests/discord_bot/test_discord_bot_api.py +++ b/backend/tests/integration/tests/discord_bot/test_discord_bot_api.py @@ -6,8 +6,10 @@ These tests hit actual API endpoints via HTTP client. import httpx import pytest -from onyx.db.discord_bot import get_discord_service_api_key -from onyx.db.discord_bot import get_or_create_discord_service_api_key +from onyx.db.discord_bot import ( + get_discord_service_api_key, + get_or_create_discord_service_api_key, +) from onyx.db.engine.sql_engine import get_session_with_current_tenant from tests.integration.common_utils.managers.discord_bot import DiscordBotManager from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/discord_bot/test_discord_bot_db.py b/backend/tests/integration/tests/discord_bot/test_discord_bot_db.py index 16ec2973a1..ec50be61d9 100644 --- a/backend/tests/integration/tests/discord_bot/test_discord_bot_db.py +++ b/backend/tests/integration/tests/discord_bot/test_discord_bot_db.py @@ -10,25 +10,26 @@ from sqlalchemy import select from sqlalchemy.orm import Session from onyx.auth.permissions import get_effective_permissions -from onyx.db.discord_bot import bulk_create_channel_configs -from onyx.db.discord_bot import create_discord_bot_config -from onyx.db.discord_bot import create_guild_config -from onyx.db.discord_bot import delete_discord_bot_config -from onyx.db.discord_bot import delete_discord_service_api_key -from onyx.db.discord_bot import delete_guild_config -from onyx.db.discord_bot import get_channel_configs -from onyx.db.discord_bot import get_discord_bot_config -from onyx.db.discord_bot import get_discord_service_api_key -from onyx.db.discord_bot import get_guild_config_by_internal_id -from onyx.db.discord_bot import get_guild_config_by_registration_key -from onyx.db.discord_bot import get_guild_configs -from onyx.db.discord_bot import get_or_create_discord_service_api_key -from onyx.db.discord_bot import sync_channel_configs -from onyx.db.discord_bot import update_discord_channel_config -from onyx.db.discord_bot import update_guild_config +from onyx.db.discord_bot import ( + bulk_create_channel_configs, + create_discord_bot_config, + create_guild_config, + delete_discord_bot_config, + delete_discord_service_api_key, + delete_guild_config, + get_channel_configs, + get_discord_bot_config, + get_discord_service_api_key, + get_guild_config_by_internal_id, + get_guild_config_by_registration_key, + get_guild_configs, + get_or_create_discord_service_api_key, + sync_channel_configs, + update_discord_channel_config, + update_guild_config, +) from onyx.db.enums import Permission -from onyx.db.models import Persona -from onyx.db.models import User +from onyx.db.models import Persona, User from onyx.db.utils import DiscordChannelView from onyx.server.manage.discord_bot.utils import generate_discord_registration_key @@ -682,8 +683,7 @@ class TestServiceApiKeyAPI: @pytest.fixture def db_session() -> Generator[Session, None, None]: """Create database session for tests.""" - from onyx.db.engine.sql_engine import get_session_with_current_tenant - from onyx.db.engine.sql_engine import SqlEngine + from onyx.db.engine.sql_engine import get_session_with_current_tenant, SqlEngine from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR SqlEngine.init_engine(pool_size=10, max_overflow=5) diff --git a/backend/tests/integration/tests/document_set/test_syncing.py b/backend/tests/integration/tests/document_set/test_syncing.py index ab31517051..fa2d9c208a 100644 --- a/backend/tests/integration/tests/document_set/test_syncing.py +++ b/backend/tests/integration/tests/document_set/test_syncing.py @@ -6,8 +6,7 @@ from tests.integration.common_utils.managers.api_key import APIKeyManager from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.document_set import DocumentSetManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser from tests.integration.common_utils.vespa import vespa_fixture diff --git a/backend/tests/integration/tests/external_apps/test_external_app_action_policies.py b/backend/tests/integration/tests/external_apps/test_external_app_action_policies.py index e401e62b2f..a4b1f947cc 100644 --- a/backend/tests/integration/tests/external_apps/test_external_app_action_policies.py +++ b/backend/tests/integration/tests/external_apps/test_external_app_action_policies.py @@ -13,8 +13,7 @@ Contract under test (observed through the admin response's ``actions`` view): - clearing an override means sending that action explicitly as ``ASK``. """ -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.external_apps.providers.slack import SlackAction from onyx.server.features.build.external_apps.models import ExternalAppAdminResponse from tests.integration.common_utils.managers.external_app import ExternalAppManager diff --git a/backend/tests/integration/tests/external_apps/test_external_apps.py b/backend/tests/integration/tests/external_apps/test_external_apps.py index 6948f1e82d..b49f301695 100644 --- a/backend/tests/integration/tests/external_apps/test_external_apps.py +++ b/backend/tests/integration/tests/external_apps/test_external_apps.py @@ -4,8 +4,10 @@ import httpx import pytest from onyx.db.enums import ExternalAppType -from onyx.server.features.build.external_apps.models import ExternalAppAdminResponse -from onyx.server.features.build.external_apps.models import ExternalAppUserResponse +from onyx.server.features.build.external_apps.models import ( + ExternalAppAdminResponse, + ExternalAppUserResponse, +) from onyx.utils.encryption import mask_credential_dict from tests.integration.common_utils.managers.external_app import ExternalAppManager from tests.integration.common_utils.managers.skill import SkillManager diff --git a/backend/tests/integration/tests/image_generation/test_image_generation_config.py b/backend/tests/integration/tests/image_generation/test_image_generation_config.py index 4288d8d6ba..d0f1d63476 100644 --- a/backend/tests/integration/tests/image_generation/test_image_generation_config.py +++ b/backend/tests/integration/tests/image_generation/test_image_generation_config.py @@ -16,8 +16,7 @@ from tests.integration.common_utils.managers.image_generation import ( from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.reset import reset_all -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser @pytest.fixture(scope="module") diff --git a/backend/tests/integration/tests/image_generation/test_image_generation_tool_visibility.py b/backend/tests/integration/tests/image_generation/test_image_generation_tool_visibility.py index 5bf328c7c2..52e0b5d680 100644 --- a/backend/tests/integration/tests/image_generation/test_image_generation_tool_visibility.py +++ b/backend/tests/integration/tests/image_generation/test_image_generation_tool_visibility.py @@ -12,8 +12,7 @@ from tests.integration.common_utils.managers.llm_provider import LLMProviderMana from tests.integration.common_utils.managers.tool import ToolManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.reset import reset_all -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser IMAGE_GENERATION_TOOL_NAME = ImageGenerationTool.NAME diff --git a/backend/tests/integration/tests/image_indexing/test_indexing_images.py b/backend/tests/integration/tests/image_indexing/test_indexing_images.py index f7972fefb0..707584a871 100644 --- a/backend/tests/integration/tests/image_indexing/test_indexing_images.py +++ b/backend/tests/integration/tests/image_indexing/test_indexing_images.py @@ -1,6 +1,5 @@ import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest @@ -15,8 +14,7 @@ from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.file import FileManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.settings import SettingsManager -from tests.integration.common_utils.test_models import DATestSettings -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestSettings, DATestUser from tests.integration.common_utils.vespa import vespa_fixture FILE_NAME = "Sample.pdf" diff --git a/backend/tests/integration/tests/indexing/conftest.py b/backend/tests/integration/tests/indexing/conftest.py index d55be5fbc0..719f91b816 100644 --- a/backend/tests/integration/tests/indexing/conftest.py +++ b/backend/tests/integration/tests/indexing/conftest.py @@ -1,8 +1,10 @@ import httpx import pytest -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT +from tests.integration.common_utils.constants import ( + MOCK_CONNECTOR_SERVER_HOST, + MOCK_CONNECTOR_SERVER_PORT, +) @pytest.fixture diff --git a/backend/tests/integration/tests/indexing/file_connector/test_file_connector_zip_metadata.py b/backend/tests/integration/tests/indexing/file_connector/test_file_connector_zip_metadata.py index a503af3711..8764a22267 100644 --- a/backend/tests/integration/tests/indexing/file_connector/test_file_connector_zip_metadata.py +++ b/backend/tests/integration/tests/indexing/file_connector/test_file_connector_zip_metadata.py @@ -1,7 +1,6 @@ import json import os -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest diff --git a/backend/tests/integration/tests/indexing/test_checkpointing.py b/backend/tests/integration/tests/indexing/test_checkpointing.py index 869272a184..8fd00f5c06 100644 --- a/backend/tests/integration/tests/indexing/test_checkpointing.py +++ b/backend/tests/integration/tests/indexing/test_checkpointing.py @@ -1,24 +1,22 @@ import uuid -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import httpx from onyx.configs.constants import DocumentSource from onyx.connectors.mock_connector.connector import MockConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import InputType +from onyx.connectors.models import ConnectorFailure, EntityFailure, InputType from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import IndexingStatus -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT +from tests.integration.common_utils.constants import ( + MOCK_CONNECTOR_SERVER_HOST, + MOCK_CONNECTOR_SERVER_PORT, +) from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager -from tests.integration.common_utils.test_document_utils import create_test_document from tests.integration.common_utils.test_document_utils import ( + create_test_document, create_test_document_failure, ) from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/indexing/test_initial_permission_sync.py b/backend/tests/integration/tests/indexing/test_initial_permission_sync.py index 74e5863e08..cef0debe94 100644 --- a/backend/tests/integration/tests/indexing/test_initial_permission_sync.py +++ b/backend/tests/integration/tests/indexing/test_initial_permission_sync.py @@ -1,32 +1,31 @@ import os import uuid -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import httpx import pytest from sqlalchemy import select from onyx.configs.constants import DocumentSource -from onyx.connectors.mock_connector.connector import EXTERNAL_USER_EMAILS -from onyx.connectors.mock_connector.connector import EXTERNAL_USER_GROUP_IDS -from onyx.connectors.mock_connector.connector import MockConnectorCheckpoint -from onyx.connectors.models import Document -from onyx.connectors.models import InputType +from onyx.connectors.mock_connector.connector import ( + EXTERNAL_USER_EMAILS, + EXTERNAL_USER_GROUP_IDS, + MockConnectorCheckpoint, +) +from onyx.connectors.models import Document, InputType from onyx.db.document import get_documents_by_ids from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import IndexingStatus -from onyx.db.enums import PermissionSyncStatus +from onyx.db.enums import AccessType, IndexingStatus, PermissionSyncStatus from onyx.db.models import DocPermissionSyncAttempt -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT +from tests.integration.common_utils.constants import ( + MOCK_CONNECTOR_SERVER_HOST, + MOCK_CONNECTOR_SERVER_PORT, +) from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager from tests.integration.common_utils.test_document_utils import create_test_document -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestCCPair, DATestUser from tests.integration.common_utils.vespa import vespa_fixture diff --git a/backend/tests/integration/tests/indexing/test_polling.py b/backend/tests/integration/tests/indexing/test_polling.py index 6bb032d254..d016266acd 100644 --- a/backend/tests/integration/tests/indexing/test_polling.py +++ b/backend/tests/integration/tests/indexing/test_polling.py @@ -1,7 +1,5 @@ import uuid -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import httpx @@ -10,8 +8,10 @@ from onyx.configs.constants import DocumentSource from onyx.connectors.mock_connector.connector import MockConnectorCheckpoint from onyx.connectors.models import InputType from onyx.db.enums import IndexingStatus -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT +from tests.integration.common_utils.constants import ( + MOCK_CONNECTOR_SERVER_HOST, + MOCK_CONNECTOR_SERVER_PORT, +) from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager from tests.integration.common_utils.test_document_utils import create_test_document diff --git a/backend/tests/integration/tests/indexing/test_repeated_error_state.py b/backend/tests/integration/tests/indexing/test_repeated_error_state.py index 91566de142..9bed382d74 100644 --- a/backend/tests/integration/tests/indexing/test_repeated_error_state.py +++ b/backend/tests/integration/tests/indexing/test_repeated_error_state.py @@ -12,8 +12,10 @@ from onyx.connectors.models import InputType from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import IndexingStatus -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST -from tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT +from tests.integration.common_utils.constants import ( + MOCK_CONNECTOR_SERVER_HOST, + MOCK_CONNECTOR_SERVER_PORT, +) from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager diff --git a/backend/tests/integration/tests/kg/test_kg_api.py b/backend/tests/integration/tests/kg/test_kg_api.py index c56ddf55ce..7eccc02fd1 100644 --- a/backend/tests/integration/tests/kg/test_kg_api.py +++ b/backend/tests/integration/tests/kg/test_kg_api.py @@ -8,15 +8,16 @@ from onyx.configs.constants import DocumentSource from onyx.connectors.models import InputType from onyx.db.connector import create_connector from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.kg_config import get_kg_config_settings -from onyx.db.kg_config import set_kg_config_settings +from onyx.db.kg_config import get_kg_config_settings, set_kg_config_settings from onyx.db.models import Connector from onyx.server.documents.models import ConnectorBase -from onyx.server.kg.models import DisableKGConfigRequest -from onyx.server.kg.models import EnableKGConfigRequest -from onyx.server.kg.models import EntityType +from onyx.server.kg.models import ( + DisableKGConfigRequest, + EnableKGConfigRequest, + EntityType, + SourceAndEntityTypeView, +) from onyx.server.kg.models import KGConfig as KGConfigAPIModel -from onyx.server.kg.models import SourceAndEntityTypeView from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.user import UserManager diff --git a/backend/tests/integration/tests/llm_provider/test_llm_provider.py b/backend/tests/integration/tests/llm_provider/test_llm_provider.py index b04dbb0da4..bc36143c45 100644 --- a/backend/tests/integration/tests/llm_provider/test_llm_provider.py +++ b/backend/tests/integration/tests/llm_provider/test_llm_provider.py @@ -5,9 +5,11 @@ import httpx import pytest from onyx.llm.constants import LlmProviderNames -from onyx.llm.model_capabilities import get_max_input_tokens -from onyx.llm.model_capabilities import litellm_thinks_model_supports_image_input -from onyx.llm.model_capabilities import model_is_reasoning_model +from onyx.llm.model_capabilities import ( + get_max_input_tokens, + litellm_thinks_model_supports_image_input, + model_is_reasoning_model, +) from onyx.llm.model_name_parser import parse_litellm_model_name from onyx.llm.well_known_providers.llm_provider_options import ( fetch_default_model_for_provider, diff --git a/backend/tests/integration/tests/llm_provider/test_llm_provider_access_control.py b/backend/tests/integration/tests/llm_provider/test_llm_provider_access_control.py index a205069916..95e4b20df0 100644 --- a/backend/tests/integration/tests/llm_provider/test_llm_provider_access_control.py +++ b/backend/tests/integration/tests/llm_provider/test_llm_provider_access_control.py @@ -5,22 +5,28 @@ from sqlalchemy import select from sqlalchemy.orm import Session from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.llm import can_user_access_llm_provider -from onyx.db.llm import fetch_user_group_ids -from onyx.db.llm import update_default_provider -from onyx.db.llm import upsert_llm_provider +from onyx.db.llm import ( + can_user_access_llm_provider, + fetch_user_group_ids, + update_default_provider, + upsert_llm_provider, +) from onyx.db.models import LLMProvider as LLMProviderModel -from onyx.db.models import LLMProvider__Persona -from onyx.db.models import LLMProvider__UserGroup -from onyx.db.models import ModelConfiguration -from onyx.db.models import Persona -from onyx.db.models import User -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import ( + LLMProvider__Persona, + LLMProvider__UserGroup, + ModelConfiguration, + Persona, + User, + User__UserGroup, + UserGroup, +) from onyx.llm.constants import LlmProviderNames from onyx.llm.factory import get_llm_for_persona -from onyx.server.manage.llm.models import LLMProviderUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest +from onyx.server.manage.llm.models import ( + LLMProviderUpsertRequest, + ModelConfigurationUpsertRequest, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.llm_provider import LLMProviderManager diff --git a/backend/tests/integration/tests/llm_workflows/test_nightly_provider_chat_workflow.py b/backend/tests/integration/tests/llm_workflows/test_nightly_provider_chat_workflow.py index 9546f67323..21e4da5b7f 100644 --- a/backend/tests/integration/tests/llm_workflows/test_nightly_provider_chat_workflow.py +++ b/backend/tests/integration/tests/llm_workflows/test_nightly_provider_chat_workflow.py @@ -4,8 +4,7 @@ import time from uuid import uuid4 import pytest -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from onyx.configs import app_configs from onyx.configs.constants import DocumentSource @@ -15,8 +14,7 @@ from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.tool import ToolManager -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import ToolName +from tests.integration.common_utils.test_models import DATestUser, ToolName _ENV_PROVIDER = "NIGHTLY_LLM_PROVIDER" _ENV_MODELS = "NIGHTLY_LLM_MODELS" diff --git a/backend/tests/integration/tests/llm_workflows/test_tool_policy_enforcement.py b/backend/tests/integration/tests/llm_workflows/test_tool_policy_enforcement.py index 4453469c7a..5caa361b8b 100644 --- a/backend/tests/integration/tests/llm_workflows/test_tool_policy_enforcement.py +++ b/backend/tests/integration/tests/llm_workflows/test_tool_policy_enforcement.py @@ -6,8 +6,7 @@ from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.persona import PersonaManager from tests.integration.common_utils.managers.tool import ToolManager -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import ToolName +from tests.integration.common_utils.test_models import DATestUser, ToolName _DUMMY_OPENAI_API_KEY = "sk-mock-tool-policy-tests" diff --git a/backend/tests/integration/tests/mcp/test_mcp_client_no_auth_flow.py b/backend/tests/integration/tests/mcp/test_mcp_client_no_auth_flow.py index 5ec071148a..b47bc9ee22 100644 --- a/backend/tests/integration/tests/mcp/test_mcp_client_no_auth_flow.py +++ b/backend/tests/integration/tests/mcp/test_mcp_client_no_auth_flow.py @@ -8,14 +8,15 @@ from pathlib import Path import pytest -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPTransport +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPTransport, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.persona import PersonaManager -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser # TODO: update mcp client tests to use constants in common_utils/constants.py # NOTE: the tests for client should be independent of the Onyx MCP server diff --git a/backend/tests/integration/tests/mcp/test_mcp_server_search.py b/backend/tests/integration/tests/mcp/test_mcp_server_search.py index cd74caefd8..bc01570515 100644 --- a/backend/tests/integration/tests/mcp/test_mcp_server_search.py +++ b/backend/tests/integration/tests/mcp/test_mcp_server_search.py @@ -5,15 +5,13 @@ from __future__ import annotations import asyncio import json import os -from collections.abc import Awaitable -from collections.abc import Callable +from collections.abc import Awaitable, Callable from typing import Any import pytest from mcp import ClientSession from mcp.client.streamable_http import streamablehttp_client -from mcp.types import CallToolResult -from mcp.types import TextContent +from mcp.types import CallToolResult, TextContent from pydantic import AnyUrl from onyx.db.enums import AccessType diff --git a/backend/tests/integration/tests/migrations/conftest.py b/backend/tests/integration/tests/migrations/conftest.py index abe59f38d1..a7811fa184 100644 --- a/backend/tests/integration/tests/migrations/conftest.py +++ b/backend/tests/integration/tests/migrations/conftest.py @@ -15,16 +15,16 @@ from collections.abc import Generator from typing import Any import pytest -from sqlalchemy import create_engine -from sqlalchemy import text +from sqlalchemy import create_engine, text from sqlalchemy.engine import Engine -from onyx.configs.app_configs import POSTGRES_HOST -from onyx.configs.app_configs import POSTGRES_PASSWORD -from onyx.configs.app_configs import POSTGRES_PORT -from onyx.configs.app_configs import POSTGRES_USER -from onyx.db.engine.sql_engine import build_connection_string -from onyx.db.engine.sql_engine import SYNC_DB_API +from onyx.configs.app_configs import ( + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, +) +from onyx.db.engine.sql_engine import build_connection_string, SYNC_DB_API from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA # Override the parent integration conftest's autouse session fixtures. diff --git a/backend/tests/integration/tests/migrations/test_alembic_main.py b/backend/tests/integration/tests/migrations/test_alembic_main.py index be1a2328f2..c58ab9e659 100644 --- a/backend/tests/integration/tests/migrations/test_alembic_main.py +++ b/backend/tests/integration/tests/migrations/test_alembic_main.py @@ -13,9 +13,11 @@ Usage: See: https://github.com/schireson/pytest-alembic """ -from pytest_alembic.tests import test_single_head_revision -from pytest_alembic.tests import test_up_down_consistency -from pytest_alembic.tests import test_upgrade +from pytest_alembic.tests import ( + test_single_head_revision, + test_up_down_consistency, + test_upgrade, +) __all__ = [ "test_single_head_revision", diff --git a/backend/tests/integration/tests/migrations/test_alembic_tenants.py b/backend/tests/integration/tests/migrations/test_alembic_tenants.py index e47bb6f9d7..90eeca9eca 100644 --- a/backend/tests/integration/tests/migrations/test_alembic_tenants.py +++ b/backend/tests/integration/tests/migrations/test_alembic_tenants.py @@ -16,18 +16,21 @@ from typing import Any import pytest from pytest_alembic import create_alembic_fixture -from pytest_alembic.tests import test_single_head_revision -from pytest_alembic.tests import test_up_down_consistency -from pytest_alembic.tests import test_upgrade +from pytest_alembic.tests import ( + test_single_head_revision, + test_up_down_consistency, + test_upgrade, +) from sqlalchemy import create_engine from sqlalchemy.engine import Engine -from onyx.configs.app_configs import POSTGRES_HOST -from onyx.configs.app_configs import POSTGRES_PASSWORD -from onyx.configs.app_configs import POSTGRES_PORT -from onyx.configs.app_configs import POSTGRES_USER -from onyx.db.engine.sql_engine import build_connection_string -from onyx.db.engine.sql_engine import SYNC_DB_API +from onyx.configs.app_configs import ( + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, +) +from onyx.db.engine.sql_engine import build_connection_string, SYNC_DB_API @pytest.fixture diff --git a/backend/tests/integration/tests/migrations/test_assistant_consolidation_migration.py b/backend/tests/integration/tests/migrations/test_assistant_consolidation_migration.py index a345803538..d4036dcf27 100644 --- a/backend/tests/integration/tests/migrations/test_assistant_consolidation_migration.py +++ b/backend/tests/integration/tests/migrations/test_assistant_consolidation_migration.py @@ -8,8 +8,7 @@ to a single default Assistant (ID 0) and the associated tool seeding. from sqlalchemy import text from onyx.db.engine.sql_engine import get_session_with_current_tenant -from tests.integration.common_utils.reset import downgrade_postgres -from tests.integration.common_utils.reset import upgrade_postgres +from tests.integration.common_utils.reset import downgrade_postgres, upgrade_postgres def test_cold_startup_default_assistant() -> None: diff --git a/backend/tests/integration/tests/migrations/test_migrations.py b/backend/tests/integration/tests/migrations/test_migrations.py index 9e068fa52d..e535ded4f2 100644 --- a/backend/tests/integration/tests/migrations/test_migrations.py +++ b/backend/tests/integration/tests/migrations/test_migrations.py @@ -7,11 +7,9 @@ import json import pytest from sqlalchemy import text -from onyx.configs.constants import ANONYMOUS_USER_UUID -from onyx.configs.constants import DEFAULT_BOOST +from onyx.configs.constants import ANONYMOUS_USER_UUID, DEFAULT_BOOST from onyx.db.engine.sql_engine import get_session_with_current_tenant -from tests.integration.common_utils.reset import downgrade_postgres -from tests.integration.common_utils.reset import upgrade_postgres +from tests.integration.common_utils.reset import downgrade_postgres, upgrade_postgres @pytest.mark.skip( diff --git a/backend/tests/integration/tests/migrations/test_tool_seeding.py b/backend/tests/integration/tests/migrations/test_tool_seeding.py index 63a82b9c54..751fea9ef8 100644 --- a/backend/tests/integration/tests/migrations/test_tool_seeding.py +++ b/backend/tests/integration/tests/migrations/test_tool_seeding.py @@ -2,8 +2,7 @@ from pydantic import BaseModel from sqlalchemy import text from onyx.db.engine.sql_engine import get_session_with_current_tenant -from tests.integration.common_utils.reset import downgrade_postgres -from tests.integration.common_utils.reset import upgrade_postgres +from tests.integration.common_utils.reset import downgrade_postgres, upgrade_postgres class ToolSeedingExpectedResult(BaseModel): diff --git a/backend/tests/integration/tests/mobile_auth/_helpers.py b/backend/tests/integration/tests/mobile_auth/_helpers.py index fa0545d9be..12bb8b421f 100644 --- a/backend/tests/integration/tests/mobile_auth/_helpers.py +++ b/backend/tests/integration/tests/mobile_auth/_helpers.py @@ -2,8 +2,7 @@ import httpx -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client diff --git a/backend/tests/integration/tests/mobile_auth/test_mobile_bearer_auth.py b/backend/tests/integration/tests/mobile_auth/test_mobile_bearer_auth.py index 57d1470464..f68ef96af1 100644 --- a/backend/tests/integration/tests/mobile_auth/test_mobile_bearer_auth.py +++ b/backend/tests/integration/tests/mobile_auth/test_mobile_bearer_auth.py @@ -9,8 +9,7 @@ from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.mobile_auth._helpers import bearer -from tests.integration.tests.mobile_auth._helpers import mobile_login +from tests.integration.tests.mobile_auth._helpers import bearer, mobile_login def test_mobile_bearer_login_refresh_logout(admin_user: DATestUser) -> None: diff --git a/backend/tests/integration/tests/mobile_auth/test_mobile_sso_exchange.py b/backend/tests/integration/tests/mobile_auth/test_mobile_sso_exchange.py index 2bfc7eb35b..bb54a4df6d 100644 --- a/backend/tests/integration/tests/mobile_auth/test_mobile_sso_exchange.py +++ b/backend/tests/integration/tests/mobile_auth/test_mobile_sso_exchange.py @@ -14,12 +14,10 @@ import httpx from onyx.auth.mobile_sso.code_store import store_sso_code from onyx.auth.users import generate_pkce_pair -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser -from tests.integration.tests.mobile_auth._helpers import bearer -from tests.integration.tests.mobile_auth._helpers import mobile_login +from tests.integration.tests.mobile_auth._helpers import bearer, mobile_login def _exchange(code: str, code_verifier: str) -> httpx.Response: diff --git a/backend/tests/integration/tests/no_vectordb/conftest.py b/backend/tests/integration/tests/no_vectordb/conftest.py index f0de3a9cc2..18f0452aaa 100644 --- a/backend/tests/integration/tests/no_vectordb/conftest.py +++ b/backend/tests/integration/tests/no_vectordb/conftest.py @@ -10,8 +10,7 @@ import os import pytest from tests.integration.common_utils.managers.llm_provider import LLMProviderManager -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser # Skip the entire module when vector DB is enabled. The env var is the # source of truth because pytestmark is evaluated at module import time diff --git a/backend/tests/integration/tests/no_vectordb/test_no_vectordb_chat.py b/backend/tests/integration/tests/no_vectordb/test_no_vectordb_chat.py index 5522ed8378..3b81e348bc 100644 --- a/backend/tests/integration/tests/no_vectordb/test_no_vectordb_chat.py +++ b/backend/tests/integration/tests/no_vectordb/test_no_vectordb_chat.py @@ -19,8 +19,7 @@ from tests.integration.common_utils.managers.file import FileManager from tests.integration.common_utils.managers.persona import PersonaManager from tests.integration.common_utils.managers.project import ProjectManager from tests.integration.common_utils.managers.tool import ToolManager -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser FILE_READER_TOOL_ID = "FileReaderTool" diff --git a/backend/tests/integration/tests/no_vectordb/test_no_vectordb_file_lifecycle.py b/backend/tests/integration/tests/no_vectordb/test_no_vectordb_file_lifecycle.py index 437d3e773c..e18dfa90a8 100644 --- a/backend/tests/integration/tests/no_vectordb/test_no_vectordb_file_lifecycle.py +++ b/backend/tests/integration/tests/no_vectordb/test_no_vectordb_file_lifecycle.py @@ -14,8 +14,7 @@ from onyx.db.enums import UserFileStatus from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.project import ProjectManager -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser POLL_INTERVAL_SECONDS = 1 POLL_TIMEOUT_SECONDS = 30 diff --git a/backend/tests/integration/tests/pat/test_pat_api.py b/backend/tests/integration/tests/pat/test_pat_api.py index 25370f9172..093ffef7a9 100644 --- a/backend/tests/integration/tests/pat/test_pat_api.py +++ b/backend/tests/integration/tests/pat/test_pat_api.py @@ -11,9 +11,7 @@ Test Suite: """ import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from onyx.auth.schemas import UserRole from tests.integration.common_utils.constants import API_SERVER_URL diff --git a/backend/tests/integration/tests/permissions/conftest.py b/backend/tests/integration/tests/permissions/conftest.py index 5bd1b45e90..69278b55c7 100644 --- a/backend/tests/integration/tests/permissions/conftest.py +++ b/backend/tests/integration/tests/permissions/conftest.py @@ -18,13 +18,14 @@ import pytest from onyx.auth.schemas import UserRole from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.pat import create_pat -from onyx.db.users import add_slack_user_if_not_exists -from onyx.db.users import batch_add_ext_perm_user_if_not_exists +from onyx.db.users import ( + add_slack_user_if_not_exists, + batch_add_ext_perm_user_if_not_exists, +) from tests.integration.common_utils.managers.api_key import APIKeyManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.reset import reset_all -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser @pytest.fixture(scope="module") diff --git a/backend/tests/integration/tests/permissions/test_admin_access.py b/backend/tests/integration/tests/permissions/test_admin_access.py index d354ec50bc..04d32bc866 100644 --- a/backend/tests/integration/tests/permissions/test_admin_access.py +++ b/backend/tests/integration/tests/permissions/test_admin_access.py @@ -12,8 +12,7 @@ import pytest from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser # Representative endpoints that use require_permission(Permission.FULL_ADMIN_PANEL_ACCESS). # One per major router file to cover breadth without redundancy. diff --git a/backend/tests/integration/tests/permissions/test_basic_access.py b/backend/tests/integration/tests/permissions/test_basic_access.py index 2e4a533eac..42f05e851a 100644 --- a/backend/tests/integration/tests/permissions/test_basic_access.py +++ b/backend/tests/integration/tests/permissions/test_basic_access.py @@ -11,8 +11,7 @@ import pytest from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser # Representative endpoints that use require_permission(Permission.BASIC_ACCESS). # One per major router file to cover breadth without redundancy. diff --git a/backend/tests/integration/tests/permissions/test_cc_pair_permissions.py b/backend/tests/integration/tests/permissions/test_cc_pair_permissions.py index 24f1a1c7d9..da9c275276 100644 --- a/backend/tests/integration/tests/permissions/test_cc_pair_permissions.py +++ b/backend/tests/integration/tests/permissions/test_cc_pair_permissions.py @@ -13,8 +13,7 @@ from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/permissions/test_chat_scopes.py b/backend/tests/integration/tests/permissions/test_chat_scopes.py index b8f9eab12c..3f18cab63b 100644 --- a/backend/tests/integration/tests/permissions/test_chat_scopes.py +++ b/backend/tests/integration/tests/permissions/test_chat_scopes.py @@ -5,8 +5,7 @@ import pytest from onyx.db.enums import Permission from tests.integration.common_utils.http_client import request_status from tests.integration.common_utils.managers.pat import PATManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser _SESSION_ID = "00000000-0000-0000-0000-000000000000" diff --git a/backend/tests/integration/tests/permissions/test_connector_permissions.py b/backend/tests/integration/tests/permissions/test_connector_permissions.py index 6f36bb57a1..0ce6ee26c1 100644 --- a/backend/tests/integration/tests/permissions/test_connector_permissions.py +++ b/backend/tests/integration/tests/permissions/test_connector_permissions.py @@ -11,8 +11,7 @@ import pytest from onyx.db.enums import AccessType from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.managers.connector import ConnectorManager -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/permissions/test_connector_read_access.py b/backend/tests/integration/tests/permissions/test_connector_read_access.py index f4386c3b79..aa00584104 100644 --- a/backend/tests/integration/tests/permissions/test_connector_read_access.py +++ b/backend/tests/integration/tests/permissions/test_connector_read_access.py @@ -12,8 +12,7 @@ import pytest from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser # Connector read endpoints gated on current_curator_or_admin_user. # The by-id path targets a non-existent id on purpose: the auth dependency runs diff --git a/backend/tests/integration/tests/permissions/test_credential_permissions.py b/backend/tests/integration/tests/permissions/test_credential_permissions.py index ffe9b476be..af748c04af 100644 --- a/backend/tests/integration/tests/permissions/test_credential_permissions.py +++ b/backend/tests/integration/tests/permissions/test_credential_permissions.py @@ -10,8 +10,7 @@ import pytest from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.managers.credential import CredentialManager -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/permissions/test_doc_set_permissions.py b/backend/tests/integration/tests/permissions/test_doc_set_permissions.py index e5dca52e0f..622813a510 100644 --- a/backend/tests/integration/tests/permissions/test_doc_set_permissions.py +++ b/backend/tests/integration/tests/permissions/test_doc_set_permissions.py @@ -7,8 +7,7 @@ from onyx.db.enums import AccessType from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.document_set import DocumentSetManager -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/permissions/test_file_connector_permissions.py b/backend/tests/integration/tests/permissions/test_file_connector_permissions.py index 56371972c0..0cd1c55579 100644 --- a/backend/tests/integration/tests/permissions/test_file_connector_permissions.py +++ b/backend/tests/integration/tests/permissions/test_file_connector_permissions.py @@ -13,8 +13,7 @@ from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/permissions/test_persona_permissions.py b/backend/tests/integration/tests/permissions/test_persona_permissions.py index 3155ebd488..e9108fd906 100644 --- a/backend/tests/integration/tests/permissions/test_persona_permissions.py +++ b/backend/tests/integration/tests/permissions/test_persona_permissions.py @@ -11,8 +11,7 @@ import httpx import pytest from tests.integration.common_utils.managers.persona import PersonaManager -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/permissions/test_user_file_permissions.py b/backend/tests/integration/tests/permissions/test_user_file_permissions.py index 4d447266ef..bb156e82ca 100644 --- a/backend/tests/integration/tests/permissions/test_user_file_permissions.py +++ b/backend/tests/integration/tests/permissions/test_user_file_permissions.py @@ -12,19 +12,15 @@ This file tests user file permissions in different scenarios: import io from typing import NamedTuple -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from onyx.configs.constants import FileOrigin from onyx.connectors.models import InputType from onyx.db.engine.sql_engine import get_session_with_current_tenant -from onyx.db.enums import AccessType -from onyx.db.enums import ChatSessionSharedStatus -from onyx.db.models import ChatSession -from onyx.db.models import Document -from onyx.db.models import ToolCall +from onyx.db.enums import AccessType, ChatSessionSharedStatus +from onyx.db.models import ChatSession, Document, ToolCall from onyx.file_store.file_store import get_default_file_store from onyx.file_store.models import FileDescriptor from onyx.server.documents.models import DocumentSource @@ -38,10 +34,12 @@ from tests.integration.common_utils.managers.file import FileManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.persona import PersonaManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestCCPair -from tests.integration.common_utils.test_models import DATestChatSession -from tests.integration.common_utils.test_models import DATestPersona -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestCCPair, + DATestChatSession, + DATestPersona, + DATestUser, +) class UserFileTestSetup(NamedTuple): diff --git a/backend/tests/integration/tests/permissions/test_user_role_permissions.py b/backend/tests/integration/tests/permissions/test_user_role_permissions.py index b39ce8d588..be5efb4f8e 100644 --- a/backend/tests/integration/tests/permissions/test_user_role_permissions.py +++ b/backend/tests/integration/tests/permissions/test_user_role_permissions.py @@ -8,8 +8,7 @@ import httpx import pytest from onyx.db.models import UserRole -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/permissions/test_whole_curator_flow.py b/backend/tests/integration/tests/permissions/test_whole_curator_flow.py index 64ac706ab1..fcdab25fbd 100644 --- a/backend/tests/integration/tests/permissions/test_whole_curator_flow.py +++ b/backend/tests/integration/tests/permissions/test_whole_curator_flow.py @@ -12,8 +12,7 @@ from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/personas/test_persona_categories.py b/backend/tests/integration/tests/personas/test_persona_categories.py index b65397d3e7..ec499a6cc7 100644 --- a/backend/tests/integration/tests/personas/test_persona_categories.py +++ b/backend/tests/integration/tests/personas/test_persona_categories.py @@ -5,8 +5,7 @@ import pytest from tests.integration.common_utils.managers.persona import PersonaLabelManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DATestPersonaLabel -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestPersonaLabel, DATestUser def test_persona_label_management(reset: None) -> None: # noqa: ARG001 diff --git a/backend/tests/integration/tests/personas/test_persona_file_context.py b/backend/tests/integration/tests/personas/test_persona_file_context.py index 9624b00472..d735ae559b 100644 --- a/backend/tests/integration/tests/personas/test_persona_file_context.py +++ b/backend/tests/integration/tests/personas/test_persona_file_context.py @@ -16,16 +16,14 @@ until files reach COMPLETED before chatting. import time from onyx.db.enums import UserFileStatus -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import MAX_DELAY +from tests.integration.common_utils.constants import API_SERVER_URL, MAX_DELAY from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.file import FileManager from tests.integration.common_utils.managers.persona import PersonaManager from tests.integration.common_utils.managers.project import ProjectManager from tests.integration.common_utils.test_file_utils import create_test_text_file -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser # --------------------------------------------------------------------------- # Helpers diff --git a/backend/tests/integration/tests/personas/test_persona_knowledge_sources.py b/backend/tests/integration/tests/personas/test_persona_knowledge_sources.py index 9d4f991474..45d3904b0f 100644 --- a/backend/tests/integration/tests/personas/test_persona_knowledge_sources.py +++ b/backend/tests/integration/tests/personas/test_persona_knowledge_sources.py @@ -13,8 +13,7 @@ from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.file import FileManager from tests.integration.common_utils.managers.persona import PersonaManager from tests.integration.common_utils.test_file_utils import create_test_text_file -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser def _get_minimal_persona( diff --git a/backend/tests/integration/tests/personas/test_persona_label_updates.py b/backend/tests/integration/tests/personas/test_persona_label_updates.py index e170d288ab..4e14685792 100644 --- a/backend/tests/integration/tests/personas/test_persona_label_updates.py +++ b/backend/tests/integration/tests/personas/test_persona_label_updates.py @@ -3,10 +3,11 @@ from uuid import uuid4 from onyx.server.features.persona.models import PersonaUpsertRequest from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.managers.persona import PersonaLabelManager -from tests.integration.common_utils.managers.persona import PersonaManager -from tests.integration.common_utils.test_models import DATestPersonaLabel -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.managers.persona import ( + PersonaLabelManager, + PersonaManager, +) +from tests.integration.common_utils.test_models import DATestPersonaLabel, DATestUser def test_update_persona_with_null_label_ids_preserves_labels( diff --git a/backend/tests/integration/tests/personas/test_persona_pagination.py b/backend/tests/integration/tests/personas/test_persona_pagination.py index 69dd2d7459..7dff2ce3d4 100644 --- a/backend/tests/integration/tests/personas/test_persona_pagination.py +++ b/backend/tests/integration/tests/personas/test_persona_pagination.py @@ -1,5 +1,7 @@ -from onyx.server.features.persona.constants import ADMIN_AGENTS_RESOURCE -from onyx.server.features.persona.constants import AGENTS_RESOURCE +from onyx.server.features.persona.constants import ( + ADMIN_AGENTS_RESOURCE, + AGENTS_RESOURCE, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.persona import PersonaManager diff --git a/backend/tests/integration/tests/projects/test_projects.py b/backend/tests/integration/tests/projects/test_projects.py index 25eec51e6b..9478a1f82f 100644 --- a/backend/tests/integration/tests/projects/test_projects.py +++ b/backend/tests/integration/tests/projects/test_projects.py @@ -7,8 +7,7 @@ from onyx.db.models import UserFile from onyx.server.features.projects.models import UserProjectSnapshot from tests.integration.common_utils.managers.project import ProjectManager from tests.integration.common_utils.reset import reset_all -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser @pytest.fixture(scope="module", autouse=True) diff --git a/backend/tests/integration/tests/pruning/test_pruning.py b/backend/tests/integration/tests/pruning/test_pruning.py index 268b75f703..2f29117a75 100644 --- a/backend/tests/integration/tests/pruning/test_pruning.py +++ b/backend/tests/integration/tests/pruning/test_pruning.py @@ -5,8 +5,7 @@ import tempfile import threading from collections.abc import Generator from contextlib import contextmanager -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from time import sleep from typing import Any diff --git a/backend/tests/integration/tests/query_history/test_query_history.py b/backend/tests/integration/tests/query_history/test_query_history.py index 45e2149041..10498de016 100644 --- a/backend/tests/integration/tests/query_history/test_query_history.py +++ b/backend/tests/integration/tests/query_history/test_query_history.py @@ -1,14 +1,11 @@ import csv import io import os -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import pytest -from onyx.configs.constants import QAFeedbackType -from onyx.configs.constants import SessionType +from onyx.configs.constants import QAFeedbackType, SessionType from tests.integration.common_utils.managers.api_key import APIKeyManager from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.chat import ChatSessionManager diff --git a/backend/tests/integration/tests/query_history/test_query_history_pagination.py b/backend/tests/integration/tests/query_history/test_query_history_pagination.py index f9cc3452de..2d03b6de02 100644 --- a/backend/tests/integration/tests/query_history/test_query_history_pagination.py +++ b/backend/tests/integration/tests/query_history/test_query_history_pagination.py @@ -5,8 +5,7 @@ import pytest from onyx.configs.constants import QAFeedbackType from tests.integration.common_utils.managers.query_history import QueryHistoryManager -from tests.integration.common_utils.test_models import DAQueryHistoryEntry -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DAQueryHistoryEntry, DATestUser from tests.integration.tests.query_history.utils import ( setup_chat_sessions_with_different_feedback, ) diff --git a/backend/tests/integration/tests/query_history/test_usage_reports.py b/backend/tests/integration/tests/query_history/test_usage_reports.py index 9195ca3fe5..1a835f2080 100644 --- a/backend/tests/integration/tests/query_history/test_usage_reports.py +++ b/backend/tests/integration/tests/query_history/test_usage_reports.py @@ -1,6 +1,4 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from ee.onyx.db.usage_export import get_all_empty_chat_message_entries from onyx.db.engine.sql_engine import get_session_with_current_tenant diff --git a/backend/tests/integration/tests/query_history/utils.py b/backend/tests/integration/tests/query_history/utils.py index d2cf98f8ac..e56510f46d 100644 --- a/backend/tests/integration/tests/query_history/utils.py +++ b/backend/tests/integration/tests/query_history/utils.py @@ -1,5 +1,4 @@ -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, ThreadPoolExecutor from onyx.configs.constants import QAFeedbackType from tests.integration.common_utils.managers.api_key import APIKeyManager @@ -8,8 +7,7 @@ from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager -from tests.integration.common_utils.test_models import DAQueryHistoryEntry -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DAQueryHistoryEntry, DATestUser def _create_chat_session_with_feedback( diff --git a/backend/tests/integration/tests/reporting/test_usage_export_api.py b/backend/tests/integration/tests/reporting/test_usage_export_api.py index 0a0cd5000c..6fa01a152a 100644 --- a/backend/tests/integration/tests/reporting/test_usage_export_api.py +++ b/backend/tests/integration/tests/reporting/test_usage_export_api.py @@ -1,11 +1,8 @@ import csv import os import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from io import BytesIO -from io import StringIO +from datetime import datetime, timedelta, timezone +from io import BytesIO, StringIO from uuid import UUID from zipfile import ZipFile diff --git a/backend/tests/integration/tests/scim/test_scim_groups.py b/backend/tests/integration/tests/scim/test_scim_groups.py index 9dda0f03e1..3010c4883e 100644 --- a/backend/tests/integration/tests/scim/test_scim_groups.py +++ b/backend/tests/integration/tests/scim/test_scim_groups.py @@ -21,15 +21,19 @@ import httpx import pytest from onyx.auth.schemas import UserRole -from tests.integration.common_utils.constants import ADMIN_USER_NAME -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import ( + ADMIN_USER_NAME, + API_SERVER_URL, + GENERAL_HEADERS, +) from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.scim_client import ScimClient from tests.integration.common_utils.managers.scim_token import ScimTokenManager -from tests.integration.common_utils.managers.user import build_email -from tests.integration.common_utils.managers.user import DEFAULT_PASSWORD -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import ( + build_email, + DEFAULT_PASSWORD, + UserManager, +) from tests.integration.common_utils.test_models import DATestUser SCIM_GROUP_SCHEMA = "urn:ietf:params:scim:schemas:core:2.0:Group" diff --git a/backend/tests/integration/tests/scim/test_scim_users.py b/backend/tests/integration/tests/scim/test_scim_users.py index 569d8280fd..d1768f00a8 100644 --- a/backend/tests/integration/tests/scim/test_scim_users.py +++ b/backend/tests/integration/tests/scim/test_scim_users.py @@ -20,30 +20,25 @@ are accepted and round-tripped correctly. Auth, revoked-token, and service-discovery tests live in test_scim_tokens.py. """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import httpx import pytest import redis -from ee.onyx.server.license.models import LicenseMetadata -from ee.onyx.server.license.models import LicenseSource -from ee.onyx.server.license.models import PlanType +from ee.onyx.server.license.models import LicenseMetadata, LicenseSource, PlanType from onyx.auth.schemas import UserRole -from onyx.configs.app_configs import REDIS_DB_NUMBER -from onyx.configs.app_configs import REDIS_HOST -from onyx.configs.app_configs import REDIS_PORT +from onyx.configs.app_configs import REDIS_DB_NUMBER, REDIS_HOST, REDIS_PORT from onyx.db.enums import AccountType from onyx.server.settings.models import ApplicationStatus -from tests.integration.common_utils.constants import ADMIN_USER_NAME -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import ADMIN_USER_NAME, GENERAL_HEADERS from tests.integration.common_utils.managers.scim_client import ScimClient from tests.integration.common_utils.managers.scim_token import ScimTokenManager -from tests.integration.common_utils.managers.user import build_email -from tests.integration.common_utils.managers.user import DEFAULT_PASSWORD -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import ( + build_email, + DEFAULT_PASSWORD, + UserManager, +) from tests.integration.common_utils.test_models import DATestUser SCIM_USER_SCHEMA = "urn:ietf:params:scim:schemas:core:2.0:User" @@ -69,11 +64,15 @@ def scim_token(idp_style: str) -> str: per IdP-style run and reuse. Uses UserManager directly to avoid fixture-scope conflicts with the function-scoped admin_user fixture. """ - from tests.integration.common_utils.constants import ADMIN_USER_NAME - from tests.integration.common_utils.constants import GENERAL_HEADERS - from tests.integration.common_utils.managers.user import build_email - from tests.integration.common_utils.managers.user import DEFAULT_PASSWORD - from tests.integration.common_utils.managers.user import UserManager + from tests.integration.common_utils.constants import ( + ADMIN_USER_NAME, + GENERAL_HEADERS, + ) + from tests.integration.common_utils.managers.user import ( + build_email, + DEFAULT_PASSWORD, + UserManager, + ) from tests.integration.common_utils.test_models import DATestUser try: diff --git a/backend/tests/integration/tests/search/test_search_api.py b/backend/tests/integration/tests/search/test_search_api.py index 184c244b24..114f8dd3b3 100644 --- a/backend/tests/integration/tests/search/test_search_api.py +++ b/backend/tests/integration/tests/search/test_search_api.py @@ -7,8 +7,7 @@ import os import httpx import pytest -from onyx.db.enums import AccessType -from onyx.db.enums import Permission +from onyx.db.enums import AccessType, Permission from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.cc_pair import CCPairManager @@ -18,9 +17,11 @@ from tests.integration.common_utils.managers.pat import PATManager from tests.integration.common_utils.managers.persona import PersonaManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import ( + DATestAPIKey, + DATestLLMProvider, + DATestUser, +) SEARCH_URL = f"{API_SERVER_URL}/search" diff --git a/backend/tests/integration/tests/search_settings/test_search_settings.py b/backend/tests/integration/tests/search_settings/test_search_settings.py index 14a9167da5..1308ffd22f 100644 --- a/backend/tests/integration/tests/search_settings/test_search_settings.py +++ b/backend/tests/integration/tests/search_settings/test_search_settings.py @@ -3,8 +3,7 @@ import httpx from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.llm_provider import LLMProviderManager -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestLLMProvider, DATestUser SEARCH_SETTINGS_URL = f"{API_SERVER_URL}/search-settings" diff --git a/backend/tests/integration/tests/skills/test_skills_admin.py b/backend/tests/integration/tests/skills/test_skills_admin.py index 52913287e8..21f50f581c 100644 --- a/backend/tests/integration/tests/skills/test_skills_admin.py +++ b/backend/tests/integration/tests/skills/test_skills_admin.py @@ -2,8 +2,7 @@ from __future__ import annotations import io import zipfile -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx import pytest @@ -13,14 +12,15 @@ from onyx.auth.schemas import UserRole from onyx.configs.constants import FileOrigin from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import SkillSharePermission -from onyx.db.models import FileRecord -from onyx.db.models import Skill +from onyx.db.models import FileRecord, Skill from onyx.file_store.file_store import get_default_file_store from onyx.server.features.skill.models import SkillPatchRequest from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.managers.skill import build_minimal_bundle -from tests.integration.common_utils.managers.skill import SkillManager +from tests.integration.common_utils.managers.skill import ( + build_minimal_bundle, + SkillManager, +) from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/skills/test_skills_management.py b/backend/tests/integration/tests/skills/test_skills_management.py index 84c172efd2..cc9fc24476 100644 --- a/backend/tests/integration/tests/skills/test_skills_management.py +++ b/backend/tests/integration/tests/skills/test_skills_management.py @@ -1,16 +1,16 @@ from __future__ import annotations -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import httpx import pytest -from onyx.db.enums import SkillAccessLevel -from onyx.db.enums import SkillSharePermission -from onyx.server.features.skill.models import SkillPatchRequest -from onyx.server.features.skill.models import SkillResponse -from onyx.server.features.skill.models import SkillUserShareRequest +from onyx.db.enums import SkillAccessLevel, SkillSharePermission +from onyx.server.features.skill.models import ( + SkillPatchRequest, + SkillResponse, + SkillUserShareRequest, +) from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.skill import SkillManager diff --git a/backend/tests/integration/tests/skills/test_skills_personal.py b/backend/tests/integration/tests/skills/test_skills_personal.py index 17bc26ec9c..84b481ab8c 100644 --- a/backend/tests/integration/tests/skills/test_skills_personal.py +++ b/backend/tests/integration/tests/skills/test_skills_personal.py @@ -19,8 +19,10 @@ from onyx.db.enums import SkillSharePermission from onyx.server.features.skill.models import SkillPatchRequest from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.managers.skill import build_minimal_bundle -from tests.integration.common_utils.managers.skill import SkillManager +from tests.integration.common_utils.managers.skill import ( + build_minimal_bundle, + SkillManager, +) from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/streaming_endpoints/test_chat_file_attachment.py b/backend/tests/integration/tests/streaming_endpoints/test_chat_file_attachment.py index d03c7bd60f..2573055378 100644 --- a/backend/tests/integration/tests/streaming_endpoints/test_chat_file_attachment.py +++ b/backend/tests/integration/tests/streaming_endpoints/test_chat_file_attachment.py @@ -6,8 +6,10 @@ from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.file import FileManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager -from tests.integration.common_utils.test_file_utils import create_test_image -from tests.integration.common_utils.test_file_utils import create_test_text_file +from tests.integration.common_utils.test_file_utils import ( + create_test_image, + create_test_text_file, +) from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/streaming_endpoints/test_resume_stream.py b/backend/tests/integration/tests/streaming_endpoints/test_resume_stream.py index 7dbf3792ed..cc9ddc4f34 100644 --- a/backend/tests/integration/tests/streaming_endpoints/test_resume_stream.py +++ b/backend/tests/integration/tests/streaming_endpoints/test_resume_stream.py @@ -19,8 +19,7 @@ from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.chat import ChatSessionManager from tests.integration.common_utils.managers.llm_provider import LLMProviderManager -from tests.integration.common_utils.test_models import DATestChatSession -from tests.integration.common_utils.test_models import DATestUser +from tests.integration.common_utils.test_models import DATestChatSession, DATestUser TERMINATED_RESPONSE_MESSAGE = ( "Response was terminated prior to completion, try regenerating." diff --git a/backend/tests/integration/tests/tools/test_force_tool_use.py b/backend/tests/integration/tests/tools/test_force_tool_use.py index 914520bc6c..d521255a15 100644 --- a/backend/tests/integration/tests/tools/test_force_tool_use.py +++ b/backend/tests/integration/tests/tools/test_force_tool_use.py @@ -9,9 +9,11 @@ from sqlalchemy import select from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.models import Tool from tests.integration.common_utils.managers.chat import ChatSessionManager -from tests.integration.common_utils.test_models import DATestImageGenerationConfig -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import ToolName +from tests.integration.common_utils.test_models import ( + DATestImageGenerationConfig, + DATestUser, + ToolName, +) def test_force_tool_use( diff --git a/backend/tests/integration/tests/tools/test_image_generation_streaming.py b/backend/tests/integration/tests/tools/test_image_generation_streaming.py index 89cbd60740..bc3af7daf7 100644 --- a/backend/tests/integration/tests/tools/test_image_generation_streaming.py +++ b/backend/tests/integration/tests/tools/test_image_generation_streaming.py @@ -12,10 +12,12 @@ from onyx.tools.tool_implementations.images.image_generation_tool import ( HEARTBEAT_INTERVAL, ) from tests.integration.common_utils.managers.chat import ChatSessionManager -from tests.integration.common_utils.test_models import DATestImageGenerationConfig -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import ToolName +from tests.integration.common_utils.test_models import ( + DATestImageGenerationConfig, + DATestLLMProvider, + DATestUser, + ToolName, +) ART_PERSONA_ID = -3 diff --git a/backend/tests/integration/tests/usergroup/test_add_users_to_group.py b/backend/tests/integration/tests/usergroup/test_add_users_to_group.py index 516a833132..a64c65a88a 100644 --- a/backend/tests/integration/tests/usergroup/test_add_users_to_group.py +++ b/backend/tests/integration/tests/usergroup/test_add_users_to_group.py @@ -7,8 +7,7 @@ from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import DATestUserGroup +from tests.integration.common_utils.test_models import DATestUser, DATestUserGroup @pytest.mark.skipif( diff --git a/backend/tests/integration/tests/usergroup/test_curator_group_scope.py b/backend/tests/integration/tests/usergroup/test_curator_group_scope.py index 0271c4d378..f80da47640 100644 --- a/backend/tests/integration/tests/usergroup/test_curator_group_scope.py +++ b/backend/tests/integration/tests/usergroup/test_curator_group_scope.py @@ -23,8 +23,7 @@ import pytest from onyx.db.models import UserRole from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client -from tests.integration.common_utils.managers.user import DATestUser -from tests.integration.common_utils.managers.user import UserManager +from tests.integration.common_utils.managers.user import DATestUser, UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager diff --git a/backend/tests/integration/tests/usergroup/test_group_membership_updates_user_permissions.py b/backend/tests/integration/tests/usergroup/test_group_membership_updates_user_permissions.py index 039af10044..70570c0cdd 100644 --- a/backend/tests/integration/tests/usergroup/test_group_membership_updates_user_permissions.py +++ b/backend/tests/integration/tests/usergroup/test_group_membership_updates_user_permissions.py @@ -6,8 +6,10 @@ from onyx.db.engine.sql_engine import get_session_with_current_tenant from onyx.db.enums import Permission from onyx.db.models import PermissionGrant from onyx.db.models import UserGroup as UserGroupModel -from onyx.db.permissions import recompute_permissions_for_group__no_commit -from onyx.db.permissions import recompute_user_permissions__no_commit +from onyx.db.permissions import ( + recompute_permissions_for_group__no_commit, + recompute_user_permissions__no_commit, +) from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/integration/tests/usergroup/test_user_group_deletion.py b/backend/tests/integration/tests/usergroup/test_user_group_deletion.py index 3eef9ca41f..dcb0f0639f 100644 --- a/backend/tests/integration/tests/usergroup/test_user_group_deletion.py +++ b/backend/tests/integration/tests/usergroup/test_user_group_deletion.py @@ -21,12 +21,14 @@ from tests.integration.common_utils.managers.llm_provider import LLMProviderMana from tests.integration.common_utils.managers.persona import PersonaManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestCredential -from tests.integration.common_utils.test_models import DATestDocumentSet -from tests.integration.common_utils.test_models import DATestLLMProvider -from tests.integration.common_utils.test_models import DATestPersona -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import DATestUserGroup +from tests.integration.common_utils.test_models import ( + DATestCredential, + DATestDocumentSet, + DATestLLMProvider, + DATestPersona, + DATestUser, + DATestUserGroup, +) from tests.integration.common_utils.vespa import vespa_fixture diff --git a/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py b/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py index 9845418d4f..8e524270b2 100644 --- a/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py +++ b/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py @@ -9,9 +9,11 @@ from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.document import DocumentManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.managers.user_group import UserGroupManager -from tests.integration.common_utils.test_models import DATestAPIKey -from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.test_models import DATestUserGroup +from tests.integration.common_utils.test_models import ( + DATestAPIKey, + DATestUser, + DATestUserGroup, +) from tests.integration.common_utils.vespa import vespa_fixture diff --git a/backend/tests/integration/tests/users/test_seat_limit.py b/backend/tests/integration/tests/users/test_seat_limit.py index cd2d7f4c3e..35e432168a 100644 --- a/backend/tests/integration/tests/users/test_seat_limit.py +++ b/backend/tests/integration/tests/users/test_seat_limit.py @@ -4,21 +4,14 @@ Verifies that when a license with a seat limit is active, new user creation (registration, invite, reactivation) is blocked with HTTP 402. """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import redis -from ee.onyx.server.license.models import LicenseMetadata -from ee.onyx.server.license.models import LicenseSource -from ee.onyx.server.license.models import PlanType -from onyx.configs.app_configs import REDIS_DB_NUMBER -from onyx.configs.app_configs import REDIS_HOST -from onyx.configs.app_configs import REDIS_PORT +from ee.onyx.server.license.models import LicenseMetadata, LicenseSource, PlanType +from onyx.configs.app_configs import REDIS_DB_NUMBER, REDIS_HOST, REDIS_PORT from onyx.server.settings.models import ApplicationStatus -from tests.integration.common_utils.constants import API_SERVER_URL -from tests.integration.common_utils.constants import GENERAL_HEADERS +from tests.integration.common_utils.constants import API_SERVER_URL, GENERAL_HEADERS from tests.integration.common_utils.http_client import client from tests.integration.common_utils.managers.user import UserManager diff --git a/backend/tests/integration/tests/users/test_slack_user_deactivation.py b/backend/tests/integration/tests/users/test_slack_user_deactivation.py index 1b0a9b3d84..eb70d8e661 100644 --- a/backend/tests/integration/tests/users/test_slack_user_deactivation.py +++ b/backend/tests/integration/tests/users/test_slack_user_deactivation.py @@ -6,19 +6,13 @@ Verifies that: - Reactivation is blocked when the seat limit is reached """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone import redis -from ee.onyx.server.license.models import LicenseMetadata -from ee.onyx.server.license.models import LicenseSource -from ee.onyx.server.license.models import PlanType +from ee.onyx.server.license.models import LicenseMetadata, LicenseSource, PlanType from onyx.auth.schemas import UserRole -from onyx.configs.app_configs import REDIS_DB_NUMBER -from onyx.configs.app_configs import REDIS_HOST -from onyx.configs.app_configs import REDIS_PORT +from onyx.configs.app_configs import REDIS_DB_NUMBER, REDIS_HOST, REDIS_PORT from onyx.server.settings.models import ApplicationStatus from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client diff --git a/backend/tests/integration/tests/web_search/test_web_search_api.py b/backend/tests/integration/tests/web_search/test_web_search_api.py index e8fb77bea5..eaf36c3dae 100644 --- a/backend/tests/integration/tests/web_search/test_web_search_api.py +++ b/backend/tests/integration/tests/web_search/test_web_search_api.py @@ -2,8 +2,7 @@ import os import pytest -from shared_configs.enums import WebContentProviderType -from shared_configs.enums import WebSearchProviderType +from shared_configs.enums import WebContentProviderType, WebSearchProviderType from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.http_client import client from tests.integration.common_utils.test_models import DATestUser diff --git a/backend/tests/regression/answer_quality/file_uploader.py b/backend/tests/regression/answer_quality/file_uploader.py index c3e900113d..06bf85c787 100644 --- a/backend/tests/regression/answer_quality/file_uploader.py +++ b/backend/tests/regression/answer_quality/file_uploader.py @@ -8,12 +8,14 @@ from types import SimpleNamespace import yaml -from tests.regression.answer_quality.api_utils import check_indexing_status -from tests.regression.answer_quality.api_utils import create_cc_pair -from tests.regression.answer_quality.api_utils import create_connector -from tests.regression.answer_quality.api_utils import create_credential -from tests.regression.answer_quality.api_utils import run_cc_once -from tests.regression.answer_quality.api_utils import upload_file +from tests.regression.answer_quality.api_utils import ( + check_indexing_status, + create_cc_pair, + create_connector, + create_credential, + run_cc_once, + upload_file, +) def unzip_and_get_file_paths(zip_file_path: str) -> list[str]: diff --git a/backend/tests/regression/answer_quality/launch_eval_env.py b/backend/tests/regression/answer_quality/launch_eval_env.py index a3e64056e7..27fbf46130 100644 --- a/backend/tests/regression/answer_quality/launch_eval_env.py +++ b/backend/tests/regression/answer_quality/launch_eval_env.py @@ -3,10 +3,12 @@ from types import SimpleNamespace import yaml -from tests.regression.answer_quality.cli_utils import manage_data_directories -from tests.regression.answer_quality.cli_utils import set_env_variables -from tests.regression.answer_quality.cli_utils import start_docker_compose -from tests.regression.answer_quality.cli_utils import switch_to_commit +from tests.regression.answer_quality.cli_utils import ( + manage_data_directories, + set_env_variables, + start_docker_compose, + switch_to_commit, +) def load_config(config_filename: str) -> SimpleNamespace: diff --git a/backend/tests/regression/search_quality/run_search_eval.py b/backend/tests/regression/search_quality/run_search_eval.py index 398a8944b9..a2e43bf344 100644 --- a/backend/tests/regression/search_quality/run_search_eval.py +++ b/backend/tests/regression/search_quality/run_search_eval.py @@ -4,13 +4,10 @@ import os import sys import time from collections import defaultdict -from concurrent.futures import as_completed -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed, ThreadPoolExecutor from datetime import datetime from pathlib import Path -from threading import Event -from threading import Lock -from threading import Semaphore +from threading import Event, Lock, Semaphore from typing import cast import matplotlib.pyplot as plt @@ -38,28 +35,33 @@ load_dotenv(env_path) # pylint: disable=E402 # flake8: noqa: E402 -from ee.onyx.server.query_and_chat.models import SearchFullResponse -from ee.onyx.server.query_and_chat.models import SendSearchQueryRequest -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW -from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import SavedSearchDoc -from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.engine.sql_engine import SqlEngine +from ee.onyx.server.query_and_chat.models import ( + SearchFullResponse, + SendSearchQueryRequest, +) +from onyx.configs.app_configs import ( + POSTGRES_API_SERVER_POOL_OVERFLOW, + POSTGRES_API_SERVER_POOL_SIZE, +) +from onyx.context.search.models import BaseFilters, SavedSearchDoc +from onyx.db.engine.sql_engine import get_session_with_tenant, SqlEngine from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE -from tests.regression.search_quality.models import AnalysisSummary -from tests.regression.search_quality.models import CombinedMetrics -from tests.regression.search_quality.models import EvalConfig -from tests.regression.search_quality.models import OneshotQAResult -from tests.regression.search_quality.models import TestQuery -from tests.regression.search_quality.utils import compute_overall_scores -from tests.regression.search_quality.utils import find_document_id -from tests.regression.search_quality.utils import get_federated_sources -from tests.regression.search_quality.utils import LazyJsonWriter -from tests.regression.search_quality.utils import ragas_evaluate -from tests.regression.search_quality.utils import search_docs_to_doc_contexts +from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE +from tests.regression.search_quality.models import ( + AnalysisSummary, + CombinedMetrics, + EvalConfig, + OneshotQAResult, + TestQuery, +) +from tests.regression.search_quality.utils import ( + compute_overall_scores, + find_document_id, + get_federated_sources, + LazyJsonWriter, + ragas_evaluate, + search_docs_to_doc_contexts, +) logger = setup_logger(__name__) diff --git a/backend/tests/regression/search_quality/utils.py b/backend/tests/regression/search_quality/utils.py index d6496f75fb..bc34c74a48 100644 --- a/backend/tests/regression/search_quality/utils.py +++ b/backend/tests/regression/search_quality/utils.py @@ -2,33 +2,35 @@ import json import re from pathlib import Path from textwrap import indent -from typing import Any -from typing import cast -from typing import TextIO +from typing import Any, cast, TextIO -from ragas import evaluate # ty: ignore[unresolved-import] -from ragas import EvaluationDataset # ty: ignore[unresolved-import] -from ragas import SingleTurnSample # ty: ignore[unresolved-import] +from ragas import ( # ty: ignore[unresolved-import] + evaluate, + EvaluationDataset, + SingleTurnSample, +) from ragas.dataset_schema import EvaluationResult # ty: ignore[unresolved-import] -from ragas.metrics import FactualCorrectness # ty: ignore[unresolved-import] -from ragas.metrics import Faithfulness # ty: ignore[unresolved-import] -from ragas.metrics import ResponseRelevancy # ty: ignore[unresolved-import] +from ragas.metrics import ( # ty: ignore[unresolved-import] + FactualCorrectness, + Faithfulness, + ResponseRelevancy, +) from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import SavedSearchDoc +from onyx.context.search.models import IndexFilters, SavedSearchDoc from onyx.db.engine.sql_engine import get_session_with_tenant -from onyx.db.models import Document -from onyx.db.models import FederatedConnector +from onyx.db.models import Document, FederatedConnector from onyx.db.search_settings import get_current_search_settings from onyx.document_index.factory import get_default_document_index from onyx.document_index.interfaces_new import DocumentSectionRequest from onyx.prompts.prompt_utils import build_doc_context_str from onyx.utils.logger import setup_logger -from tests.regression.search_quality.models import CombinedMetrics -from tests.regression.search_quality.models import GroundTruth -from tests.regression.search_quality.models import RetrievedDocument +from tests.regression.search_quality.models import ( + CombinedMetrics, + GroundTruth, + RetrievedDocument, +) logger = setup_logger(__name__) diff --git a/backend/tests/unit/background/celery/test_celery_utils.py b/backend/tests/unit/background/celery/test_celery_utils.py index 5fd5da4595..ba9e9742e9 100644 --- a/backend/tests/unit/background/celery/test_celery_utils.py +++ b/backend/tests/unit/background/celery/test_celery_utils.py @@ -8,8 +8,10 @@ import pytest from onyx.background.celery.celery_utils import extract_ids_from_runnable_connector from onyx.connectors.interfaces import SlimConnector from onyx.connectors.models import SlimDocument -from onyx.server.metrics.pruning_metrics import PRUNING_ENUMERATION_DURATION -from onyx.server.metrics.pruning_metrics import PRUNING_RATE_LIMIT_ERRORS +from onyx.server.metrics.pruning_metrics import ( + PRUNING_ENUMERATION_DURATION, + PRUNING_RATE_LIMIT_ERRORS, +) def _make_slim_connector(doc_ids: list[str]) -> SlimConnector: diff --git a/backend/tests/unit/build/test_nextjs_dev.py b/backend/tests/unit/build/test_nextjs_dev.py index a6ffe3e38f..0ce8543d22 100644 --- a/backend/tests/unit/build/test_nextjs_dev.py +++ b/backend/tests/unit/build/test_nextjs_dev.py @@ -3,8 +3,7 @@ from __future__ import annotations from pathlib import Path -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import UUID from onyx.server.features.build.sandbox import nextjs_dev diff --git a/backend/tests/unit/build/test_rewrite_asset_paths.py b/backend/tests/unit/build/test_rewrite_asset_paths.py index 764c13c0fb..15fa630e7d 100644 --- a/backend/tests/unit/build/test_rewrite_asset_paths.py +++ b/backend/tests/unit/build/test_rewrite_asset_paths.py @@ -1,8 +1,7 @@ """Unit tests for the Craft webapp proxy.""" import re -from collections.abc import AsyncGenerator -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator from pathlib import Path from types import SimpleNamespace from typing import cast @@ -10,8 +9,7 @@ from uuid import UUID import httpx import pytest -from fastapi import HTTPException -from fastapi import Request +from fastapi import HTTPException, Request from starlette.responses import StreamingResponse from onyx.db.enums import SharingScope diff --git a/backend/tests/unit/build/test_session_manager_merger.py b/backend/tests/unit/build/test_session_manager_merger.py index 394a1c21b3..c456278089 100644 --- a/backend/tests/unit/build/test_session_manager_merger.py +++ b/backend/tests/unit/build/test_session_manager_merger.py @@ -11,14 +11,15 @@ import time from collections.abc import Generator from typing import Any from unittest.mock import MagicMock -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from onyx.server.features.build.connect_app import ConnectAppRequest -from onyx.server.features.build.packets import ApprovalRequestedPacket -from onyx.server.features.build.packets import ConnectAppRequestPacket +from onyx.server.features.build.packets import ( + ApprovalRequestedPacket, + ConnectAppRequestPacket, +) from onyx.server.features.build.session import streaming as streaming_mod diff --git a/backend/tests/unit/build/test_webapp_proxy_header_stripping.py b/backend/tests/unit/build/test_webapp_proxy_header_stripping.py index 43e25cd816..42860740f8 100644 --- a/backend/tests/unit/build/test_webapp_proxy_header_stripping.py +++ b/backend/tests/unit/build/test_webapp_proxy_header_stripping.py @@ -5,9 +5,7 @@ from __future__ import annotations import asyncio from collections.abc import AsyncGenerator -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 import httpx diff --git a/backend/tests/unit/ee/onyx/db/test_license.py b/backend/tests/unit/ee/onyx/db/test_license.py index 19086a4b0f..773af59847 100644 --- a/backend/tests/unit/ee/onyx/db/test_license.py +++ b/backend/tests/unit/ee/onyx/db/test_license.py @@ -1,19 +1,16 @@ """Tests for license database CRUD operations.""" -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch -from ee.onyx.db.license import check_seat_availability -from ee.onyx.db.license import delete_license -from ee.onyx.db.license import get_license -from ee.onyx.db.license import get_used_seats -from ee.onyx.db.license import upsert_license -from ee.onyx.server.license.models import LicenseMetadata -from ee.onyx.server.license.models import LicenseSource -from ee.onyx.server.license.models import PlanType +from ee.onyx.db.license import ( + check_seat_availability, + delete_license, + get_license, + get_used_seats, + upsert_license, +) +from ee.onyx.server.license.models import LicenseMetadata, LicenseSource, PlanType from onyx.db.models import License from onyx.server.settings.models import ApplicationStatus diff --git a/backend/tests/unit/ee/onyx/db/test_user_group_audit.py b/backend/tests/unit/ee/onyx/db/test_user_group_audit.py index f1ea68c0f4..af7b686e54 100644 --- a/backend/tests/unit/ee/onyx/db/test_user_group_audit.py +++ b/backend/tests/unit/ee/onyx/db/test_user_group_audit.py @@ -7,8 +7,7 @@ must emit nothing while an add/remove of users must emit exactly one event. import json import logging from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest diff --git a/backend/tests/unit/ee/onyx/db/test_user_group_rename.py b/backend/tests/unit/ee/onyx/db/test_user_group_rename.py index b5ad5853f3..3773c92257 100644 --- a/backend/tests/unit/ee/onyx/db/test_user_group_rename.py +++ b/backend/tests/unit/ee/onyx/db/test_user_group_rename.py @@ -1,7 +1,6 @@ """Tests for user group rename DB operation.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/ee/onyx/external_permissions/box/test_box_group_sync.py b/backend/tests/unit/ee/onyx/external_permissions/box/test_box_group_sync.py index 092af73c66..1a6fc621be 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/box/test_box_group_sync.py +++ b/backend/tests/unit/ee/onyx/external_permissions/box/test_box_group_sync.py @@ -1,6 +1,5 @@ from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from box_sdk_gen import BoxClient @@ -10,9 +9,11 @@ from box_sdk_gen.schemas.user_mini import UserMini from ee.onyx.external_permissions.box import group_sync as group_sync_module from ee.onyx.external_permissions.box.group_sync import box_group_sync -from onyx.connectors.box.connector import box_all_enterprise_users_group_id -from onyx.connectors.box.connector import box_group_id -from onyx.connectors.box.connector import BoxConnector +from onyx.connectors.box.connector import ( + box_all_enterprise_users_group_id, + box_group_id, + BoxConnector, +) from onyx.db.models import ConnectorCredentialPair from tests.unit.onyx.connectors.box.fake_box_client import FakeBoxClient diff --git a/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_gdrive_domain_permission_scoping.py b/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_gdrive_domain_permission_scoping.py index a73ca38844..3b39e96bb7 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_gdrive_domain_permission_scoping.py +++ b/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_gdrive_domain_permission_scoping.py @@ -4,23 +4,21 @@ permission paths. The user-side token comes from the synced domain group's real Workspace membership, not from the user's email string.""" from types import SimpleNamespace -from typing import Any -from typing import cast +from typing import Any, cast from unittest.mock import patch from sqlalchemy.orm import Session from ee.onyx.external_permissions.google_drive.doc_sync import ( get_external_access_for_folder, -) -from ee.onyx.external_permissions.google_drive.doc_sync import ( get_external_access_for_raw_gdrive_file, ) -from ee.onyx.external_permissions.google_drive.models import GoogleDrivePermission -from ee.onyx.external_permissions.google_drive.models import PermissionType +from ee.onyx.external_permissions.google_drive.models import ( + GoogleDrivePermission, + PermissionType, +) from onyx.access.models import ExternalAccess -from onyx.access.utils import build_ext_group_name_for_onyx -from onyx.access.utils import prefix_external_group +from onyx.access.utils import build_ext_group_name_for_onyx, prefix_external_group from onyx.configs.constants import DocumentSource from onyx.connectors.google_utils.resources import GoogleDriveService from onyx.db.models import User diff --git a/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_get_drive_members_admin_403.py b/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_get_drive_members_admin_403.py index b5d605704c..df3c65f34a 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_get_drive_members_admin_403.py +++ b/backend/tests/unit/ee/onyx/external_permissions/google_drive/test_get_drive_members_admin_403.py @@ -1,5 +1,4 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from googleapiclient.errors import HttpError diff --git a/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_group_sync.py b/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_group_sync.py index 16a4c34f28..6a4c1c19a1 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_group_sync.py +++ b/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_group_sync.py @@ -5,8 +5,10 @@ import pytest from jira import JIRA from jira.exceptions import JIRAError -from ee.onyx.external_permissions.jira.group_sync import _fetch_group_member_page -from ee.onyx.external_permissions.jira.group_sync import _get_group_member_emails +from ee.onyx.external_permissions.jira.group_sync import ( + _fetch_group_member_page, + _get_group_member_emails, +) def test_get_group_member_emails_skips_deleted_group( diff --git a/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_page_access.py b/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_page_access.py index 24bfbe3ad1..f7100361c5 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_page_access.py +++ b/backend/tests/unit/ee/onyx/external_permissions/jira/test_jira_page_access.py @@ -5,8 +5,7 @@ from unittest.mock import MagicMock from pytest import LogCaptureFixture from ee.onyx.external_permissions.jira.page_access import get_project_permissions -from onyx.connectors.jira.utils import JIRA_CLOUD_API_VERSION -from onyx.connectors.jira.utils import JIRA_SERVER_API_VERSION +from onyx.connectors.jira.utils import JIRA_CLOUD_API_VERSION, JIRA_SERVER_API_VERSION PROJECT_KEY = "PROJ" diff --git a/backend/tests/unit/ee/onyx/external_permissions/salesforce/test_salesforce_perm_cache_tenant_isolation.py b/backend/tests/unit/ee/onyx/external_permissions/salesforce/test_salesforce_perm_cache_tenant_isolation.py index cd98d45da3..993d7e3b68 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/salesforce/test_salesforce_perm_cache_tenant_isolation.py +++ b/backend/tests/unit/ee/onyx/external_permissions/salesforce/test_salesforce_perm_cache_tenant_isolation.py @@ -2,8 +2,7 @@ caches: the per-tenant Salesforce client and the email->user_id map must never let one tenant's org, credentials, or resolved user id serve another tenant's access check.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import ee.onyx.external_permissions.salesforce.utils as sf_utils from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/unit/ee/onyx/external_permissions/sharepoint/test_permission_utils.py b/backend/tests/unit/ee/onyx/external_permissions/sharepoint/test_permission_utils.py index 0882e75ffe..5b3da79dc8 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/sharepoint/test_permission_utils.py +++ b/backend/tests/unit/ee/onyx/external_permissions/sharepoint/test_permission_utils.py @@ -1,33 +1,20 @@ from collections.abc import Generator from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from ee.onyx.external_permissions.sharepoint.permission_utils import ( _enumerate_ad_groups_paginated, -) -from ee.onyx.external_permissions.sharepoint.permission_utils import _get_azuread_groups -from ee.onyx.external_permissions.sharepoint.permission_utils import ( + _get_azuread_groups, _has_only_limited_access, -) -from ee.onyx.external_permissions.sharepoint.permission_utils import _is_public_item -from ee.onyx.external_permissions.sharepoint.permission_utils import ( + _is_public_item, _iter_graph_collection, -) -from ee.onyx.external_permissions.sharepoint.permission_utils import _normalize_email -from ee.onyx.external_permissions.sharepoint.permission_utils import ( + _normalize_email, AD_GROUP_ENUMERATION_THRESHOLD, -) -from ee.onyx.external_permissions.sharepoint.permission_utils import ( get_external_access_from_sharepoint, -) -from ee.onyx.external_permissions.sharepoint.permission_utils import ( get_sharepoint_external_groups, -) -from ee.onyx.external_permissions.sharepoint.permission_utils import GroupsResult -from ee.onyx.external_permissions.sharepoint.permission_utils import ( + GroupsResult, SHAREPOINT_GROUP_PRINCIPAL_TYPE, ) diff --git a/backend/tests/unit/ee/onyx/external_permissions/slack/test_grid_perm_sync.py b/backend/tests/unit/ee/onyx/external_permissions/slack/test_grid_perm_sync.py index 86b35f4d9f..58a494be80 100644 --- a/backend/tests/unit/ee/onyx/external_permissions/slack/test_grid_perm_sync.py +++ b/backend/tests/unit/ee/onyx/external_permissions/slack/test_grid_perm_sync.py @@ -1,18 +1,20 @@ """Unit tests for EE Slack perm sync on Enterprise Grid.""" -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch from ee.onyx.external_permissions.slack.channel_access import ( get_channel_access as ee_get_channel_access, ) -from ee.onyx.external_permissions.slack.doc_sync import _fetch_channel_permissions -from ee.onyx.external_permissions.slack.doc_sync import _fetch_workspace_permissions -from ee.onyx.external_permissions.slack.doc_sync import _get_slack_document_access -from ee.onyx.external_permissions.slack.utils import fetch_team_user_emails -from ee.onyx.external_permissions.slack.utils import fetch_user_id_to_email_map +from ee.onyx.external_permissions.slack.doc_sync import ( + _fetch_channel_permissions, + _fetch_workspace_permissions, + _get_slack_document_access, +) +from ee.onyx.external_permissions.slack.utils import ( + fetch_team_user_emails, + fetch_user_id_to_email_map, +) from onyx.access.models import ExternalAccess from onyx.connectors.models import SlimDocument from onyx.connectors.slack.models import ChannelType diff --git a/backend/tests/unit/ee/onyx/hooks/test_executor.py b/backend/tests/unit/ee/onyx/hooks/test_executor.py index ce4dd5485e..6da010cacd 100644 --- a/backend/tests/unit/ee/onyx/hooks/test_executor.py +++ b/backend/tests/unit/ee/onyx/hooks/test_executor.py @@ -2,20 +2,17 @@ import json from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import httpx import pytest from pydantic import BaseModel from ee.onyx.hooks.executor import _execute_hook_impl as execute_hook -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint +from onyx.db.enums import HookFailStrategy, HookPoint from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.hooks.executor import HookSkipped -from onyx.hooks.executor import HookSoftFailed +from onyx.hooks.executor import HookSkipped, HookSoftFailed from onyx.hooks.points.query_processing import QueryProcessingResponse # --------------------------------------------------------------------------- diff --git a/backend/tests/unit/ee/onyx/server/billing/conftest.py b/backend/tests/unit/ee/onyx/server/billing/conftest.py index c1ae0c6ba2..65a04941d8 100644 --- a/backend/tests/unit/ee/onyx/server/billing/conftest.py +++ b/backend/tests/unit/ee/onyx/server/billing/conftest.py @@ -1,14 +1,11 @@ """Shared fixtures and utilities for billing tests.""" -from datetime import datetime -from datetime import timezone -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock import pytest -from ee.onyx.server.license.models import LicensePayload -from ee.onyx.server.license.models import PlanType +from ee.onyx.server.license.models import LicensePayload, PlanType @pytest.fixture diff --git a/backend/tests/unit/ee/onyx/server/billing/test_billing_api.py b/backend/tests/unit/ee/onyx/server/billing/test_billing_api.py index d097658984..529ae0bb1d 100644 --- a/backend/tests/unit/ee/onyx/server/billing/test_billing_api.py +++ b/backend/tests/unit/ee/onyx/server/billing/test_billing_api.py @@ -1,16 +1,16 @@ """Tests for the unified billing API endpoints.""" -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from ee.onyx.server.billing.models import BillingInformationResponse -from ee.onyx.server.billing.models import CreateCheckoutSessionResponse -from ee.onyx.server.billing.models import CreateCustomerPortalSessionResponse -from ee.onyx.server.billing.models import SeatUpdateResponse -from ee.onyx.server.billing.models import SubscriptionStatusResponse +from ee.onyx.server.billing.models import ( + BillingInformationResponse, + CreateCheckoutSessionResponse, + CreateCustomerPortalSessionResponse, + SeatUpdateResponse, + SubscriptionStatusResponse, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/tests/unit/ee/onyx/server/billing/test_billing_cache.py b/backend/tests/unit/ee/onyx/server/billing/test_billing_cache.py index b9af6e82fb..0a3ddccb96 100644 --- a/backend/tests/unit/ee/onyx/server/billing/test_billing_cache.py +++ b/backend/tests/unit/ee/onyx/server/billing/test_billing_cache.py @@ -1,22 +1,20 @@ """Unit tests for the Redis-backed billing cache.""" from collections.abc import Generator -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch import pytest from redis.exceptions import RedisError from ee.onyx.server.billing import billing_cache as bc -from ee.onyx.server.billing.billing_cache import BILLING_CACHE_KEY -from ee.onyx.server.billing.billing_cache import cached_fetch_billing_information -from ee.onyx.server.billing.billing_cache import cached_is_tenant_on_trial -from ee.onyx.server.billing.billing_cache import invalidate_billing_cache -from ee.onyx.server.tenants.models import BillingInformation -from ee.onyx.server.tenants.models import SubscriptionStatusResponse +from ee.onyx.server.billing.billing_cache import ( + BILLING_CACHE_KEY, + cached_fetch_billing_information, + cached_is_tenant_on_trial, + invalidate_billing_cache, +) +from ee.onyx.server.tenants.models import BillingInformation, SubscriptionStatusResponse def _billing(status: str = "active") -> BillingInformation: diff --git a/backend/tests/unit/ee/onyx/server/billing/test_billing_service.py b/backend/tests/unit/ee/onyx/server/billing/test_billing_service.py index 51a94f8f6e..7ede32b063 100644 --- a/backend/tests/unit/ee/onyx/server/billing/test_billing_service.py +++ b/backend/tests/unit/ee/onyx/server/billing/test_billing_service.py @@ -1,22 +1,21 @@ """Tests for the billing service layer.""" -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from ee.onyx.server.billing.models import BillingInformationResponse -from ee.onyx.server.billing.models import CreateCheckoutSessionResponse -from ee.onyx.server.billing.models import CreateCustomerPortalSessionResponse -from ee.onyx.server.billing.models import SeatUpdateResponse -from ee.onyx.server.billing.models import SubscriptionStatusResponse +from ee.onyx.server.billing.models import ( + BillingInformationResponse, + CreateCheckoutSessionResponse, + CreateCustomerPortalSessionResponse, + SeatUpdateResponse, + SubscriptionStatusResponse, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from .conftest import make_mock_http_client -from .conftest import make_mock_response +from .conftest import make_mock_http_client, make_mock_response class TestMakeBillingRequest: diff --git a/backend/tests/unit/ee/onyx/server/billing/test_proxy.py b/backend/tests/unit/ee/onyx/server/billing/test_proxy.py index 4ba0e00e7e..b82fd6ad1b 100644 --- a/backend/tests/unit/ee/onyx/server/billing/test_proxy.py +++ b/backend/tests/unit/ee/onyx/server/billing/test_proxy.py @@ -1,17 +1,13 @@ """Tests for the billing proxy endpoints.""" -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from ee.onyx.server.license.models import LicensePayload -from .conftest import make_license_payload -from .conftest import make_mock_http_client -from .conftest import make_mock_response +from .conftest import make_license_payload, make_mock_http_client, make_mock_response class TestProxySeatUpdate: diff --git a/backend/tests/unit/ee/onyx/server/features/hooks/test_api.py b/backend/tests/unit/ee/onyx/server/features/hooks/test_api.py index 3c96916288..2ec7e904ff 100644 --- a/backend/tests/unit/ee/onyx/server/features/hooks/test_api.py +++ b/backend/tests/unit/ee/onyx/server/features/hooks/test_api.py @@ -10,19 +10,19 @@ Covers: - _raise_for_validation_failure: HookValidateStatus → OnyxError mapping """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import httpx import pytest -from ee.onyx.server.features.hooks.api import _check_ssrf_safety -from ee.onyx.server.features.hooks.api import _raise_for_validation_failure -from ee.onyx.server.features.hooks.api import _validate_endpoint +from ee.onyx.server.features.hooks.api import ( + _check_ssrf_safety, + _raise_for_validation_failure, + _validate_endpoint, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.hooks.models import HookValidateResponse -from onyx.hooks.models import HookValidateStatus +from onyx.hooks.models import HookValidateResponse, HookValidateStatus # --------------------------------------------------------------------------- # Helpers diff --git a/backend/tests/unit/ee/onyx/server/metrics/test_license_metrics.py b/backend/tests/unit/ee/onyx/server/metrics/test_license_metrics.py index 813789254c..f697abd78a 100644 --- a/backend/tests/unit/ee/onyx/server/metrics/test_license_metrics.py +++ b/backend/tests/unit/ee/onyx/server/metrics/test_license_metrics.py @@ -1,9 +1,7 @@ """Tests for the self-hosted license Prometheus collector.""" -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch from prometheus_client.core import GaugeMetricFamily diff --git a/backend/tests/unit/ee/onyx/server/middleware/test_license_enforcement.py b/backend/tests/unit/ee/onyx/server/middleware/test_license_enforcement.py index e065dcdd1a..163ded2d2b 100644 --- a/backend/tests/unit/ee/onyx/server/middleware/test_license_enforcement.py +++ b/backend/tests/unit/ee/onyx/server/middleware/test_license_enforcement.py @@ -1,10 +1,8 @@ """Tests for license enforcement middleware.""" -from collections.abc import Awaitable -from collections.abc import Callable +from collections.abc import Awaitable, Callable from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from starlette.requests import Request diff --git a/backend/tests/unit/ee/onyx/server/middleware/test_tenant_tracking.py b/backend/tests/unit/ee/onyx/server/middleware/test_tenant_tracking.py index d7ce3c298f..61f36ebd34 100644 --- a/backend/tests/unit/ee/onyx/server/middleware/test_tenant_tracking.py +++ b/backend/tests/unit/ee/onyx/server/middleware/test_tenant_tracking.py @@ -7,9 +7,7 @@ served against the wrong tenant. (Contrast license enforcement, which fails *open*; tenant routing is a data-isolation boundary and must not.) """ -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException diff --git a/backend/tests/unit/ee/onyx/server/middleware/test_tier_gate.py b/backend/tests/unit/ee/onyx/server/middleware/test_tier_gate.py index e42b56c907..39fc6bf809 100644 --- a/backend/tests/unit/ee/onyx/server/middleware/test_tier_gate.py +++ b/backend/tests/unit/ee/onyx/server/middleware/test_tier_gate.py @@ -1,10 +1,8 @@ """Tests for the unified tier_gate middleware.""" -from collections.abc import Awaitable -from collections.abc import Callable +from collections.abc import Awaitable, Callable from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from starlette.requests import Request diff --git a/backend/tests/unit/ee/onyx/server/settings/test_license_enforcement_settings.py b/backend/tests/unit/ee/onyx/server/settings/test_license_enforcement_settings.py index b23a980097..3b70668db6 100644 --- a/backend/tests/unit/ee/onyx/server/settings/test_license_enforcement_settings.py +++ b/backend/tests/unit/ee/onyx/server/settings/test_license_enforcement_settings.py @@ -1,13 +1,11 @@ """Tests for license enforcement in settings API.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from redis.exceptions import RedisError -from onyx.server.settings.models import ApplicationStatus -from onyx.server.settings.models import Settings +from onyx.server.settings.models import ApplicationStatus, Settings # Fields we assert on across all tests _ASSERT_FIELDS = { diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_billing_api.py b/backend/tests/unit/ee/onyx/server/tenants/test_billing_api.py index 64d7648534..2edd6345a3 100644 --- a/backend/tests/unit/ee/onyx/server/tenants/test_billing_api.py +++ b/backend/tests/unit/ee/onyx/server/tenants/test_billing_api.py @@ -1,8 +1,6 @@ """Tests for billing API endpoints.""" -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py b/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py index d578153bd8..5faef67b38 100644 --- a/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py +++ b/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py @@ -7,16 +7,17 @@ surfacing as ``OnyxError(SEAT_LIMIT_EXCEEDED)``). from __future__ import annotations -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import stripe -from ee.onyx.server.tenants.billing import _seat_billing_idempotency_key -from ee.onyx.server.tenants.billing import attempt_seat_billing_increase -from ee.onyx.server.tenants.billing import enforce_cloud_seat_limit -from ee.onyx.server.tenants.billing import SeatBillingDeclineReason +from ee.onyx.server.tenants.billing import ( + _seat_billing_idempotency_key, + attempt_seat_billing_increase, + enforce_cloud_seat_limit, + SeatBillingDeclineReason, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_product_gating.py b/backend/tests/unit/ee/onyx/server/tenants/test_product_gating.py index 467de680b7..362d76f16b 100644 --- a/backend/tests/unit/ee/onyx/server/tenants/test_product_gating.py +++ b/backend/tests/unit/ee/onyx/server/tenants/test_product_gating.py @@ -1,7 +1,6 @@ """Tests for product gating functions.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_provisioning.py b/backend/tests/unit/ee/onyx/server/tenants/test_provisioning.py index a6124300e7..11d2ea8759 100644 --- a/backend/tests/unit/ee/onyx/server/tenants/test_provisioning.py +++ b/backend/tests/unit/ee/onyx/server/tenants/test_provisioning.py @@ -1,7 +1,6 @@ """Tests for tenant provisioning helpers.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch _PROVISIONING_MODULE = "ee.onyx.server.tenants.provisioning" diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_proxy.py b/backend/tests/unit/ee/onyx/server/tenants/test_proxy.py index d595c1f192..ff60af26cd 100644 --- a/backend/tests/unit/ee/onyx/server/tenants/test_proxy.py +++ b/backend/tests/unit/ee/onyx/server/tenants/test_proxy.py @@ -1,25 +1,22 @@ """Tests for proxy endpoints for self-hosted data planes.""" -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from fastapi import HTTPException -from ee.onyx.server.license.models import LicensePayload -from ee.onyx.server.license.models import PlanType -from ee.onyx.server.tenants.proxy import _check_license_enforcement_enabled -from ee.onyx.server.tenants.proxy import _extract_license_from_header -from ee.onyx.server.tenants.proxy import forward_to_control_plane -from ee.onyx.server.tenants.proxy import get_license_payload -from ee.onyx.server.tenants.proxy import get_license_payload_allow_expired -from ee.onyx.server.tenants.proxy import get_optional_license_payload -from ee.onyx.server.tenants.proxy import verify_license_auth +from ee.onyx.server.license.models import LicensePayload, PlanType +from ee.onyx.server.tenants.proxy import ( + _check_license_enforcement_enabled, + _extract_license_from_header, + forward_to_control_plane, + get_license_payload, + get_license_payload_allow_expired, + get_optional_license_payload, + verify_license_auth, +) # All tests that use license auth need LICENSE_ENFORCEMENT_ENABLED=True LICENSE_ENABLED_PATCH = patch( @@ -482,8 +479,10 @@ class TestProxyCheckoutSessionWithSeats: @pytest.mark.asyncio async def test_includes_seats_in_body_when_provided(self) -> None: """Should include seats in request body when provided.""" - from ee.onyx.server.tenants.proxy import CreateCheckoutSessionRequest - from ee.onyx.server.tenants.proxy import proxy_create_checkout_session + from ee.onyx.server.tenants.proxy import ( + CreateCheckoutSessionRequest, + proxy_create_checkout_session, + ) mock_response = MagicMock() mock_response.json.return_value = {"url": "https://checkout.stripe.com/session"} @@ -527,8 +526,10 @@ class TestProxyCheckoutSessionWithSeats: @pytest.mark.asyncio async def test_excludes_seats_when_not_provided(self) -> None: """Should not include seats in request body when not provided.""" - from ee.onyx.server.tenants.proxy import CreateCheckoutSessionRequest - from ee.onyx.server.tenants.proxy import proxy_create_checkout_session + from ee.onyx.server.tenants.proxy import ( + CreateCheckoutSessionRequest, + proxy_create_checkout_session, + ) mock_response = MagicMock() mock_response.json.return_value = {"url": "https://checkout.stripe.com/session"} @@ -566,8 +567,10 @@ class TestProxyCheckoutSessionWithSeats: @pytest.mark.asyncio async def test_includes_seats_for_new_customer(self) -> None: """Should include seats for new customer without license.""" - from ee.onyx.server.tenants.proxy import CreateCheckoutSessionRequest - from ee.onyx.server.tenants.proxy import proxy_create_checkout_session + from ee.onyx.server.tenants.proxy import ( + CreateCheckoutSessionRequest, + proxy_create_checkout_session, + ) mock_response = MagicMock() mock_response.json.return_value = {"url": "https://checkout.stripe.com/session"} diff --git a/backend/tests/unit/ee/onyx/utils/test_encryption.py b/backend/tests/unit/ee/onyx/utils/test_encryption.py index 1e087907df..b8cb135ae5 100644 --- a/backend/tests/unit/ee/onyx/utils/test_encryption.py +++ b/backend/tests/unit/ee/onyx/utils/test_encryption.py @@ -8,11 +8,13 @@ from unittest.mock import patch import pytest -from ee.onyx.utils.encryption import _decrypt_bytes -from ee.onyx.utils.encryption import _encrypt_string -from ee.onyx.utils.encryption import _get_trimmed_key -from ee.onyx.utils.encryption import decrypt_bytes_to_string -from ee.onyx.utils.encryption import encrypt_string_to_bytes +from ee.onyx.utils.encryption import ( + _decrypt_bytes, + _encrypt_string, + _get_trimmed_key, + decrypt_bytes_to_string, + encrypt_string_to_bytes, +) EE_MODULE = "ee.onyx.utils.encryption" diff --git a/backend/tests/unit/ee/onyx/utils/test_license_expiry.py b/backend/tests/unit/ee/onyx/utils/test_license_expiry.py index fd59f4ac98..13c0c0b28e 100644 --- a/backend/tests/unit/ee/onyx/utils/test_license_expiry.py +++ b/backend/tests/unit/ee/onyx/utils/test_license_expiry.py @@ -1,15 +1,15 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from unittest.mock import patch import pytest -from ee.onyx.utils.license_expiry import ExpiryWarningStage -from ee.onyx.utils.license_expiry import get_expiry_warning_stage -from ee.onyx.utils.license_expiry import get_grace_days_remaining -from ee.onyx.utils.license_expiry import get_grace_period_end -from ee.onyx.utils.license_expiry import LICENSE_GRACE_PERIOD_DAYS +from ee.onyx.utils.license_expiry import ( + ExpiryWarningStage, + get_expiry_warning_stage, + get_grace_days_remaining, + get_grace_period_end, + LICENSE_GRACE_PERIOD_DAYS, +) NOW = datetime(2026, 5, 1, 12, 0, 0, tzinfo=timezone.utc) diff --git a/backend/tests/unit/ee/onyx/utils/test_license_utils.py b/backend/tests/unit/ee/onyx/utils/test_license_utils.py index 5cc70cb1a1..9aeafd689b 100644 --- a/backend/tests/unit/ee/onyx/utils/test_license_utils.py +++ b/backend/tests/unit/ee/onyx/utils/test_license_utils.py @@ -2,21 +2,19 @@ import base64 import json -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from unittest.mock import patch import pytest from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives.asymmetric import padding -from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.hazmat.primitives.asymmetric import padding, rsa -from ee.onyx.server.license.models import LicensePayload -from ee.onyx.server.license.models import PlanType -from ee.onyx.utils.license import get_license_status -from ee.onyx.utils.license import is_license_valid -from ee.onyx.utils.license import verify_license_signature +from ee.onyx.server.license.models import LicensePayload, PlanType +from ee.onyx.utils.license import ( + get_license_status, + is_license_valid, + verify_license_signature, +) from onyx.server.settings.models import ApplicationStatus diff --git a/backend/tests/unit/ee/onyx/utils/test_tier.py b/backend/tests/unit/ee/onyx/utils/test_tier.py index a5a702ceca..b8b4eb4b76 100644 --- a/backend/tests/unit/ee/onyx/utils/test_tier.py +++ b/backend/tests/unit/ee/onyx/utils/test_tier.py @@ -2,8 +2,7 @@ (cache-failure resilience) and the tier-requirement guards. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from redis.exceptions import RedisError @@ -13,8 +12,7 @@ from ee.onyx.server.license.models import CustomerTier from onyx.db.enums import AccessType from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.settings.models import ApplicationStatus -from onyx.server.settings.models import Tier +from onyx.server.settings.models import ApplicationStatus, Tier def _metadata( diff --git a/backend/tests/unit/external_apps/matching/test_approval_actions.py b/backend/tests/unit/external_apps/matching/test_approval_actions.py index e97b3bf558..df4a4875ea 100644 --- a/backend/tests/unit/external_apps/matching/test_approval_actions.py +++ b/backend/tests/unit/external_apps/matching/test_approval_actions.py @@ -1,6 +1,5 @@ from onyx.db.enums import EndpointPolicy -from onyx.external_apps.matching.engine import actions_requiring_approval -from onyx.external_apps.matching.engine import MatchedAction +from onyx.external_apps.matching.engine import actions_requiring_approval, MatchedAction def _action(action_type: str, policy: EndpointPolicy) -> MatchedAction: diff --git a/backend/tests/unit/external_apps/matching/test_engine.py b/backend/tests/unit/external_apps/matching/test_engine.py index 25f9c28881..baf096ced2 100644 --- a/backend/tests/unit/external_apps/matching/test_engine.py +++ b/backend/tests/unit/external_apps/matching/test_engine.py @@ -6,13 +6,14 @@ from __future__ import annotations import pytest from pydantic import ValidationError -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.db.models import ExternalApp -from onyx.external_apps.matching.engine import AllMatchedActions -from onyx.external_apps.matching.engine import apply_credential_gate -from onyx.external_apps.matching.engine import MatchedAction -from onyx.external_apps.matching.engine import WHOLE_DOMAIN_ACTION_TYPE +from onyx.external_apps.matching.engine import ( + AllMatchedActions, + apply_credential_gate, + MatchedAction, + WHOLE_DOMAIN_ACTION_TYPE, +) from onyx.external_apps.matching.request import ProxiedRequest diff --git a/backend/tests/unit/external_apps/presentation/test_payload_decoders.py b/backend/tests/unit/external_apps/presentation/test_payload_decoders.py index 7ac6a2ffb8..7ac9a4c604 100644 --- a/backend/tests/unit/external_apps/presentation/test_payload_decoders.py +++ b/backend/tests/unit/external_apps/presentation/test_payload_decoders.py @@ -6,8 +6,7 @@ from typing import Any from onyx.external_apps.presentation.decode import decode_payload from onyx.external_apps.presentation.payload_decoders import GmailRawMimeDecoder -from onyx.external_apps.providers.gmail import GmailAction -from onyx.external_apps.providers.gmail import GmailProvider +from onyx.external_apps.providers.gmail import GmailAction, GmailProvider # The `messages.send` decoder; draft create/update wrap the MIME under `message`. _SEND_DECODER = GmailRawMimeDecoder() diff --git a/backend/tests/unit/external_apps/test_action_policy_defaults.py b/backend/tests/unit/external_apps/test_action_policy_defaults.py index 6e856e297a..9885f62180 100644 --- a/backend/tests/unit/external_apps/test_action_policy_defaults.py +++ b/backend/tests/unit/external_apps/test_action_policy_defaults.py @@ -9,12 +9,13 @@ from __future__ import annotations import pytest -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType +from onyx.db.enums import EndpointPolicy, ExternalAppType from onyx.external_apps.providers import registry -from onyx.external_apps.providers.actions import EndpointSpec -from onyx.external_apps.providers.actions import ExternalAppAction -from onyx.external_apps.providers.actions import RestRoute +from onyx.external_apps.providers.actions import ( + EndpointSpec, + ExternalAppAction, + RestRoute, +) class _TestAction(ExternalAppAction): diff --git a/backend/tests/unit/external_apps/test_github_catalog.py b/backend/tests/unit/external_apps/test_github_catalog.py index 468d5bfb29..86dc483649 100644 --- a/backend/tests/unit/external_apps/test_github_catalog.py +++ b/backend/tests/unit/external_apps/test_github_catalog.py @@ -12,10 +12,8 @@ import json import pytest -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.matching.request import MatchContext -from onyx.external_apps.matching.request import ProxiedRequest +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.matching.request import MatchContext, ProxiedRequest from onyx.external_apps.matching.rules import rule_matches from onyx.external_apps.providers.github import GitHubAction from onyx.external_apps.providers.registry import get_endpoint_catalog diff --git a/backend/tests/unit/external_apps/test_gmail_catalog.py b/backend/tests/unit/external_apps/test_gmail_catalog.py index 77d664e17e..a1fcb1361e 100644 --- a/backend/tests/unit/external_apps/test_gmail_catalog.py +++ b/backend/tests/unit/external_apps/test_gmail_catalog.py @@ -7,10 +7,8 @@ from __future__ import annotations import pytest -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.providers.actions import path_matches -from onyx.external_apps.providers.actions import RestRoute +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.providers.actions import path_matches, RestRoute from onyx.external_apps.providers.gmail import GmailAction from onyx.external_apps.providers.registry import get_endpoint_catalog diff --git a/backend/tests/unit/external_apps/test_google_drive_provider.py b/backend/tests/unit/external_apps/test_google_drive_provider.py index 0aa16c959a..e239f23580 100644 --- a/backend/tests/unit/external_apps/test_google_drive_provider.py +++ b/backend/tests/unit/external_apps/test_google_drive_provider.py @@ -4,13 +4,13 @@ calls. Reads are auto-approved (ALWAYS); mutations default to ASK.""" from __future__ import annotations -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.providers.actions import path_matches -from onyx.external_apps.providers.actions import RestRoute +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.providers.actions import path_matches, RestRoute from onyx.external_apps.providers.base import OnyxManagedExtApp -from onyx.external_apps.providers.google_drive import GoogleDriveAction -from onyx.external_apps.providers.google_drive import GoogleDriveProvider +from onyx.external_apps.providers.google_drive import ( + GoogleDriveAction, + GoogleDriveProvider, +) from onyx.external_apps.providers.registry import PROVIDERS _READ_ACTIONS = { diff --git a/backend/tests/unit/external_apps/test_hubspot_catalog.py b/backend/tests/unit/external_apps/test_hubspot_catalog.py index 279e43103d..42ae95eb6a 100644 --- a/backend/tests/unit/external_apps/test_hubspot_catalog.py +++ b/backend/tests/unit/external_apps/test_hubspot_catalog.py @@ -9,10 +9,8 @@ from __future__ import annotations import pytest -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.matching.request import MatchContext -from onyx.external_apps.matching.request import ProxiedRequest +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.matching.request import MatchContext, ProxiedRequest from onyx.external_apps.matching.rules import rule_matches from onyx.external_apps.providers.hubspot import HubspotAction from onyx.external_apps.providers.registry import get_endpoint_catalog diff --git a/backend/tests/unit/external_apps/test_hubspot_provider.py b/backend/tests/unit/external_apps/test_hubspot_provider.py index b9011f4b12..b6399546ae 100644 --- a/backend/tests/unit/external_apps/test_hubspot_provider.py +++ b/backend/tests/unit/external_apps/test_hubspot_provider.py @@ -8,8 +8,7 @@ from __future__ import annotations from types import SimpleNamespace from typing import cast from unittest.mock import MagicMock -from urllib.parse import parse_qs -from urllib.parse import urlparse +from urllib.parse import parse_qs, urlparse import pytest diff --git a/backend/tests/unit/external_apps/test_managed_credentials.py b/backend/tests/unit/external_apps/test_managed_credentials.py index 5a362800d0..9181d53015 100644 --- a/backend/tests/unit/external_apps/test_managed_credentials.py +++ b/backend/tests/unit/external_apps/test_managed_credentials.py @@ -8,8 +8,7 @@ import pytest from onyx.db.enums import ExternalAppType from onyx.external_apps.providers.base import OnyxManagedExtApp -from onyx.external_apps.providers.registry import get_onyx_managed_provider -from onyx.external_apps.providers.registry import PROVIDERS +from onyx.external_apps.providers.registry import get_onyx_managed_provider, PROVIDERS def _gmail() -> OnyxManagedExtApp: diff --git a/backend/tests/unit/external_apps/test_notion_catalog.py b/backend/tests/unit/external_apps/test_notion_catalog.py index 6944489dd9..f592be8731 100644 --- a/backend/tests/unit/external_apps/test_notion_catalog.py +++ b/backend/tests/unit/external_apps/test_notion_catalog.py @@ -9,10 +9,8 @@ from __future__ import annotations import pytest -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.matching.request import MatchContext -from onyx.external_apps.matching.request import ProxiedRequest +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.matching.request import MatchContext, ProxiedRequest from onyx.external_apps.matching.rules import rule_matches from onyx.external_apps.providers.notion import NotionAction from onyx.external_apps.providers.registry import get_endpoint_catalog diff --git a/backend/tests/unit/external_apps/test_path_matching.py b/backend/tests/unit/external_apps/test_path_matching.py index ea5068418c..5c0680c302 100644 --- a/backend/tests/unit/external_apps/test_path_matching.py +++ b/backend/tests/unit/external_apps/test_path_matching.py @@ -3,8 +3,7 @@ import pytest from pydantic import ValidationError -from onyx.external_apps.providers.actions import path_matches -from onyx.external_apps.providers.actions import RestRoute +from onyx.external_apps.providers.actions import path_matches, RestRoute def test_non_trailing_wildcard_rejected_at_construction() -> None: diff --git a/backend/tests/unit/external_apps/test_provider_granted_scopes.py b/backend/tests/unit/external_apps/test_provider_granted_scopes.py index 303c8185ec..c4e25312af 100644 --- a/backend/tests/unit/external_apps/test_provider_granted_scopes.py +++ b/backend/tests/unit/external_apps/test_provider_granted_scopes.py @@ -17,8 +17,10 @@ import requests from onyx.db.enums import ExternalAppType from onyx.external_apps.providers import hubspot as hubspot_module -from onyx.external_apps.providers.base import OAuthExternalAppProvider -from onyx.external_apps.providers.base import parse_granted_scopes +from onyx.external_apps.providers.base import ( + OAuthExternalAppProvider, + parse_granted_scopes, +) from onyx.external_apps.providers.hubspot import HubspotProvider from onyx.external_apps.providers.registry import PROVIDERS diff --git a/backend/tests/unit/external_apps/test_slack_catalog.py b/backend/tests/unit/external_apps/test_slack_catalog.py index 348056f826..bbdc5104c0 100644 --- a/backend/tests/unit/external_apps/test_slack_catalog.py +++ b/backend/tests/unit/external_apps/test_slack_catalog.py @@ -8,14 +8,11 @@ from __future__ import annotations import pytest -from onyx.db.enums import EndpointPolicy -from onyx.db.enums import ExternalAppType -from onyx.external_apps.matching.request import MatchContext -from onyx.external_apps.matching.request import ProxiedRequest +from onyx.db.enums import EndpointPolicy, ExternalAppType +from onyx.external_apps.matching.request import MatchContext, ProxiedRequest from onyx.external_apps.matching.rules import rule_matches from onyx.external_apps.providers.registry import get_endpoint_catalog -from onyx.external_apps.providers.slack import SlackAction -from onyx.external_apps.providers.slack import SlackProvider +from onyx.external_apps.providers.slack import SlackAction, SlackProvider _CATALOG = get_endpoint_catalog(ExternalAppType.SLACK) diff --git a/backend/tests/unit/external_apps/test_token_refresh.py b/backend/tests/unit/external_apps/test_token_refresh.py index 3ba6aaa6db..a16de30bc8 100644 --- a/backend/tests/unit/external_apps/test_token_refresh.py +++ b/backend/tests/unit/external_apps/test_token_refresh.py @@ -10,8 +10,10 @@ from redis.exceptions import ConnectionError as RedisConnectionError from sqlalchemy.exc import SQLAlchemyError from onyx.external_apps import token_refresh as tr -from onyx.external_apps.providers.base import TokenRefreshTerminalError -from onyx.external_apps.providers.base import TokenRefreshTransientError +from onyx.external_apps.providers.base import ( + TokenRefreshTerminalError, + TokenRefreshTransientError, +) from onyx.external_apps.providers.google_calendar import GoogleCalendarProvider # --------------------------------------------------------------------------- diff --git a/backend/tests/unit/external_apps/test_token_utils.py b/backend/tests/unit/external_apps/test_token_utils.py index 1acf2f318a..c884fb914c 100644 --- a/backend/tests/unit/external_apps/test_token_utils.py +++ b/backend/tests/unit/external_apps/test_token_utils.py @@ -1,9 +1,6 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone -from onyx.external_apps.token_utils import needs_refresh -from onyx.external_apps.token_utils import stamp_expires_at +from onyx.external_apps.token_utils import needs_refresh, stamp_expires_at _NOW = datetime(2026, 5, 29, 12, 0, 0, tzinfo=timezone.utc) diff --git a/backend/tests/unit/fakes.py b/backend/tests/unit/fakes.py index a0906e2490..99777348ad 100644 --- a/backend/tests/unit/fakes.py +++ b/backend/tests/unit/fakes.py @@ -1,7 +1,6 @@ from __future__ import annotations -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheLock +from onyx.cache.interface import CacheBackend, CacheLock class FakeLock(CacheLock): diff --git a/backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py b/backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py index eaa4bb5b51..eeb740b3b4 100644 --- a/backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py +++ b/backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py @@ -1,6 +1,4 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import patch diff --git a/backend/tests/unit/file_store/test_file_store.py b/backend/tests/unit/file_store/test_file_store.py index dd707611ee..8446eb501c 100644 --- a/backend/tests/unit/file_store/test_file_store.py +++ b/backend/tests/unit/file_store/test_file_store.py @@ -3,25 +3,15 @@ from collections.abc import Generator from contextlib import nullcontext from io import BytesIO from typing import Any -from unittest.mock import MagicMock -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import MagicMock, Mock, patch import pytest -from sqlalchemy import create_engine -from sqlalchemy import DateTime -from sqlalchemy import Enum -from sqlalchemy import String -from sqlalchemy.orm import DeclarativeBase -from sqlalchemy.orm import Mapped -from sqlalchemy.orm import mapped_column -from sqlalchemy.orm import Session -from sqlalchemy.orm import sessionmaker +from sqlalchemy import create_engine, DateTime, Enum, String +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, Session, sessionmaker from sqlalchemy.sql import func from onyx.configs.constants import FileOrigin -from onyx.file_store.file_store import get_default_file_store -from onyx.file_store.file_store import S3BackedFileStore +from onyx.file_store.file_store import get_default_file_store, S3BackedFileStore from onyx.file_store.gcs_file_store import GCSBackedFileStore diff --git a/backend/tests/unit/file_store/test_postgres_file_store.py b/backend/tests/unit/file_store/test_postgres_file_store.py index fb78769d0e..5818583f62 100644 --- a/backend/tests/unit/file_store/test_postgres_file_store.py +++ b/backend/tests/unit/file_store/test_postgres_file_store.py @@ -4,17 +4,17 @@ These tests mock the database layer (sessions, raw connections, large objects) so they run without any external services. """ -from io import BytesIO -from io import StringIO +from io import BytesIO, StringIO from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.configs.constants import FileOrigin -from onyx.file_store.postgres_file_store import POSTGRES_BUCKET_SENTINEL -from onyx.file_store.postgres_file_store import PostgresBackedFileStore +from onyx.file_store.postgres_file_store import ( + POSTGRES_BUCKET_SENTINEL, + PostgresBackedFileStore, +) @pytest.fixture diff --git a/backend/tests/unit/model_server/test_embedding.py b/backend/tests/unit/model_server/test_embedding.py index 6d4ecb2cb0..be123371c5 100644 --- a/backend/tests/unit/model_server/test_embedding.py +++ b/backend/tests/unit/model_server/test_embedding.py @@ -1,14 +1,11 @@ import asyncio import time -from typing import Any -from typing import List -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, List +from unittest.mock import MagicMock, patch import pytest -from model_server.encoders import embed_text -from model_server.encoders import process_embed_request +from model_server.encoders import embed_text, process_embed_request from shared_configs.enums import EmbedTextType from shared_configs.model_server_models import EmbedRequest diff --git a/backend/tests/unit/onyx/access/test_user_file_access.py b/backend/tests/unit/onyx/access/test_user_file_access.py index b5fe6e61ff..c57904d409 100644 --- a/backend/tests/unit/onyx/access/test_user_file_access.py +++ b/backend/tests/unit/onyx/access/test_user_file_access.py @@ -1,11 +1,9 @@ """Tests for user file ACL computation, including shared persona access.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 -from onyx.access.access import collect_user_file_access -from onyx.access.access import get_access_for_user_files_impl +from onyx.access.access import collect_user_file_access, get_access_for_user_files_impl from onyx.access.utils import prefix_user_email from onyx.configs.constants import PUBLIC_DOC_PAT diff --git a/backend/tests/unit/onyx/auth/conftest.py b/backend/tests/unit/onyx/auth/conftest.py index b4aaee9802..6c7d5dd338 100644 --- a/backend/tests/unit/onyx/auth/conftest.py +++ b/backend/tests/unit/onyx/auth/conftest.py @@ -1,10 +1,8 @@ -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest -from onyx.db.models import OAuthAccount -from onyx.db.models import User +from onyx.db.models import OAuthAccount, User @pytest.fixture diff --git a/backend/tests/unit/onyx/auth/test_auth_audit_events.py b/backend/tests/unit/onyx/auth/test_auth_audit_events.py index a35d8c4f0a..08515ba2c7 100644 --- a/backend/tests/unit/onyx/auth/test_auth_audit_events.py +++ b/backend/tests/unit/onyx/auth/test_auth_audit_events.py @@ -9,9 +9,7 @@ pattern). Emission itself is covered in tests/unit/onyx/utils/test_audit.py. import json import logging from typing import Any -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.security import OAuth2PasswordRequestForm diff --git a/backend/tests/unit/onyx/auth/test_captcha_enterprise.py b/backend/tests/unit/onyx/auth/test_captcha_enterprise.py index 46db76cad8..e649a4e8e1 100644 --- a/backend/tests/unit/onyx/auth/test_captcha_enterprise.py +++ b/backend/tests/unit/onyx/auth/test_captcha_enterprise.py @@ -1,18 +1,17 @@ """Unit tests for the reCAPTCHA Enterprise Assessment rejection ladder.""" from collections.abc import Iterator -from datetime import datetime -from datetime import timezone -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch import pytest from onyx.auth import captcha as captcha_module -from onyx.auth.captcha import CaptchaAction -from onyx.auth.captcha import CaptchaVerificationError -from onyx.auth.captcha import verify_captcha_token +from onyx.auth.captcha import ( + CaptchaAction, + CaptchaVerificationError, + verify_captcha_token, +) def _fake_client(payload: dict) -> MagicMock: diff --git a/backend/tests/unit/onyx/auth/test_captcha_replay.py b/backend/tests/unit/onyx/auth/test_captcha_replay.py index 0967a87972..880f4db561 100644 --- a/backend/tests/unit/onyx/auth/test_captcha_replay.py +++ b/backend/tests/unit/onyx/auth/test_captcha_replay.py @@ -1,18 +1,18 @@ """Unit tests for the reCAPTCHA token replay cache.""" -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from onyx.auth import captcha as captcha_module -from onyx.auth.captcha import _replay_cache_key -from onyx.auth.captcha import _reserve_token_or_raise -from onyx.auth.captcha import CaptchaAction -from onyx.auth.captcha import CaptchaVerificationError -from onyx.auth.captcha import verify_captcha_token +from onyx.auth.captcha import ( + _replay_cache_key, + _reserve_token_or_raise, + CaptchaAction, + CaptchaVerificationError, + verify_captcha_token, +) @pytest.mark.asyncio diff --git a/backend/tests/unit/onyx/auth/test_disposable_email_validator.py b/backend/tests/unit/onyx/auth/test_disposable_email_validator.py index 652f157e87..54fc09f4d4 100644 --- a/backend/tests/unit/onyx/auth/test_disposable_email_validator.py +++ b/backend/tests/unit/onyx/auth/test_disposable_email_validator.py @@ -2,8 +2,10 @@ Tests for disposable email validation. """ -from onyx.auth.disposable_email_validator import DisposableEmailValidator -from onyx.auth.disposable_email_validator import is_disposable_email +from onyx.auth.disposable_email_validator import ( + DisposableEmailValidator, + is_disposable_email, +) class TestDisposableEmailValidator: diff --git a/backend/tests/unit/onyx/auth/test_email.py b/backend/tests/unit/onyx/auth/test_email.py index cc1e91e330..dfede61de7 100644 --- a/backend/tests/unit/onyx/auth/test_email.py +++ b/backend/tests/unit/onyx/auth/test_email.py @@ -1,7 +1,6 @@ import pytest -from onyx.auth.email_utils import build_user_email_invite -from onyx.auth.email_utils import send_email +from onyx.auth.email_utils import build_user_email_invite, send_email from onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME from onyx.db.engine.sql_engine import SqlEngine from onyx.server.runtime.onyx_runtime import OnyxRuntime diff --git a/backend/tests/unit/onyx/auth/test_email_archive_bcc.py b/backend/tests/unit/onyx/auth/test_email_archive_bcc.py index ab880c2b28..a33d8940ac 100644 --- a/backend/tests/unit/onyx/auth/test_email_archive_bcc.py +++ b/backend/tests/unit/onyx/auth/test_email_archive_bcc.py @@ -1,6 +1,5 @@ from types import SimpleNamespace -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/auth/test_jwt_provisioning.py b/backend/tests/unit/onyx/auth/test_jwt_provisioning.py index 7b930ca77e..bb1d026303 100644 --- a/backend/tests/unit/onyx/auth/test_jwt_provisioning.py +++ b/backend/tests/unit/onyx/auth/test_jwt_provisioning.py @@ -1,9 +1,7 @@ from collections.abc import Callable -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest diff --git a/backend/tests/unit/onyx/auth/test_oauth_claims_capture.py b/backend/tests/unit/onyx/auth/test_oauth_claims_capture.py index 1b3ccc2be6..6efd532649 100644 --- a/backend/tests/unit/onyx/auth/test_oauth_claims_capture.py +++ b/backend/tests/unit/onyx/auth/test_oauth_claims_capture.py @@ -1,17 +1,17 @@ import json from typing import Any -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import jwt import pytest import onyx.auth.oauth_claims_capture as claims_capture -from onyx.auth.oauth_claims_capture import capture_oauth_login_claims -from onyx.auth.oauth_claims_capture import get_captured_oauth_claims -from onyx.auth.oauth_claims_capture import get_idp_profile_fields -from onyx.auth.oauth_claims_capture import get_idp_profile_placeholder_values +from onyx.auth.oauth_claims_capture import ( + capture_oauth_login_claims, + get_captured_oauth_claims, + get_idp_profile_fields, + get_idp_profile_placeholder_values, +) @pytest.fixture(autouse=True) diff --git a/backend/tests/unit/onyx/auth/test_oauth_refresher.py b/backend/tests/unit/onyx/auth/test_oauth_refresher.py index aa462ec036..47cc30a1b9 100644 --- a/backend/tests/unit/onyx/auth/test_oauth_refresher.py +++ b/backend/tests/unit/onyx/auth/test_oauth_refresher.py @@ -1,19 +1,18 @@ -from datetime import datetime -from datetime import timezone -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch import pytest from sqlalchemy.ext.asyncio import AsyncSession from onyx.auth import oauth_refresher -from onyx.auth.oauth_refresher import _resolve_token_endpoint -from onyx.auth.oauth_refresher import _test_expire_oauth_token -from onyx.auth.oauth_refresher import check_and_refresh_oauth_tokens -from onyx.auth.oauth_refresher import check_oauth_account_has_refresh_token -from onyx.auth.oauth_refresher import get_oauth_accounts_requiring_refresh_token -from onyx.auth.oauth_refresher import refresh_oauth_token +from onyx.auth.oauth_refresher import ( + _resolve_token_endpoint, + _test_expire_oauth_token, + check_and_refresh_oauth_tokens, + check_oauth_account_has_refresh_token, + get_oauth_accounts_requiring_refresh_token, + refresh_oauth_token, +) from onyx.db.models import OAuthAccount diff --git a/backend/tests/unit/onyx/auth/test_oidc_pkce.py b/backend/tests/unit/onyx/auth/test_oidc_pkce.py index 89834e6ae9..de0c35b47e 100644 --- a/backend/tests/unit/onyx/auth/test_oidc_pkce.py +++ b/backend/tests/unit/onyx/auth/test_oidc_pkce.py @@ -1,26 +1,21 @@ -from typing import Any -from typing import cast -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch -from urllib.parse import parse_qs -from urllib.parse import urlparse +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch +from urllib.parse import parse_qs, urlparse -from fastapi import FastAPI -from fastapi import Response +from fastapi import FastAPI, Response from fastapi.testclient import TestClient -from fastapi_users.authentication import AuthenticationBackend -from fastapi_users.authentication import CookieTransport +from fastapi_users.authentication import AuthenticationBackend, CookieTransport from fastapi_users.jwt import generate_jwt -from httpx_oauth.oauth2 import BaseOAuth2 -from httpx_oauth.oauth2 import GetAccessTokenError +from httpx_oauth.oauth2 import BaseOAuth2, GetAccessTokenError -from onyx.auth.users import CSRF_TOKEN_COOKIE_NAME -from onyx.auth.users import CSRF_TOKEN_KEY -from onyx.auth.users import get_oauth_router -from onyx.auth.users import get_pkce_cookie_name -from onyx.auth.users import PKCE_COOKIE_NAME_PREFIX -from onyx.auth.users import STATE_TOKEN_AUDIENCE +from onyx.auth.users import ( + CSRF_TOKEN_COOKIE_NAME, + CSRF_TOKEN_KEY, + get_oauth_router, + get_pkce_cookie_name, + PKCE_COOKIE_NAME_PREFIX, + STATE_TOKEN_AUDIENCE, +) from onyx.error_handling.exceptions import register_onyx_exception_handlers diff --git a/backend/tests/unit/onyx/auth/test_optional_user_oauth_refresh.py b/backend/tests/unit/onyx/auth/test_optional_user_oauth_refresh.py index f0df1c43c0..04f9537520 100644 --- a/backend/tests/unit/onyx/auth/test_optional_user_oauth_refresh.py +++ b/backend/tests/unit/onyx/auth/test_optional_user_oauth_refresh.py @@ -1,6 +1,4 @@ -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from sqlalchemy.ext.asyncio import AsyncSession diff --git a/backend/tests/unit/onyx/auth/test_permissions.py b/backend/tests/unit/onyx/auth/test_permissions.py index 5f70386bc3..b17d32729f 100644 --- a/backend/tests/unit/onyx/auth/test_permissions.py +++ b/backend/tests/unit/onyx/auth/test_permissions.py @@ -9,12 +9,14 @@ from unittest.mock import MagicMock import pytest from fastapi import Request -from onyx.auth.permissions import ALL_PERMISSIONS -from onyx.auth.permissions import get_effective_permissions -from onyx.auth.permissions import IMPLIED_PERMISSIONS -from onyx.auth.permissions import NON_TOGGLEABLE_PERMISSIONS -from onyx.auth.permissions import require_permission -from onyx.auth.permissions import resolve_effective_permissions +from onyx.auth.permissions import ( + ALL_PERMISSIONS, + get_effective_permissions, + IMPLIED_PERMISSIONS, + NON_TOGGLEABLE_PERMISSIONS, + require_permission, + resolve_effective_permissions, +) from onyx.auth.users import get_anonymous_user from onyx.db.enums import Permission from onyx.error_handling.error_codes import OnyxErrorCode diff --git a/backend/tests/unit/onyx/auth/test_scoped_pat_route_gate.py b/backend/tests/unit/onyx/auth/test_scoped_pat_route_gate.py index b21e61de4b..3a6da931a6 100644 --- a/backend/tests/unit/onyx/auth/test_scoped_pat_route_gate.py +++ b/backend/tests/unit/onyx/auth/test_scoped_pat_route_gate.py @@ -14,8 +14,7 @@ from fastapi import Depends from fastapi.routing import APIRoute from onyx.auth.permissions import require_permission -from onyx.auth.users import _scoped_pat_permitted_on_route -from onyx.auth.users import scope_exempt +from onyx.auth.users import _scoped_pat_permitted_on_route, scope_exempt from onyx.db.enums import Permission diff --git a/backend/tests/unit/onyx/auth/test_seat_upgrade_helpers.py b/backend/tests/unit/onyx/auth/test_seat_upgrade_helpers.py index bce0900b9a..847781e430 100644 --- a/backend/tests/unit/onyx/auth/test_seat_upgrade_helpers.py +++ b/backend/tests/unit/onyx/auth/test_seat_upgrade_helpers.py @@ -14,16 +14,13 @@ both EE and CE builds. from __future__ import annotations -from typing import Any -from typing import Iterator -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, Iterator +from unittest.mock import MagicMock, patch import pytest from ee.onyx.db.license import user_counts_toward_seats as _canonical -from onyx.auth.users import _upgrade_will_add_seat -from onyx.auth.users import _user_currently_counts_toward_seats +from onyx.auth.users import _upgrade_will_add_seat, _user_currently_counts_toward_seats from onyx.configs.constants import ANONYMOUS_USER_EMAIL from onyx.db.enums import AccountType diff --git a/backend/tests/unit/onyx/auth/test_signup_rate_limit.py b/backend/tests/unit/onyx/auth/test_signup_rate_limit.py index e774e68715..eeb620c683 100644 --- a/backend/tests/unit/onyx/auth/test_signup_rate_limit.py +++ b/backend/tests/unit/onyx/auth/test_signup_rate_limit.py @@ -1,17 +1,17 @@ """Unit tests for the per-IP signup rate limiter.""" -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import Request from onyx.auth import signup_rate_limit as rl -from onyx.auth.signup_rate_limit import _bucket_key -from onyx.auth.signup_rate_limit import _client_ip -from onyx.auth.signup_rate_limit import _PER_IP_PER_HOUR -from onyx.auth.signup_rate_limit import enforce_signup_rate_limit +from onyx.auth.signup_rate_limit import ( + _bucket_key, + _client_ip, + _PER_IP_PER_HOUR, + enforce_signup_rate_limit, +) from onyx.error_handling.exceptions import OnyxError diff --git a/backend/tests/unit/onyx/auth/test_single_tenant_jwt_strategy.py b/backend/tests/unit/onyx/auth/test_single_tenant_jwt_strategy.py index 29a5e21dfd..f61d28e349 100644 --- a/backend/tests/unit/onyx/auth/test_single_tenant_jwt_strategy.py +++ b/backend/tests/unit/onyx/auth/test_single_tenant_jwt_strategy.py @@ -1,8 +1,6 @@ import uuid -from datetime import datetime -from datetime import timezone -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock import jwt import pytest diff --git a/backend/tests/unit/onyx/auth/test_smtp_no_auth.py b/backend/tests/unit/onyx/auth/test_smtp_no_auth.py index aa46f3a9e9..2c5790cea2 100644 --- a/backend/tests/unit/onyx/auth/test_smtp_no_auth.py +++ b/backend/tests/unit/onyx/auth/test_smtp_no_auth.py @@ -5,8 +5,7 @@ and must not call `login()` when credentials are not configured — internal IP-whitelisted relays don't support either. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/auth/test_user_default_pins.py b/backend/tests/unit/onyx/auth/test_user_default_pins.py index 3a788a84ea..fdeaa66009 100644 --- a/backend/tests/unit/onyx/auth/test_user_default_pins.py +++ b/backend/tests/unit/onyx/auth/test_user_default_pins.py @@ -1,5 +1,4 @@ -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from sqlalchemy.ext.asyncio import AsyncSession diff --git a/backend/tests/unit/onyx/auth/test_user_registration.py b/backend/tests/unit/onyx/auth/test_user_registration.py index 2fd35a0386..590fb5e1d8 100644 --- a/backend/tests/unit/onyx/auth/test_user_registration.py +++ b/backend/tests/unit/onyx/auth/test_user_registration.py @@ -9,9 +9,7 @@ Tests cover: """ from types import TracebackType -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi_users import exceptions diff --git a/backend/tests/unit/onyx/auth/test_verified_email_openid_client.py b/backend/tests/unit/onyx/auth/test_verified_email_openid_client.py index c62f8e246a..aa77d4fc1d 100644 --- a/backend/tests/unit/onyx/auth/test_verified_email_openid_client.py +++ b/backend/tests/unit/onyx/auth/test_verified_email_openid_client.py @@ -3,15 +3,16 @@ email_verified counts as unverified) and refuse discovery documents whose issuer does not own the configured endpoint.""" from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from httpx_oauth.exceptions import GetIdEmailError -from onyx.auth.oidc_client import OpenIDConfigurationIssuerMismatch -from onyx.auth.oidc_client import validate_issuer_owns_config_url -from onyx.auth.oidc_client import VerifiedEmailOpenID +from onyx.auth.oidc_client import ( + OpenIDConfigurationIssuerMismatch, + validate_issuer_owns_config_url, + VerifiedEmailOpenID, +) _ISSUER = "https://idp.companyb.com" _CONFIG_URL = f"{_ISSUER}/.well-known/openid-configuration" diff --git a/backend/tests/unit/onyx/auth/test_verify_auth_setting.py b/backend/tests/unit/onyx/auth/test_verify_auth_setting.py index 2d6de4f656..963fc91e76 100644 --- a/backend/tests/unit/onyx/auth/test_verify_auth_setting.py +++ b/backend/tests/unit/onyx/auth/test_verify_auth_setting.py @@ -3,8 +3,7 @@ from unittest.mock import MagicMock import pytest import onyx.auth.users as users -from onyx.auth.users import verify_auth_setting -from onyx.auth.users import verify_user_auth_secret +from onyx.auth.users import verify_auth_setting, verify_user_auth_secret @pytest.mark.parametrize("stale_value", ["", "basic", "cloud"]) diff --git a/backend/tests/unit/onyx/background/celery/tasks/port/test_run_port_attempt_cursor.py b/backend/tests/unit/onyx/background/celery/tasks/port/test_run_port_attempt_cursor.py index c6bf8a40c5..c9a8de595e 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/port/test_run_port_attempt_cursor.py +++ b/backend/tests/unit/onyx/background/celery/tasks/port/test_run_port_attempt_cursor.py @@ -9,8 +9,7 @@ and permanently skips the un-ported tail (silent corpus drop after the swap). from contextlib import contextmanager from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.background.celery.tasks.port import tasks as port_tasks from onyx.db.enums import PortAttemptStatus diff --git a/backend/tests/unit/onyx/background/celery/tasks/tenant_provisioning/test_check_available_tenants.py b/backend/tests/unit/onyx/background/celery/tasks/tenant_provisioning/test_check_available_tenants.py index df068cbbee..b84d6b585f 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/tenant_provisioning/test_check_available_tenants.py +++ b/backend/tests/unit/onyx/background/celery/tasks/tenant_provisioning/test_check_available_tenants.py @@ -15,8 +15,6 @@ import pytest from ee.onyx.background.celery.tasks.tenant_provisioning.tasks import ( _MAX_TENANTS_PER_RUN, -) -from ee.onyx.background.celery.tasks.tenant_provisioning.tasks import ( check_available_tenants, ) diff --git a/backend/tests/unit/onyx/background/celery/tasks/test_hierarchyfetching_queue.py b/backend/tests/unit/onyx/background/celery/tasks/test_hierarchyfetching_queue.py index 37702870fb..bf7fc35ffd 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/test_hierarchyfetching_queue.py +++ b/backend/tests/unit/onyx/background/celery/tasks/test_hierarchyfetching_queue.py @@ -1,17 +1,16 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.background.celery.tasks.hierarchyfetching.tasks import ( _connector_supports_hierarchy_fetching, -) -from onyx.background.celery.tasks.hierarchyfetching.tasks import ( check_for_hierarchy_fetching, ) from onyx.connectors.factory import ConnectorMissingException -from onyx.connectors.interfaces import BaseConnector -from onyx.connectors.interfaces import HierarchyConnector -from onyx.connectors.interfaces import HierarchyOutput -from onyx.connectors.interfaces import SecondsSinceUnixEpoch +from onyx.connectors.interfaces import ( + BaseConnector, + HierarchyConnector, + HierarchyOutput, + SecondsSinceUnixEpoch, +) TASKS_MODULE = "onyx.background.celery.tasks.hierarchyfetching.tasks" diff --git a/backend/tests/unit/onyx/background/celery/tasks/test_indexing_interrupted.py b/backend/tests/unit/onyx/background/celery/tasks/test_indexing_interrupted.py index dddffac4eb..cfb97215e9 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/test_indexing_interrupted.py +++ b/backend/tests/unit/onyx/background/celery/tasks/test_indexing_interrupted.py @@ -3,11 +3,10 @@ infrastructure (deploy / autoscaling) must not count as a failure, so it never trips the repeated-error auto-pause, while genuine consecutive FAILED still does. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.background.celery.tasks.docprocessing.utils import is_in_repeated_error_state from onyx.background.celery.tasks.docprocessing.utils import ( + is_in_repeated_error_state, NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE, ) from onyx.db.enums import IndexingStatus diff --git a/backend/tests/unit/onyx/background/celery/tasks/test_port_wiring.py b/backend/tests/unit/onyx/background/celery/tasks/test_port_wiring.py index eff9b743ce..72366b3a37 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/test_port_wiring.py +++ b/backend/tests/unit/onyx/background/celery/tasks/test_port_wiring.py @@ -15,8 +15,7 @@ from pathlib import Path from onyx.background.celery.tasks.beat_schedule import beat_task_templates from onyx.configs.app_configs import DISABLE_VECTOR_DB -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import OnyxCeleryQueues, OnyxCeleryTask def _find_supervisord_conf() -> Path: diff --git a/backend/tests/unit/onyx/background/celery/tasks/test_user_file_impl_redis_locking.py b/backend/tests/unit/onyx/background/celery/tasks/test_user_file_impl_redis_locking.py index b7506f3bd5..fc9fcde577 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/test_user_file_impl_redis_locking.py +++ b/backend/tests/unit/onyx/background/celery/tasks/test_user_file_impl_redis_locking.py @@ -6,17 +6,12 @@ Verifies that: - Both paths execute the same business logic (DB lookup, status check) """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from onyx.background.celery.tasks.user_file_processing.tasks import ( delete_user_file_impl, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( process_user_file_impl, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( project_sync_user_file_impl, ) diff --git a/backend/tests/unit/onyx/background/celery/tasks/test_user_file_processing_no_vectordb.py b/backend/tests/unit/onyx/background/celery/tasks/test_user_file_processing_no_vectordb.py index e00add3675..af29c5b1e1 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/test_user_file_processing_no_vectordb.py +++ b/backend/tests/unit/onyx/background/celery/tasks/test_user_file_processing_no_vectordb.py @@ -8,25 +8,17 @@ Verifies that when DISABLE_VECTOR_DB is True: - project_sync_user_file_impl skips vector DB metadata update """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from onyx.background.celery.tasks.user_file_processing.tasks import ( _process_user_file_without_vector_db, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( delete_user_file_impl, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( process_user_file_impl, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( project_sync_user_file_impl, ) from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, TextSection from onyx.db.enums import UserFileStatus TASKS_MODULE = "onyx.background.celery.tasks.user_file_processing.tasks" diff --git a/backend/tests/unit/onyx/background/celery/tasks/test_user_file_project_sync_queue.py b/backend/tests/unit/onyx/background/celery/tasks/test_user_file_project_sync_queue.py index e110ab8c7f..a15789b065 100644 --- a/backend/tests/unit/onyx/background/celery/tasks/test_user_file_project_sync_queue.py +++ b/backend/tests/unit/onyx/background/celery/tasks/test_user_file_project_sync_queue.py @@ -1,26 +1,21 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from onyx.background.celery.tasks.user_file_processing.tasks import ( _user_file_project_sync_queued_key, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( check_for_user_file_project_sync, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( enqueue_user_file_project_sync_task, -) -from onyx.background.celery.tasks.user_file_processing.tasks import ( process_single_user_file_project_sync, ) -from onyx.configs.constants import CELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES -from onyx.configs.constants import OnyxCeleryPriority -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask -from onyx.configs.constants import USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH +from onyx.configs.constants import ( + CELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES, + OnyxCeleryPriority, + OnyxCeleryQueues, + OnyxCeleryTask, + USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH, +) def _build_redis_mock_with_lock() -> tuple[MagicMock, MagicMock]: diff --git a/backend/tests/unit/onyx/background/celery/test_celery_redis.py b/backend/tests/unit/onyx/background/celery/test_celery_redis.py index f1c79c4d3e..01c21770fd 100644 --- a/backend/tests/unit/onyx/background/celery/test_celery_redis.py +++ b/backend/tests/unit/onyx/background/celery/test_celery_redis.py @@ -1,8 +1,7 @@ """Tests for celery_get_broker_client singleton.""" from collections.abc import Iterator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/cache/test_redis_backend.py b/backend/tests/unit/onyx/cache/test_redis_backend.py index 97c365a432..605cfbef00 100644 --- a/backend/tests/unit/onyx/cache/test_redis_backend.py +++ b/backend/tests/unit/onyx/cache/test_redis_backend.py @@ -11,8 +11,7 @@ from redis.exceptions import LockNotOwnedError from redis.lock import Lock as RedisLock from onyx.cache.interface import CacheLockLostError -from onyx.cache.redis_backend import RedisCacheBackend -from onyx.cache.redis_backend import RedisCacheLock +from onyx.cache.redis_backend import RedisCacheBackend, RedisCacheLock from onyx.redis.tenant_redis_client import TenantRedisClient diff --git a/backend/tests/unit/onyx/chat/test_argument_delta_streaming.py b/backend/tests/unit/onyx/chat/test_argument_delta_streaming.py index e9233f350c..56350c083f 100644 --- a/backend/tests/unit/onyx/chat/test_argument_delta_streaming.py +++ b/backend/tests/unit/onyx/chat/test_argument_delta_streaming.py @@ -1,6 +1,5 @@ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.chat.tool_call_args_streaming import maybe_emit_argument_delta from onyx.server.query_and_chat.placement import Placement diff --git a/backend/tests/unit/onyx/chat/test_chat_utils.py b/backend/tests/unit/onyx/chat/test_chat_utils.py index 4755f62daa..b2a5606ae0 100644 --- a/backend/tests/unit/onyx/chat/test_chat_utils.py +++ b/backend/tests/unit/onyx/chat/test_chat_utils.py @@ -2,16 +2,16 @@ from io import BytesIO from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.chat.chat_utils import _build_tool_call_response_history_message -from onyx.chat.chat_utils import _get_or_extract_plaintext -from onyx.chat.chat_utils import convert_chat_history -from onyx.chat.chat_utils import get_custom_agent_prompt +from onyx.chat.chat_utils import ( + _build_tool_call_response_history_message, + _get_or_extract_plaintext, + convert_chat_history, + get_custom_agent_prompt, +) from onyx.chat.models import ChatLoadedFile -from onyx.configs.constants import DEFAULT_PERSONA_ID -from onyx.configs.constants import MessageType +from onyx.configs.constants import DEFAULT_PERSONA_ID, MessageType from onyx.db.models import ChatMessage from onyx.file_store.models import ChatFileType from onyx.prompts.chat_prompts import TOOL_CALL_RESPONSE_CROSS_MESSAGE diff --git a/backend/tests/unit/onyx/chat/test_citation_processor.py b/backend/tests/unit/onyx/chat/test_citation_processor.py index f3d4bf61c8..dc9f83fd44 100644 --- a/backend/tests/unit/onyx/chat/test_citation_processor.py +++ b/backend/tests/unit/onyx/chat/test_citation_processor.py @@ -18,9 +18,11 @@ from datetime import datetime import pytest -from onyx.chat.citation_processor import CitationMapping -from onyx.chat.citation_processor import CitationMode -from onyx.chat.citation_processor import DynamicCitationProcessor +from onyx.chat.citation_processor import ( + CitationMapping, + CitationMode, + DynamicCitationProcessor, +) from onyx.configs.constants import DocumentSource from onyx.context.search.models import SearchDoc from onyx.server.query_and_chat.streaming_models import CitationInfo diff --git a/backend/tests/unit/onyx/chat/test_compression.py b/backend/tests/unit/onyx/chat/test_compression.py index 89d8c0126c..9e480562c7 100644 --- a/backend/tests/unit/onyx/chat/test_compression.py +++ b/backend/tests/unit/onyx/chat/test_compression.py @@ -1,26 +1,25 @@ """Unit tests for chat history compression module.""" -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch -from onyx.chat.compression import _build_llm_messages_for_summarization -from onyx.chat.compression import find_summary_for_branch -from onyx.chat.compression import generate_summary -from onyx.chat.compression import get_compression_params -from onyx.chat.compression import get_messages_to_summarize -from onyx.chat.compression import SummaryContent +from onyx.chat.compression import ( + _build_llm_messages_for_summarization, + find_summary_for_branch, + generate_summary, + get_compression_params, + get_messages_to_summarize, + SummaryContent, +) from onyx.configs.constants import MessageType -from onyx.llm.models import AssistantMessage -from onyx.llm.models import SystemMessage -from onyx.llm.models import UserMessage -from onyx.prompts.compression_prompts import PROGRESSIVE_SUMMARY_SYSTEM_PROMPT_BLOCK -from onyx.prompts.compression_prompts import PROGRESSIVE_USER_REMINDER -from onyx.prompts.compression_prompts import SUMMARIZATION_CUTOFF_MARKER -from onyx.prompts.compression_prompts import SUMMARIZATION_PROMPT -from onyx.prompts.compression_prompts import USER_REMINDER +from onyx.llm.models import AssistantMessage, SystemMessage, UserMessage +from onyx.prompts.compression_prompts import ( + PROGRESSIVE_SUMMARY_SYSTEM_PROMPT_BLOCK, + PROGRESSIVE_USER_REMINDER, + SUMMARIZATION_CUTOFF_MARKER, + SUMMARIZATION_PROMPT, + USER_REMINDER, +) # Base time for generating sequential timestamps BASE_TIME = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) diff --git a/backend/tests/unit/onyx/chat/test_context_files.py b/backend/tests/unit/onyx/chat/test_context_files.py index ba53fc6a0a..9cbfee6a36 100644 --- a/backend/tests/unit/onyx/chat/test_context_files.py +++ b/backend/tests/unit/onyx/chat/test_context_files.py @@ -6,19 +6,18 @@ Covers: - Search filter / search_usage determination in the caller """ -from unittest.mock import MagicMock -from unittest.mock import patch -from uuid import UUID -from uuid import uuid4 +from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 from onyx.chat.models import ExtractedContextFiles -from onyx.chat.process_message import determine_search_params -from onyx.chat.process_message import extract_context_files -from onyx.chat.process_message import resolve_context_user_files +from onyx.chat.process_message import ( + determine_search_params, + extract_context_files, + resolve_context_user_files, +) from onyx.configs.constants import DEFAULT_PERSONA_ID from onyx.db.models import UserFile -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import InMemoryChatFile +from onyx.file_store.models import ChatFileType, InMemoryChatFile from onyx.tools.models import SearchToolUsage # --------------------------------------------------------------------------- diff --git a/backend/tests/unit/onyx/chat/test_emitter.py b/backend/tests/unit/onyx/chat/test_emitter.py index 58c0745d78..43df71bdfc 100644 --- a/backend/tests/unit/onyx/chat/test_emitter.py +++ b/backend/tests/unit/onyx/chat/test_emitter.py @@ -8,9 +8,11 @@ import queue from onyx.chat.emitter import Emitter from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ReasoningStart +from onyx.server.query_and_chat.streaming_models import ( + OverallStop, + Packet, + ReasoningStart, +) # --------------------------------------------------------------------------- # Helpers diff --git a/backend/tests/unit/onyx/chat/test_llm_loop.py b/backend/tests/unit/onyx/chat/test_llm_loop.py index d1832680b0..19330b8458 100644 --- a/backend/tests/unit/onyx/chat/test_llm_loop.py +++ b/backend/tests/unit/onyx/chat/test_llm_loop.py @@ -5,24 +5,26 @@ from unittest.mock import Mock import pytest -from onyx.chat.llm_loop import _build_empty_llm_response_error -from onyx.chat.llm_loop import _try_fallback_tool_extraction -from onyx.chat.llm_loop import construct_message_history -from onyx.chat.llm_loop import EmptyLLMResponseError -from onyx.chat.llm_loop import select_reminder_text -from onyx.chat.models import ChatLoadedFile -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import ContextFileMetadata -from onyx.chat.models import ExtractedContextFiles -from onyx.chat.models import FileToolMetadata -from onyx.chat.models import LlmStepResult -from onyx.chat.models import ToolCallSimple +from onyx.chat.llm_loop import ( + _build_empty_llm_response_error, + _try_fallback_tool_extraction, + construct_message_history, + EmptyLLMResponseError, + select_reminder_text, +) +from onyx.chat.models import ( + ChatLoadedFile, + ChatMessageSimple, + ContextFileMetadata, + ExtractedContextFiles, + FileToolMetadata, + LlmStepResult, + ToolCallSimple, +) from onyx.configs.constants import MessageType from onyx.file_store.models import ChatFileType -from onyx.llm.interfaces import LLMConfig -from onyx.llm.interfaces import ToolChoiceOptions -from onyx.prompts.chat_prompts import IMAGE_GEN_REMINDER -from onyx.prompts.chat_prompts import OPEN_URL_REMINDER +from onyx.llm.interfaces import LLMConfig, ToolChoiceOptions +from onyx.prompts.chat_prompts import IMAGE_GEN_REMINDER, OPEN_URL_REMINDER from onyx.server.query_and_chat.placement import Placement from onyx.tools.models import ToolCallKickoff diff --git a/backend/tests/unit/onyx/chat/test_llm_step.py b/backend/tests/unit/onyx/chat/test_llm_step.py index b22bda117b..eda54f31c9 100644 --- a/backend/tests/unit/onyx/chat/test_llm_step.py +++ b/backend/tests/unit/onyx/chat/test_llm_step.py @@ -5,30 +5,27 @@ from typing import Any import pytest from onyx.chat import llm_step as llm_step_module -from onyx.chat.llm_step import _extract_tool_call_kickoffs -from onyx.chat.llm_step import _increment_turns -from onyx.chat.llm_step import _parse_tool_args_to_dict -from onyx.chat.llm_step import _resolve_tool_arguments -from onyx.chat.llm_step import _XmlToolCallContentFilter -from onyx.chat.llm_step import extract_tool_calls_from_response_text -from onyx.chat.llm_step import translate_history_to_llm_format -from onyx.chat.models import ChatLoadedFile -from onyx.chat.models import ChatMessageSimple -from onyx.chat.models import ToolCallSimple +from onyx.chat.llm_step import ( + _extract_tool_call_kickoffs, + _increment_turns, + _parse_tool_args_to_dict, + _resolve_tool_arguments, + _XmlToolCallContentFilter, + extract_tool_calls_from_response_text, + translate_history_to_llm_format, +) +from onyx.chat.models import ChatLoadedFile, ChatMessageSimple, ToolCallSimple from onyx.configs.constants import MessageType from onyx.file_store.models import ChatFileType from onyx.llm.constants import LlmProviderNames -from onyx.llm.interfaces import LLMConfig -from onyx.llm.interfaces import ToolChoiceOptions -from onyx.llm.models import AssistantMessage -from onyx.llm.models import TextContentPart -from onyx.llm.models import ToolMessage -from onyx.llm.models import UserMessage -from onyx.llm.well_known_providers.constants import AZURE_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import OPENAI_PROVIDER_NAME +from onyx.llm.interfaces import LLMConfig, ToolChoiceOptions +from onyx.llm.models import AssistantMessage, TextContentPart, ToolMessage, UserMessage +from onyx.llm.well_known_providers.constants import ( + AZURE_PROVIDER_NAME, + OPENAI_PROVIDER_NAME, +) from onyx.prompts.chat_prompts import IMAGE_DROP_REMINDER -from onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE -from onyx.prompts.constants import SYSTEM_REMINDER_TAG_OPEN +from onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE, SYSTEM_REMINDER_TAG_OPEN from onyx.server.query_and_chat.placement import Placement from onyx.utils.postgres_sanitization import sanitize_string @@ -750,9 +747,7 @@ class TestEmptyAnswerRecovery: @staticmethod def _content_stream(chunks: list[str]) -> Any: - from onyx.llm.model_response import Delta - from onyx.llm.model_response import ModelResponseStream - from onyx.llm.model_response import StreamingChoice + from onyx.llm.model_response import Delta, ModelResponseStream, StreamingChoice def _gen(*_args: Any, **_kwargs: Any) -> Any: for i, chunk in enumerate(chunks): @@ -781,8 +776,7 @@ class TestEmptyAnswerRecovery: from unittest.mock import patch from onyx.chat import llm_step as _llm_step_module - from onyx.chat.citation_processor import CitationMode - from onyx.chat.citation_processor import DynamicCitationProcessor + from onyx.chat.citation_processor import CitationMode, DynamicCitationProcessor from onyx.chat.llm_step import run_llm_step_pkt_generator llm = self._make_llm() diff --git a/backend/tests/unit/onyx/chat/test_multi_model_streaming.py b/backend/tests/unit/onyx/chat/test_multi_model_streaming.py index 373cba4e21..52343730f3 100644 --- a/backend/tests/unit/onyx/chat/test_multi_model_streaming.py +++ b/backend/tests/unit/onyx/chat/test_multi_model_streaming.py @@ -8,10 +8,8 @@ calls, so we can trigger it with lightweight mocks. import threading import time from collections.abc import Generator -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -24,10 +22,12 @@ from onyx.db.models import ChatMessage from onyx.llm.override_models import LLMOverride from onyx.server.query_and_chat.models import SendMessageRequest from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import ChatHeartbeat -from onyx.server.query_and_chat.streaming_models import OverallStop -from onyx.server.query_and_chat.streaming_models import Packet -from onyx.server.query_and_chat.streaming_models import ReasoningStart +from onyx.server.query_and_chat.streaming_models import ( + ChatHeartbeat, + OverallStop, + Packet, + ReasoningStart, +) from onyx.utils.variable_functionality import global_version diff --git a/backend/tests/unit/onyx/chat/test_multi_model_types.py b/backend/tests/unit/onyx/chat/test_multi_model_types.py index a8553c488b..985c87c79d 100644 --- a/backend/tests/unit/onyx/chat/test_multi_model_types.py +++ b/backend/tests/unit/onyx/chat/test_multi_model_types.py @@ -7,15 +7,16 @@ Tests cover: - ChatMessageDetail new fields """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from uuid import uuid4 from onyx.llm.override_models import LLMOverride -from onyx.server.query_and_chat.models import ChatMessageDetail -from onyx.server.query_and_chat.models import ModelResponseSlot -from onyx.server.query_and_chat.models import MultiModelMessageResponseIDInfo -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.server.query_and_chat.models import ( + ChatMessageDetail, + ModelResponseSlot, + MultiModelMessageResponseIDInfo, + SendMessageRequest, +) from onyx.server.query_and_chat.placement import Placement diff --git a/backend/tests/unit/onyx/chat/test_narration_tool_placement.py b/backend/tests/unit/onyx/chat/test_narration_tool_placement.py index 29d4b5c00c..31c41022d3 100644 --- a/backend/tests/unit/onyx/chat/test_narration_tool_placement.py +++ b/backend/tests/unit/onyx/chat/test_narration_tool_placement.py @@ -13,19 +13,22 @@ pre-tool answer content was emitted, so it forms its own render group. from collections.abc import Iterator from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.chat.llm_step import run_llm_step_pkt_generator from onyx.llm.interfaces import ToolChoiceOptions -from onyx.llm.model_response import ChatCompletionDeltaToolCall -from onyx.llm.model_response import Delta -from onyx.llm.model_response import FunctionCall -from onyx.llm.model_response import ModelResponseStream -from onyx.llm.model_response import StreamingChoice +from onyx.llm.model_response import ( + ChatCompletionDeltaToolCall, + Delta, + FunctionCall, + ModelResponseStream, + StreamingChoice, +) from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import AgentResponseDelta -from onyx.server.query_and_chat.streaming_models import AgentResponseStart +from onyx.server.query_and_chat.streaming_models import ( + AgentResponseDelta, + AgentResponseStart, +) def _chunk(delta: Delta) -> ModelResponseStream: diff --git a/backend/tests/unit/onyx/chat/test_process_message.py b/backend/tests/unit/onyx/chat/test_process_message.py index c8ee27dc49..e3c25f4a5d 100644 --- a/backend/tests/unit/onyx/chat/test_process_message.py +++ b/backend/tests/unit/onyx/chat/test_process_message.py @@ -1,11 +1,12 @@ import pytest -from onyx.chat.process_message import _resolve_query_processing_hook_result -from onyx.chat.process_message import remove_answer_citations +from onyx.chat.process_message import ( + _resolve_query_processing_hook_result, + remove_answer_citations, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.hooks.executor import HookSkipped -from onyx.hooks.executor import HookSoftFailed +from onyx.hooks.executor import HookSkipped, HookSoftFailed from onyx.hooks.points.query_processing import QueryProcessingResponse diff --git a/backend/tests/unit/onyx/chat/test_process_message_mock_llm.py b/backend/tests/unit/onyx/chat/test_process_message_mock_llm.py index a2ec8e0144..f4656cc840 100644 --- a/backend/tests/unit/onyx/chat/test_process_message_mock_llm.py +++ b/backend/tests/unit/onyx/chat/test_process_message_mock_llm.py @@ -3,11 +3,9 @@ from unittest.mock import Mock import pytest from onyx.chat import process_message -from onyx.chat.models import AnswerStream -from onyx.chat.models import StreamingError +from onyx.chat.models import AnswerStream, StreamingError from onyx.configs import app_configs -from onyx.server.query_and_chat.models import MessageResponseIDInfo -from onyx.server.query_and_chat.models import SendMessageRequest +from onyx.server.query_and_chat.models import MessageResponseIDInfo, SendMessageRequest def test_mock_llm_response_requires_integration_mode() -> None: diff --git a/backend/tests/unit/onyx/chat/test_save_chat.py b/backend/tests/unit/onyx/chat/test_save_chat.py index ab8d56d15c..9cf811b49c 100644 --- a/backend/tests/unit/onyx/chat/test_save_chat.py +++ b/backend/tests/unit/onyx/chat/test_save_chat.py @@ -10,8 +10,7 @@ from pytest import MonkeyPatch from onyx.chat import save_chat from onyx.chat.save_chat import _extract_referenced_file_descriptors from onyx.file_store.models import ChatFileType -from onyx.tools.models import PythonExecutionFile -from onyx.tools.models import ToolCallInfo +from onyx.tools.models import PythonExecutionFile, ToolCallInfo def _make_tool_call_info( diff --git a/backend/tests/unit/onyx/chat/test_stop_signal_checker.py b/backend/tests/unit/onyx/chat/test_stop_signal_checker.py index e388a3b330..d9ed641a3b 100644 --- a/backend/tests/unit/onyx/chat/test_stop_signal_checker.py +++ b/backend/tests/unit/onyx/chat/test_stop_signal_checker.py @@ -7,14 +7,17 @@ so no external services are needed. from uuid import uuid4 -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheLock -from onyx.chat.chat_processing_checker import is_chat_session_processing -from onyx.chat.chat_processing_checker import set_processing_status -from onyx.chat.stop_signal_checker import FENCE_TTL -from onyx.chat.stop_signal_checker import is_connected -from onyx.chat.stop_signal_checker import reset_cancel_status -from onyx.chat.stop_signal_checker import set_fence +from onyx.cache.interface import CacheBackend, CacheLock +from onyx.chat.chat_processing_checker import ( + is_chat_session_processing, + set_processing_status, +) +from onyx.chat.stop_signal_checker import ( + FENCE_TTL, + is_connected, + reset_cancel_status, + set_fence, +) class _MemoryCacheBackend(CacheBackend): diff --git a/backend/tests/unit/onyx/chat/test_stream_buffer_unit.py b/backend/tests/unit/onyx/chat/test_stream_buffer_unit.py index 0e8f4b8103..24b3199430 100644 --- a/backend/tests/unit/onyx/chat/test_stream_buffer_unit.py +++ b/backend/tests/unit/onyx/chat/test_stream_buffer_unit.py @@ -4,15 +4,16 @@ and missing-chunk gaps surfacing as non-replayable instead of broken replays.""" import os import zlib -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from onyx.chat import stream_buffer -from onyx.chat.stream_buffer import read_stream_chunks -from onyx.chat.stream_buffer import StreamBufferMeta -from onyx.chat.stream_buffer import StreamBufferWriter +from onyx.chat.stream_buffer import ( + read_stream_chunks, + StreamBufferMeta, + StreamBufferWriter, +) from tests.unit.fakes import FakeCache _RUN_ID = 42 diff --git a/backend/tests/unit/onyx/configs/test_sentry.py b/backend/tests/unit/onyx/configs/test_sentry.py index 39bd92c483..b79c529dc4 100644 --- a/backend/tests/unit/onyx/configs/test_sentry.py +++ b/backend/tests/unit/onyx/configs/test_sentry.py @@ -1,6 +1,5 @@ from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from sentry_sdk.types import Event diff --git a/backend/tests/unit/onyx/configs/test_sentry_scrubber.py b/backend/tests/unit/onyx/configs/test_sentry_scrubber.py index 3566af34c3..bdb203c210 100644 --- a/backend/tests/unit/onyx/configs/test_sentry_scrubber.py +++ b/backend/tests/unit/onyx/configs/test_sentry_scrubber.py @@ -12,8 +12,7 @@ from unittest.mock import patch from sentry_sdk.scrubber import EventScrubber from sentry_sdk.utils import AnnotatedValue -from onyx.configs.sentry import build_event_scrubber -from onyx.configs.sentry import init_sentry +from onyx.configs.sentry import build_event_scrubber, init_sentry def _frame_event(frame_vars: dict) -> dict: diff --git a/backend/tests/unit/onyx/connectors/airtable/test_airtable_index_all.py b/backend/tests/unit/onyx/connectors/airtable/test_airtable_index_all.py index 08ec552fd6..109deb46e7 100644 --- a/backend/tests/unit/onyx/connectors/airtable/test_airtable_index_all.py +++ b/backend/tests/unit/onyx/connectors/airtable/test_airtable_index_all.py @@ -1,11 +1,12 @@ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.connectors.airtable.airtable_connector import AirtableConnector -from onyx.connectors.airtable.airtable_connector import parse_airtable_url +from onyx.connectors.airtable.airtable_connector import ( + AirtableConnector, + parse_airtable_url, +) from onyx.connectors.exceptions import ConnectorValidationError from onyx.connectors.models import Document diff --git a/backend/tests/unit/onyx/connectors/asana/test_asana_connector.py b/backend/tests/unit/onyx/connectors/asana/test_asana_connector.py index a779ff1449..4a4e2037a2 100644 --- a/backend/tests/unit/onyx/connectors/asana/test_asana_connector.py +++ b/backend/tests/unit/onyx/connectors/asana/test_asana_connector.py @@ -1,10 +1,8 @@ """Tests for Asana connector configuration parsing.""" from collections.abc import Iterator -from typing import Any -from typing import NamedTuple -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, NamedTuple +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/connectors/blob/test_blob_connector_region.py b/backend/tests/unit/onyx/connectors/blob/test_blob_connector_region.py index 164847ad16..f43704163d 100644 --- a/backend/tests/unit/onyx/connectors/blob/test_blob_connector_region.py +++ b/backend/tests/unit/onyx/connectors/blob/test_blob_connector_region.py @@ -1,15 +1,15 @@ -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch import pytest from botocore.exceptions import ClientError from onyx.configs.constants import BlobType from onyx.connectors.blob.connector import BlobStorageConnector -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError +from onyx.connectors.exceptions import ( + CredentialExpiredError, + InsufficientPermissionsError, +) def _make_connector(region_name: str | None) -> BlobStorageConnector: diff --git a/backend/tests/unit/onyx/connectors/box/conftest.py b/backend/tests/unit/onyx/connectors/box/conftest.py index c94613b396..34679ac68d 100644 --- a/backend/tests/unit/onyx/connectors/box/conftest.py +++ b/backend/tests/unit/onyx/connectors/box/conftest.py @@ -1,6 +1,5 @@ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/connectors/box/fake_box_client.py b/backend/tests/unit/onyx/connectors/box/fake_box_client.py index 7de21c8ba7..0e667e625f 100644 --- a/backend/tests/unit/onyx/connectors/box/fake_box_client.py +++ b/backend/tests/unit/onyx/connectors/box/fake_box_client.py @@ -4,9 +4,7 @@ connector exercises the same attribute access paths as against the live API.""" from io import BytesIO -from box_sdk_gen.box.errors import BoxAPIError -from box_sdk_gen.box.errors import RequestInfo -from box_sdk_gen.box.errors import ResponseInfo +from box_sdk_gen.box.errors import BoxAPIError, RequestInfo, ResponseInfo from box_sdk_gen.schemas.collaboration import Collaboration from box_sdk_gen.schemas.collaborations import Collaborations from box_sdk_gen.schemas.folder_full import FolderFull diff --git a/backend/tests/unit/onyx/connectors/box/test_box_access.py b/backend/tests/unit/onyx/connectors/box/test_box_access.py index 18b9cf0461..cc979dd31d 100644 --- a/backend/tests/unit/onyx/connectors/box/test_box_access.py +++ b/backend/tests/unit/onyx/connectors/box/test_box_access.py @@ -1,12 +1,13 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import cast import pytest from box_sdk_gen import BoxClient -from box_sdk_gen.schemas.collaboration import Collaboration -from box_sdk_gen.schemas.collaboration import CollaborationRoleField -from box_sdk_gen.schemas.collaboration import CollaborationStatusField +from box_sdk_gen.schemas.collaboration import ( + Collaboration, + CollaborationRoleField, + CollaborationStatusField, +) from box_sdk_gen.schemas.file_full import FileFull from box_sdk_gen.schemas.folder_full import FolderFull from box_sdk_gen.schemas.folder_mini import FolderMini @@ -14,13 +15,17 @@ from box_sdk_gen.schemas.group_mini import GroupMini from box_sdk_gen.schemas.items import Items from box_sdk_gen.schemas.user_collaborations import UserCollaborations from box_sdk_gen.schemas.user_mini import UserMini -from box_sdk_gen.schemas.web_link import WebLinkSharedLinkEffectiveAccessField -from box_sdk_gen.schemas.web_link import WebLinkSharedLinkEffectivePermissionField -from box_sdk_gen.schemas.web_link import WebLinkSharedLinkField +from box_sdk_gen.schemas.web_link import ( + WebLinkSharedLinkEffectiveAccessField, + WebLinkSharedLinkEffectivePermissionField, + WebLinkSharedLinkField, +) -from ee.onyx.external_permissions.box.access import apply_collaborations_to_access -from ee.onyx.external_permissions.box.access import apply_shared_link_to_access -from ee.onyx.external_permissions.box.access import BoxAccessContext +from ee.onyx.external_permissions.box.access import ( + apply_collaborations_to_access, + apply_shared_link_to_access, + BoxAccessContext, +) from onyx.connectors.box.connector import BoxConnector from onyx.connectors.models import Document from tests.unit.onyx.connectors.box.fake_box_client import FakeBoxClient diff --git a/backend/tests/unit/onyx/connectors/box/test_box_connector.py b/backend/tests/unit/onyx/connectors/box/test_box_connector.py index a88bcc4729..b80cdbc59d 100644 --- a/backend/tests/unit/onyx/connectors/box/test_box_connector.py +++ b/backend/tests/unit/onyx/connectors/box/test_box_connector.py @@ -1,35 +1,40 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import cast from unittest.mock import MagicMock import pytest -from box_sdk_gen import BoxCCGAuth -from box_sdk_gen import BoxClient +from box_sdk_gen import BoxCCGAuth, BoxClient from box_sdk_gen.schemas.file_full import FileFull from box_sdk_gen.schemas.folder_full import FolderFull from box_sdk_gen.schemas.folder_mini import FolderMini from box_sdk_gen.schemas.items import Items from box_sdk_gen.schemas.user_mini import UserMini -from box_sdk_gen.schemas.web_link import WebLink -from box_sdk_gen.schemas.web_link import WebLinkSharedLinkEffectiveAccessField -from box_sdk_gen.schemas.web_link import WebLinkSharedLinkEffectivePermissionField -from box_sdk_gen.schemas.web_link import WebLinkSharedLinkField +from box_sdk_gen.schemas.web_link import ( + WebLink, + WebLinkSharedLinkEffectiveAccessField, + WebLinkSharedLinkEffectivePermissionField, + WebLinkSharedLinkField, +) from onyx.access.models import ExternalAccess -from onyx.connectors.box.connector import BoxConnector -from onyx.connectors.box.connector import parse_box_folder_id +from onyx.connectors.box.connector import BoxConnector, parse_box_folder_id from onyx.connectors.box.models import BoxFolderFrontierEntry -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TextSection -from tests.unit.onyx.connectors.box.fake_box_client import FakeBoxClient -from tests.unit.onyx.connectors.box.fake_box_client import make_box_api_error +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.models import ( + ConnectorFailure, + Document, + HierarchyNode, + TextSection, +) +from tests.unit.onyx.connectors.box.fake_box_client import ( + FakeBoxClient, + make_box_api_error, +) _OWNER = UserMini(id="u1", name="Alice", login="alice@example.com") diff --git a/backend/tests/unit/onyx/connectors/braintrust/test_braintrust_connector.py b/backend/tests/unit/onyx/connectors/braintrust/test_braintrust_connector.py index f16268b4e7..21c103cd42 100644 --- a/backend/tests/unit/onyx/connectors/braintrust/test_braintrust_connector.py +++ b/backend/tests/unit/onyx/connectors/braintrust/test_braintrust_connector.py @@ -1,22 +1,23 @@ import re import time from collections.abc import Iterator -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import IO +from datetime import datetime, timedelta, timezone +from typing import Any, IO from unittest.mock import patch import pytest -from onyx.connectors.braintrust.connector import BraintrustCheckpoint -from onyx.connectors.braintrust.connector import BraintrustConnector -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.braintrust.connector import ( + BraintrustCheckpoint, + BraintrustConnector, +) +from onyx.connectors.models import ( + ConnectorFailure, + Document, + HierarchyNode, + TabularSection, + TextSection, +) _PROJECT = {"id": "proj-1", "name": "agent-wiki"} _PROMPT = { diff --git a/backend/tests/unit/onyx/connectors/canvas/test_canvas_connector.py b/backend/tests/unit/onyx/connectors/canvas/test_canvas_connector.py index 06978b226a..5abd36436f 100644 --- a/backend/tests/unit/onyx/connectors/canvas/test_canvas_connector.py +++ b/backend/tests/unit/onyx/connectors/canvas/test_canvas_connector.py @@ -1,45 +1,51 @@ """Tests for Canvas connector — client, credentials, conversion.""" -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from ee.onyx.external_permissions.canvas.access import CanvasCoursePermissionContext -from ee.onyx.external_permissions.canvas.access import get_announcement_permissions -from ee.onyx.external_permissions.canvas.access import get_assignment_permissions -from ee.onyx.external_permissions.canvas.access import get_page_permissions +from ee.onyx.external_permissions.canvas.access import ( + CanvasCoursePermissionContext, + get_announcement_permissions, + get_assignment_permissions, + get_page_permissions, +) from ee.onyx.external_permissions.canvas.group_sync import ( _referenced_group_and_section_ids, ) from onyx.access.models import ExternalAccess from onyx.configs.constants import DocumentSource from onyx.connectors.canvas.client import CanvasApiClient -from onyx.connectors.canvas.connector import _in_time_window -from onyx.connectors.canvas.connector import _parse_canvas_dt -from onyx.connectors.canvas.connector import _unix_to_canvas_time -from onyx.connectors.canvas.connector import canvas_all_users_group_id -from onyx.connectors.canvas.connector import canvas_course_group_id -from onyx.connectors.canvas.connector import canvas_group_group_id -from onyx.connectors.canvas.connector import canvas_section_group_id -from onyx.connectors.canvas.connector import CanvasAnnouncement -from onyx.connectors.canvas.connector import CanvasAnnouncementSection -from onyx.connectors.canvas.connector import CanvasAssignment -from onyx.connectors.canvas.connector import CanvasAssignmentOverride -from onyx.connectors.canvas.connector import CanvasConnector -from onyx.connectors.canvas.connector import CanvasConnectorCheckpoint -from onyx.connectors.canvas.connector import CanvasPage -from onyx.connectors.canvas.connector import CanvasStage -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import ConnectorMissingCredentialError -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode +from onyx.connectors.canvas.connector import ( + _in_time_window, + _parse_canvas_dt, + _unix_to_canvas_time, + canvas_all_users_group_id, + canvas_course_group_id, + canvas_group_group_id, + canvas_section_group_id, + CanvasAnnouncement, + CanvasAnnouncementSection, + CanvasAssignment, + CanvasAssignmentOverride, + CanvasConnector, + CanvasConnectorCheckpoint, + CanvasPage, + CanvasStage, +) +from onyx.connectors.exceptions import ( + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.models import ( + ConnectorFailure, + ConnectorMissingCredentialError, + Document, + HierarchyNode, +) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError diff --git a/backend/tests/unit/onyx/connectors/clickup/test_clickup_connector.py b/backend/tests/unit/onyx/connectors/clickup/test_clickup_connector.py index 068e193598..53fb3aa8fa 100644 --- a/backend/tests/unit/onyx/connectors/clickup/test_clickup_connector.py +++ b/backend/tests/unit/onyx/connectors/clickup/test_clickup_connector.py @@ -1,9 +1,7 @@ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.connectors.clickup.connector import CLICKUP_API_BASE_URL -from onyx.connectors.clickup.connector import ClickupConnector +from onyx.connectors.clickup.connector import CLICKUP_API_BASE_URL, ClickupConnector def _mock_response(json_response: dict[str, Any]) -> MagicMock: diff --git a/backend/tests/unit/onyx/connectors/confluence/test_confcloud_77618_fallback.py b/backend/tests/unit/onyx/connectors/confluence/test_confcloud_77618_fallback.py index 504bbdbc97..05645a208e 100644 --- a/backend/tests/unit/onyx/connectors/confluence/test_confcloud_77618_fallback.py +++ b/backend/tests/unit/onyx/connectors/confluence/test_confcloud_77618_fallback.py @@ -19,9 +19,11 @@ from onyx.connectors.confluence.access import ( get_page_restrictions_with_per_ancestor_fetch as get_page_restrictions_with_per_ancestor_fetch_shim, ) from onyx.connectors.confluence.connector import ConfluenceConnector -from onyx.connectors.confluence.onyx_confluence import _is_confcloud_77618_response -from onyx.connectors.confluence.onyx_confluence import Confcloud77618Error -from onyx.connectors.confluence.onyx_confluence import OnyxConfluence +from onyx.connectors.confluence.onyx_confluence import ( + _is_confcloud_77618_response, + Confcloud77618Error, + OnyxConfluence, +) from onyx.connectors.interfaces import CredentialsProviderInterface _CONFCLOUD_77618_BODY = ( diff --git a/backend/tests/unit/onyx/connectors/confluence/test_confluence_checkpointing.py b/backend/tests/unit/onyx/connectors/confluence/test_confluence_checkpointing.py index a3762b0c01..c343598257 100644 --- a/backend/tests/unit/onyx/connectors/confluence/test_confluence_checkpointing.py +++ b/backend/tests/unit/onyx/connectors/confluence/test_confluence_checkpointing.py @@ -1,29 +1,32 @@ import time -from collections.abc import Callable -from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Generator +from datetime import datetime, timezone from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from requests.exceptions import HTTPError from onyx.configs.constants import DocumentSource -from onyx.connectors.confluence.connector import ConfluenceCheckpoint -from onyx.connectors.confluence.connector import ConfluenceConnector +from onyx.connectors.confluence.connector import ( + ConfluenceCheckpoint, + ConfluenceConnector, +) from onyx.connectors.confluence.onyx_confluence import OnyxConfluence -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import SlimDocument -from tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector +from onyx.connectors.exceptions import ( + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + HierarchyNode, + SlimDocument, +) from tests.unit.onyx.connectors.utils import ( + load_everything_from_checkpoint_connector, load_everything_from_checkpoint_connector_from_checkpoint, ) diff --git a/backend/tests/unit/onyx/connectors/confluence/test_onyx_confluence.py b/backend/tests/unit/onyx/connectors/confluence/test_onyx_confluence.py index 2a4a05c205..3a4f132471 100644 --- a/backend/tests/unit/onyx/connectors/confluence/test_onyx_confluence.py +++ b/backend/tests/unit/onyx/connectors/confluence/test_onyx_confluence.py @@ -7,17 +7,17 @@ import requests from requests import HTTPError from onyx.connectors.confluence import onyx_confluence as onyx_confluence_module -from onyx.connectors.confluence.onyx_confluence import _DEFAULT_PAGINATION_LIMIT -from onyx.connectors.confluence.onyx_confluence import _MINIMUM_PAGINATION_LIMIT from onyx.connectors.confluence.onyx_confluence import ( + _DEFAULT_PAGINATION_LIMIT, + _MINIMUM_PAGINATION_LIMIT, ConfluenceRestSpacePermissionsNotAvailableError, -) -from onyx.connectors.confluence.onyx_confluence import ( get_user_email_from_userkey__server, + OnyxConfluence, +) +from onyx.connectors.exceptions import ( + ConnectorValidationError, + InsufficientPermissionsError, ) -from onyx.connectors.confluence.onyx_confluence import OnyxConfluence -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import InsufficientPermissionsError from onyx.connectors.interfaces import CredentialsProviderInterface diff --git a/backend/tests/unit/onyx/connectors/cross_connector_utils/test_rate_limit.py b/backend/tests/unit/onyx/connectors/cross_connector_utils/test_rate_limit.py index 97e10485b7..a76f856fe9 100644 --- a/backend/tests/unit/onyx/connectors/cross_connector_utils/test_rate_limit.py +++ b/backend/tests/unit/onyx/connectors/cross_connector_utils/test_rate_limit.py @@ -1,12 +1,11 @@ import time -from typing import Any -from typing import cast +from typing import Any, cast import pytest import requests -from onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder from onyx.connectors.cross_connector_utils.rate_limit_wrapper import ( + rate_limit_builder, wrap_request_to_handle_ratelimiting, ) diff --git a/backend/tests/unit/onyx/connectors/cross_connector_utils/test_section_text_cap.py b/backend/tests/unit/onyx/connectors/cross_connector_utils/test_section_text_cap.py index 04b682a917..c421d5286b 100644 --- a/backend/tests/unit/onyx/connectors/cross_connector_utils/test_section_text_cap.py +++ b/backend/tests/unit/onyx/connectors/cross_connector_utils/test_section_text_cap.py @@ -5,9 +5,7 @@ must still have the text they retain per file bounded.""" from unittest.mock import patch from onyx.connectors.cross_connector_utils.section_utils import cap_sections_text -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import ImageSection, TabularSection, TextSection _CAP_PATH = ( "onyx.connectors.cross_connector_utils.section_utils." diff --git a/backend/tests/unit/onyx/connectors/cross_connector_utils/test_tabular_section_utils.py b/backend/tests/unit/onyx/connectors/cross_connector_utils/test_tabular_section_utils.py index eb589cd25c..deb3359180 100644 --- a/backend/tests/unit/onyx/connectors/cross_connector_utils/test_tabular_section_utils.py +++ b/backend/tests/unit/onyx/connectors/cross_connector_utils/test_tabular_section_utils.py @@ -1,14 +1,13 @@ import io from collections.abc import Callable -from typing import cast -from typing import IO +from typing import cast, IO import openpyxl import pytest from openpyxl.worksheet.worksheet import Worksheet -from onyx.connectors.cross_connector_utils.tabular_section_utils import is_tabular_file from onyx.connectors.cross_connector_utils.tabular_section_utils import ( + is_tabular_file, tabular_file_to_sections, ) diff --git a/backend/tests/unit/onyx/connectors/discord/test_discord_validation.py b/backend/tests/unit/onyx/connectors/discord/test_discord_validation.py index 887bd3cd27..841a6b6685 100644 --- a/backend/tests/unit/onyx/connectors/discord/test_discord_validation.py +++ b/backend/tests/unit/onyx/connectors/discord/test_discord_validation.py @@ -1,5 +1,4 @@ -from unittest.mock import AsyncMock -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from discord.errors import LoginFailure diff --git a/backend/tests/unit/onyx/connectors/fireflies/test_fireflies_poll_overlap.py b/backend/tests/unit/onyx/connectors/fireflies/test_fireflies_poll_overlap.py index 733693760a..b8b7a7e7b0 100644 --- a/backend/tests/unit/onyx/connectors/fireflies/test_fireflies_poll_overlap.py +++ b/backend/tests/unit/onyx/connectors/fireflies/test_fireflies_poll_overlap.py @@ -1,9 +1,10 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from unittest.mock import patch -from onyx.connectors.fireflies.connector import _FIREFLIES_POLL_OVERLAP_SECONDS -from onyx.connectors.fireflies.connector import FirefliesConnector +from onyx.connectors.fireflies.connector import ( + _FIREFLIES_POLL_OVERLAP_SECONDS, + FirefliesConnector, +) def _fmt(epoch_seconds: int) -> str: diff --git a/backend/tests/unit/onyx/connectors/freshdesk/test_freshdesk_due_by.py b/backend/tests/unit/onyx/connectors/freshdesk/test_freshdesk_due_by.py index f0b6addda0..de19619e64 100644 --- a/backend/tests/unit/onyx/connectors/freshdesk/test_freshdesk_due_by.py +++ b/backend/tests/unit/onyx/connectors/freshdesk/test_freshdesk_due_by.py @@ -16,8 +16,10 @@ in that recalculation window. from typing import Any -from onyx.connectors.freshdesk.connector import _create_doc_from_ticket -from onyx.connectors.freshdesk.connector import _create_metadata_from_ticket +from onyx.connectors.freshdesk.connector import ( + _create_doc_from_ticket, + _create_metadata_from_ticket, +) def _ticket(**overrides: Any) -> dict[str, Any]: diff --git a/backend/tests/unit/onyx/connectors/github/test_github_checkpointing.py b/backend/tests/unit/onyx/connectors/github/test_github_checkpointing.py index 830a1cd74d..a4765e0951 100644 --- a/backend/tests/unit/onyx/connectors/github/test_github_checkpointing.py +++ b/backend/tests/unit/onyx/connectors/github/test_github_checkpointing.py @@ -1,15 +1,11 @@ import time -from collections.abc import Callable -from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from collections.abc import Callable, Generator +from datetime import datetime, timezone from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from github import Github -from github import RateLimitExceededException +from github import Github, RateLimitExceededException from github.GithubException import GithubException from github.Issue import Issue from github.PaginatedList import PaginatedList @@ -18,15 +14,16 @@ from github.RateLimit import RateLimit from github.Repository import Repository from github.Requester import Requester -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.github.connector import GithubConnector -from onyx.connectors.github.connector import GithubConnectorStage +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) +from onyx.connectors.github.connector import GithubConnector, GithubConnectorStage from onyx.connectors.github.models import SerializedRepository from onyx.connectors.models import Document -from tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector from tests.unit.onyx.connectors.utils import ( + load_everything_from_checkpoint_connector, load_everything_from_checkpoint_connector_from_checkpoint, ) diff --git a/backend/tests/unit/onyx/connectors/github/test_github_files.py b/backend/tests/unit/onyx/connectors/github/test_github_files.py index ecaaf9f6f8..9ff91b5523 100644 --- a/backend/tests/unit/onyx/connectors/github/test_github_files.py +++ b/backend/tests/unit/onyx/connectors/github/test_github_files.py @@ -1,22 +1,17 @@ import time from collections.abc import Callable -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch import pytest from github import Github -from github.GithubException import GithubException -from github.GithubException import UnknownObjectException +from github.GithubException import GithubException, UnknownObjectException from github.RateLimit import RateLimit from github.Requester import Requester -from onyx.connectors.github.connector import _is_indexable_path -from onyx.connectors.github.connector import GithubConnector +from onyx.connectors.github.connector import _is_indexable_path, GithubConnector from onyx.connectors.github.models import SerializedRepository -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document +from onyx.connectors.models import ConnectorFailure, Document from tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector diff --git a/backend/tests/unit/onyx/connectors/github/test_github_slim_connector.py b/backend/tests/unit/onyx/connectors/github/test_github_slim_connector.py index 31f3bd17c5..a8f750ffa6 100644 --- a/backend/tests/unit/onyx/connectors/github/test_github_slim_connector.py +++ b/backend/tests/unit/onyx/connectors/github/test_github_slim_connector.py @@ -4,17 +4,14 @@ correctly, and that pruning uses the cheap slim path (no lazy loading). """ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch -from unittest.mock import PropertyMock +from unittest.mock import MagicMock, patch, PropertyMock import pytest from onyx.access.models import ExternalAccess from onyx.background.celery.celery_utils import extract_ids_from_runnable_connector from onyx.connectors.github.connector import GithubConnector -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync +from onyx.connectors.interfaces import SlimConnector, SlimConnectorWithPermSync from onyx.connectors.models import SlimDocument diff --git a/backend/tests/unit/onyx/connectors/gmail/test_connector.py b/backend/tests/unit/onyx/connectors/gmail/test_connector.py index fc0cae62e4..5e8751ccea 100644 --- a/backend/tests/unit/onyx/connectors/gmail/test_connector.py +++ b/backend/tests/unit/onyx/connectors/gmail/test_connector.py @@ -2,19 +2,18 @@ import copy import datetime import json import os -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch from onyx.access.models import ExternalAccess from onyx.configs.constants import DocumentSource -from onyx.connectors.gmail.connector import _build_time_range_query -from onyx.connectors.gmail.connector import GmailCheckpoint -from onyx.connectors.gmail.connector import GmailConnector -from onyx.connectors.gmail.connector import thread_to_document -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.gmail.connector import ( + _build_time_range_query, + GmailCheckpoint, + GmailConnector, + thread_to_document, +) +from onyx.connectors.models import Document, TextSection from tests.unit.onyx.connectors.utils import ( load_everything_from_checkpoint_connector_from_checkpoint, ) diff --git a/backend/tests/unit/onyx/connectors/gong/test_gong_checkpoint.py b/backend/tests/unit/onyx/connectors/gong/test_gong_checkpoint.py index acd18dc5cd..27843bdc15 100644 --- a/backend/tests/unit/onyx/connectors/gong/test_gong_checkpoint.py +++ b/backend/tests/unit/onyx/connectors/gong/test_gong_checkpoint.py @@ -1,14 +1,11 @@ import time from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.connectors.gong.connector import GongConnector -from onyx.connectors.gong.connector import GongConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document +from onyx.connectors.gong.connector import GongConnector, GongConnectorCheckpoint +from onyx.connectors.models import ConnectorFailure, Document def _make_transcript(call_id: str) -> dict[str, Any]: diff --git a/backend/tests/unit/onyx/connectors/google_drive/test_drive_conversion_batching.py b/backend/tests/unit/onyx/connectors/google_drive/test_drive_conversion_batching.py index 2a4f997cbd..a19f2a1c06 100644 --- a/backend/tests/unit/onyx/connectors/google_drive/test_drive_conversion_batching.py +++ b/backend/tests/unit/onyx/connectors/google_drive/test_drive_conversion_batching.py @@ -2,20 +2,18 @@ documents materialize at most one DRIVE_CONVERSION_BATCH_SIZE chunk at a time, and hierarchy-deferred files wait for their ancestor node before converting.""" -from collections.abc import Callable -from collections.abc import Iterator +from collections.abc import Callable, Iterator from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.connectors.google_drive.connector import GoogleDriveConnector -from onyx.connectors.google_drive.models import DriveRetrievalStage -from onyx.connectors.google_drive.models import GoogleDriveCheckpoint -from onyx.connectors.google_drive.models import RetrievedDriveFile -from onyx.connectors.models import Document -from onyx.connectors.models import HierarchyNode -from onyx.utils.threadpool_concurrency import ThreadSafeDict -from onyx.utils.threadpool_concurrency import ThreadSafeSet +from onyx.connectors.google_drive.models import ( + DriveRetrievalStage, + GoogleDriveCheckpoint, + RetrievedDriveFile, +) +from onyx.connectors.models import Document, HierarchyNode +from onyx.utils.threadpool_concurrency import ThreadSafeDict, ThreadSafeSet _BATCH = 10 _CONN_MODULE = "onyx.connectors.google_drive.connector" diff --git a/backend/tests/unit/onyx/connectors/google_drive/test_drive_shortcut_following.py b/backend/tests/unit/onyx/connectors/google_drive/test_drive_shortcut_following.py index bcc71f48b7..9fb593183f 100644 --- a/backend/tests/unit/onyx/connectors/google_drive/test_drive_shortcut_following.py +++ b/backend/tests/unit/onyx/connectors/google_drive/test_drive_shortcut_following.py @@ -1,21 +1,25 @@ from collections.abc import Iterator -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch import pytest from googleapiclient.discovery import Resource -from onyx.connectors.google_drive.constants import DRIVE_FOLDER_TYPE -from onyx.connectors.google_drive.constants import DRIVE_SHORTCUT_TYPE -from onyx.connectors.google_drive.doc_conversion import convert_drive_item_to_document -from onyx.connectors.google_drive.doc_conversion import download_request -from onyx.connectors.google_drive.file_retrieval import _get_files_in_parent -from onyx.connectors.google_drive.file_retrieval import crawl_folders_for_files -from onyx.connectors.google_drive.file_retrieval import DRIVE_RESOURCE_KEY_FIELD -from onyx.connectors.google_drive.file_retrieval import DRIVE_RESOURCE_KEY_HEADER -from onyx.connectors.google_drive.file_retrieval import DriveFileFieldType +from onyx.connectors.google_drive.constants import ( + DRIVE_FOLDER_TYPE, + DRIVE_SHORTCUT_TYPE, +) +from onyx.connectors.google_drive.doc_conversion import ( + convert_drive_item_to_document, + download_request, +) +from onyx.connectors.google_drive.file_retrieval import ( + _get_files_in_parent, + crawl_folders_for_files, + DRIVE_RESOURCE_KEY_FIELD, + DRIVE_RESOURCE_KEY_HEADER, + DriveFileFieldType, +) from onyx.connectors.google_drive.models import DriveRetrievalStage _FILE_RETRIEVAL_MODULE = "onyx.connectors.google_drive.file_retrieval" diff --git a/backend/tests/unit/onyx/connectors/google_drive/test_slim_retrieval.py b/backend/tests/unit/onyx/connectors/google_drive/test_slim_retrieval.py index 769e4f2177..726987140e 100644 --- a/backend/tests/unit/onyx/connectors/google_drive/test_slim_retrieval.py +++ b/backend/tests/unit/onyx/connectors/google_drive/test_slim_retrieval.py @@ -7,23 +7,22 @@ Verifies that: - celery_utils routing picks retrieve_all_slim_docs() for GoogleDriveConnector """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from google.auth.exceptions import RefreshError from onyx.background.celery.celery_utils import extract_ids_from_runnable_connector from onyx.connectors.google_drive.connector import GoogleDriveConnector from onyx.connectors.google_drive.file_retrieval import DriveFileFieldType -from onyx.connectors.google_drive.models import DriveRetrievalStage -from onyx.connectors.google_drive.models import GoogleDriveCheckpoint -from onyx.connectors.google_drive.models import StageCompletion +from onyx.connectors.google_drive.models import ( + DriveRetrievalStage, + GoogleDriveCheckpoint, + StageCompletion, +) from onyx.connectors.google_utils.resources import ImpersonationError -from onyx.connectors.interfaces import SlimConnector -from onyx.connectors.interfaces import SlimConnectorWithPermSync +from onyx.connectors.interfaces import SlimConnector, SlimConnectorWithPermSync from onyx.connectors.models import SlimDocument -from onyx.utils.threadpool_concurrency import ThreadSafeDict -from onyx.utils.threadpool_concurrency import ThreadSafeSet +from onyx.utils.threadpool_concurrency import ThreadSafeDict, ThreadSafeSet def _make_done_checkpoint() -> GoogleDriveCheckpoint: diff --git a/backend/tests/unit/onyx/connectors/google_utils/test_google_credential_storage.py b/backend/tests/unit/onyx/connectors/google_utils/test_google_credential_storage.py index 0b49e328f1..583fd7f57c 100644 --- a/backend/tests/unit/onyx/connectors/google_utils/test_google_credential_storage.py +++ b/backend/tests/unit/onyx/connectors/google_utils/test_google_credential_storage.py @@ -3,44 +3,36 @@ the app cred from the row, reconstruct it from the row's token blob when absent, and restore the PKCE verifier for the callback token exchange.""" import json -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from typing import Any -from typing import cast +from datetime import datetime, timedelta, timezone +from typing import Any, cast import pytest from sqlalchemy.orm import Session -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import KV_CRED_KEY -from onyx.connectors.google_utils.google_kv import build_service_account_creds -from onyx.connectors.google_utils.google_kv import get_auth_url -from onyx.connectors.google_utils.google_kv import update_credential_access_tokens -from onyx.connectors.google_utils.google_kv import verify_csrf +from onyx.configs.constants import DocumentSource, KV_CRED_KEY +from onyx.connectors.google_utils.google_kv import ( + build_service_account_creds, + get_auth_url, + update_credential_access_tokens, + verify_csrf, +) from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_AUTHENTICATION_METHOD, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_APP_CREDENTIAL_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY, -) -from onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY -from onyx.connectors.google_utils.shared_constants import ( + DB_CREDENTIALS_DICT_TOKEN_KEY, DB_CREDENTIALS_PRIMARY_ADMIN_KEY, -) -from onyx.connectors.google_utils.shared_constants import ( GoogleOAuthAuthenticationMethod, ) from onyx.db.models import User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.key_value_store.interface import KvKeyNotFoundError -from onyx.server.documents.models import GoogleAppCredentials -from onyx.server.documents.models import GoogleAppWebCredentials -from onyx.server.documents.models import GoogleServiceAccountKey +from onyx.server.documents.models import ( + GoogleAppCredentials, + GoogleAppWebCredentials, + GoogleServiceAccountKey, +) class _StubCredentialJson: diff --git a/backend/tests/unit/onyx/connectors/google_utils/test_impersonation_guard.py b/backend/tests/unit/onyx/connectors/google_utils/test_impersonation_guard.py index 1f157cd29d..ab545c87b5 100644 --- a/backend/tests/unit/onyx/connectors/google_utils/test_impersonation_guard.py +++ b/backend/tests/unit/onyx/connectors/google_utils/test_impersonation_guard.py @@ -3,9 +3,11 @@ from unittest.mock import MagicMock from google.auth.exceptions import RefreshError -from onyx.connectors.google_utils.resources import ImpersonationError -from onyx.connectors.google_utils.resources import make_user_removal_checker -from onyx.connectors.google_utils.resources import UserRemovedError +from onyx.connectors.google_utils.resources import ( + ImpersonationError, + make_user_removal_checker, + UserRemovedError, +) from onyx.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR diff --git a/backend/tests/unit/onyx/connectors/hubspot/test_hubspot_inline_associations.py b/backend/tests/unit/onyx/connectors/hubspot/test_hubspot_inline_associations.py index 2b5a2ea441..8202cf0196 100644 --- a/backend/tests/unit/onyx/connectors/hubspot/test_hubspot_inline_associations.py +++ b/backend/tests/unit/onyx/connectors/hubspot/test_hubspot_inline_associations.py @@ -1,10 +1,8 @@ from __future__ import annotations -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.connectors.hubspot.connector import HubSpotConnector diff --git a/backend/tests/unit/onyx/connectors/jira/conftest.py b/backend/tests/unit/onyx/connectors/jira/conftest.py index efbf66cb7d..71af1ead98 100644 --- a/backend/tests/unit/onyx/connectors/jira/conftest.py +++ b/backend/tests/unit/onyx/connectors/jira/conftest.py @@ -1,6 +1,5 @@ from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from jira import JIRA diff --git a/backend/tests/unit/onyx/connectors/jira/test_jira_bulk_fetch.py b/backend/tests/unit/onyx/connectors/jira/test_jira_bulk_fetch.py index 3a7328dc39..5a052bb4b0 100644 --- a/backend/tests/unit/onyx/connectors/jira/test_jira_bulk_fetch.py +++ b/backend/tests/unit/onyx/connectors/jira/test_jira_bulk_fetch.py @@ -6,8 +6,7 @@ import requests from jira import JIRA from jira.resources import Issue -from onyx.connectors.jira.connector import _JIRA_BULK_FETCH_LIMIT -from onyx.connectors.jira.connector import bulk_fetch_issues +from onyx.connectors.jira.connector import _JIRA_BULK_FETCH_LIMIT, bulk_fetch_issues def _make_raw_issue(issue_id: str) -> dict[str, Any]: diff --git a/backend/tests/unit/onyx/connectors/jira/test_jira_checkpointing.py b/backend/tests/unit/onyx/connectors/jira/test_jira_checkpointing.py index 6af326ccff..9964d896fd 100644 --- a/backend/tests/unit/onyx/connectors/jira/test_jira_checkpointing.py +++ b/backend/tests/unit/onyx/connectors/jira/test_jira_checkpointing.py @@ -1,29 +1,23 @@ import time -from collections.abc import Callable -from collections.abc import Generator -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from collections.abc import Callable, Generator +from datetime import datetime, timezone +from typing import Any, cast +from unittest.mock import MagicMock, patch import pytest -from jira import JIRA -from jira import JIRAError +from jira import JIRA, JIRAError from jira.resources import Issue from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.jira.connector import JiraConnector -from onyx.connectors.jira.connector import JiraConnectorCheckpoint +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.jira.connector import JiraConnector, JiraConnectorCheckpoint from onyx.connectors.jira.utils import JIRA_SERVER_API_VERSION -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import SlimDocument +from onyx.connectors.models import ConnectorFailure, Document, SlimDocument from onyx.utils.logger import setup_logger from tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector diff --git a/backend/tests/unit/onyx/connectors/jira/test_jira_error_handling.py b/backend/tests/unit/onyx/connectors/jira/test_jira_error_handling.py index c935e4ba66..e9aabab674 100644 --- a/backend/tests/unit/onyx/connectors/jira/test_jira_error_handling.py +++ b/backend/tests/unit/onyx/connectors/jira/test_jira_error_handling.py @@ -4,12 +4,13 @@ import time from unittest.mock import MagicMock import pytest -from jira import JIRA -from jira import JIRAError +from jira import JIRA, JIRAError -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) from onyx.connectors.jira.connector import JiraConnector from tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector diff --git a/backend/tests/unit/onyx/connectors/jira/test_jira_large_ticket_handling.py b/backend/tests/unit/onyx/connectors/jira/test_jira_large_ticket_handling.py index eef27a6bec..0733620a81 100644 --- a/backend/tests/unit/onyx/connectors/jira/test_jira_large_ticket_handling.py +++ b/backend/tests/unit/onyx/connectors/jira/test_jira_large_ticket_handling.py @@ -1,14 +1,12 @@ from collections.abc import Generator from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from jira.resources import Issue from pytest_mock import MockFixture -from onyx.connectors.jira.connector import _perform_jql_search -from onyx.connectors.jira.connector import process_jira_issue +from onyx.connectors.jira.connector import _perform_jql_search, process_jira_issue @pytest.fixture diff --git a/backend/tests/unit/onyx/connectors/jira/test_jira_permission_sync.py b/backend/tests/unit/onyx/connectors/jira/test_jira_permission_sync.py index 87859e35b6..2b7ee717ae 100644 --- a/backend/tests/unit/onyx/connectors/jira/test_jira_permission_sync.py +++ b/backend/tests/unit/onyx/connectors/jira/test_jira_permission_sync.py @@ -1,7 +1,5 @@ -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/connectors/jira/test_jira_slim_retrieval.py b/backend/tests/unit/onyx/connectors/jira/test_jira_slim_retrieval.py index 433f88f096..eb98030e5c 100644 --- a/backend/tests/unit/onyx/connectors/jira/test_jira_slim_retrieval.py +++ b/backend/tests/unit/onyx/connectors/jira/test_jira_slim_retrieval.py @@ -1,12 +1,10 @@ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from jira.resources import Issue from onyx.access.models import ExternalAccess -from onyx.connectors.jira.connector import JiraConnector -from onyx.connectors.jira.connector import JiraConnectorCheckpoint +from onyx.connectors.jira.connector import JiraConnector, JiraConnectorCheckpoint from onyx.connectors.models import SlimDocument diff --git a/backend/tests/unit/onyx/connectors/linear/test_linear_load_credentials.py b/backend/tests/unit/onyx/connectors/linear/test_linear_load_credentials.py index d88c8f3b90..be78d8cd73 100644 --- a/backend/tests/unit/onyx/connectors/linear/test_linear_load_credentials.py +++ b/backend/tests/unit/onyx/connectors/linear/test_linear_load_credentials.py @@ -4,8 +4,7 @@ from __future__ import annotations import time from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_client.py b/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_client.py index 95cf7ccc1b..96c45f335c 100644 --- a/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_client.py +++ b/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_client.py @@ -7,8 +7,7 @@ from unittest.mock import MagicMock import pytest import requests -from onyx.connectors.lumapps.client import LumAppsClientError -from onyx.connectors.lumapps.client import OnyxLumApps +from onyx.connectors.lumapps.client import LumAppsClientError, OnyxLumApps def _client_with_token() -> tuple[OnyxLumApps, MagicMock]: diff --git a/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_connector.py b/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_connector.py index 9cc30954a9..aad147de4f 100644 --- a/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_connector.py +++ b/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_connector.py @@ -1,15 +1,12 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any from unittest.mock import MagicMock import pytest from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.lumapps.connector import _parse_dt -from onyx.connectors.lumapps.connector import LumAppsConnector -from onyx.connectors.models import Document -from onyx.connectors.models import SlimDocument +from onyx.connectors.lumapps.connector import _parse_dt, LumAppsConnector +from onyx.connectors.models import Document, SlimDocument def _make_connector(list_response: dict[str, Any]) -> LumAppsConnector: diff --git a/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_utils.py b/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_utils.py index e1b8a6d472..dbad64af56 100644 --- a/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_utils.py +++ b/backend/tests/unit/onyx/connectors/lumapps/test_lumapps_utils.py @@ -1,9 +1,11 @@ """Unit tests for the LumApps connector's pure helpers (no API / services).""" -from onyx.connectors.lumapps.utils import extract_body_text -from onyx.connectors.lumapps.utils import pick_lang -from onyx.connectors.lumapps.utils import resolve_metadata_labels -from onyx.connectors.lumapps.utils import slugify_family_key +from onyx.connectors.lumapps.utils import ( + extract_body_text, + pick_lang, + resolve_metadata_labels, + slugify_family_key, +) def test_slugify_family_key() -> None: diff --git a/backend/tests/unit/onyx/connectors/notion/test_notion_block_recursion.py b/backend/tests/unit/onyx/connectors/notion/test_notion_block_recursion.py index 63ccfd46dc..29f537e149 100644 --- a/backend/tests/unit/onyx/connectors/notion/test_notion_block_recursion.py +++ b/backend/tests/unit/onyx/connectors/notion/test_notion_block_recursion.py @@ -10,9 +10,11 @@ from typing import Any from unittest.mock import patch from onyx.connectors.models import Document -from onyx.connectors.notion.connector import BlockReadOutput -from onyx.connectors.notion.connector import NotionConnector -from onyx.connectors.notion.connector import NotionPage +from onyx.connectors.notion.connector import ( + BlockReadOutput, + NotionConnector, + NotionPage, +) # Pinned per test for determinism; _DEEP is well past it to force the overflow. _PROD_RECURSION_LIMIT = 1000 diff --git a/backend/tests/unit/onyx/connectors/notion/test_notion_database_parent_hierarchy.py b/backend/tests/unit/onyx/connectors/notion/test_notion_database_parent_hierarchy.py index d25f3fb3fe..2e68e284d7 100644 --- a/backend/tests/unit/onyx/connectors/notion/test_notion_database_parent_hierarchy.py +++ b/backend/tests/unit/onyx/connectors/notion/test_notion_database_parent_hierarchy.py @@ -7,13 +7,14 @@ emits a hierarchy node for that page even if its BlockReadOutput is empty upsert can be promoted. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.connectors.models import HierarchyNode -from onyx.connectors.notion.connector import BlockReadOutput -from onyx.connectors.notion.connector import NotionConnector -from onyx.connectors.notion.connector import NotionPage +from onyx.connectors.notion.connector import ( + BlockReadOutput, + NotionConnector, + NotionPage, +) from onyx.db.enums import HierarchyNodeType diff --git a/backend/tests/unit/onyx/connectors/notion/test_notion_datasource.py b/backend/tests/unit/onyx/connectors/notion/test_notion_datasource.py index 0d35770619..d0207af175 100644 --- a/backend/tests/unit/onyx/connectors/notion/test_notion_datasource.py +++ b/backend/tests/unit/onyx/connectors/notion/test_notion_datasource.py @@ -4,14 +4,15 @@ Tests the new data source discovery + querying flow and the data_source_id -> database_id parent resolution. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from requests.exceptions import HTTPError -from onyx.connectors.notion.connector import NotionConnector -from onyx.connectors.notion.connector import NotionDataSource -from onyx.connectors.notion.connector import NotionPage +from onyx.connectors.notion.connector import ( + NotionConnector, + NotionDataSource, + NotionPage, +) def _make_connector() -> NotionConnector: diff --git a/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_custom_config.py b/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_custom_config.py index 9c634bf552..fc44d2fdf0 100644 --- a/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_custom_config.py +++ b/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_custom_config.py @@ -9,10 +9,11 @@ exactly which fields and associations (child objects) to retrieve for each objec import json from typing import Any -from onyx.connectors.salesforce.connector import _validate_custom_query_config -from onyx.connectors.salesforce.connector import SalesforceConnector -from onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE -from onyx.connectors.salesforce.utils import MODIFIED_FIELD +from onyx.connectors.salesforce.connector import ( + _validate_custom_query_config, + SalesforceConnector, +) +from onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE, MODIFIED_FIELD def test_custom_query_config() -> None: diff --git a/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_sqlite.py b/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_sqlite.py index 77a0e2a617..b6403a29a9 100644 --- a/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_sqlite.py +++ b/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_sqlite.py @@ -5,8 +5,7 @@ import shutil import tempfile import time from collections import defaultdict -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from pathlib import Path from typing import cast @@ -14,21 +13,21 @@ import pytest from onyx.configs.constants import DocumentSource from onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import Document -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TextSection -from onyx.connectors.salesforce.doc_conversion import _extract_section -from onyx.connectors.salesforce.doc_conversion import ID_PREFIX +from onyx.connectors.models import BasicExpertInfo, Document, ImageSection, TextSection +from onyx.connectors.salesforce.doc_conversion import _extract_section, ID_PREFIX from onyx.connectors.salesforce.onyx_salesforce import OnyxSalesforce -from onyx.connectors.salesforce.salesforce_calls import _bulk_retrieve_from_salesforce -from onyx.connectors.salesforce.salesforce_calls import _make_time_filter_for_sf_type -from onyx.connectors.salesforce.salesforce_calls import _make_time_filtered_query -from onyx.connectors.salesforce.salesforce_calls import get_object_by_id_query +from onyx.connectors.salesforce.salesforce_calls import ( + _bulk_retrieve_from_salesforce, + _make_time_filter_for_sf_type, + _make_time_filtered_query, + get_object_by_id_query, +) from onyx.connectors.salesforce.sqlite_functions import OnyxSalesforceSQLite -from onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE -from onyx.connectors.salesforce.utils import MODIFIED_FIELD -from onyx.connectors.salesforce.utils import USER_OBJECT_TYPE +from onyx.connectors.salesforce.utils import ( + ACCOUNT_OBJECT_TYPE, + MODIFIED_FIELD, + USER_OBJECT_TYPE, +) from onyx.utils.logger import setup_logger # from onyx.connectors.salesforce.onyx_salesforce_type import OnyxSalesforceType diff --git a/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_utils.py b/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_utils.py index ae1bdf8497..a835ffc4e5 100644 --- a/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_utils.py +++ b/backend/tests/unit/onyx/connectors/salesforce/test_salesforce_utils.py @@ -1,7 +1,9 @@ import pytest -from onyx.connectors.salesforce.utils import is_valid_sf_identifier -from onyx.connectors.salesforce.utils import validate_sf_identifier +from onyx.connectors.salesforce.utils import ( + is_valid_sf_identifier, + validate_sf_identifier, +) class TestSfIdentifier: diff --git a/backend/tests/unit/onyx/connectors/salesforce/test_yield_doc_batches.py b/backend/tests/unit/onyx/connectors/salesforce/test_yield_doc_batches.py index 4fa4a9e34a..9468c51e2e 100644 --- a/backend/tests/unit/onyx/connectors/salesforce/test_yield_doc_batches.py +++ b/backend/tests/unit/onyx/connectors/salesforce/test_yield_doc_batches.py @@ -1,19 +1,22 @@ """Unit tests for _yield_doc_batches and metadata type conversion in SalesforceConnector.""" from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.configs.constants import DocumentSource from onyx.connectors.models import Document -from onyx.connectors.salesforce.connector import _convert_to_metadata_value -from onyx.connectors.salesforce.connector import SalesforceConnector -from onyx.connectors.salesforce.utils import ID_FIELD -from onyx.connectors.salesforce.utils import MODIFIED_FIELD -from onyx.connectors.salesforce.utils import NAME_FIELD -from onyx.connectors.salesforce.utils import SalesforceObject +from onyx.connectors.salesforce.connector import ( + _convert_to_metadata_value, + SalesforceConnector, +) +from onyx.connectors.salesforce.utils import ( + ID_FIELD, + MODIFIED_FIELD, + NAME_FIELD, + SalesforceObject, +) class TestConvertToMetadataValue: diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_delta_checkpointing.py b/backend/tests/unit/onyx/connectors/sharepoint/test_delta_checkpointing.py index 52122eeb3f..ccee9d8a4c 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_delta_checkpointing.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_delta_checkpointing.py @@ -13,20 +13,23 @@ from __future__ import annotations from collections import deque from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import pytest -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import TextSection -from onyx.connectors.sharepoint.connector import DriveItemData -from onyx.connectors.sharepoint.connector import SharepointConnector -from onyx.connectors.sharepoint.connector import SharepointConnectorCheckpoint -from onyx.connectors.sharepoint.connector import SiteDescriptor +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentSource, + TextSection, +) +from onyx.connectors.sharepoint.connector import ( + DriveItemData, + SharepointConnector, + SharepointConnectorCheckpoint, + SiteDescriptor, +) # --------------------------------------------------------------------------- # Constants diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_denylist.py b/backend/tests/unit/onyx/connectors/sharepoint/test_denylist.py index 940248b838..2b7a07aded 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_denylist.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_denylist.py @@ -2,12 +2,14 @@ from __future__ import annotations import pytest -from onyx.connectors.sharepoint.connector import _build_item_relative_path -from onyx.connectors.sharepoint.connector import _is_path_excluded -from onyx.connectors.sharepoint.connector import _is_site_excluded -from onyx.connectors.sharepoint.connector import DriveItemData -from onyx.connectors.sharepoint.connector import SharepointConnector -from onyx.connectors.sharepoint.connector import SiteDescriptor +from onyx.connectors.sharepoint.connector import ( + _build_item_relative_path, + _is_path_excluded, + _is_site_excluded, + DriveItemData, + SharepointConnector, + SiteDescriptor, +) class TestIsSiteExcluded: diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_drive_matching.py b/backend/tests/unit/onyx/connectors/sharepoint/test_drive_matching.py index 0a6c687ba8..fbc073d22c 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_drive_matching.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_drive_matching.py @@ -1,22 +1,20 @@ from __future__ import annotations from collections import deque -from collections.abc import Generator -from collections.abc import Sequence -from datetime import datetime -from datetime import timezone +from collections.abc import Generator, Sequence +from datetime import datetime, timezone from typing import Any import pytest -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import TextSection -from onyx.connectors.sharepoint.connector import DriveItemData -from onyx.connectors.sharepoint.connector import SHARED_DOCUMENTS_MAP -from onyx.connectors.sharepoint.connector import SharepointConnector -from onyx.connectors.sharepoint.connector import SharepointConnectorCheckpoint -from onyx.connectors.sharepoint.connector import SiteDescriptor +from onyx.connectors.models import Document, DocumentSource, TextSection +from onyx.connectors.sharepoint.connector import ( + DriveItemData, + SHARED_DOCUMENTS_MAP, + SharepointConnector, + SharepointConnectorCheckpoint, + SiteDescriptor, +) class _FakeQuery: diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_fetch_site_pages.py b/backend/tests/unit/onyx/connectors/sharepoint/test_fetch_site_pages.py index 5db35768cb..1a1eb247c2 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_fetch_site_pages.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_fetch_site_pages.py @@ -18,11 +18,13 @@ from office365.runtime.client_request_exception import ClientRequestException from requests import Response from requests.exceptions import HTTPError -from onyx.connectors.sharepoint.connector import _is_per_site_graph_failure -from onyx.connectors.sharepoint.connector import GRAPH_INVALID_REQUEST_CODE -from onyx.connectors.sharepoint.connector import PER_SITE_GRAPH_FAILURE_STATUSES -from onyx.connectors.sharepoint.connector import SharepointConnector -from onyx.connectors.sharepoint.connector import SiteDescriptor +from onyx.connectors.sharepoint.connector import ( + _is_per_site_graph_failure, + GRAPH_INVALID_REQUEST_CODE, + PER_SITE_GRAPH_FAILURE_STATUSES, + SharepointConnector, + SiteDescriptor, +) SITE_URL = "https://tenant.sharepoint.com/sites/ClassicSite" FAKE_SITE_ID = "tenant.sharepoint.com,abc123,def456" diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_load_checkpoint_resilience.py b/backend/tests/unit/onyx/connectors/sharepoint/test_load_checkpoint_resilience.py index fbba86766a..dd7f2ae8b4 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_load_checkpoint_resilience.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_load_checkpoint_resilience.py @@ -13,22 +13,25 @@ from __future__ import annotations from collections import deque from collections.abc import Generator -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from typing import Any import pytest -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import EntityFailure -from onyx.connectors.models import TextSection -from onyx.connectors.sharepoint.connector import DriveItemData -from onyx.connectors.sharepoint.connector import SharepointConnector -from onyx.connectors.sharepoint.connector import SharepointConnectorCheckpoint -from onyx.connectors.sharepoint.connector import SiteDescriptor +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + DocumentSource, + EntityFailure, + TextSection, +) +from onyx.connectors.sharepoint.connector import ( + DriveItemData, + SharepointConnector, + SharepointConnectorCheckpoint, + SiteDescriptor, +) SITE_URL = "https://example.sharepoint.com/sites/sample" DRIVE_WEB_URL = f"{SITE_URL}/Shared Documents" diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_load_credentials.py b/backend/tests/unit/onyx/connectors/sharepoint/test_load_credentials.py index 7e088a7327..d542c4132b 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_load_credentials.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_load_credentials.py @@ -3,8 +3,7 @@ from __future__ import annotations import base64 -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.connectors.sharepoint.connector import SharepointConnector diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_rest_client_context_caching.py b/backend/tests/unit/onyx/connectors/sharepoint/test_rest_client_context_caching.py index 87fb4edc6a..8b2e52977a 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_rest_client_context_caching.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_rest_client_context_caching.py @@ -2,11 +2,12 @@ from __future__ import annotations -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.connectors.sharepoint.connector import _REST_CTX_MAX_AGE_S -from onyx.connectors.sharepoint.connector import SharepointConnector +from onyx.connectors.sharepoint.connector import ( + _REST_CTX_MAX_AGE_S, + SharepointConnector, +) SITE_A = "https://tenant.sharepoint.com/sites/SiteA" SITE_B = "https://tenant.sharepoint.com/sites/SiteB" diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_slim_and_validation.py b/backend/tests/unit/onyx/connectors/sharepoint/test_slim_and_validation.py index c29a55df9c..ae88c3c6fa 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_slim_and_validation.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_slim_and_validation.py @@ -3,8 +3,7 @@ validate_connector_settings RoleAssignments permission probe.""" from __future__ import annotations -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -197,8 +196,7 @@ def test_retrieve_all_slim_docs_does_not_fetch_permissions( ) -> None: """retrieve_all_slim_docs (pruning path) never calls _create_rest_client_context and returns SlimDocuments with empty ExternalAccess.""" - from onyx.connectors.models import ExternalAccess - from onyx.connectors.models import SlimDocument + from onyx.connectors.models import ExternalAccess, SlimDocument from onyx.connectors.sharepoint.connector import DriveItemData connector = _make_connector() diff --git a/backend/tests/unit/onyx/connectors/sharepoint/test_streaming_download_retry.py b/backend/tests/unit/onyx/connectors/sharepoint/test_streaming_download_retry.py index d968336a31..2c9d6c4e30 100644 --- a/backend/tests/unit/onyx/connectors/sharepoint/test_streaming_download_retry.py +++ b/backend/tests/unit/onyx/connectors/sharepoint/test_streaming_download_retry.py @@ -10,17 +10,18 @@ per-document failure (which then trips the indexing failure threshold). from __future__ import annotations from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import requests from onyx.connectors.sharepoint import connector as sp_connector -from onyx.connectors.sharepoint.connector import _download_via_graph_api -from onyx.connectors.sharepoint.connector import _download_with_cap -from onyx.connectors.sharepoint.connector import _redact_url_for_logging -from onyx.connectors.sharepoint.connector import SizeCapExceeded +from onyx.connectors.sharepoint.connector import ( + _download_via_graph_api, + _download_with_cap, + _redact_url_for_logging, + SizeCapExceeded, +) CAP = 10 * 1024 * 1024 # 10 MiB cap; well above the byte payloads used in tests diff --git a/backend/tests/unit/onyx/connectors/slack/test_enterprise_grid.py b/backend/tests/unit/onyx/connectors/slack/test_enterprise_grid.py index c99f5e1473..5561b4bd96 100644 --- a/backend/tests/unit/onyx/connectors/slack/test_enterprise_grid.py +++ b/backend/tests/unit/onyx/connectors/slack/test_enterprise_grid.py @@ -4,23 +4,22 @@ These tests exercise the Grid-specific code paths added to the Slack connector using a mocked Slack ``WebClient``. They do not hit the real Slack API. """ -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch import pytest from slack_sdk.errors import SlackApiError -from onyx.connectors.slack.connector import _channel_team_id -from onyx.connectors.slack.connector import _channel_to_hierarchy_node -from onyx.connectors.slack.connector import channel_team_ids -from onyx.connectors.slack.connector import fetch_team_url -from onyx.connectors.slack.connector import get_channels_across_teams -from onyx.connectors.slack.connector import list_grid_team_ids -from onyx.connectors.slack.connector import SlackConnector -from onyx.connectors.slack.models import ChannelType -from onyx.connectors.slack.models import MessageType +from onyx.connectors.slack.connector import ( + _channel_team_id, + _channel_to_hierarchy_node, + channel_team_ids, + fetch_team_url, + get_channels_across_teams, + list_grid_team_ids, + SlackConnector, +) +from onyx.connectors.slack.models import ChannelType, MessageType from onyx.connectors.slack.utils import get_message_link diff --git a/backend/tests/unit/onyx/connectors/slack/test_message_filtering.py b/backend/tests/unit/onyx/connectors/slack/test_message_filtering.py index 78b5c8745f..b18b66535d 100644 --- a/backend/tests/unit/onyx/connectors/slack/test_message_filtering.py +++ b/backend/tests/unit/onyx/connectors/slack/test_message_filtering.py @@ -1,9 +1,11 @@ import pytest -from onyx.connectors.slack.connector import _bot_inclusive_msg_filter -from onyx.connectors.slack.connector import default_msg_filter -from onyx.connectors.slack.connector import SlackConnector -from onyx.connectors.slack.connector import SlackMessageFilterReason +from onyx.connectors.slack.connector import ( + _bot_inclusive_msg_filter, + default_msg_filter, + SlackConnector, + SlackMessageFilterReason, +) from onyx.connectors.slack.models import MessageType # -- default_msg_filter tests -- diff --git a/backend/tests/unit/onyx/connectors/slack/test_slack_filter_channels.py b/backend/tests/unit/onyx/connectors/slack/test_slack_filter_channels.py index 12770eb024..de1af6c8bd 100644 --- a/backend/tests/unit/onyx/connectors/slack/test_slack_filter_channels.py +++ b/backend/tests/unit/onyx/connectors/slack/test_slack_filter_channels.py @@ -2,14 +2,12 @@ after includes, support exact-name and full-match regex modes, and never fail on names absent from the workspace.""" -from typing import Any -from typing import cast +from typing import Any, cast import pytest from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.slack.connector import _validate_channel_regexes -from onyx.connectors.slack.connector import filter_channels +from onyx.connectors.slack.connector import _validate_channel_regexes, filter_channels from onyx.connectors.slack.models import ChannelType diff --git a/backend/tests/unit/onyx/connectors/teams/test_collect_teams.py b/backend/tests/unit/onyx/connectors/teams/test_collect_teams.py index f2cc8dd129..ea67b7fd36 100644 --- a/backend/tests/unit/onyx/connectors/teams/test_collect_teams.py +++ b/backend/tests/unit/onyx/connectors/teams/test_collect_teams.py @@ -74,9 +74,11 @@ def test_single_quote_escaping() -> None: def test_helper_functions() -> None: """Test the helper functions for team name processing.""" - from onyx.connectors.teams.connector import _can_use_odata_filter - from onyx.connectors.teams.connector import _escape_odata_string - from onyx.connectors.teams.connector import _has_odata_incompatible_chars + from onyx.connectors.teams.connector import ( + _can_use_odata_filter, + _escape_odata_string, + _has_odata_incompatible_chars, + ) # Test OData string escaping assert _escape_odata_string("Team's Group") == "Team''s Group" diff --git a/backend/tests/unit/onyx/connectors/teams/test_execute_query_with_retry.py b/backend/tests/unit/onyx/connectors/teams/test_execute_query_with_retry.py index 80830346f0..1f1c136102 100644 --- a/backend/tests/unit/onyx/connectors/teams/test_execute_query_with_retry.py +++ b/backend/tests/unit/onyx/connectors/teams/test_execute_query_with_retry.py @@ -7,15 +7,16 @@ gateway/server hiccups). A live daily-connector test cannot reliably induce a """ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from office365.runtime.client_request_exception import ClientRequestException -from onyx.connectors.teams.utils import _backoff_seconds -from onyx.connectors.teams.utils import execute_query_with_retry -from onyx.connectors.teams.utils import GRAPH_API_RETRYABLE_STATUSES +from onyx.connectors.teams.utils import ( + _backoff_seconds, + execute_query_with_retry, + GRAPH_API_RETRYABLE_STATUSES, +) def _client_request_exception( diff --git a/backend/tests/unit/onyx/connectors/teams/test_retry_request.py b/backend/tests/unit/onyx/connectors/teams/test_retry_request.py index 37e9783f46..96a76bcff1 100644 --- a/backend/tests/unit/onyx/connectors/teams/test_retry_request.py +++ b/backend/tests/unit/onyx/connectors/teams/test_retry_request.py @@ -8,14 +8,12 @@ here in isolation. """ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from requests.exceptions import HTTPError -from onyx.connectors.teams.utils import _retry -from onyx.connectors.teams.utils import GRAPH_API_RETRYABLE_STATUSES +from onyx.connectors.teams.utils import _retry, GRAPH_API_RETRYABLE_STATUSES def _response( diff --git a/backend/tests/unit/onyx/connectors/test_connector_factory.py b/backend/tests/unit/onyx/connectors/test_connector_factory.py index 6bc773e147..bd7ce1680a 100644 --- a/backend/tests/unit/onyx/connectors/test_connector_factory.py +++ b/backend/tests/unit/onyx/connectors/test_connector_factory.py @@ -8,25 +8,26 @@ Unit tests for lazy loading connector factory to validate: import importlib from types import SimpleNamespace -from unittest.mock import MagicMock -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import MagicMock, Mock, patch import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import UnexpectedValidationError -from onyx.connectors.factory import _connector_cache -from onyx.connectors.factory import _load_connector_class -from onyx.connectors.factory import ConnectorMissingException -from onyx.connectors.factory import identify_connector_class -from onyx.connectors.factory import instantiate_connector -from onyx.connectors.factory import validate_ccpair_for_user +from onyx.connectors.exceptions import ( + ConnectorValidationError, + UnexpectedValidationError, +) +from onyx.connectors.factory import ( + _connector_cache, + _load_connector_class, + ConnectorMissingException, + identify_connector_class, + instantiate_connector, + validate_ccpair_for_user, +) from onyx.connectors.interfaces import BaseConnector from onyx.connectors.models import InputType -from onyx.connectors.registry import CONNECTOR_CLASS_MAP -from onyx.connectors.registry import ConnectorMapping +from onyx.connectors.registry import CONNECTOR_CLASS_MAP, ConnectorMapping from onyx.db.enums import AccessType diff --git a/backend/tests/unit/onyx/connectors/test_document_metadata_coercion.py b/backend/tests/unit/onyx/connectors/test_document_metadata_coercion.py index 700c44b28b..1bc7659ed5 100644 --- a/backend/tests/unit/onyx/connectors/test_document_metadata_coercion.py +++ b/backend/tests/unit/onyx/connectors/test_document_metadata_coercion.py @@ -1,7 +1,5 @@ from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentBase -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, DocumentBase, TextSection def _minimal_doc_kwargs(metadata: dict) -> dict: diff --git a/backend/tests/unit/onyx/connectors/utils.py b/backend/tests/unit/onyx/connectors/utils.py index 4f372146dc..5f54895158 100644 --- a/backend/tests/unit/onyx/connectors/utils.py +++ b/backend/tests/unit/onyx/connectors/utils.py @@ -1,14 +1,10 @@ -from typing import Generic -from typing import TypeVar +from typing import Generic, TypeVar from pydantic import BaseModel from onyx.connectors.connector_runner import CheckpointOutputWrapper -from onyx.connectors.interfaces import CheckpointedConnector -from onyx.connectors.interfaces import SecondsSinceUnixEpoch -from onyx.connectors.models import ConnectorCheckpoint -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document +from onyx.connectors.interfaces import CheckpointedConnector, SecondsSinceUnixEpoch +from onyx.connectors.models import ConnectorCheckpoint, ConnectorFailure, Document _ITERATION_LIMIT = 100_000 diff --git a/backend/tests/unit/onyx/connectors/web/test_doc_updated_at.py b/backend/tests/unit/onyx/connectors/web/test_doc_updated_at.py index bb9e5a3d94..ae7f01135a 100644 --- a/backend/tests/unit/onyx/connectors/web/test_doc_updated_at.py +++ b/backend/tests/unit/onyx/connectors/web/test_doc_updated_at.py @@ -10,14 +10,12 @@ rely on the content hash instead. from __future__ import annotations -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.connectors.models import Document -from onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS -from onyx.connectors.web.connector import WebConnector +from onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS, WebConnector from onyx.file_processing.html_utils import ParsedHTML diff --git a/backend/tests/unit/onyx/connectors/web/test_slim_retrieval.py b/backend/tests/unit/onyx/connectors/web/test_slim_retrieval.py index aeb7af43a0..64c1262716 100644 --- a/backend/tests/unit/onyx/connectors/web/test_slim_retrieval.py +++ b/backend/tests/unit/onyx/connectors/web/test_slim_retrieval.py @@ -3,14 +3,12 @@ from __future__ import annotations from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.connectors.models import SlimDocument -from onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS -from onyx.connectors.web.connector import WebConnector +from onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS, WebConnector BASE_URL = "http://example.com" diff --git a/backend/tests/unit/onyx/connectors/web/test_ssrf_protection.py b/backend/tests/unit/onyx/connectors/web/test_ssrf_protection.py index c608b874aa..1c6eebb450 100644 --- a/backend/tests/unit/onyx/connectors/web/test_ssrf_protection.py +++ b/backend/tests/unit/onyx/connectors/web/test_ssrf_protection.py @@ -14,10 +14,12 @@ import pytest from onyx.connectors.exceptions import ConnectorValidationError from onyx.connectors.web import connector as web_connector -from onyx.connectors.web.connector import check_internet_connection -from onyx.connectors.web.connector import extract_urls_from_sitemap -from onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS -from onyx.connectors.web.connector import WebConnector +from onyx.connectors.web.connector import ( + check_internet_connection, + extract_urls_from_sitemap, + WEB_CONNECTOR_VALID_SETTINGS, + WebConnector, +) from onyx.server.security.models import SSRFProtectionLevel LOOPBACK_URL = "http://127.0.0.1:9999/ssrf-poc" diff --git a/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_checkpointing.py b/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_checkpointing.py index 4a8a191cfe..0da8ea39a1 100644 --- a/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_checkpointing.py +++ b/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_checkpointing.py @@ -1,21 +1,19 @@ import time -from collections.abc import Callable -from collections.abc import Generator -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from collections.abc import Callable, Generator +from typing import Any, cast +from unittest.mock import MagicMock, patch import pytest from requests.exceptions import HTTPError from onyx.configs.constants import DocumentSource -from onyx.connectors.exceptions import ConnectorValidationError -from onyx.connectors.exceptions import CredentialExpiredError -from onyx.connectors.exceptions import InsufficientPermissionsError +from onyx.connectors.exceptions import ( + ConnectorValidationError, + CredentialExpiredError, + InsufficientPermissionsError, +) from onyx.connectors.models import Document -from onyx.connectors.zendesk.connector import ZendeskClient -from onyx.connectors.zendesk.connector import ZendeskConnector +from onyx.connectors.zendesk.connector import ZendeskClient, ZendeskConnector from tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector diff --git a/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_rate_limit.py b/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_rate_limit.py index c95d95616b..0bd3eb6243 100644 --- a/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_rate_limit.py +++ b/backend/tests/unit/onyx/connectors/zendesk/test_zendesk_rate_limit.py @@ -1,8 +1,7 @@ from __future__ import annotations import types -from typing import Any -from typing import Dict +from typing import Any, Dict import pytest diff --git a/backend/tests/unit/onyx/context/search/federated/test_build_thread_text.py b/backend/tests/unit/onyx/context/search/federated/test_build_thread_text.py index 9ceaeb7def..770412d993 100644 --- a/backend/tests/unit/onyx/context/search/federated/test_build_thread_text.py +++ b/backend/tests/unit/onyx/context/search/federated/test_build_thread_text.py @@ -1,7 +1,6 @@ """Tests for _build_thread_text function.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.context.search.federated.slack_search import _build_thread_text diff --git a/backend/tests/unit/onyx/context/search/federated/test_slack_query_construction.py b/backend/tests/unit/onyx/context/search/federated/test_slack_query_construction.py index f16a22f8d6..e96fee76e2 100644 --- a/backend/tests/unit/onyx/context/search/federated/test_slack_query_construction.py +++ b/backend/tests/unit/onyx/context/search/federated/test_slack_query_construction.py @@ -1,7 +1,9 @@ from unittest.mock import MagicMock -from onyx.context.search.federated.slack_search_utils import build_channel_query_filter -from onyx.context.search.federated.slack_search_utils import matches_exclude_pattern +from onyx.context.search.federated.slack_search_utils import ( + build_channel_query_filter, + matches_exclude_pattern, +) from onyx.onyxbot.slack.models import ChannelType diff --git a/backend/tests/unit/onyx/context/search/federated/test_slack_thread_context.py b/backend/tests/unit/onyx/context/search/federated/test_slack_thread_context.py index f2b9255a1f..fee0f7a0ce 100644 --- a/backend/tests/unit/onyx/context/search/federated/test_slack_thread_context.py +++ b/backend/tests/unit/onyx/context/search/federated/test_slack_thread_context.py @@ -1,19 +1,18 @@ """Tests for Slack thread context fetching with rate limit handling.""" from datetime import datetime -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from slack_sdk.errors import SlackApiError from onyx.context.search.federated.models import SlackMessage -from onyx.context.search.federated.slack_search import _fetch_thread_context from onyx.context.search.federated.slack_search import ( + _fetch_thread_context, fetch_thread_contexts_with_rate_limit_handling, + SlackRateLimitError, + ThreadContextResult, ) -from onyx.context.search.federated.slack_search import SlackRateLimitError -from onyx.context.search.federated.slack_search import ThreadContextResult def _create_mock_message( diff --git a/backend/tests/unit/onyx/context/search/federated/test_url_override.py b/backend/tests/unit/onyx/context/search/federated/test_url_override.py index 88a7322ac0..0ccb65eaf8 100644 --- a/backend/tests/unit/onyx/context/search/federated/test_url_override.py +++ b/backend/tests/unit/onyx/context/search/federated/test_url_override.py @@ -1,7 +1,6 @@ """Tests for Slack URL parsing and direct thread fetch via URL override.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.context.search.federated.models import DirectThreadFetch from onyx.context.search.federated.slack_search import _fetch_thread_from_url diff --git a/backend/tests/unit/onyx/context/search/test_base_filters_time_compat.py b/backend/tests/unit/onyx/context/search/test_base_filters_time_compat.py index 1e5664d6c4..2722ac80a4 100644 --- a/backend/tests/unit/onyx/context/search/test_base_filters_time_compat.py +++ b/backend/tests/unit/onyx/context/search/test_base_filters_time_compat.py @@ -5,11 +5,9 @@ Older API clients send time_cutoff in request bodies it must keep filtering on document update time via updated_at_range. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import TimeRange +from onyx.context.search.models import BaseFilters, TimeRange _START = datetime(2025, 11, 24, tzinfo=timezone.utc) diff --git a/backend/tests/unit/onyx/db/engine/test_engine_disposal.py b/backend/tests/unit/onyx/db/engine/test_engine_disposal.py index c0617fb69d..b35cd4198d 100644 --- a/backend/tests/unit/onyx/db/engine/test_engine_disposal.py +++ b/backend/tests/unit/onyx/db/engine/test_engine_disposal.py @@ -16,9 +16,7 @@ the cached engine references are released. from __future__ import annotations from contextlib import ExitStack -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/db/engine/test_postgres_ssl.py b/backend/tests/unit/onyx/db/engine/test_postgres_ssl.py index 9f1509c833..d395aedf18 100644 --- a/backend/tests/unit/onyx/db/engine/test_postgres_ssl.py +++ b/backend/tests/unit/onyx/db/engine/test_postgres_ssl.py @@ -8,8 +8,7 @@ from unittest.mock import patch import certifi import pytest from cryptography import x509 -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID diff --git a/backend/tests/unit/onyx/db/test_assign_default_groups.py b/backend/tests/unit/onyx/db/test_assign_default_groups.py index 690de7f555..7c8046b12f 100644 --- a/backend/tests/unit/onyx/db/test_assign_default_groups.py +++ b/backend/tests/unit/onyx/db/test_assign_default_groups.py @@ -17,8 +17,7 @@ import pytest from sqlalchemy.exc import IntegrityError from onyx.db.enums import AccountType -from onyx.db.models import User__UserGroup -from onyx.db.models import UserGroup +from onyx.db.models import User__UserGroup, UserGroup from onyx.db.users import assign_user_to_default_groups__no_commit diff --git a/backend/tests/unit/onyx/db/test_chat_message_cleanup.py b/backend/tests/unit/onyx/db/test_chat_message_cleanup.py index c9f3ba5787..edc5c8f833 100644 --- a/backend/tests/unit/onyx/db/test_chat_message_cleanup.py +++ b/backend/tests/unit/onyx/db/test_chat_message_cleanup.py @@ -4,9 +4,7 @@ Verifies that user-owned files (those with user_file_id) are never deleted during chat session cleanup — only chat-only files should be removed. """ -from unittest.mock import call -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import call, MagicMock, patch from uuid import uuid4 from onyx.db.chat import delete_messages_and_files_from_chat_session diff --git a/backend/tests/unit/onyx/db/test_chat_sessions.py b/backend/tests/unit/onyx/db/test_chat_sessions.py index 1395ba8aa2..6edf658744 100644 --- a/backend/tests/unit/onyx/db/test_chat_sessions.py +++ b/backend/tests/unit/onyx/db/test_chat_sessions.py @@ -5,12 +5,9 @@ correctly filtered out while preserving recently created sessions, matching the behavior specified in PR #7233. """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session diff --git a/backend/tests/unit/onyx/db/test_dal.py b/backend/tests/unit/onyx/db/test_dal.py index fdd69ec46a..e5159ec2d0 100644 --- a/backend/tests/unit/onyx/db/test_dal.py +++ b/backend/tests/unit/onyx/db/test_dal.py @@ -1,5 +1,4 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/db/test_delete_user.py b/backend/tests/unit/onyx/db/test_delete_user.py index acb5198729..06e4e1a08d 100644 --- a/backend/tests/unit/onyx/db/test_delete_user.py +++ b/backend/tests/unit/onyx/db/test_delete_user.py @@ -1,15 +1,15 @@ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch -from uuid import UUID -from uuid import uuid4 +from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 -from onyx.db.models import DocumentSet -from onyx.db.models import DocumentSet__User -from onyx.db.models import Persona -from onyx.db.models import Persona__User -from onyx.db.models import SamlAccount -from onyx.db.models import User__UserGroup +from onyx.db.models import ( + DocumentSet, + DocumentSet__User, + Persona, + Persona__User, + SamlAccount, + User__UserGroup, +) from onyx.db.users import delete_user_from_db diff --git a/backend/tests/unit/onyx/db/test_external_app_auth_template.py b/backend/tests/unit/onyx/db/test_external_app_auth_template.py index d147ae32d3..d36daf6a34 100644 --- a/backend/tests/unit/onyx/db/test_external_app_auth_template.py +++ b/backend/tests/unit/onyx/db/test_external_app_auth_template.py @@ -7,8 +7,7 @@ from typing import Any import pytest -from onyx.db.external_app import resolve_masked_credentials -from onyx.db.external_app import validate_auth_template +from onyx.db.external_app import resolve_masked_credentials, validate_auth_template from onyx.error_handling.exceptions import OnyxError from onyx.utils.encryption import mask_string from onyx.utils.sensitive import SensitiveValue diff --git a/backend/tests/unit/onyx/db/test_index_attempt_errors.py b/backend/tests/unit/onyx/db/test_index_attempt_errors.py index 7e9b972bc3..f65e520579 100644 --- a/backend/tests/unit/onyx/db/test_index_attempt_errors.py +++ b/backend/tests/unit/onyx/db/test_index_attempt_errors.py @@ -1,7 +1,6 @@ """Tests for get_index_attempt_errors_across_connectors.""" -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone from unittest.mock import MagicMock from onyx.db.index_attempt import get_index_attempt_errors_across_connectors diff --git a/backend/tests/unit/onyx/db/test_llm_sync.py b/backend/tests/unit/onyx/db/test_llm_sync.py index ebe172578d..fcc33e9a87 100644 --- a/backend/tests/unit/onyx/db/test_llm_sync.py +++ b/backend/tests/unit/onyx/db/test_llm_sync.py @@ -1,7 +1,6 @@ """Tests for LLM provider model sync functionality.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/db/test_projects_upload_task_expiry.py b/backend/tests/unit/onyx/db/test_projects_upload_task_expiry.py index 3d6174b331..7af244b810 100644 --- a/backend/tests/unit/onyx/db/test_projects_upload_task_expiry.py +++ b/backend/tests/unit/onyx/db/test_projects_upload_task_expiry.py @@ -6,13 +6,14 @@ on every send_task call to prevent phantom task accumulation if the worker is down or slow. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 -from onyx.configs.constants import CELERY_USER_FILE_PROCESSING_TASK_EXPIRES -from onyx.configs.constants import OnyxCeleryQueues -from onyx.configs.constants import OnyxCeleryTask +from onyx.configs.constants import ( + CELERY_USER_FILE_PROCESSING_TASK_EXPIRES, + OnyxCeleryQueues, + OnyxCeleryTask, +) from onyx.db.models import UserFile from onyx.db.projects import upload_files_to_user_files_with_indexing diff --git a/backend/tests/unit/onyx/db/test_scim_dal.py b/backend/tests/unit/onyx/db/test_scim_dal.py index 4adb53f61c..849bd3a7d5 100644 --- a/backend/tests/unit/onyx/db/test_scim_dal.py +++ b/backend/tests/unit/onyx/db/test_scim_dal.py @@ -5,9 +5,7 @@ from uuid import uuid4 import pytest from ee.onyx.db.scim import ScimDAL -from onyx.db.models import ScimGroupMapping -from onyx.db.models import ScimToken -from onyx.db.models import ScimUserMapping +from onyx.db.models import ScimGroupMapping, ScimToken, ScimUserMapping from tests.unit.onyx.db.conftest import model_attrs diff --git a/backend/tests/unit/onyx/db/test_usage.py b/backend/tests/unit/onyx/db/test_usage.py index a956bcac3a..b3323604ea 100644 --- a/backend/tests/unit/onyx/db/test_usage.py +++ b/backend/tests/unit/onyx/db/test_usage.py @@ -1,20 +1,20 @@ """Unit tests for tenant usage tracking and limits.""" -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch import pytest -from onyx.db.usage import check_usage_limit -from onyx.db.usage import get_current_window_start -from onyx.db.usage import get_or_create_tenant_usage -from onyx.db.usage import get_tenant_usage_stats -from onyx.db.usage import increment_usage -from onyx.db.usage import TenantUsageStats -from onyx.db.usage import UsageLimitExceededError -from onyx.db.usage import UsageType +from onyx.db.usage import ( + check_usage_limit, + get_current_window_start, + get_or_create_tenant_usage, + get_tenant_usage_stats, + increment_usage, + TenantUsageStats, + UsageLimitExceededError, + UsageType, +) class TestGetCurrentWindowStart: diff --git a/backend/tests/unit/onyx/db/test_voice.py b/backend/tests/unit/onyx/db/test_voice.py index 5c3243c790..32c209aab1 100644 --- a/backend/tests/unit/onyx/db/test_voice.py +++ b/backend/tests/unit/onyx/db/test_voice.py @@ -6,20 +6,22 @@ from uuid import uuid4 import pytest from onyx.db.models import VoiceProvider -from onyx.db.voice import deactivate_stt_provider -from onyx.db.voice import deactivate_tts_provider -from onyx.db.voice import delete_voice_provider -from onyx.db.voice import fetch_default_stt_provider -from onyx.db.voice import fetch_default_tts_provider -from onyx.db.voice import fetch_voice_provider_by_id -from onyx.db.voice import fetch_voice_provider_by_type -from onyx.db.voice import fetch_voice_providers -from onyx.db.voice import MAX_VOICE_PLAYBACK_SPEED -from onyx.db.voice import MIN_VOICE_PLAYBACK_SPEED -from onyx.db.voice import set_default_stt_provider -from onyx.db.voice import set_default_tts_provider -from onyx.db.voice import update_user_voice_settings -from onyx.db.voice import upsert_voice_provider +from onyx.db.voice import ( + deactivate_stt_provider, + deactivate_tts_provider, + delete_voice_provider, + fetch_default_stt_provider, + fetch_default_tts_provider, + fetch_voice_provider_by_id, + fetch_voice_provider_by_type, + fetch_voice_providers, + MAX_VOICE_PLAYBACK_SPEED, + MIN_VOICE_PLAYBACK_SPEED, + set_default_stt_provider, + set_default_tts_provider, + update_user_voice_settings, + upsert_voice_provider, +) from onyx.error_handling.exceptions import OnyxError diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_document_chunk_serialization.py b/backend/tests/unit/onyx/document_index/opensearch/test_document_chunk_serialization.py index beef03f239..5498d58bfb 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_document_chunk_serialization.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_document_chunk_serialization.py @@ -5,12 +5,13 @@ pin that the pydantic model serializes tz-aware datetimes to epoch ints, parses them back to UTC datetimes, and omits the fields entirely when unset. """ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from onyx.document_index.opensearch.schema import CREATED_AT_FIELD_NAME -from onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors -from onyx.document_index.opensearch.schema import LAST_UPDATED_FIELD_NAME +from onyx.document_index.opensearch.schema import ( + CREATED_AT_FIELD_NAME, + DocumentChunkWithoutVectors, + LAST_UPDATED_FIELD_NAME, +) def _make_chunk( diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_index_pair_secondary_missing.py b/backend/tests/unit/onyx/document_index/opensearch/test_index_pair_secondary_missing.py index fd4e1f5c60..665c032ea9 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_index_pair_secondary_missing.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_index_pair_secondary_missing.py @@ -18,15 +18,19 @@ from unittest.mock import MagicMock import pytest -from onyx.document_index.interfaces_new import MetadataUpdateRequest -from onyx.document_index.interfaces_new import SecondaryIndexDocumentMissingError -from onyx.document_index.interfaces_new import TenantState -from onyx.document_index.opensearch.client import OpenSearchDocumentMissingError -from onyx.document_index.opensearch.client import OpenSearchUpdateError +from onyx.document_index.interfaces_new import ( + MetadataUpdateRequest, + SecondaryIndexDocumentMissingError, + TenantState, +) +from onyx.document_index.opensearch.client import ( + OpenSearchDocumentMissingError, + OpenSearchUpdateError, +) from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, + OpenSearchIndexPair, ) -from onyx.document_index.opensearch.opensearch_document_index import OpenSearchIndexPair from onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA 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 03a7d9be78..dd9dd91252 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 @@ -1,7 +1,6 @@ import importlib import os -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from opensearchpy import Urllib3AWSV4SignerAuth diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_batch_flush.py b/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_batch_flush.py index 1b6de729c6..b717206f15 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_batch_flush.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_opensearch_batch_flush.py @@ -1,17 +1,13 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.access.models import DocumentAccess from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import TenantState +from onyx.connectors.models import Document, TextSection +from onyx.document_index.interfaces_new import IndexingMetadata, TenantState from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, ) -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocMetadataAwareIndexChunk +from onyx.indexing.models import ChunkEmbedding, DocMetadataAwareIndexChunk def _make_chunk( diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_port_copy.py b/backend/tests/unit/onyx/document_index/opensearch/test_port_copy.py index b04fd4fa0e..27232a2570 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_port_copy.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_port_copy.py @@ -5,8 +5,7 @@ create-only write (after the slow re-embed) so a doc deleted while the batch was being read/embedded is not resurrected into the FUTURE index. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.document_index.opensearch.port_copy import copy_present_chunks_to_future from onyx.indexing.port_reembed import ReembedStrategy diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_time_cutoff_filter.py b/backend/tests/unit/onyx/document_index/opensearch/test_time_cutoff_filter.py index a8260ebc39..cbc34852c5 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_time_cutoff_filter.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_time_cutoff_filter.py @@ -4,16 +4,16 @@ when undated documents are included, via `DocumentQuery._get_search_filters` document_index/FILTER_SEMANTICS.md ("Time filtering"). """ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from typing import Any from onyx.context.search.models import TimeRange from onyx.document_index.interfaces_new import TenantState from onyx.document_index.opensearch.constants import ASSUMED_DOCUMENT_AGE_DAYS -from onyx.document_index.opensearch.schema import CREATED_AT_FIELD_NAME -from onyx.document_index.opensearch.schema import LAST_UPDATED_FIELD_NAME +from onyx.document_index.opensearch.schema import ( + CREATED_AT_FIELD_NAME, + LAST_UPDATED_FIELD_NAME, +) from onyx.document_index.opensearch.search import DocumentQuery from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA diff --git a/backend/tests/unit/onyx/document_index/test_disabled_document_index.py b/backend/tests/unit/onyx/document_index/test_disabled_document_index.py index a89dd6a9dd..8061f22f99 100644 --- a/backend/tests/unit/onyx/document_index/test_disabled_document_index.py +++ b/backend/tests/unit/onyx/document_index/test_disabled_document_index.py @@ -12,10 +12,8 @@ import pytest from onyx.context.search.enums import QueryType from onyx.context.search.models import IndexFilters from onyx.db.enums import EmbeddingPrecision -from onyx.document_index.disabled import DisabledDocumentIndex -from onyx.document_index.disabled import VECTOR_DB_DISABLED_ERROR -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import MetadataUpdateRequest +from onyx.document_index.disabled import DisabledDocumentIndex, VECTOR_DB_DISABLED_ERROR +from onyx.document_index.interfaces_new import IndexingMetadata, MetadataUpdateRequest ESCAPED_ERROR = re.escape(VECTOR_DB_DISABLED_ERROR) diff --git a/backend/tests/unit/onyx/document_index/vespa/test_vespa_batch_flush.py b/backend/tests/unit/onyx/document_index/vespa/test_vespa_batch_flush.py index 028d13547a..5af8d8b8b8 100644 --- a/backend/tests/unit/onyx/document_index/vespa/test_vespa_batch_flush.py +++ b/backend/tests/unit/onyx/document_index/vespa/test_vespa_batch_flush.py @@ -5,20 +5,15 @@ the streaming logic, ID cleaning/mapping, and DocumentInsertionRecord construction. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.access.models import DocumentAccess from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection -from onyx.document_index.interfaces_new import IndexingMetadata -from onyx.document_index.interfaces_new import TenantState +from onyx.connectors.models import Document, TextSection +from onyx.document_index.interfaces_new import IndexingMetadata, TenantState from onyx.document_index.vespa.internal_types import EnrichedDocumentIndexingInfo from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import IndexChunk +from onyx.indexing.models import ChunkEmbedding, DocMetadataAwareIndexChunk, IndexChunk def _make_chunk( diff --git a/backend/tests/unit/onyx/error_handling/test_exceptions.py b/backend/tests/unit/onyx/error_handling/test_exceptions.py index 45e1b19601..96d4b187c9 100644 --- a/backend/tests/unit/onyx/error_handling/test_exceptions.py +++ b/backend/tests/unit/onyx/error_handling/test_exceptions.py @@ -5,8 +5,7 @@ from fastapi import FastAPI from fastapi.testclient import TestClient from onyx.error_handling.error_codes import OnyxErrorCode -from onyx.error_handling.exceptions import OnyxError -from onyx.error_handling.exceptions import register_onyx_exception_handlers +from onyx.error_handling.exceptions import OnyxError, register_onyx_exception_handlers class TestOnyxError: diff --git a/backend/tests/unit/onyx/federated_connectors/test_federated_connector_factory.py b/backend/tests/unit/onyx/federated_connectors/test_federated_connector_factory.py index cc9de0fc3c..7662573b43 100644 --- a/backend/tests/unit/onyx/federated_connectors/test_federated_connector_factory.py +++ b/backend/tests/unit/onyx/federated_connectors/test_federated_connector_factory.py @@ -7,20 +7,22 @@ Unit tests for federated connector lazy loading factory to validate: """ import importlib -from unittest.mock import MagicMock -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import MagicMock, Mock, patch import pytest from onyx.configs.constants import FederatedConnectorSource -from onyx.federated_connectors.factory import _federated_connector_cache -from onyx.federated_connectors.factory import _load_federated_connector_class -from onyx.federated_connectors.factory import FederatedConnectorMissingException -from onyx.federated_connectors.factory import get_federated_connector_cls +from onyx.federated_connectors.factory import ( + _federated_connector_cache, + _load_federated_connector_class, + FederatedConnectorMissingException, + get_federated_connector_cls, +) from onyx.federated_connectors.interfaces import FederatedConnector -from onyx.federated_connectors.registry import FEDERATED_CONNECTOR_CLASS_MAP -from onyx.federated_connectors.registry import FederatedConnectorMapping +from onyx.federated_connectors.registry import ( + FEDERATED_CONNECTOR_CLASS_MAP, + FederatedConnectorMapping, +) class TestFederatedConnectorMappingValidation: diff --git a/backend/tests/unit/onyx/federated_connectors/test_oauth_utils.py b/backend/tests/unit/onyx/federated_connectors/test_oauth_utils.py index 7850825668..6aa8bc3bd2 100644 --- a/backend/tests/unit/onyx/federated_connectors/test_oauth_utils.py +++ b/backend/tests/unit/onyx/federated_connectors/test_oauth_utils.py @@ -9,12 +9,13 @@ from unittest.mock import patch import pytest -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheLock -from onyx.federated_connectors.oauth_utils import generate_oauth_state -from onyx.federated_connectors.oauth_utils import OAUTH_STATE_TTL -from onyx.federated_connectors.oauth_utils import OAuthSession -from onyx.federated_connectors.oauth_utils import verify_oauth_state +from onyx.cache.interface import CacheBackend, CacheLock +from onyx.federated_connectors.oauth_utils import ( + generate_oauth_state, + OAUTH_STATE_TTL, + OAuthSession, + verify_oauth_state, +) class _MemoryCacheBackend(CacheBackend): diff --git a/backend/tests/unit/onyx/file_processing/test_image_summarization_errors.py b/backend/tests/unit/onyx/file_processing/test_image_summarization_errors.py index 527542d344..b414a275fc 100644 --- a/backend/tests/unit/onyx/file_processing/test_image_summarization_errors.py +++ b/backend/tests/unit/onyx/file_processing/test_image_summarization_errors.py @@ -7,13 +7,14 @@ Verifies that: 3. The ValueError raised on LLM failure preserves the original exception """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.file_processing.image_summarization import _summarize_image -from onyx.file_processing.image_summarization import summarize_image_with_error_handling +from onyx.file_processing.image_summarization import ( + _summarize_image, + summarize_image_with_error_handling, +) class TestSummarizeImageErrorMessage: diff --git a/backend/tests/unit/onyx/file_processing/test_image_utils.py b/backend/tests/unit/onyx/file_processing/test_image_utils.py index 2a1da38aba..b8ad4666ee 100644 --- a/backend/tests/unit/onyx/file_processing/test_image_utils.py +++ b/backend/tests/unit/onyx/file_processing/test_image_utils.py @@ -1,9 +1,6 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection +from onyx.connectors.models import ImageSection, TabularSection, TextSection from onyx.file_processing.image_utils import make_image_callback # Minimal valid file headers for testing magic-byte detection diff --git a/backend/tests/unit/onyx/file_processing/test_pdf.py b/backend/tests/unit/onyx/file_processing/test_pdf.py index 70da4b982c..3f9b9a8263 100644 --- a/backend/tests/unit/onyx/file_processing/test_pdf.py +++ b/backend/tests/unit/onyx/file_processing/test_pdf.py @@ -16,9 +16,11 @@ import pytest from pypdfium2 import PdfiumError from onyx.file_processing import extract_file_text -from onyx.file_processing.extract_file_text import count_pdf_embedded_images -from onyx.file_processing.extract_file_text import pdf_to_text -from onyx.file_processing.extract_file_text import read_pdf_file +from onyx.file_processing.extract_file_text import ( + count_pdf_embedded_images, + pdf_to_text, + read_pdf_file, +) from onyx.file_processing.password_validation import is_pdf_protected from onyx.utils.process_isolation import IsolatedProcessCrashed diff --git a/backend/tests/unit/onyx/file_processing/test_pptx_to_text.py b/backend/tests/unit/onyx/file_processing/test_pptx_to_text.py index 147ebcc045..d47f2fbbb9 100644 --- a/backend/tests/unit/onyx/file_processing/test_pptx_to_text.py +++ b/backend/tests/unit/onyx/file_processing/test_pptx_to_text.py @@ -5,9 +5,11 @@ from pptx.chart.data import CategoryChartData from pptx.enum.chart import XL_CHART_TYPE from pptx.util import Inches -from onyx.file_processing.extract_file_text import extract_pptx_images -from onyx.file_processing.extract_file_text import pptx_to_text -from onyx.file_processing.extract_file_text import read_pptx_file +from onyx.file_processing.extract_file_text import ( + extract_pptx_images, + pptx_to_text, + read_pptx_file, +) def _make_1x1_png() -> bytes: diff --git a/backend/tests/unit/onyx/file_processing/test_xlsx_to_text.py b/backend/tests/unit/onyx/file_processing/test_xlsx_to_text.py index 9e6846f0f7..322787e41b 100644 --- a/backend/tests/unit/onyx/file_processing/test_xlsx_to_text.py +++ b/backend/tests/unit/onyx/file_processing/test_xlsx_to_text.py @@ -5,9 +5,11 @@ from unittest.mock import patch import openpyxl from openpyxl.worksheet.worksheet import Worksheet -from onyx.file_processing.extract_file_text import _sheet_to_csv -from onyx.file_processing.extract_file_text import xlsx_sheet_extraction -from onyx.file_processing.extract_file_text import xlsx_to_text +from onyx.file_processing.extract_file_text import ( + _sheet_to_csv, + xlsx_sheet_extraction, + xlsx_to_text, +) def _make_xlsx(sheets: dict[str, list[list[str]]]) -> io.BytesIO: diff --git a/backend/tests/unit/onyx/file_store/test_delete_file.py b/backend/tests/unit/onyx/file_store/test_delete_file.py index cfba1657d5..18b16de897 100644 --- a/backend/tests/unit/onyx/file_store/test_delete_file.py +++ b/backend/tests/unit/onyx/file_store/test_delete_file.py @@ -1,7 +1,6 @@ """Tests for FileStore.delete_file error_on_missing behavior.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/file_store/test_document_batch_storage.py b/backend/tests/unit/onyx/file_store/test_document_batch_storage.py index 1cde71e1b2..87c1c34a83 100644 --- a/backend/tests/unit/onyx/file_store/test_document_batch_storage.py +++ b/backend/tests/unit/onyx/file_store/test_document_batch_storage.py @@ -1,7 +1,6 @@ """Tests for FileStoreDocumentBatchStorage.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.file_store.document_batch_storage import FileStoreDocumentBatchStorage from onyx.file_store.file_store import S3BackedFileStore diff --git a/backend/tests/unit/onyx/file_store/test_store_plaintext.py b/backend/tests/unit/onyx/file_store/test_store_plaintext.py index 7a5213bb17..e93e84486a 100644 --- a/backend/tests/unit/onyx/file_store/test_store_plaintext.py +++ b/backend/tests/unit/onyx/file_store/test_store_plaintext.py @@ -1,13 +1,11 @@ """Tests for store_plaintext caching behavior in file_store.utils.""" from io import BytesIO -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from onyx.file_store.models import ChatFileType -from onyx.file_store.utils import load_user_file -from onyx.file_store.utils import store_plaintext +from onyx.file_store.utils import load_user_file, store_plaintext _UTILS_MODULE = "onyx.file_store.utils" diff --git a/backend/tests/unit/onyx/hooks/test_models.py b/backend/tests/unit/onyx/hooks/test_models.py index 0fb53b786e..a878c6e7fa 100644 --- a/backend/tests/unit/onyx/hooks/test_models.py +++ b/backend/tests/unit/onyx/hooks/test_models.py @@ -1,10 +1,8 @@ import pytest from pydantic import ValidationError -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint -from onyx.hooks.models import HookCreateRequest -from onyx.hooks.models import HookUpdateRequest +from onyx.db.enums import HookFailStrategy, HookPoint +from onyx.hooks.models import HookCreateRequest, HookUpdateRequest def test_hook_update_request_rejects_empty() -> None: diff --git a/backend/tests/unit/onyx/hooks/test_query_processing_spec.py b/backend/tests/unit/onyx/hooks/test_query_processing_spec.py index c5bec46499..53bf87b3e6 100644 --- a/backend/tests/unit/onyx/hooks/test_query_processing_spec.py +++ b/backend/tests/unit/onyx/hooks/test_query_processing_spec.py @@ -1,5 +1,4 @@ -from onyx.db.enums import HookFailStrategy -from onyx.db.enums import HookPoint +from onyx.db.enums import HookFailStrategy, HookPoint from onyx.hooks.points.query_processing import QueryProcessingSpec diff --git a/backend/tests/unit/onyx/hooks/test_registry.py b/backend/tests/unit/onyx/hooks/test_registry.py index 8e877ab3d0..84e3fa8774 100644 --- a/backend/tests/unit/onyx/hooks/test_registry.py +++ b/backend/tests/unit/onyx/hooks/test_registry.py @@ -2,9 +2,7 @@ import pytest from onyx.db.enums import HookPoint from onyx.hooks import registry as registry_module -from onyx.hooks.registry import get_all_specs -from onyx.hooks.registry import get_hook_point_spec -from onyx.hooks.registry import validate_registry +from onyx.hooks.registry import get_all_specs, get_hook_point_spec, validate_registry def test_registry_covers_all_hook_points() -> None: diff --git a/backend/tests/unit/onyx/image_gen/test_generation.py b/backend/tests/unit/onyx/image_gen/test_generation.py index 320d61461b..5b0319e066 100644 --- a/backend/tests/unit/onyx/image_gen/test_generation.py +++ b/backend/tests/unit/onyx/image_gen/test_generation.py @@ -1,14 +1,14 @@ from types import SimpleNamespace -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.image_gen.exceptions import ImageGenerationNotConfiguredError -from onyx.image_gen.generation import generate_images_with_default_config -from onyx.image_gen.generation import is_image_generation_configured -from onyx.image_gen.interfaces import ImageShape -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.generation import ( + generate_images_with_default_config, + is_image_generation_configured, +) +from onyx.image_gen.interfaces import ImageShape, ReferenceImage _HELPER = "onyx.image_gen.generation" diff --git a/backend/tests/unit/onyx/image_gen/test_provider_building.py b/backend/tests/unit/onyx/image_gen/test_provider_building.py index f88d9694f7..a8f305cc48 100644 --- a/backend/tests/unit/onyx/image_gen/test_provider_building.py +++ b/backend/tests/unit/onyx/image_gen/test_provider_building.py @@ -5,8 +5,7 @@ import pytest from onyx.image_gen.exceptions import ImageProviderCredentialsError from onyx.image_gen.factory import get_image_generation_provider -from onyx.image_gen.interfaces import ImageGenerationProviderCredentials -from onyx.image_gen.interfaces import ReferenceImage +from onyx.image_gen.interfaces import ImageGenerationProviderCredentials, ReferenceImage from onyx.image_gen.providers.azure_img_gen import AzureImageGenerationProvider from onyx.image_gen.providers.openai_img_gen import OpenAIImageGenerationProvider from onyx.image_gen.providers.vertex_img_gen import VertexImageGenerationProvider diff --git a/backend/tests/unit/onyx/indexing/test_censoring.py b/backend/tests/unit/onyx/indexing/test_censoring.py index e14dd4d88b..0de902970d 100644 --- a/backend/tests/unit/onyx/indexing/test_censoring.py +++ b/backend/tests/unit/onyx/indexing/test_censoring.py @@ -1,6 +1,5 @@ import os -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/indexing/test_chunker.py b/backend/tests/unit/onyx/indexing/test_chunker.py index b61365c2a4..7192092f5c 100644 --- a/backend/tests/unit/onyx/indexing/test_chunker.py +++ b/backend/tests/unit/onyx/indexing/test_chunker.py @@ -3,11 +3,9 @@ from unittest.mock import Mock import pytest -from onyx.configs.app_configs import USE_CHUNK_SUMMARY -from onyx.configs.app_configs import USE_DOCUMENT_SUMMARY +from onyx.configs.app_configs import USE_CHUNK_SUMMARY, USE_DOCUMENT_SUMMARY from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, TextSection from onyx.indexing.chunker import Chunker from onyx.indexing.embedder import DefaultIndexingEmbedder from onyx.indexing.indexing_pipeline import process_image_sections diff --git a/backend/tests/unit/onyx/indexing/test_document_chunker.py b/backend/tests/unit/onyx/indexing/test_document_chunker.py index 1e02c5c7b9..e7b5b3cf93 100644 --- a/backend/tests/unit/onyx/indexing/test_document_chunker.py +++ b/backend/tests/unit/onyx/indexing/test_document_chunker.py @@ -3,12 +3,13 @@ import io import pytest from chonkie import SentenceChunker -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import SECTION_SEPARATOR -from onyx.connectors.models import IndexingDocument -from onyx.connectors.models import Section -from onyx.connectors.models import SectionType -from onyx.connectors.models import TabularSection +from onyx.configs.constants import DocumentSource, SECTION_SEPARATOR +from onyx.connectors.models import ( + IndexingDocument, + Section, + SectionType, + TabularSection, +) from onyx.indexing.chunking import DocumentChunker from onyx.indexing.chunking import text_section_chunker as text_chunker_module from onyx.natural_language_processing.utils import BaseTokenizer diff --git a/backend/tests/unit/onyx/indexing/test_document_push.py b/backend/tests/unit/onyx/indexing/test_document_push.py index 5c4be0fe5a..b60abe7f6f 100644 --- a/backend/tests/unit/onyx/indexing/test_document_push.py +++ b/backend/tests/unit/onyx/indexing/test_document_push.py @@ -1,15 +1,16 @@ """Unit tests for the config-driven document push (non-EE).""" from collections.abc import Iterator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import httpx import pytest -from onyx.indexing.document_push import DocumentPushPayload -from onyx.indexing.document_push import get_document_push_config -from onyx.indexing.document_push import push_document_via_config +from onyx.indexing.document_push import ( + DocumentPushPayload, + get_document_push_config, + push_document_via_config, +) _MODULE = "onyx.indexing.document_push" diff --git a/backend/tests/unit/onyx/indexing/test_embed_chunks_in_batches.py b/backend/tests/unit/onyx/indexing/test_embed_chunks_in_batches.py index 2a9c0d5bc2..ce069b7faa 100644 --- a/backend/tests/unit/onyx/indexing/test_embed_chunks_in_batches.py +++ b/backend/tests/unit/onyx/indexing/test_embed_chunks_in_batches.py @@ -11,19 +11,18 @@ Tests cover: """ from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.connectors.models import ConnectorFailure -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentFailure -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + ConnectorFailure, + Document, + DocumentFailure, + DocumentSource, + TextSection, +) from onyx.indexing.chunk_batch_store import ChunkBatchStore from onyx.indexing.indexing_pipeline import _embed_chunks_to_store -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import IndexChunk +from onyx.indexing.models import ChunkEmbedding, DocAwareChunk, IndexChunk def _make_doc(doc_id: str) -> Document: diff --git a/backend/tests/unit/onyx/indexing/test_embedder.py b/backend/tests/unit/onyx/indexing/test_embedder.py index 5b4717378a..462d5775a7 100644 --- a/backend/tests/unit/onyx/indexing/test_embedder.py +++ b/backend/tests/unit/onyx/indexing/test_embedder.py @@ -1,18 +1,13 @@ from collections.abc import Generator -from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection +from onyx.connectors.models import Document, TextSection from onyx.indexing.embedder import DefaultIndexingEmbedder -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocAwareChunk -from onyx.indexing.models import IndexChunk -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType +from onyx.indexing.models import ChunkEmbedding, DocAwareChunk, IndexChunk +from shared_configs.enums import EmbeddingProvider, EmbedTextType @pytest.fixture diff --git a/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py b/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py index 9a81cd66bc..3e18695eca 100644 --- a/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py +++ b/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py @@ -1,38 +1,36 @@ import random import threading import time -from datetime import datetime -from datetime import timezone -from typing import Any -from typing import cast -from typing import List -from unittest.mock import MagicMock -from unittest.mock import Mock -from unittest.mock import patch +from datetime import datetime, timezone +from typing import Any, cast, List +from unittest.mock import MagicMock, Mock, patch import pytest -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import ImageSection -from onyx.connectors.models import TabularSection -from onyx.connectors.models import TextSection -from onyx.hooks.executor import HookSkipped -from onyx.hooks.executor import HookSoftFailed -from onyx.hooks.points.document_ingestion import DocumentIngestionResponse -from onyx.hooks.points.document_ingestion import DocumentIngestionSection +from onyx.connectors.models import ( + Document, + DocumentSource, + ImageSection, + TabularSection, + TextSection, +) +from onyx.hooks.executor import HookSkipped, HookSoftFailed +from onyx.hooks.points.document_ingestion import ( + DocumentIngestionResponse, + DocumentIngestionSection, +) from onyx.indexing.chunker import Chunker from onyx.indexing.embedder import DefaultIndexingEmbedder -from onyx.indexing.indexing_pipeline import _apply_document_ingestion_hook -from onyx.indexing.indexing_pipeline import add_contextual_summaries -from onyx.indexing.indexing_pipeline import filter_documents -from onyx.indexing.indexing_pipeline import get_docs_to_update -from onyx.indexing.indexing_pipeline import process_image_sections +from onyx.indexing.indexing_pipeline import ( + _apply_document_ingestion_hook, + add_contextual_summaries, + filter_documents, + get_docs_to_update, + process_image_sections, +) from onyx.llm.constants import LlmProviderNames from onyx.llm.model_capabilities import get_max_input_tokens -from onyx.llm.model_response import Choice -from onyx.llm.model_response import Message -from onyx.llm.model_response import ModelResponse +from onyx.llm.model_response import Choice, Message, ModelResponse def create_test_document( diff --git a/backend/tests/unit/onyx/indexing/test_persistent_indexing_helpers.py b/backend/tests/unit/onyx/indexing/test_persistent_indexing_helpers.py index 14f41e1e37..69ce594acf 100644 --- a/backend/tests/unit/onyx/indexing/test_persistent_indexing_helpers.py +++ b/backend/tests/unit/onyx/indexing/test_persistent_indexing_helpers.py @@ -1,16 +1,16 @@ """Unit tests for the PERSISTENT_INDEXING catch-all helpers in `onyx.indexing.persistent_indexing`.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.configs.constants import DocumentSource -from onyx.connectors.models import Document -from onyx.connectors.models import TextSection -from onyx.indexing.persistent_indexing import build_generic_connector_failure -from onyx.indexing.persistent_indexing import record_generic_failure +from onyx.connectors.models import Document, TextSection +from onyx.indexing.persistent_indexing import ( + build_generic_connector_failure, + record_generic_failure, +) def _make_doc( diff --git a/backend/tests/unit/onyx/indexing/test_personas_in_chunks.py b/backend/tests/unit/onyx/indexing/test_personas_in_chunks.py index 19214216d6..79c1de7c81 100644 --- a/backend/tests/unit/onyx/indexing/test_personas_in_chunks.py +++ b/backend/tests/unit/onyx/indexing/test_personas_in_chunks.py @@ -4,17 +4,12 @@ Covers Phase 1 (schema plumbing) and Phase 2 (write at index time) of the unify-assistant-project-files plan. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from onyx.access.models import DocumentAccess -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import TextSection -from onyx.indexing.models import ChunkEmbedding -from onyx.indexing.models import DocMetadataAwareIndexChunk -from onyx.indexing.models import IndexChunk +from onyx.connectors.models import Document, DocumentSource, TextSection +from onyx.indexing.models import ChunkEmbedding, DocMetadataAwareIndexChunk, IndexChunk def _make_index_chunk( diff --git a/backend/tests/unit/onyx/indexing/test_tabular_section_chunker.py b/backend/tests/unit/onyx/indexing/test_tabular_section_chunker.py index c5f57e6815..514b652d3b 100644 --- a/backend/tests/unit/onyx/indexing/test_tabular_section_chunker.py +++ b/backend/tests/unit/onyx/indexing/test_tabular_section_chunker.py @@ -16,8 +16,7 @@ from collections.abc import Iterator import pytest -from onyx.connectors.models import Section -from onyx.connectors.models import TabularSection +from onyx.connectors.models import Section, TabularSection from onyx.indexing.chunking.section_chunker import AccumulatorState from onyx.indexing.chunking.tabular_section_chunker import TabularChunker from onyx.indexing.chunking.tabular_section_chunker.analysis import analyze_sheet @@ -26,13 +25,10 @@ from onyx.indexing.chunking.tabular_section_chunker.sheet_descriptor import ( ) from onyx.indexing.chunking.tabular_section_chunker.total_descriptor import ( build_total_descriptor_chunks, -) -from onyx.indexing.chunking.tabular_section_chunker.total_descriptor import ( TOTALS_HEADER, ) from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.utils.csv_utils import parse_csv_string -from onyx.utils.csv_utils import read_csv_header +from onyx.utils.csv_utils import parse_csv_string, read_csv_header class CharTokenizer(BaseTokenizer): diff --git a/backend/tests/unit/onyx/llm/test_factory.py b/backend/tests/unit/onyx/llm/test_factory.py index 08c41ca542..057b5e9b64 100644 --- a/backend/tests/unit/onyx/llm/test_factory.py +++ b/backend/tests/unit/onyx/llm/test_factory.py @@ -1,12 +1,9 @@ from unittest.mock import patch from onyx.llm.constants import LlmProviderNames -from onyx.llm.factory import _build_provider_extra_headers -from onyx.llm.factory import get_llm -from onyx.llm.factory import llm_from_provider +from onyx.llm.factory import _build_provider_extra_headers, get_llm, llm_from_provider from onyx.llm.well_known_providers.constants import OLLAMA_API_KEY_CONFIG_KEY -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import ModelConfigurationView +from onyx.server.manage.llm.models import LLMProviderView, ModelConfigurationView def test_build_provider_extra_headers_adds_bearer_for_ollama_api_key() -> None: diff --git a/backend/tests/unit/onyx/llm/test_litellm_monkey_patches.py b/backend/tests/unit/onyx/llm/test_litellm_monkey_patches.py index 12f141a22f..406896aed7 100644 --- a/backend/tests/unit/onyx/llm/test_litellm_monkey_patches.py +++ b/backend/tests/unit/onyx/llm/test_litellm_monkey_patches.py @@ -6,13 +6,14 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.ollama.chat.transformation import OllamaChatCompletionResponseIterator -from litellm.types.llms.openai import ResponseAPIUsage -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.utils import ModelResponse from openai.types.responses.response_output_message import ResponseOutputMessage from openai.types.responses.response_output_text import ResponseOutputText -from openai.types.responses.response_reasoning_item import ResponseReasoningItem -from openai.types.responses.response_reasoning_item import Summary +from openai.types.responses.response_reasoning_item import ( + ResponseReasoningItem, + Summary, +) from onyx.llm.litellm_singleton.monkey_patches import apply_monkey_patches diff --git a/backend/tests/unit/onyx/llm/test_llm_exception_classification.py b/backend/tests/unit/onyx/llm/test_llm_exception_classification.py index 3413a71c3f..742646bb71 100644 --- a/backend/tests/unit/onyx/llm/test_llm_exception_classification.py +++ b/backend/tests/unit/onyx/llm/test_llm_exception_classification.py @@ -4,10 +4,12 @@ BadRequestError and must be matched first, or context overflow is mislabeled BAD_REQUEST instead of CONTEXT_TOO_LONG. """ -from litellm.exceptions import APIError -from litellm.exceptions import BadRequestError -from litellm.exceptions import ContentPolicyViolationError -from litellm.exceptions import ContextWindowExceededError +from litellm.exceptions import ( + APIError, + BadRequestError, + ContentPolicyViolationError, + ContextWindowExceededError, +) from onyx.llm.utils import litellm_exception_to_error_msg diff --git a/backend/tests/unit/onyx/llm/test_llm_provider_options.py b/backend/tests/unit/onyx/llm/test_llm_provider_options.py index df4c0c2f50..6396a43cf5 100644 --- a/backend/tests/unit/onyx/llm/test_llm_provider_options.py +++ b/backend/tests/unit/onyx/llm/test_llm_provider_options.py @@ -1,12 +1,15 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone import pytest -from onyx.llm.well_known_providers.auto_update_models import LLMProviderRecommendation -from onyx.llm.well_known_providers.auto_update_models import LLMRecommendations -from onyx.llm.well_known_providers.constants import OPENAI_PROVIDER_NAME -from onyx.llm.well_known_providers.constants import VERTEXAI_PROVIDER_NAME +from onyx.llm.well_known_providers.auto_update_models import ( + LLMProviderRecommendation, + LLMRecommendations, +) +from onyx.llm.well_known_providers.constants import ( + OPENAI_PROVIDER_NAME, + VERTEXAI_PROVIDER_NAME, +) from onyx.llm.well_known_providers.llm_provider_options import ( model_configurations_for_provider, ) diff --git a/backend/tests/unit/onyx/llm/test_llm_utils_scrub.py b/backend/tests/unit/onyx/llm/test_llm_utils_scrub.py index ad623ec695..170a7b33c4 100644 --- a/backend/tests/unit/onyx/llm/test_llm_utils_scrub.py +++ b/backend/tests/unit/onyx/llm/test_llm_utils_scrub.py @@ -6,11 +6,12 @@ from echoing API keys back through error messages. from collections.abc import Iterable from typing import Any -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMConfig -from onyx.llm.utils import collect_llm_credential_values -from onyx.llm.utils import is_sensitive_custom_config_key -from onyx.llm.utils import scrub_sensitive_values +from onyx.llm.interfaces import LLM, LLMConfig +from onyx.llm.utils import ( + collect_llm_credential_values, + is_sensitive_custom_config_key, + scrub_sensitive_values, +) from onyx.llm.utils import ( test_llm as run_test_llm, ) # aliased to avoid pytest collection diff --git a/backend/tests/unit/onyx/llm/test_model_map.py b/backend/tests/unit/onyx/llm/test_model_map.py index e4227cc7fa..4570f398fe 100644 --- a/backend/tests/unit/onyx/llm/test_model_map.py +++ b/backend/tests/unit/onyx/llm/test_model_map.py @@ -4,9 +4,11 @@ import litellm from onyx.configs.model_configs import GEN_AI_MODEL_FALLBACK_MAX_TOKENS from onyx.llm.constants import LlmProviderNames -from onyx.llm.model_capabilities import find_model_obj -from onyx.llm.model_capabilities import get_model_map -from onyx.llm.model_capabilities import model_is_reasoning_model +from onyx.llm.model_capabilities import ( + find_model_obj, + get_model_map, + model_is_reasoning_model, +) def test_partial_match_in_model_map() -> None: diff --git a/backend/tests/unit/onyx/llm/test_model_response.py b/backend/tests/unit/onyx/llm/test_model_response.py index 45c64ecf60..fe85fdb8b1 100644 --- a/backend/tests/unit/onyx/llm/test_model_response.py +++ b/backend/tests/unit/onyx/llm/test_model_response.py @@ -1,16 +1,17 @@ from __future__ import annotations -from typing import cast -from typing import TYPE_CHECKING +from typing import cast, TYPE_CHECKING import pytest -from onyx.llm.model_response import ChatCompletionDeltaToolCall -from onyx.llm.model_response import from_litellm_model_response -from onyx.llm.model_response import from_litellm_model_response_stream -from onyx.llm.model_response import FunctionCall -from onyx.llm.model_response import ModelResponse -from onyx.llm.model_response import ModelResponseStream +from onyx.llm.model_response import ( + ChatCompletionDeltaToolCall, + from_litellm_model_response, + from_litellm_model_response_stream, + FunctionCall, + ModelResponse, + ModelResponseStream, +) if TYPE_CHECKING: from litellm.types.utils import ModelResponse as LiteLLMModelResponse diff --git a/backend/tests/unit/onyx/llm/test_multi_llm.py b/backend/tests/unit/onyx/llm/test_multi_llm.py index 6440b67d4f..30eea18ba3 100644 --- a/backend/tests/unit/onyx/llm/test_multi_llm.py +++ b/backend/tests/unit/onyx/llm/test_multi_llm.py @@ -2,13 +2,11 @@ import os import threading import time from typing import Any -from unittest.mock import ANY -from unittest.mock import patch +from unittest.mock import ANY, patch import litellm import pytest -from litellm.types.utils import ChatCompletionDeltaToolCall -from litellm.types.utils import Delta +from litellm.types.utils import ChatCompletionDeltaToolCall, Delta from litellm.types.utils import Function as LiteLLMFunction import onyx.llm.models @@ -16,18 +14,21 @@ from onyx.configs.app_configs import MOCK_LLM_RESPONSE from onyx.llm.constants import LlmProviderNames from onyx.llm.interfaces import LLMUserIdentity from onyx.llm.model_capabilities import get_max_input_tokens -from onyx.llm.model_response import ModelResponse -from onyx.llm.model_response import ModelResponseStream -from onyx.llm.models import AssistantMessage -from onyx.llm.models import FunctionCall -from onyx.llm.models import LanguageModelInput -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import ToolCall -from onyx.llm.models import ToolChoiceOptions -from onyx.llm.models import UserMessage -from onyx.llm.multi_llm import _parse_anthropic_model_version -from onyx.llm.multi_llm import LitellmLLM -from onyx.llm.multi_llm import temporary_env_and_lock +from onyx.llm.model_response import ModelResponse, ModelResponseStream +from onyx.llm.models import ( + AssistantMessage, + FunctionCall, + LanguageModelInput, + ReasoningEffort, + ToolCall, + ToolChoiceOptions, + UserMessage, +) +from onyx.llm.multi_llm import ( + _parse_anthropic_model_version, + LitellmLLM, + temporary_env_and_lock, +) VERTEX_OPUS_MODELS_REJECTING_OUTPUT_CONFIG = [ "claude-opus-4-5@20251101", diff --git a/backend/tests/unit/onyx/llm/test_multi_llm_stream_retry.py b/backend/tests/unit/onyx/llm/test_multi_llm_stream_retry.py index 84cd5755b0..6aea4c8990 100644 --- a/backend/tests/unit/onyx/llm/test_multi_llm_stream_retry.py +++ b/backend/tests/unit/onyx/llm/test_multi_llm_stream_retry.py @@ -1,14 +1,11 @@ from collections.abc import Iterator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from litellm.exceptions import Timeout as LiteLLMTimeout from onyx.llm.interfaces import LanguageModelInput -from onyx.llm.model_response import Delta -from onyx.llm.model_response import ModelResponseStream -from onyx.llm.model_response import StreamingChoice +from onyx.llm.model_response import Delta, ModelResponseStream, StreamingChoice from onyx.llm.models import UserMessage from onyx.llm.multi_llm import LitellmLLM diff --git a/backend/tests/unit/onyx/llm/test_reasoning_effort_mapping.py b/backend/tests/unit/onyx/llm/test_reasoning_effort_mapping.py index b5c9a1c42d..5f49cb5215 100644 --- a/backend/tests/unit/onyx/llm/test_reasoning_effort_mapping.py +++ b/backend/tests/unit/onyx/llm/test_reasoning_effort_mapping.py @@ -1,5 +1,4 @@ -from onyx.llm.models import OPENAI_REASONING_EFFORT -from onyx.llm.models import ReasoningEffort +from onyx.llm.models import OPENAI_REASONING_EFFORT, ReasoningEffort # Valid OpenAI reasoning effort values per the API documentation # https://platform.openai.com/docs/api-reference/responses diff --git a/backend/tests/unit/onyx/llm/test_request_context.py b/backend/tests/unit/onyx/llm/test_request_context.py index 27a53b1378..610ea1fb1a 100644 --- a/backend/tests/unit/onyx/llm/test_request_context.py +++ b/backend/tests/unit/onyx/llm/test_request_context.py @@ -1,8 +1,10 @@ import contextvars -from onyx.llm.request_context import get_llm_mock_response -from onyx.llm.request_context import reset_llm_mock_response -from onyx.llm.request_context import set_llm_mock_response +from onyx.llm.request_context import ( + get_llm_mock_response, + reset_llm_mock_response, + set_llm_mock_response, +) def test_reset_llm_mock_response_same_context() -> None: diff --git a/backend/tests/unit/onyx/llm/test_token_limit_lookups.py b/backend/tests/unit/onyx/llm/test_token_limit_lookups.py index 81ae4292cc..91f8d64294 100644 --- a/backend/tests/unit/onyx/llm/test_token_limit_lookups.py +++ b/backend/tests/unit/onyx/llm/test_token_limit_lookups.py @@ -4,9 +4,11 @@ from unittest.mock import patch from onyx.configs.model_configs import GEN_AI_MODEL_FALLBACK_MAX_TOKENS -from onyx.llm.model_capabilities import get_llm_max_output_tokens -from onyx.llm.model_capabilities import get_max_input_tokens -from onyx.llm.model_capabilities import llm_max_input_tokens +from onyx.llm.model_capabilities import ( + get_llm_max_output_tokens, + get_max_input_tokens, + llm_max_input_tokens, +) class TestLlmMaxInputTokens: diff --git a/backend/tests/unit/onyx/llm/test_tracing_wrap.py b/backend/tests/unit/onyx/llm/test_tracing_wrap.py index e24720a552..aa76d86ca6 100644 --- a/backend/tests/unit/onyx/llm/test_tracing_wrap.py +++ b/backend/tests/unit/onyx/llm/test_tracing_wrap.py @@ -18,41 +18,41 @@ declared but unused — hence the `ARG002` suppression at the file level. from __future__ import annotations -from collections.abc import Callable -from collections.abc import Generator -from collections.abc import Iterator -from typing import Any -from typing import cast +from collections.abc import Callable, Generator, Iterator +from typing import Any, cast from unittest.mock import patch import pytest -from onyx.llm.interfaces import LLM -from onyx.llm.interfaces import LLMConfig -from onyx.llm.interfaces import LLMUserIdentity -from onyx.llm.model_response import ChatCompletionDeltaToolCall -from onyx.llm.model_response import Choice -from onyx.llm.model_response import Delta +from onyx.llm.interfaces import LLM, LLMConfig, LLMUserIdentity +from onyx.llm.model_response import ( + ChatCompletionDeltaToolCall, + Choice, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoice, + Usage, +) from onyx.llm.model_response import FunctionCall as DeltaFunctionCall -from onyx.llm.model_response import Message -from onyx.llm.model_response import ModelResponse -from onyx.llm.model_response import ModelResponseStream -from onyx.llm.model_response import StreamingChoice -from onyx.llm.model_response import Usage -from onyx.llm.models import LanguageModelInput -from onyx.llm.models import ReasoningEffort -from onyx.llm.models import ToolChoiceOptions -from onyx.llm.models import UserMessage -from onyx.llm.tracing_wrap import _ALREADY_WRAPPED_ATTR -from onyx.llm.tracing_wrap import _extract_prompt_and_tools -from onyx.llm.tracing_wrap import _finalize_tool_calls -from onyx.llm.tracing_wrap import _merge_tool_call_delta -from onyx.llm.tracing_wrap import _outer_generation_span_active -from onyx.llm.tracing_wrap import _validate_prompt_param -from onyx.llm.tracing_wrap import wrap_invoke -from onyx.llm.tracing_wrap import wrap_stream -from onyx.tracing.framework.create import generation_span -from onyx.tracing.framework.create import trace +from onyx.llm.models import ( + LanguageModelInput, + ReasoningEffort, + ToolChoiceOptions, + UserMessage, +) +from onyx.llm.tracing_wrap import ( + _ALREADY_WRAPPED_ATTR, + _extract_prompt_and_tools, + _finalize_tool_calls, + _merge_tool_call_delta, + _outer_generation_span_active, + _validate_prompt_param, + wrap_invoke, + wrap_stream, +) +from onyx.tracing.framework.create import generation_span, trace _TEST_MODEL_RESPONSE = ModelResponse( id="test-id", diff --git a/backend/tests/unit/onyx/llm/test_true_openai_model.py b/backend/tests/unit/onyx/llm/test_true_openai_model.py index b25cc8dbce..33f348b2fc 100644 --- a/backend/tests/unit/onyx/llm/test_true_openai_model.py +++ b/backend/tests/unit/onyx/llm/test_true_openai_model.py @@ -1,6 +1,5 @@ from onyx.llm.constants import LlmProviderNames -from onyx.llm.model_capabilities import get_model_map -from onyx.llm.model_capabilities import is_true_openai_model +from onyx.llm.model_capabilities import get_model_map, is_true_openai_model class TestIsTrueOpenAIModel: diff --git a/backend/tests/unit/onyx/llm/test_vision_model_selection_logging.py b/backend/tests/unit/onyx/llm/test_vision_model_selection_logging.py index 36f9ec8e96..7c8d44cc14 100644 --- a/backend/tests/unit/onyx/llm/test_vision_model_selection_logging.py +++ b/backend/tests/unit/onyx/llm/test_vision_model_selection_logging.py @@ -7,8 +7,7 @@ Verifies that operators get clear feedback about: 3. When no vision-capable model exists at all """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.llm.factory import get_default_llm_with_vision diff --git a/backend/tests/unit/onyx/natural_language_processing/test_search_nlp_models.py b/backend/tests/unit/onyx/natural_language_processing/test_search_nlp_models.py index cf42cabe56..7e66a89215 100644 --- a/backend/tests/unit/onyx/natural_language_processing/test_search_nlp_models.py +++ b/backend/tests/unit/onyx/natural_language_processing/test_search_nlp_models.py @@ -1,10 +1,7 @@ from collections.abc import AsyncGenerator from threading import Lock -from typing import Any -from typing import cast -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest from httpx import AsyncClient @@ -12,13 +9,13 @@ from litellm.exceptions import RateLimitError from tenacity import wait_none from onyx.llm.constants import LlmProviderNames -from onyx.natural_language_processing.search_nlp_models import clean_model_name -from onyx.natural_language_processing.search_nlp_models import CloudEmbedding -from onyx.natural_language_processing.search_nlp_models import EmbeddingModel -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType -from shared_configs.model_server_models import EmbedRequest -from shared_configs.model_server_models import EmbedResponse +from onyx.natural_language_processing.search_nlp_models import ( + clean_model_name, + CloudEmbedding, + EmbeddingModel, +) +from shared_configs.enums import EmbeddingProvider, EmbedTextType +from shared_configs.model_server_models import EmbedRequest, EmbedResponse @pytest.fixture diff --git a/backend/tests/unit/onyx/onyxbot/discord/conftest.py b/backend/tests/unit/onyx/onyxbot/discord/conftest.py index 6ace16b7de..9588087f17 100644 --- a/backend/tests/unit/onyx/onyxbot/discord/conftest.py +++ b/backend/tests/unit/onyx/onyxbot/discord/conftest.py @@ -3,8 +3,7 @@ import random from collections.abc import Callable from typing import Any -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import discord import pytest diff --git a/backend/tests/unit/onyx/onyxbot/discord/test_api_client.py b/backend/tests/unit/onyx/onyxbot/discord/test_api_client.py index b4e19398fe..1718f76deb 100644 --- a/backend/tests/unit/onyx/onyxbot/discord/test_api_client.py +++ b/backend/tests/unit/onyx/onyxbot/discord/test_api_client.py @@ -4,9 +4,7 @@ Tests for OnyxAPIClient class functionality. """ from typing import Any -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import aiohttp import pytest @@ -14,9 +12,11 @@ import pytest from onyx.chat.models import ChatFullResponse from onyx.onyxbot.discord.api_client import OnyxAPIClient from onyx.onyxbot.discord.constants import API_REQUEST_TIMEOUT -from onyx.onyxbot.discord.exceptions import APIConnectionError -from onyx.onyxbot.discord.exceptions import APIResponseError -from onyx.onyxbot.discord.exceptions import APITimeoutError +from onyx.onyxbot.discord.exceptions import ( + APIConnectionError, + APIResponseError, + APITimeoutError, +) class MockAsyncContextManager: diff --git a/backend/tests/unit/onyx/onyxbot/discord/test_cache_manager.py b/backend/tests/unit/onyx/onyxbot/discord/test_cache_manager.py index 81caed9832..16a0bac7d4 100644 --- a/backend/tests/unit/onyx/onyxbot/discord/test_cache_manager.py +++ b/backend/tests/unit/onyx/onyxbot/discord/test_cache_manager.py @@ -5,13 +5,11 @@ Tests for DiscordCacheManager class functionality. import asyncio import time -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.onyxbot.discord.cache import DiscordCacheManager -from onyx.onyxbot.discord.cache import TenantDiscordData +from onyx.onyxbot.discord.cache import DiscordCacheManager, TenantDiscordData class TestCacheInitialization: diff --git a/backend/tests/unit/onyx/onyxbot/discord/test_context_builders.py b/backend/tests/unit/onyx/onyxbot/discord/test_context_builders.py index 4c42ce6fba..fa53db8933 100644 --- a/backend/tests/unit/onyx/onyxbot/discord/test_context_builders.py +++ b/backend/tests/unit/onyx/onyxbot/discord/test_context_builders.py @@ -4,20 +4,20 @@ Tests the thread and reply context building logic with mocked Discord API. """ from typing import Any -from unittest.mock import AsyncMock -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import discord import pytest from onyx.onyxbot.discord.constants import MAX_CONTEXT_MESSAGES -from onyx.onyxbot.discord.handle_message import _build_conversation_context -from onyx.onyxbot.discord.handle_message import _build_reply_chain_context -from onyx.onyxbot.discord.handle_message import _build_thread_context -from onyx.onyxbot.discord.handle_message import _format_messages_as_context -from onyx.onyxbot.discord.handle_message import format_message_content -from tests.unit.onyx.onyxbot.discord.conftest import AsyncIteratorMock -from tests.unit.onyx.onyxbot.discord.conftest import mock_message +from onyx.onyxbot.discord.handle_message import ( + _build_conversation_context, + _build_reply_chain_context, + _build_thread_context, + _format_messages_as_context, + format_message_content, +) +from tests.unit.onyx.onyxbot.discord.conftest import AsyncIteratorMock, mock_message class TestThreadContextBuilder: diff --git a/backend/tests/unit/onyx/onyxbot/discord/test_discord_utils.py b/backend/tests/unit/onyx/onyxbot/discord/test_discord_utils.py index d638a3a4f3..23821ee74b 100644 --- a/backend/tests/unit/onyx/onyxbot/discord/test_discord_utils.py +++ b/backend/tests/unit/onyx/onyxbot/discord/test_discord_utils.py @@ -5,13 +5,14 @@ Tests for: - Registration key parsing (parse_discord_registration_key, generate_discord_registration_key) """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.onyxbot.discord.utils import get_bot_token -from onyx.server.manage.discord_bot.utils import generate_discord_registration_key -from onyx.server.manage.discord_bot.utils import parse_discord_registration_key -from onyx.server.manage.discord_bot.utils import REGISTRATION_KEY_PREFIX +from onyx.server.manage.discord_bot.utils import ( + generate_discord_registration_key, + parse_discord_registration_key, + REGISTRATION_KEY_PREFIX, +) class TestGetBotToken: diff --git a/backend/tests/unit/onyx/onyxbot/discord/test_message_utils.py b/backend/tests/unit/onyx/onyxbot/discord/test_message_utils.py index 01ba98690f..03f2600b07 100644 --- a/backend/tests/unit/onyx/onyxbot/discord/test_message_utils.py +++ b/backend/tests/unit/onyx/onyxbot/discord/test_message_utils.py @@ -9,8 +9,7 @@ from unittest.mock import MagicMock from onyx.chat.models import ChatFullResponse from onyx.onyxbot.discord.constants import MAX_MESSAGE_LENGTH -from onyx.onyxbot.discord.handle_message import _append_citations -from onyx.onyxbot.discord.handle_message import _split_message +from onyx.onyxbot.discord.handle_message import _append_citations, _split_message class TestSplitMessage: diff --git a/backend/tests/unit/onyx/onyxbot/discord/test_should_respond.py b/backend/tests/unit/onyx/onyxbot/discord/test_should_respond.py index 72ac3f309d..048425ef6b 100644 --- a/backend/tests/unit/onyx/onyxbot/discord/test_should_respond.py +++ b/backend/tests/unit/onyx/onyxbot/discord/test_should_respond.py @@ -3,15 +3,15 @@ Tests the decision tree for when the bot should respond to messages. """ -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import discord import pytest -from onyx.onyxbot.discord.handle_message import check_implicit_invocation -from onyx.onyxbot.discord.handle_message import should_respond +from onyx.onyxbot.discord.handle_message import ( + check_implicit_invocation, + should_respond, +) class TestBasicShouldRespond: diff --git a/backend/tests/unit/onyx/onyxbot/test_handle_regular_answer.py b/backend/tests/unit/onyx/onyxbot/test_handle_regular_answer.py index a78e19df52..015e90dc30 100644 --- a/backend/tests/unit/onyx/onyxbot/test_handle_regular_answer.py +++ b/backend/tests/unit/onyx/onyxbot/test_handle_regular_answer.py @@ -2,8 +2,7 @@ in handle_regular_answer.py.""" from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from slack_sdk.errors import SlackApiError @@ -11,15 +10,17 @@ from slack_sdk.errors import SlackApiError from onyx.chat.models import ChatBasicResponse from onyx.context.search.models import Tag from onyx.onyxbot.slack.constants import SLACK_CHANNEL_REF_PATTERN -from onyx.onyxbot.slack.handlers.handle_regular_answer import handle_regular_answer -from onyx.onyxbot.slack.handlers.handle_regular_answer import resolve_channel_references from onyx.onyxbot.slack.handlers.handle_regular_answer import ( + handle_regular_answer, + resolve_channel_references, SLACK_PERSONA_ACCESS_DENIED_MESSAGE, ) -from onyx.onyxbot.slack.models import ChannelType -from onyx.onyxbot.slack.models import SlackContext -from onyx.onyxbot.slack.models import SlackMessageInfo -from onyx.onyxbot.slack.models import ThreadMessage +from onyx.onyxbot.slack.models import ( + ChannelType, + SlackContext, + SlackMessageInfo, + ThreadMessage, +) _HANDLE_REGULAR_ANSWER = "onyx.onyxbot.slack.handlers.handle_regular_answer" diff --git a/backend/tests/unit/onyx/onyxbot/test_slack_channel_config.py b/backend/tests/unit/onyx/onyxbot/test_slack_channel_config.py index ad158f5aac..ea88462ba9 100644 --- a/backend/tests/unit/onyx/onyxbot/test_slack_channel_config.py +++ b/backend/tests/unit/onyx/onyxbot/test_slack_channel_config.py @@ -1,5 +1,4 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.db.slack_channel_config import create_slack_channel_persona diff --git a/backend/tests/unit/onyx/onyxbot/test_slack_formatting.py b/backend/tests/unit/onyx/onyxbot/test_slack_formatting.py index d65871f3fb..adf54891ac 100644 --- a/backend/tests/unit/onyx/onyxbot/test_slack_formatting.py +++ b/backend/tests/unit/onyx/onyxbot/test_slack_formatting.py @@ -1,8 +1,10 @@ -from onyx.onyxbot.slack.formatting import _convert_slack_links_to_markdown -from onyx.onyxbot.slack.formatting import _normalize_link_destinations -from onyx.onyxbot.slack.formatting import _sanitize_html -from onyx.onyxbot.slack.formatting import _transform_outside_code_blocks -from onyx.onyxbot.slack.formatting import format_slack_message +from onyx.onyxbot.slack.formatting import ( + _convert_slack_links_to_markdown, + _normalize_link_destinations, + _sanitize_html, + _transform_outside_code_blocks, + format_slack_message, +) from onyx.onyxbot.slack.utils import remove_slack_text_interactions from onyx.utils.text_processing import decode_escapes diff --git a/backend/tests/unit/onyx/onyxbot/test_slack_gating.py b/backend/tests/unit/onyx/onyxbot/test_slack_gating.py index e6a1e3a0e0..95bde01fc2 100644 --- a/backend/tests/unit/onyx/onyxbot/test_slack_gating.py +++ b/backend/tests/unit/onyx/onyxbot/test_slack_gating.py @@ -2,8 +2,7 @@ from collections.abc import Generator from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/onyxbot/test_slack_icons.py b/backend/tests/unit/onyx/onyxbot/test_slack_icons.py index 3363b0b199..cf1803ad6c 100644 --- a/backend/tests/unit/onyx/onyxbot/test_slack_icons.py +++ b/backend/tests/unit/onyx/onyxbot/test_slack_icons.py @@ -3,10 +3,12 @@ from pathlib import Path import pytest from onyx.configs.constants import DocumentSource -from onyx.onyxbot.slack.icons import _DEFAULT_SOURCE_IMAGE_FILENAME -from onyx.onyxbot.slack.icons import _PUBLIC_SOURCE_IMAGE_BASE_URL -from onyx.onyxbot.slack.icons import _SOURCE_IMAGE_FILENAMES -from onyx.onyxbot.slack.icons import source_to_github_img_link +from onyx.onyxbot.slack.icons import ( + _DEFAULT_SOURCE_IMAGE_FILENAME, + _PUBLIC_SOURCE_IMAGE_BASE_URL, + _SOURCE_IMAGE_FILENAMES, + source_to_github_img_link, +) SLACK_SUPPORTED_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif"} SOURCE_ICON_DIR = ( diff --git a/backend/tests/unit/onyx/prompts/test_prompt_utils.py b/backend/tests/unit/onyx/prompts/test_prompt_utils.py index 22451912c5..ffae21a232 100644 --- a/backend/tests/unit/onyx/prompts/test_prompt_utils.py +++ b/backend/tests/unit/onyx/prompts/test_prompt_utils.py @@ -1,14 +1,17 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.prompts.chat_prompts import CITATION_GUIDANCE_REPLACEMENT_PAT -from onyx.prompts.chat_prompts import DATETIME_REPLACEMENT_PAT -from onyx.prompts.chat_prompts import REQUIRE_CITATION_GUIDANCE +from onyx.prompts.chat_prompts import ( + CITATION_GUIDANCE_REPLACEMENT_PAT, + DATETIME_REPLACEMENT_PAT, + REQUIRE_CITATION_GUIDANCE, +) from onyx.prompts.constants import REMINDER_TAG_DESCRIPTION -from onyx.prompts.prompt_utils import apply_prompt_placeholders -from onyx.prompts.prompt_utils import replace_current_datetime_tag -from onyx.prompts.prompt_utils import replace_reminder_tag -from onyx.prompts.prompt_utils import substitute_user_placeholders +from onyx.prompts.prompt_utils import ( + apply_prompt_placeholders, + replace_current_datetime_tag, + replace_reminder_tag, + substitute_user_placeholders, +) def test_replace_reminder_tag_pattern() -> None: diff --git a/backend/tests/unit/onyx/secondary_llm_flows/test_document_filter.py b/backend/tests/unit/onyx/secondary_llm_flows/test_document_filter.py index f1e4956c9b..f9cf87746e 100644 --- a/backend/tests/unit/onyx/secondary_llm_flows/test_document_filter.py +++ b/backend/tests/unit/onyx/secondary_llm_flows/test_document_filter.py @@ -1,17 +1,20 @@ from collections.abc import Iterator from contextlib import contextmanager -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.configs.chat_configs import SECONDARY_LLM_FLOW_TIMEOUT_S from onyx.configs.constants import DocumentSource -from onyx.context.search.models import ContextExpansionType -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection +from onyx.context.search.models import ( + ContextExpansionType, + InferenceChunk, + InferenceSection, +) from onyx.llm.interfaces import LLM from onyx.llm.multi_llm import LLMTimeoutError -from onyx.secondary_llm_flows.document_filter import classify_section_relevance -from onyx.secondary_llm_flows.document_filter import select_sections_for_expansion +from onyx.secondary_llm_flows.document_filter import ( + classify_section_relevance, + select_sections_for_expansion, +) @contextmanager diff --git a/backend/tests/unit/onyx/secondary_llm_flows/test_source_filter.py b/backend/tests/unit/onyx/secondary_llm_flows/test_source_filter.py index 827098cd62..930a947aa7 100644 --- a/backend/tests/unit/onyx/secondary_llm_flows/test_source_filter.py +++ b/backend/tests/unit/onyx/secondary_llm_flows/test_source_filter.py @@ -1,14 +1,11 @@ from __future__ import annotations from contextlib import nullcontext -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MessageType +from onyx.configs.constants import DocumentSource, MessageType from onyx.llm.models import UserMessage -from onyx.secondary_llm_flows.source_filter import decide_search_scope -from onyx.secondary_llm_flows.source_filter import SearchCycle +from onyx.secondary_llm_flows.source_filter import decide_search_scope, SearchCycle from onyx.tools.models import ChatMinimalTextMessage A = DocumentSource.ZENDESK diff --git a/backend/tests/unit/onyx/secondary_llm_flows/test_time_filter.py b/backend/tests/unit/onyx/secondary_llm_flows/test_time_filter.py index 9e1744a934..905958d239 100644 --- a/backend/tests/unit/onyx/secondary_llm_flows/test_time_filter.py +++ b/backend/tests/unit/onyx/secondary_llm_flows/test_time_filter.py @@ -1,19 +1,18 @@ from __future__ import annotations from contextlib import nullcontext -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch from onyx.configs.constants import MessageType -from onyx.context.search.models import BaseFilters -from onyx.context.search.models import TimeRange +from onyx.context.search.models import BaseFilters, TimeRange from onyx.llm.models import UserMessage -from onyx.secondary_llm_flows.time_filter import _parse_time_decision -from onyx.secondary_llm_flows.time_filter import decide_time_filter -from onyx.secondary_llm_flows.time_filter import DocumentTimeField -from onyx.secondary_llm_flows.time_filter import TimeFilter +from onyx.secondary_llm_flows.time_filter import ( + _parse_time_decision, + decide_time_filter, + DocumentTimeField, + TimeFilter, +) from onyx.tools.models import ChatMinimalTextMessage diff --git a/backend/tests/unit/onyx/server/auth/test_captcha_api.py b/backend/tests/unit/onyx/server/auth/test_captcha_api.py index 41e2c336aa..61015f5026 100644 --- a/backend/tests/unit/onyx/server/auth/test_captcha_api.py +++ b/backend/tests/unit/onyx/server/auth/test_captcha_api.py @@ -1,7 +1,6 @@ """Unit tests for the reCAPTCHA OAuth verify endpoint + cookie middleware.""" -from unittest.mock import AsyncMock -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from fastapi import FastAPI diff --git a/backend/tests/unit/onyx/server/auth/test_login_captcha_middleware.py b/backend/tests/unit/onyx/server/auth/test_login_captcha_middleware.py index 44dce8cb00..724e46ae30 100644 --- a/backend/tests/unit/onyx/server/auth/test_login_captcha_middleware.py +++ b/backend/tests/unit/onyx/server/auth/test_login_captcha_middleware.py @@ -1,13 +1,11 @@ """Unit tests for LoginCaptchaMiddleware.""" -from unittest.mock import AsyncMock -from unittest.mock import patch +from unittest.mock import AsyncMock, patch from fastapi import FastAPI from fastapi.testclient import TestClient -from onyx.auth.captcha import CaptchaAction -from onyx.auth.captcha import CaptchaVerificationError +from onyx.auth.captcha import CaptchaAction, CaptchaVerificationError from onyx.error_handling.exceptions import register_onyx_exception_handlers from onyx.server.auth import captcha_api as captcha_api_module from onyx.server.auth.captcha_api import LoginCaptchaMiddleware diff --git a/backend/tests/unit/onyx/server/features/build/external_apps/test_api_bundle_cleanup.py b/backend/tests/unit/onyx/server/features/build/external_apps/test_api_bundle_cleanup.py index dbfe06ec88..b69b7fc4cf 100644 --- a/backend/tests/unit/onyx/server/features/build/external_apps/test_api_bundle_cleanup.py +++ b/backend/tests/unit/onyx/server/features/build/external_apps/test_api_bundle_cleanup.py @@ -2,8 +2,7 @@ from __future__ import annotations import io from types import SimpleNamespace -from typing import Any -from typing import cast +from typing import Any, cast from unittest.mock import MagicMock import pytest diff --git a/backend/tests/unit/onyx/server/features/build/external_apps/test_api_user_response.py b/backend/tests/unit/onyx/server/features/build/external_apps/test_api_user_response.py index a0cf23ac93..3a420a929d 100644 --- a/backend/tests/unit/onyx/server/features/build/external_apps/test_api_user_response.py +++ b/backend/tests/unit/onyx/server/features/build/external_apps/test_api_user_response.py @@ -2,12 +2,10 @@ from __future__ import annotations import json from types import SimpleNamespace -from typing import Any -from typing import cast +from typing import Any, cast from onyx.db.enums import ExternalAppType -from onyx.db.models import ExternalApp -from onyx.db.models import ExternalAppUserCredential +from onyx.db.models import ExternalApp, ExternalAppUserCredential from onyx.server.features.build.external_apps.api import _to_user_response from onyx.utils.encryption import mask_string from onyx.utils.sensitive import SensitiveValue diff --git a/backend/tests/unit/onyx/server/features/build/interactive_turns/test_executor.py b/backend/tests/unit/onyx/server/features/build/interactive_turns/test_executor.py index 96c457941f..4a6c378a63 100644 --- a/backend/tests/unit/onyx/server/features/build/interactive_turns/test_executor.py +++ b/backend/tests/unit/onyx/server/features/build/interactive_turns/test_executor.py @@ -1,33 +1,33 @@ from __future__ import annotations import time -from collections.abc import Callable -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager from types import SimpleNamespace -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from onyx.server.features.build.interactive_turns import executor -from onyx.server.features.build.interactive_turns.state import claim_turn_for_runner -from onyx.server.features.build.interactive_turns.state import create_interactive_turn -from onyx.server.features.build.interactive_turns.state import get_active_turn -from onyx.server.features.build.interactive_turns.state import get_turn -from onyx.server.features.build.interactive_turns.state import InteractiveTurn -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_CANCELLED -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_FAILED -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_RUNNING -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_SUCCEEDED -from onyx.server.features.build.sandbox.event_schema import ActivityTimeoutError +from onyx.server.features.build.interactive_turns.state import ( + claim_turn_for_runner, + create_interactive_turn, + get_active_turn, + get_turn, + InteractiveTurn, + TURN_STATUS_CANCELLED, + TURN_STATUS_FAILED, + TURN_STATUS_RUNNING, + TURN_STATUS_SUCCEEDED, +) +from onyx.server.features.build.sandbox.event_schema import ( + ActivityTimeoutError, + PromptResponse, + TURN_ERROR_CODE_TIMEOUT, +) from onyx.server.features.build.sandbox.event_schema import Error as SandboxError -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TIMEOUT from onyx.server.features.build.sandbox.serve_transport import ( PROMPT_SLOT_FAST_FAIL_ACQUIRE_SECONDS, -) -from onyx.server.features.build.sandbox.serve_transport import ( PROMPT_SLOT_WAIT_OUT_ORPHAN_SECONDS, ) from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/unit/onyx/server/features/build/interactive_turns/test_state.py b/backend/tests/unit/onyx/server/features/build/interactive_turns/test_state.py index 4b88517e0b..3d0891f232 100644 --- a/backend/tests/unit/onyx/server/features/build/interactive_turns/test_state.py +++ b/backend/tests/unit/onyx/server/features/build/interactive_turns/test_state.py @@ -1,22 +1,23 @@ from __future__ import annotations from collections.abc import Callable -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 -from onyx.server.features.build.interactive_turns.state import acquire_active_turn_lock -from onyx.server.features.build.interactive_turns.state import claim_turn_for_runner -from onyx.server.features.build.interactive_turns.state import create_interactive_turn -from onyx.server.features.build.interactive_turns.state import finish_turn -from onyx.server.features.build.interactive_turns.state import get_active_turn -from onyx.server.features.build.interactive_turns.state import get_turn -from onyx.server.features.build.interactive_turns.state import get_turn_for_request -from onyx.server.features.build.interactive_turns.state import InteractiveTurn -from onyx.server.features.build.interactive_turns.state import REQUEST_ID_TTL_SECONDS -from onyx.server.features.build.interactive_turns.state import touch_turn -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_FAILED -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_QUEUED -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_RUNNING +from onyx.server.features.build.interactive_turns.state import ( + acquire_active_turn_lock, + claim_turn_for_runner, + create_interactive_turn, + finish_turn, + get_active_turn, + get_turn, + get_turn_for_request, + InteractiveTurn, + REQUEST_ID_TTL_SECONDS, + touch_turn, + TURN_STATUS_FAILED, + TURN_STATUS_QUEUED, + TURN_STATUS_RUNNING, +) from tests.unit.fakes import FakeCache diff --git a/backend/tests/unit/onyx/server/features/build/interactive_turns/test_turns_api.py b/backend/tests/unit/onyx/server/features/build/interactive_turns/test_turns_api.py index db07ca5814..eb171983bd 100644 --- a/backend/tests/unit/onyx/server/features/build/interactive_turns/test_turns_api.py +++ b/backend/tests/unit/onyx/server/features/build/interactive_turns/test_turns_api.py @@ -5,8 +5,7 @@ from collections.abc import Iterator from contextlib import contextmanager from types import SimpleNamespace from typing import cast -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session @@ -14,11 +13,13 @@ from sqlalchemy.orm import Session from onyx.db.models import User from onyx.error_handling.exceptions import OnyxError from onyx.server.features.build.interactive_turns import api as turns_api -from onyx.server.features.build.interactive_turns.state import claim_turn_for_runner -from onyx.server.features.build.interactive_turns.state import create_interactive_turn -from onyx.server.features.build.interactive_turns.state import finish_turn -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_FAILED -from onyx.server.features.build.interactive_turns.state import TURN_STATUS_SUCCEEDED +from onyx.server.features.build.interactive_turns.state import ( + claim_turn_for_runner, + create_interactive_turn, + finish_turn, + TURN_STATUS_FAILED, + TURN_STATUS_SUCCEEDED, +) from tests.unit.fakes import FakeCache diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/sandbox_daemon/test_sandbox_daemon.py b/backend/tests/unit/onyx/server/features/build/sandbox/sandbox_daemon/test_sandbox_daemon.py index cd9e2208a7..4757526c48 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/sandbox_daemon/test_sandbox_daemon.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/sandbox_daemon/test_sandbox_daemon.py @@ -27,35 +27,18 @@ from uuid import UUID import httpx import pytest from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey -from cryptography.hazmat.primitives.serialization import Encoding -from cryptography.hazmat.primitives.serialization import PublicFormat +from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat from fastapi.testclient import TestClient from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( FilesystemListRequest, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_FILESYSTEM_LIST_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_HEALTH_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_CREATE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_MARK_RESTORED_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_RESTORE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_READY_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_SNAPSHOT_CREATE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( sidecar_snapshot_restore_path, ) diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_bus_closed_property.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_bus_closed_property.py index 4ee786c560..5786788712 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_bus_closed_property.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_bus_closed_property.py @@ -15,8 +15,10 @@ replaces the dead bus on the next access. These tests pin both halves. from __future__ import annotations -from onyx.server.features.build.sandbox.opencode.event_bus import BUS_CLOSED_SENTINEL -from onyx.server.features.build.sandbox.opencode.event_bus import PodEventBus +from onyx.server.features.build.sandbox.opencode.event_bus import ( + BUS_CLOSED_SENTINEL, + PodEventBus, +) def test_closed_property_starts_false() -> None: diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_cold_pod_retry.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_cold_pod_retry.py index 9e32b09acd..5bad6cb2b9 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_cold_pod_retry.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_cold_pod_retry.py @@ -19,8 +19,10 @@ from typing import Any import httpx -from onyx.server.features.build.sandbox.opencode.serve_client import ClientTimeouts -from onyx.server.features.build.sandbox.opencode.serve_client import OpencodeServeClient +from onyx.server.features.build.sandbox.opencode.serve_client import ( + ClientTimeouts, + OpencodeServeClient, +) _STALE_ID = "ses_stale_001" _FRESH_ID = "ses_fresh_002" diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_manager_config.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_manager_config.py index 5894c6e841..c6c16c9e3a 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_manager_config.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_manager_config.py @@ -21,54 +21,28 @@ import onyx.server.features.build.sandbox.docker.docker_sandbox_manager as dsm from onyx.server.features.build.configs import SANDBOX_PROXY_INJECTED_PLACEHOLDER from onyx.server.features.build.sandbox.docker.dev_mode_serve import ( OPENCODE_SERVE_CONTAINER_PORT, -) -from onyx.server.features.build.sandbox.docker.dev_mode_serve import ( OPENCODE_SERVE_HOST_BIND_IP, ) from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( _sandbox_container_name, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( _sandbox_volume_name, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( _sanitize_relative_path, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( _validate_strict_path, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( build_container_create_kwargs, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( build_sandbox_labels, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( ContainerCreateKwargs, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( LABEL_COMPONENT, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( LABEL_COMPONENT_VALUE, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( LABEL_SANDBOX_ID, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( LABEL_TENANT_ID, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( LABEL_USER_ID, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( SANDBOX_TMP_PATH, -) -from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( SANDBOX_TMPFS_OPTIONS, ) -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY -from onyx.server.features.build.sandbox.labels import LABEL_K8S_MANAGED_BY_ONYX +from onyx.server.features.build.sandbox.labels import ( + LABEL_K8S_MANAGED_BY, + LABEL_K8S_MANAGED_BY_ONYX, +) SANDBOX_ID = UUID("12345678-1234-1234-1234-1234567890ab") USER_ID = UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_multi_provider_config.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_multi_provider_config.py index 7a93155dfa..4ad2e58de2 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_multi_provider_config.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_multi_provider_config.py @@ -12,8 +12,10 @@ from __future__ import annotations import json import onyx.server.features.build.sandbox.docker.docker_sandbox_manager as dsm -from onyx.server.features.build.configs import BUILD_MODE_NOT_CONFIGURED_API_KEY -from onyx.server.features.build.configs import SANDBOX_PROXY_INJECTED_PLACEHOLDER +from onyx.server.features.build.configs import ( + BUILD_MODE_NOT_CONFIGURED_API_KEY, + SANDBOX_PROXY_INJECTED_PLACEHOLDER, +) from onyx.server.features.build.sandbox.models import LLMProviderConfig from onyx.server.features.build.sandbox.util.opencode_config import ( build_multi_provider_opencode_config, diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_serve_plumbing.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_serve_plumbing.py index 06a677d1f7..53cfd280f4 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_serve_plumbing.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_docker_serve_plumbing.py @@ -23,8 +23,10 @@ import pytest import onyx.server.features.build.sandbox.docker.dev_mode_serve as dev_mode_serve import onyx.server.features.build.sandbox.docker.docker_sandbox_manager as dsm -from onyx.server.features.build.configs import OPENCODE_SERVE_PORT -from onyx.server.features.build.configs import OPENCODE_SERVER_PASSWORD +from onyx.server.features.build.configs import ( + OPENCODE_SERVE_PORT, + OPENCODE_SERVER_PASSWORD, +) from onyx.server.features.build.sandbox.docker.docker_sandbox_manager import ( DockerSandboxManager, ) diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_opencode_session.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_opencode_session.py index f5f8ed29f3..2ab93d11b6 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_opencode_session.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_opencode_session.py @@ -2,8 +2,7 @@ from __future__ import annotations from collections.abc import Callable from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_session.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_session.py index 884d5f8d4e..b1b3ec0b41 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_session.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_ensure_session.py @@ -28,8 +28,10 @@ from typing import Any import httpx import pytest -from onyx.server.features.build.sandbox.opencode.serve_client import ClientTimeouts -from onyx.server.features.build.sandbox.opencode.serve_client import OpencodeServeClient +from onyx.server.features.build.sandbox.opencode.serve_client import ( + ClientTimeouts, + OpencodeServeClient, +) _STALE_ID = "ses_stale_old_id_001" _FRESH_ID = "ses_fresh_new_id_002" diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_event_bus.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_event_bus.py index d00458f9fe..969a2af86e 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_event_bus.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_event_bus.py @@ -21,11 +21,13 @@ import httpx import pytest from onyx.server.features.build.sandbox.opencode import event_bus as event_bus_mod -from onyx.server.features.build.sandbox.opencode.event_bus import _extract_session_id -from onyx.server.features.build.sandbox.opencode.event_bus import _parse_sse_block -from onyx.server.features.build.sandbox.opencode.event_bus import _Subscription -from onyx.server.features.build.sandbox.opencode.event_bus import BUS_CLOSED_SENTINEL -from onyx.server.features.build.sandbox.opencode.event_bus import PodEventBus +from onyx.server.features.build.sandbox.opencode.event_bus import ( + _extract_session_id, + _parse_sse_block, + _Subscription, + BUS_CLOSED_SENTINEL, + PodEventBus, +) # --------------------------------------------------------------------------- # Fixtures diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_k8s_push_error_mapping.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_k8s_push_error_mapping.py index a9bc89540d..fe34770bf9 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_k8s_push_error_mapping.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_k8s_push_error_mapping.py @@ -14,50 +14,43 @@ import hashlib import io from collections.abc import Callable from types import SimpleNamespace -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch -from uuid import UUID -from uuid import uuid4 +from typing import Any, cast +from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 import httpx import pytest -from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey -from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey -from cryptography.hazmat.primitives.serialization import Encoding -from cryptography.hazmat.primitives.serialization import NoEncryption -from cryptography.hazmat.primitives.serialization import PrivateFormat +from cryptography.hazmat.primitives.asymmetric.ed25519 import ( + Ed25519PrivateKey, + Ed25519PublicKey, +) +from cryptography.hazmat.primitives.serialization import ( + Encoding, + NoEncryption, + PrivateFormat, +) from kubernetes import client from kubernetes.client.rest import ApiException from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_MARK_RESTORED_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_OPENCODE_HISTORY_RESTORE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_PUSH_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_SNAPSHOT_CREATE_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( sidecar_snapshot_restore_path, ) from onyx.server.features.build.sandbox.kubernetes import sidecar_client from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( _build_targz, -) -from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, ) from onyx.server.features.build.sandbox.kubernetes.sidecar_client import SidecarClient -from onyx.server.features.build.sandbox.models import FatalWriteError -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import RetriableWriteError +from onyx.server.features.build.sandbox.models import ( + FatalWriteError, + FileSet, + LLMProviderConfig, + RetriableWriteError, +) # Path to httpx.Client as imported inside the sidecar transport module. _HTTPX_CLIENT_PATH = ( diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_pod_spec.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_pod_spec.py index 224910737b..48ac0525a2 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_pod_spec.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_pod_spec.py @@ -30,9 +30,11 @@ from pathlib import Path import pytest import yaml from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey -from cryptography.hazmat.primitives.serialization import Encoding -from cryptography.hazmat.primitives.serialization import NoEncryption -from cryptography.hazmat.primitives.serialization import PrivateFormat +from cryptography.hazmat.primitives.serialization import ( + Encoding, + NoEncryption, + PrivateFormat, +) from kubernetes import client import onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager as ksm @@ -40,11 +42,7 @@ import onyx.server.features.build.sandbox.kubernetes.sidecar_client as sidecar from onyx.server.features.build.configs import SANDBOX_PROXY_INJECTED_PLACEHOLDER from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( PUSH_DAEMON_PORT, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_HEALTH_PATH, -) -from onyx.server.features.build.sandbox.image.sandbox_daemon.contract import ( SIDECAR_READY_PATH, ) from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_push_orchestration.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_push_orchestration.py index 00910338a7..08a98d830d 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_push_orchestration.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_push_orchestration.py @@ -6,9 +6,11 @@ from uuid import UUID import pytest -from onyx.server.features.build.sandbox.models import FatalWriteError -from onyx.server.features.build.sandbox.models import FileSet -from onyx.server.features.build.sandbox.models import RetriableWriteError +from onyx.server.features.build.sandbox.models import ( + FatalWriteError, + FileSet, + RetriableWriteError, +) from tests.common.craft.stubs import StubSandboxManager SB_1 = UUID("00000000-0000-0000-0000-000000000001") diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_timeout_plumbing.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_timeout_plumbing.py index 3d9e8c2a9e..7348a8c768 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_timeout_plumbing.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_timeout_plumbing.py @@ -2,15 +2,14 @@ from __future__ import annotations -from collections.abc import Callable -from collections.abc import Generator +from collections.abc import Callable, Generator from typing import Any from uuid import uuid4 from onyx.server.features.build.configs import ( OPENCODE_PROMPT_INACTIVITY_TIMEOUT_SECONDS, + PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS, ) -from onyx.server.features.build.configs import PROMPT_SLOT_KEEP_ALIVE_MAX_SECONDS from onyx.server.features.build.sandbox.event_schema import PromptResponse from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_with_bus.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_with_bus.py index 93623841fe..961937b5ec 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_with_bus.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_send_message_with_bus.py @@ -28,17 +28,21 @@ from typing import Any import httpx import pytest -from onyx.server.features.build.sandbox.event_schema import ActivityTimeoutError -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import AgentThoughtChunk -from onyx.server.features.build.sandbox.event_schema import Error -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TIMEOUT -from onyx.server.features.build.sandbox.event_schema import TURN_ERROR_CODE_TRANSPORT +from onyx.server.features.build.sandbox.event_schema import ( + ActivityTimeoutError, + AgentMessageChunk, + AgentThoughtChunk, + Error, + PromptResponse, + TURN_ERROR_CODE_TIMEOUT, + TURN_ERROR_CODE_TRANSPORT, +) from onyx.server.features.build.sandbox.opencode import serve_client from onyx.server.features.build.sandbox.opencode.event_bus import PodEventBus -from onyx.server.features.build.sandbox.opencode.serve_client import ClientTimeouts -from onyx.server.features.build.sandbox.opencode.serve_client import OpencodeServeClient +from onyx.server.features.build.sandbox.opencode.serve_client import ( + ClientTimeouts, + OpencodeServeClient, +) from onyx.server.features.build.sandbox.sse import SSEKeepalive _SESSION = "ses_test_123" diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_readiness_probe.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_readiness_probe.py index ff472211a9..bde3059781 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_readiness_probe.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_readiness_probe.py @@ -14,8 +14,10 @@ from uuid import UUID import pytest from onyx.server.features.build.sandbox.opencode.serve_client import OpencodeServeClient -from onyx.server.features.build.sandbox.serve_transport import _ServeMixin -from onyx.server.features.build.sandbox.serve_transport import ServeConnectionInfo +from onyx.server.features.build.sandbox.serve_transport import ( + _ServeMixin, + ServeConnectionInfo, +) _SBX = UUID("12345678-1234-1234-1234-1234567890ab") _POD_IP_URL = "http://10.244.0.97:4096" diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_transport_live_coalescing.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_transport_live_coalescing.py index 6a5e5080ae..0de840d934 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_transport_live_coalescing.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_serve_transport_live_coalescing.py @@ -12,8 +12,10 @@ import pytest from onyx.server.features.build.packets import SubagentStartedPacket from onyx.server.features.build.sandbox import serve_transport -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import PromptResponse +from onyx.server.features.build.sandbox.event_schema import ( + AgentMessageChunk, + PromptResponse, +) from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( KubernetesSandboxManager, ) diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_snapshot_manager_streams.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_snapshot_manager_streams.py index c3dbeb5076..c679cff0e0 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_snapshot_manager_streams.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_snapshot_manager_streams.py @@ -9,9 +9,7 @@ must stay stable, so we assert it here against a fake ``FileStore``. from __future__ import annotations import io -from typing import Any -from typing import cast -from typing import IO +from typing import Any, cast, IO import pytest diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_tarball.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_tarball.py index 5d45323134..f711e983b5 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_tarball.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_tarball.py @@ -9,8 +9,7 @@ import pytest from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import ( _build_targz, ) -from onyx.server.features.build.sandbox.models import FatalWriteError -from onyx.server.features.build.sandbox.models import FileSet +from onyx.server.features.build.sandbox.models import FatalWriteError, FileSet def test_round_trip() -> None: diff --git a/backend/tests/unit/onyx/server/features/build/sandbox/test_translate_opencode_event.py b/backend/tests/unit/onyx/server/features/build/sandbox/test_translate_opencode_event.py index 5f2acad167..49679ce6ee 100644 --- a/backend/tests/unit/onyx/server/features/build/sandbox/test_translate_opencode_event.py +++ b/backend/tests/unit/onyx/server/features/build/sandbox/test_translate_opencode_event.py @@ -11,22 +11,24 @@ from typing import Any import pytest -from onyx.server.features.build.packets import CompactionPacket -from onyx.server.features.build.packets import ContextUsagePacket -from onyx.server.features.build.packets import SubagentStartedPacket -from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk -from onyx.server.features.build.sandbox.event_schema import AgentThoughtChunk -from onyx.server.features.build.sandbox.event_schema import Error -from onyx.server.features.build.sandbox.event_schema import PromptResponse -from onyx.server.features.build.sandbox.event_schema import ToolCallProgress -from onyx.server.features.build.sandbox.event_schema import ToolCallStart +from onyx.server.features.build.packets import ( + CompactionPacket, + ContextUsagePacket, + SubagentStartedPacket, +) +from onyx.server.features.build.sandbox.event_schema import ( + AgentMessageChunk, + AgentThoughtChunk, + Error, + PromptResponse, + ToolCallProgress, + ToolCallStart, +) from onyx.server.features.build.sandbox.opencode.serve_client import ( _synthesize_tool_content, -) -from onyx.server.features.build.sandbox.opencode.serve_client import _tool_status -from onyx.server.features.build.sandbox.opencode.serve_client import _TurnState -from onyx.server.features.build.sandbox.opencode.serve_client import _wrap_raw_output -from onyx.server.features.build.sandbox.opencode.serve_client import ( + _tool_status, + _TurnState, + _wrap_raw_output, translate_opencode_event, ) diff --git a/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_executor.py b/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_executor.py index b1cd2b804a..bb6d15f5b8 100644 --- a/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_executor.py +++ b/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_executor.py @@ -1,5 +1,7 @@ -from onyx.server.features.build.scheduled_tasks.executor import _clip_summary -from onyx.server.features.build.scheduled_tasks.executor import SUMMARY_MAX_CHARS +from onyx.server.features.build.scheduled_tasks.executor import ( + _clip_summary, + SUMMARY_MAX_CHARS, +) def test_clip_summary_short_text_passthrough() -> None: diff --git a/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_schedule.py b/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_schedule.py index 27293c0c24..2349efa242 100644 --- a/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_schedule.py +++ b/backend/tests/unit/onyx/server/features/build/scheduled_tasks/test_schedule.py @@ -1,9 +1,10 @@ -from datetime import datetime -from datetime import timezone +from datetime import datetime, timezone -from onyx.server.features.build.scheduled_tasks.schedule import compute_next_run_at -from onyx.server.features.build.scheduled_tasks.schedule import human_readable -from onyx.server.features.build.scheduled_tasks.schedule import next_n_fires +from onyx.server.features.build.scheduled_tasks.schedule import ( + compute_next_run_at, + human_readable, + next_n_fires, +) def test_compute_next_run_at_uses_utc_cron_fields() -> None: diff --git a/backend/tests/unit/onyx/server/features/build/session/test_get_all_build_mode_llm_configs.py b/backend/tests/unit/onyx/server/features/build/session/test_get_all_build_mode_llm_configs.py index 596fb648ec..dd8ca5735c 100644 --- a/backend/tests/unit/onyx/server/features/build/session/test_get_all_build_mode_llm_configs.py +++ b/backend/tests/unit/onyx/server/features/build/session/test_get_all_build_mode_llm_configs.py @@ -13,25 +13,23 @@ from collections.abc import Iterator from contextlib import AbstractContextManager from datetime import datetime from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session from onyx.db.enums import SandboxStatus -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import Sandbox, User from onyx.error_handling.exceptions import OnyxError -from onyx.server.features.build.configs import BUILD_MODE_ALLOWED_PROVIDER_TYPES -from onyx.server.features.build.configs import BUILD_MODE_NOT_CONFIGURED_API_KEY -from onyx.server.features.build.sandbox.models import LLMProviderConfig -from onyx.server.features.build.sandbox.models import SandboxInfo +from onyx.server.features.build.configs import ( + BUILD_MODE_ALLOWED_PROVIDER_TYPES, + BUILD_MODE_NOT_CONFIGURED_API_KEY, +) +from onyx.server.features.build.sandbox.models import LLMProviderConfig, SandboxInfo from onyx.server.features.build.session.llm_config import get_all_build_mode_llm_configs from onyx.server.features.build.session.manager import SessionManager -from onyx.server.manage.llm.models import LLMProviderView -from onyx.server.manage.llm.models import ModelConfigurationView +from onyx.server.manage.llm.models import LLMProviderView, ModelConfigurationView # The recommended default model per provider type is sourced from the shared # recommended-models config (GitHub-or-bundled) at runtime. Stub it here so diff --git a/backend/tests/unit/onyx/server/features/build/session/test_messages_background_turn.py b/backend/tests/unit/onyx/server/features/build/session/test_messages_background_turn.py index c7cdf3cab5..75482be111 100644 --- a/backend/tests/unit/onyx/server/features/build/session/test_messages_background_turn.py +++ b/backend/tests/unit/onyx/server/features/build/session/test_messages_background_turn.py @@ -3,8 +3,7 @@ from __future__ import annotations from types import SimpleNamespace from typing import cast from unittest.mock import MagicMock -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from sqlalchemy.orm import Session diff --git a/backend/tests/unit/onyx/server/features/build/session/test_prewarm_opencode_session.py b/backend/tests/unit/onyx/server/features/build/session/test_prewarm_opencode_session.py index 4dff0aad47..1117a68d16 100644 --- a/backend/tests/unit/onyx/server/features/build/session/test_prewarm_opencode_session.py +++ b/backend/tests/unit/onyx/server/features/build/session/test_prewarm_opencode_session.py @@ -1,12 +1,10 @@ from __future__ import annotations -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest -from onyx.db.models import BuildSession -from onyx.db.models import Sandbox +from onyx.db.models import BuildSession, Sandbox from onyx.server.features.build.session.manager import SessionManager diff --git a/backend/tests/unit/onyx/server/features/build/session/test_snapshot_api.py b/backend/tests/unit/onyx/server/features/build/session/test_snapshot_api.py index ab4035a3e8..8144122b1a 100644 --- a/backend/tests/unit/onyx/server/features/build/session/test_snapshot_api.py +++ b/backend/tests/unit/onyx/server/features/build/session/test_snapshot_api.py @@ -11,8 +11,7 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.db.models import Sandbox -from onyx.db.models import User +from onyx.db.models import Sandbox, User from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.server.features.build.sandbox.models import SnapshotResult diff --git a/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py b/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py index 5eb5e663d8..d006a0ced7 100644 --- a/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py +++ b/backend/tests/unit/onyx/server/features/build/session/test_streaming_state.py @@ -5,16 +5,13 @@ Tests for chunk accumulation and finalize semantics — no DB required. from __future__ import annotations -from typing import Any -from typing import cast -from uuid import UUID -from uuid import uuid4 +from typing import Any, cast +from uuid import UUID, uuid4 import pytest from onyx.db.models import BuildSession -from onyx.server.features.build.packets import CompactionPacket -from onyx.server.features.build.packets import ContextUsagePacket +from onyx.server.features.build.packets import CompactionPacket, ContextUsagePacket from onyx.server.features.build.sandbox.event_schema import AgentMessageChunk from onyx.server.features.build.session import streaming from onyx.server.features.build.session.streaming import BuildStreamingState diff --git a/backend/tests/unit/onyx/server/features/build/test_feature_gate.py b/backend/tests/unit/onyx/server/features/build/test_feature_gate.py index 37725d92eb..a217738a71 100644 --- a/backend/tests/unit/onyx/server/features/build/test_feature_gate.py +++ b/backend/tests/unit/onyx/server/features/build/test_feature_gate.py @@ -11,14 +11,12 @@ from __future__ import annotations from typing import Any from unittest.mock import MagicMock -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from onyx.db.enums import AccountType -from onyx.feature_flags.interface import FeatureFlagProvider -from onyx.feature_flags.interface import NoOpFeatureFlagProvider +from onyx.feature_flags.interface import FeatureFlagProvider, NoOpFeatureFlagProvider from onyx.server.features.build import utils as build_utils from onyx.server.features.build.utils import is_craft_enabled_for_user from onyx.server.settings.models import Settings diff --git a/backend/tests/unit/onyx/server/features/build/test_upload_validation.py b/backend/tests/unit/onyx/server/features/build/test_upload_validation.py index e2b119c116..e252e2a67a 100644 --- a/backend/tests/unit/onyx/server/features/build/test_upload_validation.py +++ b/backend/tests/unit/onyx/server/features/build/test_upload_validation.py @@ -13,8 +13,7 @@ from __future__ import annotations from pathlib import Path from onyx.server.features.build.configs import MAX_UPLOAD_FILE_SIZE_BYTES -from onyx.server.features.build.utils import sanitize_filename -from onyx.server.features.build.utils import validate_file +from onyx.server.features.build.utils import sanitize_filename, validate_file def test_sanitize_filename_strips_path_components() -> None: diff --git a/backend/tests/unit/onyx/server/features/hierarchy/test_user_access_info.py b/backend/tests/unit/onyx/server/features/hierarchy/test_user_access_info.py index 6c1055cacd..acbecb7464 100644 --- a/backend/tests/unit/onyx/server/features/hierarchy/test_user_access_info.py +++ b/backend/tests/unit/onyx/server/features/hierarchy/test_user_access_info.py @@ -3,8 +3,7 @@ These tests mock all database operations and don't require a real database. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from sqlalchemy.orm import Session diff --git a/backend/tests/unit/onyx/server/features/image_generation/test_image_generation_endpoint.py b/backend/tests/unit/onyx/server/features/image_generation/test_image_generation_endpoint.py index 7f17585037..bc09a49267 100644 --- a/backend/tests/unit/onyx/server/features/image_generation/test_image_generation_endpoint.py +++ b/backend/tests/unit/onyx/server/features/image_generation/test_image_generation_endpoint.py @@ -8,9 +8,11 @@ from onyx.error_handling.exceptions import OnyxError from onyx.image_gen.exceptions import ImageGenerationNotConfiguredError from onyx.image_gen.generation import GeneratedImageData from onyx.image_gen.interfaces import ImageShape -from onyx.server.features.image_generation.api import generate_image -from onyx.server.features.image_generation.api import ImageGenerationRequest -from onyx.server.features.image_generation.api import ReferenceImagePayload +from onyx.server.features.image_generation.api import ( + generate_image, + ImageGenerationRequest, + ReferenceImagePayload, +) _HELPER = ( "onyx.server.features.image_generation.api.generate_images_with_default_config" diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py index 8602689250..2f1138631d 100644 --- a/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_refresh.py @@ -11,8 +11,7 @@ import time from collections.abc import Iterator from contextlib import contextmanager from types import SimpleNamespace -from typing import Any -from typing import cast +from typing import Any, cast from urllib.parse import parse_qs import httpx diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py index 9e161ce5a4..779b72a960 100644 --- a/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py @@ -11,14 +11,12 @@ import httpx import pytest from onyx.auth import oauth_token_manager -from onyx.auth.oauth_token_manager import exchange_oauth_code_for_token -from onyx.auth.oauth_token_manager import OAuthFlowParams +from onyx.auth.oauth_token_manager import exchange_oauth_code_for_token, OAuthFlowParams from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.server.features.mcp import api from onyx.server.features.mcp import ssrf as mcp_ssrf -from onyx.server.security.models import SecuritySettings -from onyx.server.security.models import SSRFProtectionLevel +from onyx.server.security.models import SecuritySettings, SSRFProtectionLevel from onyx.server.security.store import _build_env_defaults from onyx.utils.url import SSRFException diff --git a/backend/tests/unit/onyx/server/features/mcp/test_oauth_credentials_resolver.py b/backend/tests/unit/onyx/server/features/mcp/test_oauth_credentials_resolver.py index c36c27fee5..9bc935fcde 100644 --- a/backend/tests/unit/onyx/server/features/mcp/test_oauth_credentials_resolver.py +++ b/backend/tests/unit/onyx/server/features/mcp/test_oauth_credentials_resolver.py @@ -13,9 +13,11 @@ import pytest from mcp.shared.auth import OAuthClientInformationFull from pydantic import AnyUrl -from onyx.server.features.mcp.api import _build_oauth_admin_config_data -from onyx.server.features.mcp.api import _build_oauth_admin_config_data_for_update -from onyx.server.features.mcp.api import _resolve_oauth_credentials +from onyx.server.features.mcp.api import ( + _build_oauth_admin_config_data, + _build_oauth_admin_config_data_for_update, + _resolve_oauth_credentials, +) from onyx.server.features.mcp.models import MCPOAuthKeys from onyx.utils.encryption import mask_string diff --git a/backend/tests/unit/onyx/server/features/mcp/test_per_user_listing_invariants.py b/backend/tests/unit/onyx/server/features/mcp/test_per_user_listing_invariants.py index 875ee67543..8035ee1723 100644 --- a/backend/tests/unit/onyx/server/features/mcp/test_per_user_listing_invariants.py +++ b/backend/tests/unit/onyx/server/features/mcp/test_per_user_listing_invariants.py @@ -15,15 +15,16 @@ for them regardless of caller role or the contents of the shared row. """ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.auth.schemas import UserRole -from onyx.db.enums import MCPAuthenticationPerformer -from onyx.db.enums import MCPAuthenticationType -from onyx.db.enums import MCPOAuthProviderMode -from onyx.db.enums import MCPServerStatus -from onyx.db.enums import MCPTransport +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPOAuthProviderMode, + MCPServerStatus, + MCPTransport, +) from onyx.server.features.mcp.api import _db_mcp_server_to_api_mcp_server diff --git a/backend/tests/unit/onyx/server/features/persona/test_knowledge_sources.py b/backend/tests/unit/onyx/server/features/persona/test_knowledge_sources.py index 811552b822..a5ecdc00a6 100644 --- a/backend/tests/unit/onyx/server/features/persona/test_knowledge_sources.py +++ b/backend/tests/unit/onyx/server/features/persona/test_knowledge_sources.py @@ -1,10 +1,8 @@ """Unit tests for MinimalPersonaSnapshot.from_model knowledge_sources aggregation.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import FederatedConnectorSource +from onyx.configs.constants import DocumentSource, FederatedConnectorSource from onyx.server.features.document_set.models import DocumentSetSummary from onyx.server.features.persona.models import MinimalPersonaSnapshot diff --git a/backend/tests/unit/onyx/server/manage/embedding/test_embedding_api.py b/backend/tests/unit/onyx/server/manage/embedding/test_embedding_api.py index 366f8cd74b..4efacfb37a 100644 --- a/backend/tests/unit/onyx/server/manage/embedding/test_embedding_api.py +++ b/backend/tests/unit/onyx/server/manage/embedding/test_embedding_api.py @@ -1,13 +1,16 @@ from types import SimpleNamespace -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.db.models import SearchSettings -from onyx.server.manage.embedding.api import list_embedding_models -from onyx.server.manage.embedding.api import list_embedding_providers -from onyx.utils.encryption import decrypt_bytes_to_string -from onyx.utils.encryption import encrypt_string_to_bytes -from onyx.utils.encryption import mask_string +from onyx.server.manage.embedding.api import ( + list_embedding_models, + list_embedding_providers, +) +from onyx.utils.encryption import ( + decrypt_bytes_to_string, + encrypt_string_to_bytes, + mask_string, +) from onyx.utils.sensitive import SensitiveValue from shared_configs.enums import EmbeddingProvider diff --git a/backend/tests/unit/onyx/server/manage/llm/test_fetch_models_api.py b/backend/tests/unit/onyx/server/manage/llm/test_fetch_models_api.py index 09503f8014..08be1d8df2 100644 --- a/backend/tests/unit/onyx/server/manage/llm/test_fetch_models_api.py +++ b/backend/tests/unit/onyx/server/manage/llm/test_fetch_models_api.py @@ -7,26 +7,27 @@ sync-to-DB behavior when provider_id is specified. import os from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import httpx import pytest from onyx.db.enums import LLMModelFlowType from onyx.error_handling.exceptions import OnyxError -from onyx.server.manage.llm.models import BedrockModelsRequest -from onyx.server.manage.llm.models import BifrostFinalModelResponse -from onyx.server.manage.llm.models import BifrostModelsRequest -from onyx.server.manage.llm.models import LitellmFinalModelResponse -from onyx.server.manage.llm.models import LitellmModelsRequest -from onyx.server.manage.llm.models import LMStudioFinalModelResponse -from onyx.server.manage.llm.models import LMStudioModelsRequest -from onyx.server.manage.llm.models import OllamaFinalModelResponse -from onyx.server.manage.llm.models import OllamaModelsRequest -from onyx.server.manage.llm.models import OpenAICompatibleModelsRequest -from onyx.server.manage.llm.models import OpenRouterFinalModelResponse -from onyx.server.manage.llm.models import OpenRouterModelsRequest +from onyx.server.manage.llm.models import ( + BedrockModelsRequest, + BifrostFinalModelResponse, + BifrostModelsRequest, + LitellmFinalModelResponse, + LitellmModelsRequest, + LMStudioFinalModelResponse, + LMStudioModelsRequest, + OllamaFinalModelResponse, + OllamaModelsRequest, + OpenAICompatibleModelsRequest, + OpenRouterFinalModelResponse, + OpenRouterModelsRequest, +) class TestGetOllamaAvailableModels: @@ -1666,8 +1667,10 @@ class TestGetBedrockAvailableModels: """Editing an existing provider sends the masked bearer token. The endpoint must swap it back for the stored value so the per-session token provider carries real credentials, not the masked placeholder.""" - from onyx.server.manage.llm.api import _mask_string - from onyx.server.manage.llm.api import get_bedrock_available_models + from onyx.server.manage.llm.api import ( + _mask_string, + get_bedrock_available_models, + ) real_token = "real-bearer-token-secret" masked_token = _mask_string(real_token) diff --git a/backend/tests/unit/onyx/server/manage/llm/test_llm_provider_utils.py b/backend/tests/unit/onyx/server/manage/llm/test_llm_provider_utils.py index 579c3a8ad3..781201eaa9 100644 --- a/backend/tests/unit/onyx/server/manage/llm/test_llm_provider_utils.py +++ b/backend/tests/unit/onyx/server/manage/llm/test_llm_provider_utils.py @@ -1,11 +1,13 @@ """Tests for LLM provider utilities.""" -from onyx.server.manage.llm.utils import generate_bedrock_display_name -from onyx.server.manage.llm.utils import generate_ollama_display_name -from onyx.server.manage.llm.utils import is_embedding_model -from onyx.server.manage.llm.utils import is_reasoning_model -from onyx.server.manage.llm.utils import is_valid_bedrock_model -from onyx.server.manage.llm.utils import strip_openrouter_vendor_prefix +from onyx.server.manage.llm.utils import ( + generate_bedrock_display_name, + generate_ollama_display_name, + is_embedding_model, + is_reasoning_model, + is_valid_bedrock_model, + strip_openrouter_vendor_prefix, +) class TestGenerateBedrockDisplayName: diff --git a/backend/tests/unit/onyx/server/manage/llm/test_model_configuration.py b/backend/tests/unit/onyx/server/manage/llm/test_model_configuration.py index 8d6381df4d..428aeaa421 100644 --- a/backend/tests/unit/onyx/server/manage/llm/test_model_configuration.py +++ b/backend/tests/unit/onyx/server/manage/llm/test_model_configuration.py @@ -7,13 +7,14 @@ These tests verify the flow plumbing: from the stored flow. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.db.enums import LLMModelFlowType -from onyx.server.manage.llm.models import LLMProviderDescriptor -from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest -from onyx.server.manage.llm.models import ModelConfigurationView +from onyx.server.manage.llm.models import ( + LLMProviderDescriptor, + ModelConfigurationUpsertRequest, + ModelConfigurationView, +) # ModelConfigurationView.from_model — dynamic provider branch diff --git a/backend/tests/unit/onyx/server/manage/llm/test_provider_cache.py b/backend/tests/unit/onyx/server/manage/llm/test_provider_cache.py index 808475e090..c84ea80cb5 100644 --- a/backend/tests/unit/onyx/server/manage/llm/test_provider_cache.py +++ b/backend/tests/unit/onyx/server/manage/llm/test_provider_cache.py @@ -2,9 +2,11 @@ import pytest from redis.exceptions import RedisError from onyx.server.manage.llm import provider_cache -from onyx.server.manage.llm.models import DefaultModel -from onyx.server.manage.llm.models import LLMProviderDescriptor -from onyx.server.manage.llm.models import LLMProviderResponse +from onyx.server.manage.llm.models import ( + DefaultModel, + LLMProviderDescriptor, + LLMProviderResponse, +) class _FakeCache: diff --git a/backend/tests/unit/onyx/server/manage/test_bulk_invite_limit.py b/backend/tests/unit/onyx/server/manage/test_bulk_invite_limit.py index 2ad97e425b..4ecfc16ffc 100644 --- a/backend/tests/unit/onyx/server/manage/test_bulk_invite_limit.py +++ b/backend/tests/unit/onyx/server/manage/test_bulk_invite_limit.py @@ -5,17 +5,14 @@ import logging from collections.abc import Iterator from contextlib import contextmanager from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.manage.models import EmailInviteStatus -from onyx.server.manage.models import UserByEmail -from onyx.server.manage.users import bulk_invite_users -from onyx.server.manage.users import remove_invited_user +from onyx.server.manage.models import EmailInviteStatus, UserByEmail +from onyx.server.manage.users import bulk_invite_users, remove_invited_user def _make_shared_session_mock(next_total: int) -> MagicMock: diff --git a/backend/tests/unit/onyx/server/manage/test_invite_rate_limit.py b/backend/tests/unit/onyx/server/manage/test_invite_rate_limit.py index 45381046f0..6276b0dbd9 100644 --- a/backend/tests/unit/onyx/server/manage/test_invite_rate_limit.py +++ b/backend/tests/unit/onyx/server/manage/test_invite_rate_limit.py @@ -10,8 +10,10 @@ from redis.exceptions import ConnectionError as RedisConnectionError from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.redis.tenant_redis_client import TenantRedisClient -from onyx.server.manage.invite_rate_limit import enforce_invite_rate_limit -from onyx.server.manage.invite_rate_limit import enforce_remove_invited_rate_limit +from onyx.server.manage.invite_rate_limit import ( + enforce_invite_rate_limit, + enforce_remove_invited_rate_limit, +) class _StubRedis: diff --git a/backend/tests/unit/onyx/server/manage/test_reindex_instant_backfill_guard.py b/backend/tests/unit/onyx/server/manage/test_reindex_instant_backfill_guard.py index 392654fed0..5f5e0a8135 100644 --- a/backend/tests/unit/onyx/server/manage/test_reindex_instant_backfill_guard.py +++ b/backend/tests/unit/onyx/server/manage/test_reindex_instant_backfill_guard.py @@ -1,8 +1,7 @@ """Guard that blocks a new reindex while an INSTANT swap is still backfilling the live index (superseding it would abandon that backfill).""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/onyx/server/manage/test_users_directory_visibility.py b/backend/tests/unit/onyx/server/manage/test_users_directory_visibility.py index babc1b6c25..ccc04ba82f 100644 --- a/backend/tests/unit/onyx/server/manage/test_users_directory_visibility.py +++ b/backend/tests/unit/onyx/server/manage/test_users_directory_visibility.py @@ -1,15 +1,12 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest -from onyx.db.enums import AccountType -from onyx.db.enums import Permission +from onyx.db.enums import AccountType, Permission from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.server.manage.users import list_all_users_basic_info -from onyx.server.manage.users import verify_user_logged_in +from onyx.server.manage.users import list_all_users_basic_info, verify_user_logged_in from onyx.server.security.models import SecuritySettings from onyx.server.security.store import _build_env_defaults diff --git a/backend/tests/unit/onyx/server/manage/voice/test_chunked_transcriber.py b/backend/tests/unit/onyx/server/manage/voice/test_chunked_transcriber.py index d9bd5f2c1a..a749dd9a75 100644 --- a/backend/tests/unit/onyx/server/manage/voice/test_chunked_transcriber.py +++ b/backend/tests/unit/onyx/server/manage/voice/test_chunked_transcriber.py @@ -9,12 +9,14 @@ from unittest.mock import AsyncMock import pytest -from onyx.server.manage.voice.websocket_api import ChunkedTranscriber -from onyx.server.manage.voice.websocket_api import pcm16_rms -from onyx.server.manage.voice.websocket_api import PCM_BYTES_PER_SECOND -from onyx.server.manage.voice.websocket_api import PCM_SAMPLE_RATE -from onyx.server.manage.voice.websocket_api import SILENCE_RMS_THRESHOLD -from onyx.server.manage.voice.websocket_api import trim_pcm16_silence +from onyx.server.manage.voice.websocket_api import ( + ChunkedTranscriber, + pcm16_rms, + PCM_BYTES_PER_SECOND, + PCM_SAMPLE_RATE, + SILENCE_RMS_THRESHOLD, + trim_pcm16_silence, +) def _silence(seconds: float) -> bytes: diff --git a/backend/tests/unit/onyx/server/query_and_chat/test_any_rate_limit_exists_tenant_scope.py b/backend/tests/unit/onyx/server/query_and_chat/test_any_rate_limit_exists_tenant_scope.py index 665c20b052..710d510c05 100644 --- a/backend/tests/unit/onyx/server/query_and_chat/test_any_rate_limit_exists_tenant_scope.py +++ b/backend/tests/unit/onyx/server/query_and_chat/test_any_rate_limit_exists_tenant_scope.py @@ -7,8 +7,7 @@ bypassed.""" from collections.abc import Iterator from contextlib import contextmanager -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import onyx.server.query_and_chat.token_limit as token_limit from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR diff --git a/backend/tests/unit/onyx/server/scim/conftest.py b/backend/tests/unit/onyx/server/scim/conftest.py index 6a19fae7bd..3f62036044 100644 --- a/backend/tests/unit/onyx/server/scim/conftest.py +++ b/backend/tests/unit/onyx/server/scim/conftest.py @@ -5,8 +5,7 @@ from __future__ import annotations import json from collections.abc import Generator from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -14,18 +13,16 @@ from fastapi.responses import JSONResponse from sqlalchemy.orm import Session from ee.onyx.server.scim.api import ScimJSONResponse -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimListResponse -from ee.onyx.server.scim.models import ScimName -from ee.onyx.server.scim.models import ScimUserResource +from ee.onyx.server.scim.models import ( + ScimGroupResource, + ScimListResponse, + ScimName, + ScimUserResource, +) from ee.onyx.server.scim.providers.base import ScimProvider from ee.onyx.server.scim.providers.entra import EntraProvider from ee.onyx.server.scim.providers.okta import OktaProvider -from onyx.db.models import ScimToken -from onyx.db.models import ScimUserMapping -from onyx.db.models import User -from onyx.db.models import UserGroup -from onyx.db.models import UserRole +from onyx.db.models import ScimToken, ScimUserMapping, User, UserGroup, UserRole # Every supported SCIM provider must appear here so that all endpoint tests # run against it. When adding a new provider, add its class to this list. diff --git a/backend/tests/unit/onyx/server/scim/test_admin.py b/backend/tests/unit/onyx/server/scim/test_admin.py index 59d05b63bd..3e139eded2 100644 --- a/backend/tests/unit/onyx/server/scim/test_admin.py +++ b/backend/tests/unit/onyx/server/scim/test_admin.py @@ -1,19 +1,19 @@ """Tests for SCIM admin token management endpoints.""" from datetime import datetime -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session from ee.onyx.db.scim import ScimDAL -from ee.onyx.server.enterprise_settings.api import create_scim_token -from ee.onyx.server.enterprise_settings.api import get_active_scim_token +from ee.onyx.server.enterprise_settings.api import ( + create_scim_token, + get_active_scim_token, +) from ee.onyx.server.scim.models import ScimTokenCreate -from onyx.db.models import ScimToken -from onyx.db.models import User +from onyx.db.models import ScimToken, User @pytest.fixture diff --git a/backend/tests/unit/onyx/server/scim/test_auth.py b/backend/tests/unit/onyx/server/scim/test_auth.py index ac645c843d..8c23acbbda 100644 --- a/backend/tests/unit/onyx/server/scim/test_auth.py +++ b/backend/tests/unit/onyx/server/scim/test_auth.py @@ -2,11 +2,13 @@ from unittest.mock import MagicMock import pytest -from ee.onyx.server.scim.auth import _hash_scim_token -from ee.onyx.server.scim.auth import generate_scim_token -from ee.onyx.server.scim.auth import SCIM_TOKEN_PREFIX -from ee.onyx.server.scim.auth import ScimAuthError -from ee.onyx.server.scim.auth import verify_scim_token +from ee.onyx.server.scim.auth import ( + _hash_scim_token, + generate_scim_token, + SCIM_TOKEN_PREFIX, + ScimAuthError, + verify_scim_token, +) class TestGenerateScimToken: diff --git a/backend/tests/unit/onyx/server/scim/test_entra.py b/backend/tests/unit/onyx/server/scim/test_entra.py index 18bd0c39de..89b22c0bd4 100644 --- a/backend/tests/unit/onyx/server/scim/test_entra.py +++ b/backend/tests/unit/onyx/server/scim/test_entra.py @@ -9,48 +9,53 @@ PATCH value dicts). from __future__ import annotations import json -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from fastapi import Response -from ee.onyx.server.scim.api import create_user -from ee.onyx.server.scim.api import delete_user -from ee.onyx.server.scim.api import get_group -from ee.onyx.server.scim.api import get_resource_types -from ee.onyx.server.scim.api import get_schemas -from ee.onyx.server.scim.api import get_service_provider_config -from ee.onyx.server.scim.api import get_user -from ee.onyx.server.scim.api import list_groups -from ee.onyx.server.scim.api import list_users -from ee.onyx.server.scim.api import patch_group -from ee.onyx.server.scim.api import patch_user -from ee.onyx.server.scim.api import replace_user -from ee.onyx.server.scim.api import ScimJSONResponse -from ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA -from ee.onyx.server.scim.models import SCIM_USER_SCHEMA -from ee.onyx.server.scim.models import ScimEnterpriseExtension -from ee.onyx.server.scim.models import ScimGroupMember -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimManagerRef -from ee.onyx.server.scim.models import ScimMappingFields -from ee.onyx.server.scim.models import ScimName -from ee.onyx.server.scim.models import ScimPatchOperation -from ee.onyx.server.scim.models import ScimPatchOperationType -from ee.onyx.server.scim.models import ScimPatchRequest -from ee.onyx.server.scim.models import ScimPatchResourceValue -from ee.onyx.server.scim.models import ScimUserResource +from ee.onyx.server.scim.api import ( + create_user, + delete_user, + get_group, + get_resource_types, + get_schemas, + get_service_provider_config, + get_user, + list_groups, + list_users, + patch_group, + patch_user, + replace_user, + ScimJSONResponse, +) +from ee.onyx.server.scim.models import ( + SCIM_ENTERPRISE_USER_SCHEMA, + SCIM_USER_SCHEMA, + ScimEnterpriseExtension, + ScimGroupMember, + ScimGroupResource, + ScimManagerRef, + ScimMappingFields, + ScimName, + ScimPatchOperation, + ScimPatchOperationType, + ScimPatchRequest, + ScimPatchResourceValue, + ScimUserResource, +) from ee.onyx.server.scim.providers.base import ScimProvider from ee.onyx.server.scim.providers.entra import EntraProvider -from tests.unit.onyx.server.scim.conftest import make_db_group -from tests.unit.onyx.server.scim.conftest import make_db_user -from tests.unit.onyx.server.scim.conftest import make_scim_user -from tests.unit.onyx.server.scim.conftest import make_user_mapping -from tests.unit.onyx.server.scim.conftest import parse_scim_group -from tests.unit.onyx.server.scim.conftest import parse_scim_list -from tests.unit.onyx.server.scim.conftest import parse_scim_user +from tests.unit.onyx.server.scim.conftest import ( + make_db_group, + make_db_user, + make_scim_user, + make_user_mapping, + parse_scim_group, + parse_scim_list, + parse_scim_user, +) @pytest.fixture diff --git a/backend/tests/unit/onyx/server/scim/test_filtering.py b/backend/tests/unit/onyx/server/scim/test_filtering.py index 13d70ba641..866fce2dad 100644 --- a/backend/tests/unit/onyx/server/scim/test_filtering.py +++ b/backend/tests/unit/onyx/server/scim/test_filtering.py @@ -1,8 +1,10 @@ import pytest -from ee.onyx.server.scim.filtering import parse_scim_filter -from ee.onyx.server.scim.filtering import ScimFilter -from ee.onyx.server.scim.filtering import ScimFilterOperator +from ee.onyx.server.scim.filtering import ( + parse_scim_filter, + ScimFilter, + ScimFilterOperator, +) class TestParseScimFilter: diff --git a/backend/tests/unit/onyx/server/scim/test_group_endpoints.py b/backend/tests/unit/onyx/server/scim/test_group_endpoints.py index ea6f7feb72..9331ee85b7 100644 --- a/backend/tests/unit/onyx/server/scim/test_group_endpoints.py +++ b/backend/tests/unit/onyx/server/scim/test_group_endpoints.py @@ -2,30 +2,35 @@ from __future__ import annotations -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from fastapi import Response -from ee.onyx.server.scim.api import create_group -from ee.onyx.server.scim.api import delete_group -from ee.onyx.server.scim.api import get_group -from ee.onyx.server.scim.api import list_groups -from ee.onyx.server.scim.api import patch_group -from ee.onyx.server.scim.api import replace_group -from ee.onyx.server.scim.models import ScimGroupMember -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimPatchOperation -from ee.onyx.server.scim.models import ScimPatchOperationType -from ee.onyx.server.scim.models import ScimPatchRequest +from ee.onyx.server.scim.api import ( + create_group, + delete_group, + get_group, + list_groups, + patch_group, + replace_group, +) +from ee.onyx.server.scim.models import ( + ScimGroupMember, + ScimGroupResource, + ScimPatchOperation, + ScimPatchOperationType, + ScimPatchRequest, +) from ee.onyx.server.scim.patch import ScimPatchError from ee.onyx.server.scim.providers.base import ScimProvider -from tests.unit.onyx.server.scim.conftest import assert_scim_error -from tests.unit.onyx.server.scim.conftest import make_db_group -from tests.unit.onyx.server.scim.conftest import make_scim_group -from tests.unit.onyx.server.scim.conftest import parse_scim_group -from tests.unit.onyx.server.scim.conftest import parse_scim_list +from tests.unit.onyx.server.scim.conftest import ( + assert_scim_error, + make_db_group, + make_scim_group, + parse_scim_group, + parse_scim_list, +) class TestListGroups: diff --git a/backend/tests/unit/onyx/server/scim/test_patch.py b/backend/tests/unit/onyx/server/scim/test_patch.py index 1103c4525e..ac62c07dce 100644 --- a/backend/tests/unit/onyx/server/scim/test_patch.py +++ b/backend/tests/unit/onyx/server/scim/test_patch.py @@ -1,19 +1,23 @@ import pytest -from ee.onyx.server.scim.models import ScimEmail -from ee.onyx.server.scim.models import ScimGroupMember -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimMeta -from ee.onyx.server.scim.models import ScimName -from ee.onyx.server.scim.models import ScimPatchOperation -from ee.onyx.server.scim.models import ScimPatchOperationType -from ee.onyx.server.scim.models import ScimPatchRequest -from ee.onyx.server.scim.models import ScimPatchResourceValue -from ee.onyx.server.scim.models import ScimPatchValue -from ee.onyx.server.scim.models import ScimUserResource -from ee.onyx.server.scim.patch import apply_group_patch -from ee.onyx.server.scim.patch import apply_user_patch -from ee.onyx.server.scim.patch import ScimPatchError +from ee.onyx.server.scim.models import ( + ScimEmail, + ScimGroupMember, + ScimGroupResource, + ScimMeta, + ScimName, + ScimPatchOperation, + ScimPatchOperationType, + ScimPatchRequest, + ScimPatchResourceValue, + ScimPatchValue, + ScimUserResource, +) +from ee.onyx.server.scim.patch import ( + apply_group_patch, + apply_user_patch, + ScimPatchError, +) from ee.onyx.server.scim.providers.entra import EntraProvider from ee.onyx.server.scim.providers.okta import OktaProvider diff --git a/backend/tests/unit/onyx/server/scim/test_providers.py b/backend/tests/unit/onyx/server/scim/test_providers.py index 79228386ac..0dde0db3b4 100644 --- a/backend/tests/unit/onyx/server/scim/test_providers.py +++ b/backend/tests/unit/onyx/server/scim/test_providers.py @@ -1,20 +1,25 @@ from unittest.mock import MagicMock -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 -from ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA -from ee.onyx.server.scim.models import SCIM_USER_SCHEMA -from ee.onyx.server.scim.models import ScimEmail -from ee.onyx.server.scim.models import ScimGroupMember -from ee.onyx.server.scim.models import ScimGroupResource -from ee.onyx.server.scim.models import ScimMeta -from ee.onyx.server.scim.models import ScimName -from ee.onyx.server.scim.models import ScimUserGroupRef -from ee.onyx.server.scim.models import ScimUserResource -from ee.onyx.server.scim.providers.base import COMMON_IGNORED_PATCH_PATHS -from ee.onyx.server.scim.providers.base import get_default_provider -from ee.onyx.server.scim.providers.entra import _ENTRA_IGNORED_PATCH_PATHS -from ee.onyx.server.scim.providers.entra import EntraProvider +from ee.onyx.server.scim.models import ( + SCIM_ENTERPRISE_USER_SCHEMA, + SCIM_USER_SCHEMA, + ScimEmail, + ScimGroupMember, + ScimGroupResource, + ScimMeta, + ScimName, + ScimUserGroupRef, + ScimUserResource, +) +from ee.onyx.server.scim.providers.base import ( + COMMON_IGNORED_PATCH_PATHS, + get_default_provider, +) +from ee.onyx.server.scim.providers.entra import ( + _ENTRA_IGNORED_PATCH_PATHS, + EntraProvider, +) from ee.onyx.server.scim.providers.okta import OktaProvider diff --git a/backend/tests/unit/onyx/server/scim/test_user_endpoints.py b/backend/tests/unit/onyx/server/scim/test_user_endpoints.py index 6273b9049b..21c5853ee2 100644 --- a/backend/tests/unit/onyx/server/scim/test_user_endpoints.py +++ b/backend/tests/unit/onyx/server/scim/test_user_endpoints.py @@ -2,38 +2,43 @@ from __future__ import annotations -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from fastapi import Response from sqlalchemy.exc import IntegrityError from ee.onyx.db.license import seat_lock_id_for_tenant -from ee.onyx.server.scim.api import _check_seat_availability -from ee.onyx.server.scim.api import _scim_name_to_str -from ee.onyx.server.scim.api import create_user -from ee.onyx.server.scim.api import delete_user -from ee.onyx.server.scim.api import get_user -from ee.onyx.server.scim.api import list_users -from ee.onyx.server.scim.api import patch_user -from ee.onyx.server.scim.api import replace_user -from ee.onyx.server.scim.models import ScimMappingFields -from ee.onyx.server.scim.models import ScimName -from ee.onyx.server.scim.models import ScimPatchOperation -from ee.onyx.server.scim.models import ScimPatchOperationType -from ee.onyx.server.scim.models import ScimPatchRequest -from ee.onyx.server.scim.models import ScimUserResource +from ee.onyx.server.scim.api import ( + _check_seat_availability, + _scim_name_to_str, + create_user, + delete_user, + get_user, + list_users, + patch_user, + replace_user, +) +from ee.onyx.server.scim.models import ( + ScimMappingFields, + ScimName, + ScimPatchOperation, + ScimPatchOperationType, + ScimPatchRequest, + ScimUserResource, +) from ee.onyx.server.scim.patch import ScimPatchError from ee.onyx.server.scim.providers.base import ScimProvider from onyx.db.enums import AccountType from onyx.db.models import UserRole -from tests.unit.onyx.server.scim.conftest import assert_scim_error -from tests.unit.onyx.server.scim.conftest import make_db_user -from tests.unit.onyx.server.scim.conftest import make_scim_user -from tests.unit.onyx.server.scim.conftest import make_user_mapping -from tests.unit.onyx.server.scim.conftest import parse_scim_list -from tests.unit.onyx.server.scim.conftest import parse_scim_user +from tests.unit.onyx.server.scim.conftest import ( + assert_scim_error, + make_db_user, + make_scim_user, + make_user_mapping, + parse_scim_list, + parse_scim_user, +) class TestListUsers: diff --git a/backend/tests/unit/onyx/server/security/test_models.py b/backend/tests/unit/onyx/server/security/test_models.py index ee84d11742..c0689ae9fa 100644 --- a/backend/tests/unit/onyx/server/security/test_models.py +++ b/backend/tests/unit/onyx/server/security/test_models.py @@ -5,15 +5,17 @@ from typing import Any import pytest from pydantic import ValidationError -from onyx.server.security.models import _derive_operator_locked_fields -from onyx.server.security.models import _operator_locked -from onyx.server.security.models import _tenant_editable -from onyx.server.security.models import OPERATOR_LOCKED_FIELDS -from onyx.server.security.models import PASSWORD_LENGTH_CAP -from onyx.server.security.models import PASSWORD_MAX_LENGTH_FLOOR -from onyx.server.security.models import SecuritySettings -from onyx.server.security.models import SecuritySettingsOverrides -from onyx.server.security.models import SSRFProtectionLevel +from onyx.server.security.models import ( + _derive_operator_locked_fields, + _operator_locked, + _tenant_editable, + OPERATOR_LOCKED_FIELDS, + PASSWORD_LENGTH_CAP, + PASSWORD_MAX_LENGTH_FLOOR, + SecuritySettings, + SecuritySettingsOverrides, + SSRFProtectionLevel, +) _VALID_EFFECTIVE_KWARGS: dict[str, Any] = { "user_directory_admin_only": False, diff --git a/backend/tests/unit/onyx/server/settings/test_admin_put_settings.py b/backend/tests/unit/onyx/server/settings/test_admin_put_settings.py index 200daf3e19..26978d9deb 100644 --- a/backend/tests/unit/onyx/server/settings/test_admin_put_settings.py +++ b/backend/tests/unit/onyx/server/settings/test_admin_put_settings.py @@ -5,8 +5,7 @@ import pytest from pydantic import ValidationError from onyx.server.settings import api as settings_api -from onyx.server.settings.models import CRAFT_INSTRUCTIONS_MAX_LENGTH -from onyx.server.settings.models import Settings +from onyx.server.settings.models import CRAFT_INSTRUCTIONS_MAX_LENGTH, Settings def _put_settings( diff --git a/backend/tests/unit/onyx/server/test_admin_audit_events.py b/backend/tests/unit/onyx/server/test_admin_audit_events.py index a9693b56a0..c2974c5337 100644 --- a/backend/tests/unit/onyx/server/test_admin_audit_events.py +++ b/backend/tests/unit/onyx/server/test_admin_audit_events.py @@ -10,14 +10,15 @@ tests/unit/onyx/utils/test_audit.py. import json import logging from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.server.api_key.api import create_api_key -from onyx.server.api_key.api import delete_api_key -from onyx.server.api_key.api import regenerate_existing_api_key +from onyx.server.api_key.api import ( + create_api_key, + delete_api_key, + regenerate_existing_api_key, +) from onyx.server.manage.llm.api import delete_llm_provider diff --git a/backend/tests/unit/onyx/server/test_full_user_snapshot.py b/backend/tests/unit/onyx/server/test_full_user_snapshot.py index e3b729979b..bee6b0347d 100644 --- a/backend/tests/unit/onyx/server/test_full_user_snapshot.py +++ b/backend/tests/unit/onyx/server/test_full_user_snapshot.py @@ -4,8 +4,7 @@ from uuid import uuid4 from onyx.auth.schemas import UserRole from onyx.db.enums import AccountType -from onyx.server.models import FullUserSnapshot -from onyx.server.models import UserGroupInfo +from onyx.server.models import FullUserSnapshot, UserGroupInfo def _mock_user( diff --git a/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py b/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py index 51c4470c7c..0de41086e9 100644 --- a/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py +++ b/backend/tests/unit/onyx/server/test_mcp_known_provider_oauth_helpers.py @@ -2,32 +2,31 @@ import asyncio import time from contextlib import contextmanager from types import SimpleNamespace -from typing import cast -from typing import Iterator -from urllib.parse import parse_qs -from urllib.parse import urlparse +from typing import cast, Iterator +from urllib.parse import parse_qs, urlparse import httpx import pytest from mcp.client.auth import OAuthClientProvider -from mcp.shared.auth import OAuthClientInformationFull -from mcp.shared.auth import OAuthMetadata -from mcp.shared.auth import OAuthToken -from pydantic import AnyHttpUrl -from pydantic import AnyUrl +from mcp.shared.auth import OAuthClientInformationFull, OAuthMetadata, OAuthToken +from pydantic import AnyHttpUrl, AnyUrl import onyx.server.features.mcp.oauth as mcp_oauth -from onyx.auth.oauth_token_manager import build_oauth_authorization_url -from onyx.auth.oauth_token_manager import exchange_oauth_code_for_token +from onyx.auth.oauth_token_manager import ( + build_oauth_authorization_url, + exchange_oauth_code_for_token, +) from onyx.db.enums import MCPOAuthProviderMode from onyx.db.models import MCPServer as DbMCPServer from onyx.error_handling.exceptions import OnyxError from onyx.server.features.mcp.api import _mcp_known_provider_flow_params from onyx.server.features.mcp.models import MCPOAuthKeys -from onyx.server.features.mcp.oauth import _absolute_token_expiry -from onyx.server.features.mcp.oauth import _known_provider_oauth_metadata -from onyx.server.features.mcp.oauth import _token_dict_with_preserved_refresh -from onyx.server.features.mcp.oauth import make_oauth_provider +from onyx.server.features.mcp.oauth import ( + _absolute_token_expiry, + _known_provider_oauth_metadata, + _token_dict_with_preserved_refresh, + make_oauth_provider, +) def _make_mcp_server_stub( diff --git a/backend/tests/unit/onyx/server/test_oidc_multi.py b/backend/tests/unit/onyx/server/test_oidc_multi.py index daad9e98f2..74787c10ea 100644 --- a/backend/tests/unit/onyx/server/test_oidc_multi.py +++ b/backend/tests/unit/onyx/server/test_oidc_multi.py @@ -4,17 +4,18 @@ per-provider client construction, and OAuth state/CSRF validation. No DB, no network, no live IdP.""" from types import SimpleNamespace -from typing import Any -from typing import cast +from typing import Any, cast import pytest from sqlalchemy.orm import Session -from onyx.auth.users import CSRF_TOKEN_COOKIE_NAME -from onyx.auth.users import CSRF_TOKEN_KEY -from onyx.auth.users import decode_and_validate_oauth_state -from onyx.auth.users import generate_csrf_token -from onyx.auth.users import generate_state_token +from onyx.auth.users import ( + CSRF_TOKEN_COOKIE_NAME, + CSRF_TOKEN_KEY, + decode_and_validate_oauth_state, + generate_csrf_token, + generate_state_token, +) from onyx.db.enums import SSOProviderType from onyx.db.models import SSOProvider from onyx.error_handling.exceptions import OnyxError diff --git a/backend/tests/unit/onyx/server/test_pool_metrics.py b/backend/tests/unit/onyx/server/test_pool_metrics.py index dd6c54713e..1b14c10678 100644 --- a/backend/tests/unit/onyx/server/test_pool_metrics.py +++ b/backend/tests/unit/onyx/server/test_pool_metrics.py @@ -1,19 +1,17 @@ """Unit tests for SQLAlchemy connection pool Prometheus metrics.""" from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from fastapi import FastAPI from sqlalchemy.pool import NullPool -from onyx.server.metrics.postgres_connection_pool import _register_pool_events -from onyx.server.metrics.postgres_connection_pool import PoolStateCollector from onyx.server.metrics.postgres_connection_pool import ( + _register_pool_events, + PoolStateCollector, setup_postgres_connection_pool_metrics, ) -from onyx.utils.middleware import _build_route_map -from onyx.utils.middleware import _match_route +from onyx.utils.middleware import _build_route_map, _match_route # --- PoolStateCollector tests --- diff --git a/backend/tests/unit/onyx/server/test_projects_file_utils.py b/backend/tests/unit/onyx/server/test_projects_file_utils.py index 47c3dc262e..88db92df84 100644 --- a/backend/tests/unit/onyx/server/test_projects_file_utils.py +++ b/backend/tests/unit/onyx/server/test_projects_file_utils.py @@ -5,8 +5,7 @@ import pytest from fastapi import UploadFile from onyx.natural_language_processing import utils as nlp_utils -from onyx.natural_language_processing.utils import BaseTokenizer -from onyx.natural_language_processing.utils import count_tokens +from onyx.natural_language_processing.utils import BaseTokenizer, count_tokens from onyx.server.features.projects import projects_file_utils as utils from onyx.server.settings.models import Settings diff --git a/backend/tests/unit/onyx/server/test_prometheus_instrumentation.py b/backend/tests/unit/onyx/server/test_prometheus_instrumentation.py index dd4b11cc76..d6efbc7168 100644 --- a/backend/tests/unit/onyx/server/test_prometheus_instrumentation.py +++ b/backend/tests/unit/onyx/server/test_prometheus_instrumentation.py @@ -2,13 +2,11 @@ import threading from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from fastapi import FastAPI from fastapi.testclient import TestClient -from prometheus_client import CollectorRegistry -from prometheus_client import Gauge +from prometheus_client import CollectorRegistry, Gauge from onyx.server.metrics.per_tenant import per_tenant_request_callback from onyx.server.metrics.prometheus_setup import setup_prometheus_metrics diff --git a/backend/tests/unit/onyx/server/test_saml_multi.py b/backend/tests/unit/onyx/server/test_saml_multi.py index 4f29cdd1a0..dc568796a4 100644 --- a/backend/tests/unit/onyx/server/test_saml_multi.py +++ b/backend/tests/unit/onyx/server/test_saml_multi.py @@ -5,8 +5,7 @@ gate, and email extraction from SAML attributes. No DB or live IdP.""" import base64 from types import SimpleNamespace -from typing import Any -from typing import cast +from typing import Any, cast import pytest from onelogin.saml2.auth import OneLogin_Saml2_Auth diff --git a/backend/tests/unit/onyx/server/test_settings_store.py b/backend/tests/unit/onyx/server/test_settings_store.py index 1453a037a9..5717186bd2 100644 --- a/backend/tests/unit/onyx/server/test_settings_store.py +++ b/backend/tests/unit/onyx/server/test_settings_store.py @@ -5,9 +5,9 @@ from onyx.key_value_store.interface import KvKeyNotFoundError from onyx.server.settings import store as settings_store from onyx.server.settings.models import ( DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_NO_VECTOR_DB, + DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_VECTOR_DB, + Settings, ) -from onyx.server.settings.models import DEFAULT_FILE_TOKEN_COUNT_THRESHOLD_K_VECTOR_DB -from onyx.server.settings.models import Settings class _FakeKvStore: diff --git a/backend/tests/unit/onyx/server/test_token_rate_limit_budget_models.py b/backend/tests/unit/onyx/server/test_token_rate_limit_budget_models.py index b31cb3d512..56904a0cb8 100644 --- a/backend/tests/unit/onyx/server/test_token_rate_limit_budget_models.py +++ b/backend/tests/unit/onyx/server/test_token_rate_limit_budget_models.py @@ -6,8 +6,10 @@ from pydantic import ValidationError from onyx.configs.constants import TokenRateLimitScope from onyx.db.models import TokenRateLimit from onyx.server.query_and_chat.token_limit import _is_rate_limited -from onyx.server.token_rate_limits.models import TokenRateLimitArgs -from onyx.server.token_rate_limits.models import TokenRateLimitDisplay +from onyx.server.token_rate_limits.models import ( + TokenRateLimitArgs, + TokenRateLimitDisplay, +) def _rate_limit( diff --git a/backend/tests/unit/onyx/server/test_upload_files.py b/backend/tests/unit/onyx/server/test_upload_files.py index c213362bf0..ad912034da 100644 --- a/backend/tests/unit/onyx/server/test_upload_files.py +++ b/backend/tests/unit/onyx/server/test_upload_files.py @@ -1,7 +1,6 @@ import io import zipfile -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from zipfile import BadZipFile import pytest diff --git a/backend/tests/unit/onyx/skills/test_bundle_files.py b/backend/tests/unit/onyx/skills/test_bundle_files.py index 2839c944b5..048c1bdb28 100644 --- a/backend/tests/unit/onyx/skills/test_bundle_files.py +++ b/backend/tests/unit/onyx/skills/test_bundle_files.py @@ -4,10 +4,12 @@ import zipfile import pytest from onyx.error_handling.exceptions import OnyxError -from onyx.skills.bundle import build_single_file_bundle -from onyx.skills.bundle import build_skill_md -from onyx.skills.bundle import inspect_custom_bundle -from onyx.skills.bundle import update_custom_bundle_files +from onyx.skills.bundle import ( + build_single_file_bundle, + build_skill_md, + inspect_custom_bundle, + update_custom_bundle_files, +) def _zip(entries: list[tuple[str, bytes]]) -> bytes: diff --git a/backend/tests/unit/onyx/skills/test_bundle_safety.py b/backend/tests/unit/onyx/skills/test_bundle_safety.py index 0e24089302..24547d2158 100644 --- a/backend/tests/unit/onyx/skills/test_bundle_safety.py +++ b/backend/tests/unit/onyx/skills/test_bundle_safety.py @@ -9,8 +9,7 @@ import zipfile import pytest from onyx.error_handling.exceptions import OnyxError -from onyx.skills.bundle import _ZIP_UNIX_CREATE_SYSTEM -from onyx.skills.bundle import normalize_custom_bundle +from onyx.skills.bundle import _ZIP_UNIX_CREATE_SYSTEM, normalize_custom_bundle def _build_zip( diff --git a/backend/tests/unit/onyx/skills/test_bundle_validation.py b/backend/tests/unit/onyx/skills/test_bundle_validation.py index 1697ce9f58..6e5d8f3028 100644 --- a/backend/tests/unit/onyx/skills/test_bundle_validation.py +++ b/backend/tests/unit/onyx/skills/test_bundle_validation.py @@ -10,12 +10,14 @@ import pytest from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError -from onyx.skills.bundle import _ZIP_UNIX_CREATE_SYSTEM -from onyx.skills.bundle import compute_bundle_sha256 -from onyx.skills.bundle import normalize_custom_bundle -from onyx.skills.bundle import read_custom_bundle_instructions -from onyx.skills.bundle import rewrite_custom_bundle_skill_md -from onyx.skills.bundle import strip_skill_md_frontmatter +from onyx.skills.bundle import ( + _ZIP_UNIX_CREATE_SYSTEM, + compute_bundle_sha256, + normalize_custom_bundle, + read_custom_bundle_instructions, + rewrite_custom_bundle_skill_md, + strip_skill_md_frontmatter, +) from onyx.skills.metadata import parse_skill_document diff --git a/backend/tests/unit/onyx/skills/test_content.py b/backend/tests/unit/onyx/skills/test_content.py index a06fa1b865..3068bba65a 100644 --- a/backend/tests/unit/onyx/skills/test_content.py +++ b/backend/tests/unit/onyx/skills/test_content.py @@ -9,8 +9,10 @@ from onyx.db.models import Skill from onyx.file_store.file_store import FileStore from onyx.skills import built_in as built_in_module from onyx.skills.built_in import BuiltInSkillDefinition -from onyx.skills.content import read_builtin_skill_instructions -from onyx.skills.content import read_custom_skill_bundle_instructions +from onyx.skills.content import ( + read_builtin_skill_instructions, + read_custom_skill_bundle_instructions, +) def _build_zip(entries: list[tuple[str, bytes]]) -> bytes: diff --git a/backend/tests/unit/onyx/skills/test_ingest.py b/backend/tests/unit/onyx/skills/test_ingest.py index 9429abb7fd..fde6e350d0 100644 --- a/backend/tests/unit/onyx/skills/test_ingest.py +++ b/backend/tests/unit/onyx/skills/test_ingest.py @@ -8,9 +8,11 @@ import pytest from onyx.error_handling.exceptions import OnyxError from onyx.file_store.file_store import FileStore -from onyx.skills.ingest import ingest_skill_bundle -from onyx.skills.ingest import ingested_skill_bundle -from onyx.skills.ingest import IngestedBundle +from onyx.skills.ingest import ( + ingest_skill_bundle, + ingested_skill_bundle, + IngestedBundle, +) def test_ingest_normalizes_wrapped_bundle_before_hashing_and_storage() -> None: diff --git a/backend/tests/unit/onyx/skills/test_metadata.py b/backend/tests/unit/onyx/skills/test_metadata.py index de1a13730d..38417000fa 100644 --- a/backend/tests/unit/onyx/skills/test_metadata.py +++ b/backend/tests/unit/onyx/skills/test_metadata.py @@ -5,8 +5,7 @@ from __future__ import annotations import pytest from onyx.error_handling.exceptions import OnyxError -from onyx.skills.metadata import parse_skill_document -from onyx.skills.metadata import serialize_skill_md +from onyx.skills.metadata import parse_skill_document, serialize_skill_md def _skill_md(frontmatter: str, body: str = "Do the work.") -> bytes: diff --git a/backend/tests/unit/onyx/skills/test_pptx_lint.py b/backend/tests/unit/onyx/skills/test_pptx_lint.py index 9bf7f1cb43..b93f2d525e 100644 --- a/backend/tests/unit/onyx/skills/test_pptx_lint.py +++ b/backend/tests/unit/onyx/skills/test_pptx_lint.py @@ -17,8 +17,10 @@ pytest.importorskip("PIL") from pptx.dml.color import RGBColor # noqa: E402 from pptx.enum.text import MSO_AUTO_SIZE # noqa: E402 -from pptx.util import Inches # noqa: E402 -from pptx.util import Pt # noqa: E402 +from pptx.util import ( # noqa: E402 + Inches, + Pt, +) LINT_PATH = ( Path(__file__).parents[4] diff --git a/backend/tests/unit/onyx/skills/test_push.py b/backend/tests/unit/onyx/skills/test_push.py index 2d5b20a10a..f8dcc55b53 100644 --- a/backend/tests/unit/onyx/skills/test_push.py +++ b/backend/tests/unit/onyx/skills/test_push.py @@ -10,8 +10,7 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session -from onyx.db.models import Skill -from onyx.db.models import User +from onyx.db.models import Skill, User from onyx.file_store.file_store import FileStore from onyx.skills import push diff --git a/backend/tests/unit/onyx/skills/test_skill_access.py b/backend/tests/unit/onyx/skills/test_skill_access.py index 03150db797..1c55c465e3 100644 --- a/backend/tests/unit/onyx/skills/test_skill_access.py +++ b/backend/tests/unit/onyx/skills/test_skill_access.py @@ -3,12 +3,8 @@ from uuid import uuid4 import pytest from onyx.auth.schemas import UserRole -from onyx.db.enums import SkillAccessLevel -from onyx.db.enums import SkillSharePermission -from onyx.db.models import Skill -from onyx.db.models import Skill__User -from onyx.db.models import Skill__UserGroup -from onyx.db.models import User +from onyx.db.enums import SkillAccessLevel, SkillSharePermission +from onyx.db.models import Skill, Skill__User, Skill__UserGroup, User from onyx.error_handling.exceptions import OnyxError from onyx.server.features.skill.api import _ensure_can_edit_org_visibility from onyx.server.features.skill.response_helpers import user_permission_for_skill diff --git a/backend/tests/unit/onyx/tools/custom/test_custom_tools.py b/backend/tests/unit/onyx/tools/custom/test_custom_tools.py index d6fe7ea98d..495e0074ad 100644 --- a/backend/tests/unit/onyx/tools/custom/test_custom_tools.py +++ b/backend/tests/unit/onyx/tools/custom/test_custom_tools.py @@ -6,13 +6,12 @@ from unittest.mock import patch import pytest from onyx.server.query_and_chat.placement import Placement -from onyx.tools.models import DynamicSchemaInfo -from onyx.tools.models import ToolResponse +from onyx.tools.models import DynamicSchemaInfo, ToolResponse from onyx.tools.tool_implementations.custom.custom_tool import ( build_custom_tools_from_openapi_schema_and_headers, + CustomToolCallSummary, + validate_openapi_schema, ) -from onyx.tools.tool_implementations.custom.custom_tool import CustomToolCallSummary -from onyx.tools.tool_implementations.custom.custom_tool import validate_openapi_schema from onyx.tools.tool_implementations.custom.openapi_parsing import ( openapi_to_method_specs, ) diff --git a/backend/tests/unit/onyx/tools/test_construct_tools_no_vectordb.py b/backend/tests/unit/onyx/tools/test_construct_tools_no_vectordb.py index a6eb023c50..7fe7f8df3f 100644 --- a/backend/tests/unit/onyx/tools/test_construct_tools_no_vectordb.py +++ b/backend/tests/unit/onyx/tools/test_construct_tools_no_vectordb.py @@ -7,8 +7,7 @@ Verifies that: - FileReaderTool.is_available() returns True when vector DB is disabled """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.tools.tool_implementations.file_reader.file_reader_tool import FileReaderTool diff --git a/backend/tests/unit/onyx/tools/test_file_reader_tool.py b/backend/tests/unit/onyx/tools/test_file_reader_tool.py index e725b9da33..b990d8d31e 100644 --- a/backend/tests/unit/onyx/tools/test_file_reader_tool.py +++ b/backend/tests/unit/onyx/tools/test_file_reader_tool.py @@ -8,21 +8,19 @@ Verifies: - is_available() reflects DISABLE_VECTOR_DB """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest -from onyx.file_store.models import ChatFileType -from onyx.file_store.models import InMemoryChatFile +from onyx.file_store.models import ChatFileType, InMemoryChatFile from onyx.server.query_and_chat.placement import Placement from onyx.tools.models import ToolCallException -from onyx.tools.tool_implementations.file_reader.file_reader_tool import FILE_ID_FIELD -from onyx.tools.tool_implementations.file_reader.file_reader_tool import FileReaderTool -from onyx.tools.tool_implementations.file_reader.file_reader_tool import MAX_NUM_CHARS -from onyx.tools.tool_implementations.file_reader.file_reader_tool import NUM_CHARS_FIELD from onyx.tools.tool_implementations.file_reader.file_reader_tool import ( + FILE_ID_FIELD, + FileReaderTool, + MAX_NUM_CHARS, + NUM_CHARS_FIELD, START_CHAR_FIELD, ) diff --git a/backend/tests/unit/onyx/tools/test_image_generation_reference_resolution.py b/backend/tests/unit/onyx/tools/test_image_generation_reference_resolution.py index a3fd780ec2..67f72e18c8 100644 --- a/backend/tests/unit/onyx/tools/test_image_generation_reference_resolution.py +++ b/backend/tests/unit/onyx/tools/test_image_generation_reference_resolution.py @@ -9,8 +9,7 @@ don't re-validate against an allow-list in the tool itself. """ from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -18,8 +17,6 @@ from onyx.image_gen.interfaces import ImageShape from onyx.tools.models import ToolCallException from onyx.tools.tool_implementations.images.image_generation_tool import ( ImageGenerationTool, -) -from onyx.tools.tool_implementations.images.image_generation_tool import ( REFERENCE_IMAGE_FILE_IDS_FIELD, ) diff --git a/backend/tests/unit/onyx/tools/test_mcp_tool_schema.py b/backend/tests/unit/onyx/tools/test_mcp_tool_schema.py index 0e59377ea7..0f04d09ef1 100644 --- a/backend/tests/unit/onyx/tools/test_mcp_tool_schema.py +++ b/backend/tests/unit/onyx/tools/test_mcp_tool_schema.py @@ -11,8 +11,7 @@ a JSON-Schema-valid `parameters` dict with a `properties` key. """ from typing import Any -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -21,8 +20,10 @@ from onyx.server.query_and_chat.placement import Placement from onyx.tools.interface import Tool from onyx.tools.models import ToolResponse from onyx.tools.tool_constructor import _disambiguate_mcp_tool_names -from onyx.tools.tool_implementations.mcp.mcp_tool import _normalize_parameters_schema -from onyx.tools.tool_implementations.mcp.mcp_tool import MCPTool +from onyx.tools.tool_implementations.mcp.mcp_tool import ( + _normalize_parameters_schema, + MCPTool, +) class TestNormalizeParametersSchema: diff --git a/backend/tests/unit/onyx/tools/test_no_vectordb.py b/backend/tests/unit/onyx/tools/test_no_vectordb.py index 239bd1241c..4e4a74ed17 100644 --- a/backend/tests/unit/onyx/tools/test_no_vectordb.py +++ b/backend/tests/unit/onyx/tools/test_no_vectordb.py @@ -4,8 +4,7 @@ Verifies that SearchTool is unavailable, OpenURLTool stays available (crawl-only), and FileReaderTool remains available. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from sqlalchemy.orm import Session diff --git a/backend/tests/unit/onyx/tools/test_python_tool_availability.py b/backend/tests/unit/onyx/tools/test_python_tool_availability.py index 5870bbacaa..c2c39fa9e1 100644 --- a/backend/tests/unit/onyx/tools/test_python_tool_availability.py +++ b/backend/tests/unit/onyx/tools/test_python_tool_availability.py @@ -8,8 +8,7 @@ Verifies that PythonTool reports itself as unavailable when either: Also verifies that the health check result is cached with a TTL. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from sqlalchemy.orm import Session @@ -174,8 +173,6 @@ def test_health_check_cached_on_second_call() -> None: def test_health_check_refreshed_after_ttl_expires(mock_time: MagicMock) -> None: from onyx.tools.tool_implementations.python.code_interpreter_client import ( _HEALTH_CACHE_TTL_SECONDS, - ) - from onyx.tools.tool_implementations.python.code_interpreter_client import ( CodeInterpreterClient, ) diff --git a/backend/tests/unit/onyx/tools/test_search_llm_json.py b/backend/tests/unit/onyx/tools/test_search_llm_json.py index cedaa7f0bb..bfa9dc7717 100644 --- a/backend/tests/unit/onyx/tools/test_search_llm_json.py +++ b/backend/tests/unit/onyx/tools/test_search_llm_json.py @@ -5,13 +5,12 @@ import json import pytest from onyx.configs.constants import DocumentSource -from onyx.context.search.models import InferenceChunk -from onyx.context.search.models import InferenceSection +from onyx.context.search.models import InferenceChunk, InferenceSection from onyx.context.search.utils import sandbox_filename_for_document from onyx.tools.tool_implementations.utils import ( convert_inference_sections_to_llm_string, + FILE_ASSOCIATED_GUIDANCE, ) -from onyx.tools.tool_implementations.utils import FILE_ASSOCIATED_GUIDANCE FID = "550e8400-e29b-41d4-a716-446655440000" diff --git a/backend/tests/unit/onyx/tools/test_search_tool_time_filter_cache.py b/backend/tests/unit/onyx/tools/test_search_tool_time_filter_cache.py index a9417b6855..e7976cc7cd 100644 --- a/backend/tests/unit/onyx/tools/test_search_tool_time_filter_cache.py +++ b/backend/tests/unit/onyx/tools/test_search_tool_time_filter_cache.py @@ -1,13 +1,13 @@ from __future__ import annotations -from datetime import datetime -from datetime import timezone -from unittest.mock import MagicMock -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch -from onyx.secondary_llm_flows.time_filter import decide_time_filter -from onyx.secondary_llm_flows.time_filter import DocumentTimeField -from onyx.secondary_llm_flows.time_filter import TimeFilter +from onyx.secondary_llm_flows.time_filter import ( + decide_time_filter, + DocumentTimeField, + TimeFilter, +) from onyx.tools.tool_implementations.search.search_tool import SearchTool diff --git a/backend/tests/unit/onyx/tools/test_tool_runner_chat_files.py b/backend/tests/unit/onyx/tools/test_tool_runner_chat_files.py index 58178b2d37..76b440e22a 100644 --- a/backend/tests/unit/onyx/tools/test_tool_runner_chat_files.py +++ b/backend/tests/unit/onyx/tools/test_tool_runner_chat_files.py @@ -10,8 +10,7 @@ from unittest.mock import patch import pytest from onyx.db.models import UserFile -from onyx.tools.models import ChatFile -from onyx.tools.models import PythonToolOverrideKwargs +from onyx.tools.models import ChatFile, PythonToolOverrideKwargs class TestChatFilesPassingToPythonTool: diff --git a/backend/tests/unit/onyx/tools/tool_implementations/bash/test_bash_tool.py b/backend/tests/unit/onyx/tools/tool_implementations/bash/test_bash_tool.py index 472e42cb32..4130ee954b 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/bash/test_bash_tool.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/bash/test_bash_tool.py @@ -14,20 +14,22 @@ Covers: import json from collections.abc import Iterator from contextlib import contextmanager -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.configs.app_configs import CODE_INTERPRETER_DEFAULT_TIMEOUT_MS -from onyx.configs.app_configs import CODE_INTERPRETER_MAX_OUTPUT_LENGTH +from onyx.configs.app_configs import ( + CODE_INTERPRETER_DEFAULT_TIMEOUT_MS, + CODE_INTERPRETER_MAX_OUTPUT_LENGTH, +) from onyx.server.query_and_chat.placement import Placement -from onyx.server.query_and_chat.streaming_models import BashToolDelta -from onyx.server.query_and_chat.streaming_models import BashToolStart +from onyx.server.query_and_chat.streaming_models import BashToolDelta, BashToolStart from onyx.tools.models import ToolCallException -from onyx.tools.tool_implementations.bash.bash_tool import BashTool -from onyx.tools.tool_implementations.bash.bash_tool import BashToolOverrideKwargs -from onyx.tools.tool_implementations.bash.bash_tool import CMD_FIELD +from onyx.tools.tool_implementations.bash.bash_tool import ( + BashTool, + BashToolOverrideKwargs, + CMD_FIELD, +) from onyx.tools.tool_implementations.python.code_interpreter_client import ( BashExecResponse, ) diff --git a/backend/tests/unit/onyx/tools/tool_implementations/coding_agent/test_coding_agent_tool.py b/backend/tests/unit/onyx/tools/tool_implementations/coding_agent/test_coding_agent_tool.py index 20176c7bae..77f4b08839 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/coding_agent/test_coding_agent_tool.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/coding_agent/test_coding_agent_tool.py @@ -5,8 +5,7 @@ Coverage is intentionally narrow — the heavy lifting lives in file exists to lock in the wiring this Tool wrapper is responsible for. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch from onyx.tools.tool_implementations.coding_agent.coding_agent_tool import ( CodingAgentTool, diff --git a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler.py b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler.py index f892f03585..c95bbf79c2 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler.py @@ -2,8 +2,7 @@ from __future__ import annotations import time from collections.abc import Callable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from pydantic import BaseModel @@ -13,11 +12,9 @@ from onyx.server.security.models import SSRFProtectionLevel from onyx.server.security.store import _build_env_defaults from onyx.tools.tool_implementations.open_url.onyx_web_crawler import ( DEFAULT_CONNECT_TIMEOUT_SECONDS, -) -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import ( DEFAULT_READ_TIMEOUT_SECONDS, + OnyxWebCrawler, ) -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import OnyxWebCrawler class FakeResponse(BaseModel): diff --git a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler_playwright_fallback.py b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler_playwright_fallback.py index a5d1ff7d8e..47c4c6a73e 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler_playwright_fallback.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_onyx_web_crawler_playwright_fallback.py @@ -8,12 +8,13 @@ one-shot headless render and re-parse. from __future__ import annotations -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import onyx.tools.tool_implementations.open_url.onyx_web_crawler as crawler_module -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import FailureReason -from onyx.tools.tool_implementations.open_url.onyx_web_crawler import OnyxWebCrawler +from onyx.tools.tool_implementations.open_url.onyx_web_crawler import ( + FailureReason, + OnyxWebCrawler, +) from onyx.utils.playwright_fetch import RenderedPage SUCCESS_HTML = "Real Page

Hello world, this is real content from the page after rendering.

" diff --git a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_open_url_failure_messages.py b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_open_url_failure_messages.py index a88d020191..b7eee5085d 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_open_url_failure_messages.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_open_url_failure_messages.py @@ -10,8 +10,8 @@ from onyx.tools.tool_implementations.open_url.models import FailedFetch from onyx.tools.tool_implementations.open_url.onyx_web_crawler import FailureReason from onyx.tools.tool_implementations.open_url.open_url_tool import ( _build_failure_message, + _format_failed_url, ) -from onyx.tools.tool_implementations.open_url.open_url_tool import _format_failed_url def test_format_failed_url_with_reason() -> None: diff --git a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_snippet_matcher.py b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_snippet_matcher.py index 8fa5f01e6d..a98ebe8036 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_snippet_matcher.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_snippet_matcher.py @@ -5,8 +5,7 @@ import unicodedata # used to verify NFC expansion test preconditions from pathlib import Path import pytest -from pydantic import BaseModel -from pydantic import field_validator +from pydantic import BaseModel, field_validator from onyx.tools.tool_implementations.open_url.snippet_matcher import ( find_snippet_in_content, diff --git a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_url_normalization.py b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_url_normalization.py index 3c4aaa1ada..572c4f90cf 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_url_normalization.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/open_url/test_url_normalization.py @@ -6,9 +6,7 @@ from onyx.configs.constants import DocumentSource from onyx.tools.tool_implementations.open_url.open_url_tool import _url_lookup_variants from onyx.tools.tool_implementations.open_url.url_normalization import ( _detect_source_type, -) -from onyx.tools.tool_implementations.open_url.url_normalization import normalize_url -from onyx.tools.tool_implementations.open_url.url_normalization import ( + normalize_url, normalize_url_candidates, ) diff --git a/backend/tests/unit/onyx/tools/tool_implementations/python/test_code_interpreter_client.py b/backend/tests/unit/onyx/tools/tool_implementations/python/test_code_interpreter_client.py index 1a1a3b38a0..89ced13d34 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/python/test_code_interpreter_client.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/python/test_code_interpreter_client.py @@ -10,8 +10,7 @@ from __future__ import annotations import time from collections.abc import Generator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import requests @@ -19,24 +18,12 @@ import requests from onyx.tools.tool_implementations.python import code_interpreter_client as cic from onyx.tools.tool_implementations.python.code_interpreter_client import ( BashExecResponse, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import ( CodeInterpreterClient, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import ( CodeInterpreterVersionError, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import ( CreateSessionResponse, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import FileInput -from onyx.tools.tool_implementations.python.code_interpreter_client import ( + FileInput, HealthResponse, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import ( StreamOutputEvent, -) -from onyx.tools.tool_implementations.python.code_interpreter_client import ( StreamResultEvent, ) diff --git a/backend/tests/unit/onyx/tools/tool_implementations/python/test_python_tool_upload_cache.py b/backend/tests/unit/onyx/tools/tool_implementations/python/test_python_tool_upload_cache.py index 3d6f862c1f..7d07a8df51 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/python/test_python_tool_upload_cache.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/python/test_python_tool_upload_cache.py @@ -6,19 +6,18 @@ on every agent loop iteration. """ import json -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch -from onyx.tools.models import ChatFile -from onyx.tools.models import PythonToolOverrideKwargs -from onyx.tools.models import ToolResponse +from onyx.tools.models import ChatFile, PythonToolOverrideKwargs, ToolResponse from onyx.tools.tool_implementations.python.code_interpreter_client import ( StreamResultEvent, ) -from onyx.tools.tool_implementations.python.python_tool import _build_staging_notice -from onyx.tools.tool_implementations.python.python_tool import _code_references_file -from onyx.tools.tool_implementations.python.python_tool import _select_files_for_staging -from onyx.tools.tool_implementations.python.python_tool import PythonTool +from onyx.tools.tool_implementations.python.python_tool import ( + _build_staging_notice, + _code_references_file, + _select_files_for_staging, + PythonTool, +) TOOL_MODULE = "onyx.tools.tool_implementations.python.python_tool" diff --git a/backend/tests/unit/onyx/tools/tool_implementations/search/test_search_tool_run.py b/backend/tests/unit/onyx/tools/tool_implementations/search/test_search_tool_run.py index fcadca6744..dac6766805 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/search/test_search_tool_run.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/search/test_search_tool_run.py @@ -1,17 +1,13 @@ from __future__ import annotations -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import MessageType +from onyx.configs.constants import DocumentSource, MessageType from onyx.context.search.models import BaseFilters from onyx.server.query_and_chat.placement import Placement from onyx.server.query_and_chat.streaming_models import SearchToolFilterDelta -from onyx.tools.models import ChatMinimalTextMessage -from onyx.tools.models import SearchToolOverrideKwargs +from onyx.tools.models import ChatMinimalTextMessage, SearchToolOverrideKwargs from onyx.tools.tool_implementations.search.search_tool import SearchTool MODULE = "onyx.tools.tool_implementations.search.search_tool" diff --git a/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_brave_client.py b/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_brave_client.py index d87e4a1142..9854ebc929 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_brave_client.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_brave_client.py @@ -1,7 +1,6 @@ from __future__ import annotations -from typing import Any -from typing import cast +from typing import Any, cast import pytest import requests diff --git a/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_providers.py b/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_providers.py index 3388ab4830..57a823c845 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_providers.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_providers.py @@ -3,8 +3,6 @@ import pytest from onyx.tools.tool_implementations.web_search.clients.brave_client import BraveClient from onyx.tools.tool_implementations.web_search.providers import ( build_search_provider_from_config, -) -from onyx.tools.tool_implementations.web_search.providers import ( provider_requires_api_key, ) from shared_configs.enums import WebSearchProviderType diff --git a/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_tool_run.py b/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_tool_run.py index beeb9b8ac6..187d5d126d 100644 --- a/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_tool_run.py +++ b/backend/tests/unit/onyx/tools/tool_implementations/websearch/test_web_search_tool_run.py @@ -1,20 +1,17 @@ from __future__ import annotations -from typing import Any -from typing import cast -from unittest.mock import MagicMock -from unittest.mock import patch +from typing import Any, cast +from unittest.mock import MagicMock, patch import pytest from onyx.server.query_and_chat.placement import Placement -from onyx.tools.models import ToolCallException -from onyx.tools.models import WebSearchToolOverrideKwargs +from onyx.tools.models import ToolCallException, WebSearchToolOverrideKwargs from onyx.tools.tool_implementations.web_search.models import WebSearchResult from onyx.tools.tool_implementations.web_search.web_search_tool import ( _normalize_queries_input, + WebSearchTool, ) -from onyx.tools.tool_implementations.web_search.web_search_tool import WebSearchTool def _make_result( diff --git a/backend/tests/unit/onyx/tracing/test_dynamic_processor.py b/backend/tests/unit/onyx/tracing/test_dynamic_processor.py index a619fccb82..426fa8a29f 100644 --- a/backend/tests/unit/onyx/tracing/test_dynamic_processor.py +++ b/backend/tests/unit/onyx/tracing/test_dynamic_processor.py @@ -1,15 +1,13 @@ """Unit tests for the live-reload behavior of DynamicTracingProcessor.""" -from typing import Any -from typing import cast +from typing import Any, cast from unittest.mock import patch from onyx.tracing.dynamic_processor import DynamicTracingProcessor from onyx.tracing.framework.processor_interface import TracingProcessor from onyx.tracing.framework.spans import Span from onyx.tracing.framework.traces import Trace -from onyx.tracing.provider_config import BraintrustConfig -from onyx.tracing.provider_config import EffectiveTracingConfig +from onyx.tracing.provider_config import BraintrustConfig, EffectiveTracingConfig RESOLVE = "onyx.tracing.dynamic_processor.resolve_effective_tracing_config" BUILD = "onyx.tracing.dynamic_processor.build_delegates" diff --git a/backend/tests/unit/onyx/tracing/test_tracing_setup.py b/backend/tests/unit/onyx/tracing/test_tracing_setup.py index 026a28cb65..8417f1692f 100644 --- a/backend/tests/unit/onyx/tracing/test_tracing_setup.py +++ b/backend/tests/unit/onyx/tracing/test_tracing_setup.py @@ -4,9 +4,11 @@ from unittest.mock import patch from onyx.tracing import setup as tracing_setup from onyx.tracing.dynamic_processor import DynamicTracingProcessor -from onyx.tracing.provider_config import BraintrustConfig -from onyx.tracing.provider_config import EffectiveTracingConfig -from onyx.tracing.provider_config import LangfuseConfig +from onyx.tracing.provider_config import ( + BraintrustConfig, + EffectiveTracingConfig, + LangfuseConfig, +) RESOLVE = "onyx.tracing.dynamic_processor.resolve_effective_tracing_config" BUILD = "onyx.tracing.dynamic_processor.build_delegates" diff --git a/backend/tests/unit/onyx/utils/test_audit.py b/backend/tests/unit/onyx/utils/test_audit.py index 60603dcaf4..65f6712040 100644 --- a/backend/tests/unit/onyx/utils/test_audit.py +++ b/backend/tests/unit/onyx/utils/test_audit.py @@ -8,23 +8,23 @@ and Redis-backed dedup that degrades safely when Redis is unavailable. import io import json import logging -from typing import Any -from typing import cast -from typing import TextIO +from typing import Any, cast, TextIO from unittest.mock import MagicMock import pytest from onyx.utils import audit -from onyx.utils.audit import _OCSF_CLASS_BY_ACTION -from onyx.utils.audit import actor_from_user -from onyx.utils.audit import AUDIT_HANDLER_NAME -from onyx.utils.audit import AUDIT_LOGGER_ROOT -from onyx.utils.audit import AuditAction -from onyx.utils.audit import AuditActor -from onyx.utils.audit import AuditOutcome -from onyx.utils.audit import emit_audit_event -from onyx.utils.audit import OCSFEventClass +from onyx.utils.audit import ( + _OCSF_CLASS_BY_ACTION, + actor_from_user, + AUDIT_HANDLER_NAME, + AUDIT_LOGGER_ROOT, + AuditAction, + AuditActor, + AuditOutcome, + emit_audit_event, + OCSFEventClass, +) def _capture(caplog: pytest.LogCaptureFixture) -> list[dict[str, Any]]: diff --git a/backend/tests/unit/onyx/utils/test_client_ip.py b/backend/tests/unit/onyx/utils/test_client_ip.py index cfce28e71b..c667788d9b 100644 --- a/backend/tests/unit/onyx/utils/test_client_ip.py +++ b/backend/tests/unit/onyx/utils/test_client_ip.py @@ -2,13 +2,10 @@ from unittest.mock import MagicMock -from fastapi import FastAPI -from fastapi import Request +from fastapi import FastAPI, Request from fastapi.testclient import TestClient -from onyx.utils.client_ip import ClientIPMiddleware -from onyx.utils.client_ip import current_client_ip -from onyx.utils.client_ip import get_client_ip +from onyx.utils.client_ip import ClientIPMiddleware, current_client_ip, get_client_ip def _fake_request( diff --git a/backend/tests/unit/onyx/utils/test_csv_utils.py b/backend/tests/unit/onyx/utils/test_csv_utils.py index 1382d1d3b0..7a9c9faf84 100644 --- a/backend/tests/unit/onyx/utils/test_csv_utils.py +++ b/backend/tests/unit/onyx/utils/test_csv_utils.py @@ -1,8 +1,10 @@ import pytest -from onyx.utils.csv_utils import sanitize_csv_cell -from onyx.utils.csv_utils import sanitize_csv_cell_or_none -from onyx.utils.csv_utils import sanitize_csv_row +from onyx.utils.csv_utils import ( + sanitize_csv_cell, + sanitize_csv_cell_or_none, + sanitize_csv_row, +) @pytest.mark.parametrize( diff --git a/backend/tests/unit/onyx/utils/test_external_endpoint.py b/backend/tests/unit/onyx/utils/test_external_endpoint.py index 79870b395d..8bb88a24e3 100644 --- a/backend/tests/unit/onyx/utils/test_external_endpoint.py +++ b/backend/tests/unit/onyx/utils/test_external_endpoint.py @@ -1,13 +1,11 @@ """Unit tests for the generic external endpoint client.""" -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import httpx from pydantic import BaseModel -from onyx.utils.external_endpoint import ExternalEndpointConfig -from onyx.utils.external_endpoint import post_json_to_endpoint +from onyx.utils.external_endpoint import ExternalEndpointConfig, post_json_to_endpoint _CONFIG = ExternalEndpointConfig( endpoint_url="https://endpoint.example.com/x", timeout_seconds=5.0 diff --git a/backend/tests/unit/onyx/utils/test_github.py b/backend/tests/unit/onyx/utils/test_github.py index f02ab9eb42..49b43d13f8 100644 --- a/backend/tests/unit/onyx/utils/test_github.py +++ b/backend/tests/unit/onyx/utils/test_github.py @@ -1,15 +1,16 @@ """Unit tests for onyx.utils.github.""" from collections.abc import Iterable -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import requests -from onyx.utils.github import download_github_repo -from onyx.utils.github import GITHUB_TARBALL_URL -from onyx.utils.github import parse_github_repo +from onyx.utils.github import ( + download_github_repo, + GITHUB_TARBALL_URL, + parse_github_repo, +) class TestParseGithubRepo: diff --git a/backend/tests/unit/onyx/utils/test_gpu_utils.py b/backend/tests/unit/onyx/utils/test_gpu_utils.py index a7dad4860e..b8842541a8 100644 --- a/backend/tests/unit/onyx/utils/test_gpu_utils.py +++ b/backend/tests/unit/onyx/utils/test_gpu_utils.py @@ -4,8 +4,7 @@ Test cases for onyx/utils/gpu_utils.py with DISABLE_MODEL_SERVER environment var import os from unittest import TestCase -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import requests diff --git a/backend/tests/unit/onyx/utils/test_json_river.py b/backend/tests/unit/onyx/utils/test_json_river.py index d4db746a46..83e6bba20c 100644 --- a/backend/tests/unit/onyx/utils/test_json_river.py +++ b/backend/tests/unit/onyx/utils/test_json_river.py @@ -4,8 +4,7 @@ import json import pytest -from onyx.utils.jsonriver import JsonValue -from onyx.utils.jsonriver import Parser +from onyx.utils.jsonriver import JsonValue, Parser def _all_deltas(chunks: list[str]) -> list[JsonValue]: diff --git a/backend/tests/unit/onyx/utils/test_postgres_sanitization.py b/backend/tests/unit/onyx/utils/test_postgres_sanitization.py index 737eb56504..051c2f5ced 100644 --- a/backend/tests/unit/onyx/utils/test_postgres_sanitization.py +++ b/backend/tests/unit/onyx/utils/test_postgres_sanitization.py @@ -1,18 +1,22 @@ from pytest import MonkeyPatch from onyx.access.models import ExternalAccess -from onyx.connectors.models import BasicExpertInfo -from onyx.connectors.models import Document -from onyx.connectors.models import DocumentSource -from onyx.connectors.models import HierarchyNode -from onyx.connectors.models import IndexAttemptMetadata -from onyx.connectors.models import TextSection +from onyx.connectors.models import ( + BasicExpertInfo, + Document, + DocumentSource, + HierarchyNode, + IndexAttemptMetadata, + TextSection, +) from onyx.db.enums import HierarchyNodeType from onyx.indexing import indexing_pipeline -from onyx.utils.postgres_sanitization import sanitize_document_for_postgres -from onyx.utils.postgres_sanitization import sanitize_hierarchy_node_for_postgres -from onyx.utils.postgres_sanitization import sanitize_json_like -from onyx.utils.postgres_sanitization import sanitize_string +from onyx.utils.postgres_sanitization import ( + sanitize_document_for_postgres, + sanitize_hierarchy_node_for_postgres, + sanitize_json_like, + sanitize_string, +) # ---- sanitize_string tests ---- diff --git a/backend/tests/unit/onyx/utils/test_process_isolation.py b/backend/tests/unit/onyx/utils/test_process_isolation.py index b5de242cce..5f4a398321 100644 --- a/backend/tests/unit/onyx/utils/test_process_isolation.py +++ b/backend/tests/unit/onyx/utils/test_process_isolation.py @@ -10,9 +10,11 @@ import time import pytest -from onyx.utils.process_isolation import IsolatedProcessCrashed -from onyx.utils.process_isolation import IsolatedProcessTimeout -from onyx.utils.process_isolation import run_in_isolated_process +from onyx.utils.process_isolation import ( + IsolatedProcessCrashed, + IsolatedProcessTimeout, + run_in_isolated_process, +) def test_returns_value() -> None: diff --git a/backend/tests/unit/onyx/utils/test_retry_after.py b/backend/tests/unit/onyx/utils/test_retry_after.py index fadab44d34..230f1fe382 100644 --- a/backend/tests/unit/onyx/utils/test_retry_after.py +++ b/backend/tests/unit/onyx/utils/test_retry_after.py @@ -1,6 +1,4 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from email.utils import format_datetime import pytest diff --git a/backend/tests/unit/onyx/utils/test_sensitive.py b/backend/tests/unit/onyx/utils/test_sensitive.py index c598c6e040..4d38e0969c 100644 --- a/backend/tests/unit/onyx/utils/test_sensitive.py +++ b/backend/tests/unit/onyx/utils/test_sensitive.py @@ -6,8 +6,7 @@ from typing import Any import pytest from onyx.utils.encryption import mask_env_value_for_logging -from onyx.utils.sensitive import SensitiveAccessError -from onyx.utils.sensitive import SensitiveValue +from onyx.utils.sensitive import SensitiveAccessError, SensitiveValue def _encrypt_string(value: str) -> bytes: diff --git a/backend/tests/unit/onyx/utils/test_sensitive_typing.py b/backend/tests/unit/onyx/utils/test_sensitive_typing.py index c00a19dbe5..f860a9542e 100644 --- a/backend/tests/unit/onyx/utils/test_sensitive_typing.py +++ b/backend/tests/unit/onyx/utils/test_sensitive_typing.py @@ -14,8 +14,7 @@ from typing import Any def demonstrate_correct_usage() -> None: """Shows correct patterns that pass type checking.""" - from onyx.utils.encryption import decrypt_bytes_to_string - from onyx.utils.encryption import encrypt_string_to_bytes + from onyx.utils.encryption import decrypt_bytes_to_string, encrypt_string_to_bytes from onyx.utils.sensitive import SensitiveValue # Create a SensitiveValue diff --git a/backend/tests/unit/onyx/utils/test_threadpool_concurrency.py b/backend/tests/unit/onyx/utils/test_threadpool_concurrency.py index cbd15a1fa1..b355de2dec 100644 --- a/backend/tests/unit/onyx/utils/test_threadpool_concurrency.py +++ b/backend/tests/unit/onyx/utils/test_threadpool_concurrency.py @@ -1,17 +1,18 @@ import contextvars import threading import time -from collections.abc import Generator -from collections.abc import Iterator +from collections.abc import Generator, Iterator from concurrent.futures import ThreadPoolExecutor import pytest -from onyx.utils.threadpool_concurrency import parallel_yield -from onyx.utils.threadpool_concurrency import run_in_background -from onyx.utils.threadpool_concurrency import run_with_timeout -from onyx.utils.threadpool_concurrency import ThreadSafeDict -from onyx.utils.threadpool_concurrency import wait_on_background +from onyx.utils.threadpool_concurrency import ( + parallel_yield, + run_in_background, + run_with_timeout, + ThreadSafeDict, + wait_on_background, +) # Create a context variable for testing test_context_var = contextvars.ContextVar("test_var", default="default") diff --git a/backend/tests/unit/onyx/utils/test_threadpool_contextvars.py b/backend/tests/unit/onyx/utils/test_threadpool_contextvars.py index ffdbcc4ae0..c087259e17 100644 --- a/backend/tests/unit/onyx/utils/test_threadpool_contextvars.py +++ b/backend/tests/unit/onyx/utils/test_threadpool_contextvars.py @@ -1,13 +1,15 @@ import contextvars import time -from onyx.utils.threadpool_concurrency import FunctionCall -from onyx.utils.threadpool_concurrency import run_functions_in_parallel -from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel -from onyx.utils.threadpool_concurrency import run_in_background -from onyx.utils.threadpool_concurrency import run_with_timeout -from onyx.utils.threadpool_concurrency import start_thread_with_context -from onyx.utils.threadpool_concurrency import wait_on_background +from onyx.utils.threadpool_concurrency import ( + FunctionCall, + run_functions_in_parallel, + run_functions_tuples_in_parallel, + run_in_background, + run_with_timeout, + start_thread_with_context, + wait_on_background, +) # Create a test contextvar test_var = contextvars.ContextVar("test_var", default="default") diff --git a/backend/tests/unit/onyx/utils/test_tls.py b/backend/tests/unit/onyx/utils/test_tls.py index 6dddbd34f8..ea05ab75ec 100644 --- a/backend/tests/unit/onyx/utils/test_tls.py +++ b/backend/tests/unit/onyx/utils/test_tls.py @@ -5,8 +5,7 @@ from pathlib import Path import certifi import pytest from cryptography import x509 -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID diff --git a/backend/tests/unit/onyx/utils/test_url_ssrf.py b/backend/tests/unit/onyx/utils/test_url_ssrf.py index 12328d125b..d78f3a6610 100644 --- a/backend/tests/unit/onyx/utils/test_url_ssrf.py +++ b/backend/tests/unit/onyx/utils/test_url_ssrf.py @@ -5,16 +5,17 @@ These tests verify that the SSRF protection correctly blocks requests to internal/private IP addresses and other potentially dangerous destinations. """ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.utils.url import _is_ip_private_or_reserved -from onyx.utils.url import _validate_and_resolve_url -from onyx.utils.url import ssrf_safe_get -from onyx.utils.url import SSRFException -from onyx.utils.url import validate_outbound_http_url +from onyx.utils.url import ( + _is_ip_private_or_reserved, + _validate_and_resolve_url, + ssrf_safe_get, + SSRFException, + validate_outbound_http_url, +) class TestIsIpPrivateOrReserved: diff --git a/backend/tests/unit/onyx/utils/test_vespa_query.py b/backend/tests/unit/onyx/utils/test_vespa_query.py index f0bd6191e3..935a9ad2e8 100644 --- a/backend/tests/unit/onyx/utils/test_vespa_query.py +++ b/backend/tests/unit/onyx/utils/test_vespa_query.py @@ -1,23 +1,20 @@ -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone -from onyx.configs.constants import DocumentSource -from onyx.configs.constants import INDEX_SEPARATOR -from onyx.context.search.models import IndexFilters -from onyx.context.search.models import Tag -from onyx.context.search.models import TimeRange +from onyx.configs.constants import DocumentSource, INDEX_SEPARATOR +from onyx.context.search.models import IndexFilters, Tag, TimeRange from onyx.document_index.vespa.shared_utils.vespa_request_builders import ( build_vespa_filters, ) -from onyx.document_index.vespa_constants import DOC_UPDATED_AT -from onyx.document_index.vespa_constants import DOCUMENT_SETS -from onyx.document_index.vespa_constants import HIDDEN -from onyx.document_index.vespa_constants import METADATA_LIST -from onyx.document_index.vespa_constants import PERSONAS -from onyx.document_index.vespa_constants import SOURCE_TYPE -from onyx.document_index.vespa_constants import TENANT_ID -from onyx.document_index.vespa_constants import USER_PROJECT +from onyx.document_index.vespa_constants import ( + DOC_UPDATED_AT, + DOCUMENT_SETS, + HIDDEN, + METADATA_LIST, + PERSONAS, + SOURCE_TYPE, + TENANT_ID, + USER_PROJECT, +) from shared_configs.configs import MULTI_TENANT diff --git a/backend/tests/unit/onyx/voice/providers/test_elevenlabs_provider.py b/backend/tests/unit/onyx/voice/providers/test_elevenlabs_provider.py index d495baba41..35c6cd3bc2 100644 --- a/backend/tests/unit/onyx/voice/providers/test_elevenlabs_provider.py +++ b/backend/tests/unit/onyx/voice/providers/test_elevenlabs_provider.py @@ -1,9 +1,11 @@ import struct -from onyx.voice.providers.elevenlabs import _http_to_ws_url -from onyx.voice.providers.elevenlabs import DEFAULT_ELEVENLABS_API_BASE -from onyx.voice.providers.elevenlabs import ElevenLabsSTTMessageType -from onyx.voice.providers.elevenlabs import ElevenLabsVoiceProvider +from onyx.voice.providers.elevenlabs import ( + _http_to_ws_url, + DEFAULT_ELEVENLABS_API_BASE, + ElevenLabsSTTMessageType, + ElevenLabsVoiceProvider, +) # --- _http_to_ws_url --- diff --git a/backend/tests/unit/onyx/voice/providers/test_openai_provider.py b/backend/tests/unit/onyx/voice/providers/test_openai_provider.py index 1ad4de3688..338d97f2d2 100644 --- a/backend/tests/unit/onyx/voice/providers/test_openai_provider.py +++ b/backend/tests/unit/onyx/voice/providers/test_openai_provider.py @@ -2,10 +2,12 @@ import io import struct import wave -from onyx.voice.providers.openai import _create_wav_header -from onyx.voice.providers.openai import _http_to_ws_url -from onyx.voice.providers.openai import OpenAIRealtimeMessageType -from onyx.voice.providers.openai import OpenAIVoiceProvider +from onyx.voice.providers.openai import ( + _create_wav_header, + _http_to_ws_url, + OpenAIRealtimeMessageType, + OpenAIVoiceProvider, +) # --- _http_to_ws_url --- diff --git a/backend/tests/unit/sandbox_proxy/conftest.py b/backend/tests/unit/sandbox_proxy/conftest.py index 01ef6bf063..4aa82fe48c 100644 --- a/backend/tests/unit/sandbox_proxy/conftest.py +++ b/backend/tests/unit/sandbox_proxy/conftest.py @@ -12,20 +12,19 @@ from __future__ import annotations from typing import Any from unittest.mock import MagicMock -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 from mitmproxy import http from onyx.db.enums import EndpointPolicy -from onyx.external_apps.matching.engine import AllMatchedActions -from onyx.external_apps.matching.engine import MatchedAction +from onyx.external_apps.matching.engine import AllMatchedActions, MatchedAction from onyx.sandbox_proxy.addons.gate import _IdentityResolver -from onyx.sandbox_proxy.credential_injection import CredentialResolver -from onyx.sandbox_proxy.credential_injection import InjectionContext -from onyx.sandbox_proxy.identity import ResolvedSandbox -from onyx.sandbox_proxy.identity import SandboxIdentity -from onyx.sandbox_proxy.identity import SandboxIPLookup +from onyx.sandbox_proxy.credential_injection import CredentialResolver, InjectionContext +from onyx.sandbox_proxy.identity import ( + ResolvedSandbox, + SandboxIdentity, + SandboxIPLookup, +) _SANDBOX_ID = UUID("11111111-1111-1111-1111-111111111111") diff --git a/backend/tests/unit/sandbox_proxy/test_approval_cache.py b/backend/tests/unit/sandbox_proxy/test_approval_cache.py index 8aa1cd21a7..b220b04d26 100644 --- a/backend/tests/unit/sandbox_proxy/test_approval_cache.py +++ b/backend/tests/unit/sandbox_proxy/test_approval_cache.py @@ -2,13 +2,14 @@ from uuid import uuid4 import pytest -from onyx.cache.interface import CacheBackend -from onyx.cache.interface import CacheLock +from onyx.cache.interface import CacheBackend, CacheLock from onyx.db.enums import ApprovalDecision -from onyx.sandbox_proxy.approval_cache import _wake_key -from onyx.sandbox_proxy.approval_cache import cache_session_grant_actions -from onyx.sandbox_proxy.approval_cache import cached_session_grants_cover -from onyx.sandbox_proxy.approval_cache import wait_for_wake +from onyx.sandbox_proxy.approval_cache import ( + _wake_key, + cache_session_grant_actions, + cached_session_grants_cover, + wait_for_wake, +) class _MemoryCache(CacheBackend): diff --git a/backend/tests/unit/sandbox_proxy/test_backend_dispatch.py b/backend/tests/unit/sandbox_proxy/test_backend_dispatch.py index 73972532f7..89fe21db09 100644 --- a/backend/tests/unit/sandbox_proxy/test_backend_dispatch.py +++ b/backend/tests/unit/sandbox_proxy/test_backend_dispatch.py @@ -1,5 +1,4 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/sandbox_proxy/test_ca_bootstrap.py b/backend/tests/unit/sandbox_proxy/test_ca_bootstrap.py index 8827518a24..01eeccd7e6 100644 --- a/backend/tests/unit/sandbox_proxy/test_ca_bootstrap.py +++ b/backend/tests/unit/sandbox_proxy/test_ca_bootstrap.py @@ -4,14 +4,11 @@ from pathlib import Path import pytest from cryptography import x509 -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID -from onyx.sandbox_proxy.ca import CABootstrap -from onyx.sandbox_proxy.ca import CAStore -from onyx.sandbox_proxy.ca import CAStoreConflictError +from onyx.sandbox_proxy.ca import CABootstrap, CAStore, CAStoreConflictError class _InMemoryStore(CAStore): diff --git a/backend/tests/unit/sandbox_proxy/test_credential_injection.py b/backend/tests/unit/sandbox_proxy/test_credential_injection.py index abc30185f2..ae5f62ea28 100644 --- a/backend/tests/unit/sandbox_proxy/test_credential_injection.py +++ b/backend/tests/unit/sandbox_proxy/test_credential_injection.py @@ -9,15 +9,19 @@ from __future__ import annotations from unittest.mock import MagicMock from onyx.external_apps.matching.engine import AllMatchedActions -from onyx.sandbox_proxy.credential_injection import CredentialInjectionDispatcher -from onyx.sandbox_proxy.credential_injection import CredentialResolver -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionContext -from onyx.sandbox_proxy.credential_injection import InjectionOutcome -from tests.unit.sandbox_proxy.conftest import make_flow -from tests.unit.sandbox_proxy.conftest import make_matched_actions -from tests.unit.sandbox_proxy.conftest import make_resolved_sandbox -from tests.unit.sandbox_proxy.conftest import RecordingCredentialResolver +from onyx.sandbox_proxy.credential_injection import ( + CredentialInjectionDispatcher, + CredentialResolver, + CredentialUnavailableError, + InjectionContext, + InjectionOutcome, +) +from tests.unit.sandbox_proxy.conftest import ( + make_flow, + make_matched_actions, + make_resolved_sandbox, + RecordingCredentialResolver, +) def _ctx(*, matched_actions: AllMatchedActions | None = None) -> InjectionContext: diff --git a/backend/tests/unit/sandbox_proxy/test_errors.py b/backend/tests/unit/sandbox_proxy/test_errors.py index a0596e71b6..c9f0b78a59 100644 --- a/backend/tests/unit/sandbox_proxy/test_errors.py +++ b/backend/tests/unit/sandbox_proxy/test_errors.py @@ -9,8 +9,7 @@ import json import pytest -from onyx.sandbox_proxy.errors import http_403 -from onyx.sandbox_proxy.errors import SandboxProxyError +from onyx.sandbox_proxy.errors import http_403, SandboxProxyError @pytest.mark.parametrize("code", list(SandboxProxyError)) diff --git a/backend/tests/unit/sandbox_proxy/test_external_app_resolver.py b/backend/tests/unit/sandbox_proxy/test_external_app_resolver.py index d3f3ef191a..62f73b8ced 100644 --- a/backend/tests/unit/sandbox_proxy/test_external_app_resolver.py +++ b/backend/tests/unit/sandbox_proxy/test_external_app_resolver.py @@ -16,13 +16,17 @@ from unittest.mock import MagicMock import pytest from onyx.external_apps.matching.engine import AllMatchedActions -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionContext +from onyx.sandbox_proxy.credential_injection import ( + CredentialUnavailableError, + InjectionContext, +) from onyx.sandbox_proxy.resolvers import external_app from onyx.sandbox_proxy.resolvers.external_app import ExternalAppResolver -from tests.unit.sandbox_proxy.conftest import make_flow -from tests.unit.sandbox_proxy.conftest import make_matched_actions -from tests.unit.sandbox_proxy.conftest import make_resolved_sandbox +from tests.unit.sandbox_proxy.conftest import ( + make_flow, + make_matched_actions, + make_resolved_sandbox, +) @pytest.fixture(autouse=True) diff --git a/backend/tests/unit/sandbox_proxy/test_gate.py b/backend/tests/unit/sandbox_proxy/test_gate.py index fd33e284d2..007ef862b1 100644 --- a/backend/tests/unit/sandbox_proxy/test_gate.py +++ b/backend/tests/unit/sandbox_proxy/test_gate.py @@ -18,39 +18,35 @@ import logging from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass -from typing import AbstractSet -from typing import Any +from typing import AbstractSet, Any from unittest.mock import MagicMock -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest -from mitmproxy import connection -from mitmproxy import http +from mitmproxy import connection, http from mitmproxy.proxy import server_hooks from redis.exceptions import RedisError -from onyx.db.enums import ApprovalDecidedVia -from onyx.db.enums import ApprovalDecision -from onyx.db.enums import EndpointPolicy -from onyx.external_apps.matching.engine import AllMatchedActions -from onyx.external_apps.matching.engine import MatchedAction +from onyx.db.enums import ApprovalDecidedVia, ApprovalDecision, EndpointPolicy +from onyx.external_apps.matching.engine import AllMatchedActions, MatchedAction from onyx.sandbox_proxy.addons import gate -from onyx.sandbox_proxy.addons.gate import GateAddon -from onyx.sandbox_proxy.addons.gate import ParkedApprovals -from onyx.sandbox_proxy.credential_injection import CredentialInjectionDispatcher -from onyx.sandbox_proxy.credential_injection import CredentialResolver -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionOutcome +from onyx.sandbox_proxy.addons.gate import GateAddon, ParkedApprovals +from onyx.sandbox_proxy.credential_injection import ( + CredentialInjectionDispatcher, + CredentialResolver, + CredentialUnavailableError, + InjectionOutcome, +) from onyx.sandbox_proxy.errors import SandboxProxyError -from onyx.sandbox_proxy.identity import ResolvedSandbox -from onyx.sandbox_proxy.identity import SessionContext +from onyx.sandbox_proxy.identity import ResolvedSandbox, SessionContext from onyx.sandbox_proxy.request_evaluator import RequestEvaluator -from tests.unit.sandbox_proxy.conftest import make_flow -from tests.unit.sandbox_proxy.conftest import make_matched_actions -from tests.unit.sandbox_proxy.conftest import make_resolved_sandbox -from tests.unit.sandbox_proxy.conftest import RecordingCredentialResolver -from tests.unit.sandbox_proxy.conftest import StubResolver +from tests.unit.sandbox_proxy.conftest import ( + make_flow, + make_matched_actions, + make_resolved_sandbox, + RecordingCredentialResolver, + StubResolver, +) # --------------------------------------------------------------------------- # Stubs diff --git a/backend/tests/unit/sandbox_proxy/test_healthz.py b/backend/tests/unit/sandbox_proxy/test_healthz.py index c30e30a324..666a877d2d 100644 --- a/backend/tests/unit/sandbox_proxy/test_healthz.py +++ b/backend/tests/unit/sandbox_proxy/test_healthz.py @@ -6,8 +6,7 @@ from http.server import HTTPServer import pytest from onyx.sandbox_proxy.identity import SandboxIPLookup -from onyx.sandbox_proxy.server import _build_healthz_handler -from onyx.sandbox_proxy.server import _Readiness +from onyx.sandbox_proxy.server import _build_healthz_handler, _Readiness class _FakeLookup(SandboxIPLookup): diff --git a/backend/tests/unit/sandbox_proxy/test_identity.py b/backend/tests/unit/sandbox_proxy/test_identity.py index 5cc7544a2c..b89db5c564 100644 --- a/backend/tests/unit/sandbox_proxy/test_identity.py +++ b/backend/tests/unit/sandbox_proxy/test_identity.py @@ -1,15 +1,16 @@ from collections.abc import Iterator from contextlib import contextmanager from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from onyx.sandbox_proxy import identity as identity_mod -from onyx.sandbox_proxy.identity import IdentityResolver -from onyx.sandbox_proxy.identity import ResolvedSandbox -from onyx.sandbox_proxy.identity import SandboxIdentity +from onyx.sandbox_proxy.identity import ( + IdentityResolver, + ResolvedSandbox, + SandboxIdentity, +) from tests.unit.sandbox_proxy.conftest import StaticLookup diff --git a/backend/tests/unit/sandbox_proxy/test_identity_docker.py b/backend/tests/unit/sandbox_proxy/test_identity_docker.py index 11e11cbb73..dca8f5c2f7 100644 --- a/backend/tests/unit/sandbox_proxy/test_identity_docker.py +++ b/backend/tests/unit/sandbox_proxy/test_identity_docker.py @@ -6,8 +6,10 @@ from unittest.mock import MagicMock import pytest from docker import DockerClient -from onyx.sandbox_proxy.identity_docker import _identity_from_container -from onyx.sandbox_proxy.identity_docker import DockerEventsLookup +from onyx.sandbox_proxy.identity_docker import ( + _identity_from_container, + DockerEventsLookup, +) _DEFAULT_NETWORK = "onyx_craft_sandbox" diff --git a/backend/tests/unit/sandbox_proxy/test_informer.py b/backend/tests/unit/sandbox_proxy/test_informer.py index eca7d16562..bb89834d43 100644 --- a/backend/tests/unit/sandbox_proxy/test_informer.py +++ b/backend/tests/unit/sandbox_proxy/test_informer.py @@ -1,11 +1,9 @@ -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from kubernetes import client -from onyx.sandbox_proxy.identity_k8s import _identity_from_pod -from onyx.sandbox_proxy.identity_k8s import K8sInformerLookup +from onyx.sandbox_proxy.identity_k8s import _identity_from_pod, K8sInformerLookup def _make_pod( diff --git a/backend/tests/unit/sandbox_proxy/test_llm_provider_key_resolver.py b/backend/tests/unit/sandbox_proxy/test_llm_provider_key_resolver.py index c5d3a20962..237c8493f1 100644 --- a/backend/tests/unit/sandbox_proxy/test_llm_provider_key_resolver.py +++ b/backend/tests/unit/sandbox_proxy/test_llm_provider_key_resolver.py @@ -16,14 +16,15 @@ from unittest.mock import MagicMock import pytest -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionContext +from onyx.sandbox_proxy.credential_injection import ( + CredentialUnavailableError, + InjectionContext, +) from onyx.sandbox_proxy.resolvers import llm_provider_key from onyx.sandbox_proxy.resolvers.llm_provider_key import LLMProviderKeyResolver from onyx.server.features.build.configs import BUILD_MODE_ALLOWED_PROVIDER_TYPES from onyx.utils.sensitive import make_mock_sensitive_value -from tests.unit.sandbox_proxy.conftest import make_flow -from tests.unit.sandbox_proxy.conftest import make_resolved_sandbox +from tests.unit.sandbox_proxy.conftest import make_flow, make_resolved_sandbox class _FakeProvider: diff --git a/backend/tests/unit/sandbox_proxy/test_onyx_pat_resolver.py b/backend/tests/unit/sandbox_proxy/test_onyx_pat_resolver.py index 06562f4c3a..29271564b6 100644 --- a/backend/tests/unit/sandbox_proxy/test_onyx_pat_resolver.py +++ b/backend/tests/unit/sandbox_proxy/test_onyx_pat_resolver.py @@ -15,14 +15,14 @@ from unittest.mock import MagicMock import pytest -from onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME -from onyx.auth.constants import API_KEY_HEADER_NAME -from onyx.sandbox_proxy.credential_injection import CredentialUnavailableError -from onyx.sandbox_proxy.credential_injection import InjectionContext +from onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME, API_KEY_HEADER_NAME +from onyx.sandbox_proxy.credential_injection import ( + CredentialUnavailableError, + InjectionContext, +) from onyx.sandbox_proxy.resolvers import onyx_pat from onyx.sandbox_proxy.resolvers.onyx_pat import OnyxPatResolver -from tests.unit.sandbox_proxy.conftest import make_flow -from tests.unit.sandbox_proxy.conftest import make_resolved_sandbox +from tests.unit.sandbox_proxy.conftest import make_flow, make_resolved_sandbox _API_URL = "https://api.onyx.example.com" _API_HOST = "api.onyx.example.com" diff --git a/backend/tests/unit/sandbox_proxy/test_response_streaming.py b/backend/tests/unit/sandbox_proxy/test_response_streaming.py index 8d5926956b..189caff303 100644 --- a/backend/tests/unit/sandbox_proxy/test_response_streaming.py +++ b/backend/tests/unit/sandbox_proxy/test_response_streaming.py @@ -18,14 +18,11 @@ import http.client import socket import threading import time -from collections.abc import Callable -from collections.abc import Iterator -from http.server import BaseHTTPRequestHandler -from http.server import ThreadingHTTPServer +from collections.abc import Callable, Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any -from uuid import UUID -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest from mitmproxy import http as mitm_http @@ -33,8 +30,7 @@ from mitmproxy.options import Options from mitmproxy.tools.dump import DumpMaster from onyx.sandbox_proxy.addons import gate -from onyx.sandbox_proxy.addons.gate import _IdentityResolver -from onyx.sandbox_proxy.addons.gate import GateAddon +from onyx.sandbox_proxy.addons.gate import _IdentityResolver, GateAddon from onyx.sandbox_proxy.credential_injection import CredentialInjectionDispatcher from onyx.sandbox_proxy.identity import ResolvedSandbox from onyx.sandbox_proxy.request_evaluator import RequestEvaluator diff --git a/backend/tests/unit/scripts/test_env_inventory.py b/backend/tests/unit/scripts/test_env_inventory.py index 3d63ff2c21..df5c72c462 100644 --- a/backend/tests/unit/scripts/test_env_inventory.py +++ b/backend/tests/unit/scripts/test_env_inventory.py @@ -9,12 +9,14 @@ import ast from pathlib import Path import pytest -from scripts.env_inventory import classify_var -from scripts.env_inventory import diff_baseline -from scripts.env_inventory import EnvVisitor -from scripts.env_inventory import format_baseline -from scripts.env_inventory import is_sensitive -from scripts.env_inventory import read_baseline +from scripts.env_inventory import ( + classify_var, + diff_baseline, + EnvVisitor, + format_baseline, + is_sensitive, + read_baseline, +) @pytest.mark.parametrize( diff --git a/backend/tests/unit/server/documents/test_stage_metric_residual.py b/backend/tests/unit/server/documents/test_stage_metric_residual.py index 8fd6302b6e..51f1084e0f 100644 --- a/backend/tests/unit/server/documents/test_stage_metric_residual.py +++ b/backend/tests/unit/server/documents/test_stage_metric_residual.py @@ -5,11 +5,12 @@ tested without a DB. Covers the residual math, the in-span component set, the clamp, and the absent-BATCH_TOTAL case. """ -from onyx.db.index_attempt_metrics_models import IndexAttemptStage -from onyx.db.index_attempt_metrics_models import STAGE_SCOPE -from onyx.server.documents.models import _BATCH_TOTAL_COMPONENT_STAGES -from onyx.server.documents.models import IndexAttemptStageMetricSnapshot -from onyx.server.documents.models import synthesize_unaccounted +from onyx.db.index_attempt_metrics_models import IndexAttemptStage, STAGE_SCOPE +from onyx.server.documents.models import ( + _BATCH_TOTAL_COMPONENT_STAGES, + IndexAttemptStageMetricSnapshot, + synthesize_unaccounted, +) def _snap( diff --git a/backend/tests/unit/server/metrics/test_celery_task_metrics.py b/backend/tests/unit/server/metrics/test_celery_task_metrics.py index 5d98ae9839..0b5a9ad1ec 100644 --- a/backend/tests/unit/server/metrics/test_celery_task_metrics.py +++ b/backend/tests/unit/server/metrics/test_celery_task_metrics.py @@ -7,14 +7,16 @@ from unittest.mock import MagicMock import pytest from onyx.background.celery.apps.app_base import on_before_task_publish -from onyx.server.metrics.celery_task_metrics import _task_start_times -from onyx.server.metrics.celery_task_metrics import on_celery_task_postrun -from onyx.server.metrics.celery_task_metrics import on_celery_task_prerun -from onyx.server.metrics.celery_task_metrics import TASK_COMPLETED -from onyx.server.metrics.celery_task_metrics import TASK_DURATION -from onyx.server.metrics.celery_task_metrics import TASK_QUEUE_WAIT -from onyx.server.metrics.celery_task_metrics import TASK_STARTED -from onyx.server.metrics.celery_task_metrics import TASKS_ACTIVE +from onyx.server.metrics.celery_task_metrics import ( + _task_start_times, + on_celery_task_postrun, + on_celery_task_prerun, + TASK_COMPLETED, + TASK_DURATION, + TASK_QUEUE_WAIT, + TASK_STARTED, + TASKS_ACTIVE, +) @pytest.fixture(autouse=True) diff --git a/backend/tests/unit/server/metrics/test_connector_state_metrics.py b/backend/tests/unit/server/metrics/test_connector_state_metrics.py index f517cba29e..2792e287cc 100644 --- a/backend/tests/unit/server/metrics/test_connector_state_metrics.py +++ b/backend/tests/unit/server/metrics/test_connector_state_metrics.py @@ -1,11 +1,8 @@ """Unit tests for the DB-snapshot connector state metrics collector.""" from contextlib import nullcontext -from datetime import datetime -from datetime import timedelta -from datetime import timezone -from threading import Event -from threading import Thread +from datetime import datetime, timedelta, timezone +from threading import Event, Thread from unittest.mock import MagicMock import pytest @@ -14,12 +11,12 @@ from prometheus_client import CollectorRegistry import onyx.server.metrics.connector_state_metrics as csm from onyx.configs.constants import DocumentSource from onyx.db.connector_credential_pair import ConnectorStateSnapshot -from onyx.db.enums import AccessType -from onyx.db.enums import ConnectorCredentialPairStatus -from onyx.db.enums import IndexingMode -from onyx.server.metrics.connector_state_metrics import _enum_label -from onyx.server.metrics.connector_state_metrics import _to_unix_ts -from onyx.server.metrics.connector_state_metrics import ConnectorStateMetricsCollector +from onyx.db.enums import AccessType, ConnectorCredentialPairStatus, IndexingMode +from onyx.server.metrics.connector_state_metrics import ( + _enum_label, + _to_unix_ts, + ConnectorStateMetricsCollector, +) def _snapshot() -> ConnectorStateSnapshot: diff --git a/backend/tests/unit/server/metrics/test_deletion_metrics.py b/backend/tests/unit/server/metrics/test_deletion_metrics.py index 3d57b3dd99..0f71f6fecf 100644 --- a/backend/tests/unit/server/metrics/test_deletion_metrics.py +++ b/backend/tests/unit/server/metrics/test_deletion_metrics.py @@ -2,16 +2,18 @@ import pytest -from onyx.server.metrics.deletion_metrics import DELETION_BLOCKED -from onyx.server.metrics.deletion_metrics import DELETION_COMPLETED -from onyx.server.metrics.deletion_metrics import DELETION_FENCE_RESET -from onyx.server.metrics.deletion_metrics import DELETION_STARTED -from onyx.server.metrics.deletion_metrics import DELETION_TASKSET_DURATION -from onyx.server.metrics.deletion_metrics import inc_deletion_blocked -from onyx.server.metrics.deletion_metrics import inc_deletion_completed -from onyx.server.metrics.deletion_metrics import inc_deletion_fence_reset -from onyx.server.metrics.deletion_metrics import inc_deletion_started -from onyx.server.metrics.deletion_metrics import observe_deletion_taskset_duration +from onyx.server.metrics.deletion_metrics import ( + DELETION_BLOCKED, + DELETION_COMPLETED, + DELETION_FENCE_RESET, + DELETION_STARTED, + DELETION_TASKSET_DURATION, + inc_deletion_blocked, + inc_deletion_completed, + inc_deletion_fence_reset, + inc_deletion_started, + observe_deletion_taskset_duration, +) class TestIncDeletionStarted: diff --git a/backend/tests/unit/server/metrics/test_embedding_metrics.py b/backend/tests/unit/server/metrics/test_embedding_metrics.py index 0805471708..2b86aea020 100644 --- a/backend/tests/unit/server/metrics/test_embedding_metrics.py +++ b/backend/tests/unit/server/metrics/test_embedding_metrics.py @@ -2,19 +2,20 @@ from unittest.mock import patch -from onyx.server.metrics.embedding import _client_duration -from onyx.server.metrics.embedding import _embedding_input_chars_total -from onyx.server.metrics.embedding import _embedding_requests_total -from onyx.server.metrics.embedding import _embedding_texts_total -from onyx.server.metrics.embedding import _embeddings_in_progress -from onyx.server.metrics.embedding import LOCAL_PROVIDER_LABEL -from onyx.server.metrics.embedding import observe_embedding_client -from onyx.server.metrics.embedding import provider_label -from onyx.server.metrics.embedding import PROVIDER_LABEL_NAME -from onyx.server.metrics.embedding import TEXT_TYPE_LABEL_NAME -from onyx.server.metrics.embedding import track_embedding_in_progress -from shared_configs.enums import EmbeddingProvider -from shared_configs.enums import EmbedTextType +from onyx.server.metrics.embedding import ( + _client_duration, + _embedding_input_chars_total, + _embedding_requests_total, + _embedding_texts_total, + _embeddings_in_progress, + LOCAL_PROVIDER_LABEL, + observe_embedding_client, + provider_label, + PROVIDER_LABEL_NAME, + TEXT_TYPE_LABEL_NAME, + track_embedding_in_progress, +) +from shared_configs.enums import EmbeddingProvider, EmbedTextType class TestProviderLabel: diff --git a/backend/tests/unit/server/metrics/test_indexing_pipeline_collectors.py b/backend/tests/unit/server/metrics/test_indexing_pipeline_collectors.py index e8bb649585..46853888a5 100644 --- a/backend/tests/unit/server/metrics/test_indexing_pipeline_collectors.py +++ b/backend/tests/unit/server/metrics/test_indexing_pipeline_collectors.py @@ -1,8 +1,7 @@ """Tests for indexing pipeline Prometheus collectors.""" from collections.abc import Iterator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest diff --git a/backend/tests/unit/server/metrics/test_indexing_pipeline_setup.py b/backend/tests/unit/server/metrics/test_indexing_pipeline_setup.py index d1e573ebde..da202d2af0 100644 --- a/backend/tests/unit/server/metrics/test_indexing_pipeline_setup.py +++ b/backend/tests/unit/server/metrics/test_indexing_pipeline_setup.py @@ -2,8 +2,10 @@ from unittest.mock import MagicMock -from onyx.server.metrics.indexing_pipeline import QueueDepthCollector -from onyx.server.metrics.indexing_pipeline import RedisHealthCollector +from onyx.server.metrics.indexing_pipeline import ( + QueueDepthCollector, + RedisHealthCollector, +) class TestCollectorCeleryAppSetup: diff --git a/backend/tests/unit/server/metrics/test_indexing_task_metrics.py b/backend/tests/unit/server/metrics/test_indexing_task_metrics.py index 757e137432..8c82336fb1 100644 --- a/backend/tests/unit/server/metrics/test_indexing_task_metrics.py +++ b/backend/tests/unit/server/metrics/test_indexing_task_metrics.py @@ -1,19 +1,20 @@ """Tests for per-connector indexing task Prometheus metrics.""" from collections.abc import Iterator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.server.metrics.indexing_task_metrics import _connector_cache -from onyx.server.metrics.indexing_task_metrics import _indexing_start_times -from onyx.server.metrics.indexing_task_metrics import ConnectorInfo -from onyx.server.metrics.indexing_task_metrics import INDEXING_TASK_COMPLETED -from onyx.server.metrics.indexing_task_metrics import INDEXING_TASK_DURATION -from onyx.server.metrics.indexing_task_metrics import INDEXING_TASK_STARTED -from onyx.server.metrics.indexing_task_metrics import on_indexing_task_postrun -from onyx.server.metrics.indexing_task_metrics import on_indexing_task_prerun +from onyx.server.metrics.indexing_task_metrics import ( + _connector_cache, + _indexing_start_times, + ConnectorInfo, + INDEXING_TASK_COMPLETED, + INDEXING_TASK_DURATION, + INDEXING_TASK_STARTED, + on_indexing_task_postrun, + on_indexing_task_prerun, +) @pytest.fixture(autouse=True) diff --git a/backend/tests/unit/server/metrics/test_mcp_metrics.py b/backend/tests/unit/server/metrics/test_mcp_metrics.py index 5043b5ff03..b51017d3cd 100644 --- a/backend/tests/unit/server/metrics/test_mcp_metrics.py +++ b/backend/tests/unit/server/metrics/test_mcp_metrics.py @@ -1,7 +1,5 @@ import asyncio -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest @@ -11,11 +9,9 @@ from onyx.db.enums import MCPAuthenticationType from onyx.mcp_server.api import create_mcp_fastapi_app from onyx.mcp_server.auth import OnyxTokenVerifier from onyx.mcp_server.tools import search -from onyx.server.metrics import mcp_client -from onyx.server.metrics import metrics_auth +from onyx.server.metrics import mcp_client, metrics_auth from onyx.server.metrics.mcp_common import MCPToolCallStatus -from onyx.server.metrics.mcp_server import MCPAuthResult -from onyx.server.metrics.mcp_server import MCPServerToolName +from onyx.server.metrics.mcp_server import MCPAuthResult, MCPServerToolName from onyx.server.query_and_chat.placement import Placement from onyx.tools.tool_implementations.mcp.mcp_tool import MCPTool diff --git a/backend/tests/unit/server/metrics/test_metrics_server.py b/backend/tests/unit/server/metrics/test_metrics_server.py index 28436be819..32443bb9b0 100644 --- a/backend/tests/unit/server/metrics/test_metrics_server.py +++ b/backend/tests/unit/server/metrics/test_metrics_server.py @@ -1,13 +1,11 @@ """Tests for the Prometheus metrics server module.""" from collections.abc import Iterator -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest -from onyx.server.metrics.metrics_server import _DEFAULT_PORTS -from onyx.server.metrics.metrics_server import start_metrics_server +from onyx.server.metrics.metrics_server import _DEFAULT_PORTS, start_metrics_server @pytest.fixture(autouse=True) diff --git a/backend/tests/unit/server/metrics/test_opensearch_search_metrics.py b/backend/tests/unit/server/metrics/test_opensearch_search_metrics.py index 51adf60faf..2609934a72 100644 --- a/backend/tests/unit/server/metrics/test_opensearch_search_metrics.py +++ b/backend/tests/unit/server/metrics/test_opensearch_search_metrics.py @@ -5,15 +5,17 @@ from unittest.mock import patch import pytest from onyx.document_index.opensearch.constants import OpenSearchSearchType -from onyx.server.metrics.opensearch_search import _client_duration -from onyx.server.metrics.opensearch_search import _client_server_overhead -from onyx.server.metrics.opensearch_search import _search_errors -from onyx.server.metrics.opensearch_search import _search_total -from onyx.server.metrics.opensearch_search import _searches_in_progress -from onyx.server.metrics.opensearch_search import _server_duration -from onyx.server.metrics.opensearch_search import observe_opensearch_search -from onyx.server.metrics.opensearch_search import record_opensearch_search_error -from onyx.server.metrics.opensearch_search import track_opensearch_search +from onyx.server.metrics.opensearch_search import ( + _client_duration, + _client_server_overhead, + _search_errors, + _search_total, + _searches_in_progress, + _server_duration, + observe_opensearch_search, + record_opensearch_search_error, + track_opensearch_search, +) class TestRecordOpenSearchSearchError: diff --git a/backend/tests/unit/server/metrics/test_perm_sync_metrics.py b/backend/tests/unit/server/metrics/test_perm_sync_metrics.py index 6faf7b22e9..7f59f09eaa 100644 --- a/backend/tests/unit/server/metrics/test_perm_sync_metrics.py +++ b/backend/tests/unit/server/metrics/test_perm_sync_metrics.py @@ -2,26 +2,26 @@ import pytest -from onyx.server.metrics.perm_sync_metrics import DOC_PERM_SYNC_DB_UPDATE_DURATION -from onyx.server.metrics.perm_sync_metrics import DOC_PERM_SYNC_DOCS_PROCESSED -from onyx.server.metrics.perm_sync_metrics import DOC_PERM_SYNC_DURATION -from onyx.server.metrics.perm_sync_metrics import DOC_PERM_SYNC_ERRORS -from onyx.server.metrics.perm_sync_metrics import GROUP_SYNC_DURATION -from onyx.server.metrics.perm_sync_metrics import GROUP_SYNC_ERRORS -from onyx.server.metrics.perm_sync_metrics import GROUP_SYNC_GROUPS_PROCESSED -from onyx.server.metrics.perm_sync_metrics import GROUP_SYNC_UPSERT_DURATION -from onyx.server.metrics.perm_sync_metrics import GROUP_SYNC_USERS_PROCESSED -from onyx.server.metrics.perm_sync_metrics import inc_doc_perm_sync_docs_processed -from onyx.server.metrics.perm_sync_metrics import inc_doc_perm_sync_errors -from onyx.server.metrics.perm_sync_metrics import inc_group_sync_errors -from onyx.server.metrics.perm_sync_metrics import inc_group_sync_groups_processed -from onyx.server.metrics.perm_sync_metrics import inc_group_sync_users_processed from onyx.server.metrics.perm_sync_metrics import ( + DOC_PERM_SYNC_DB_UPDATE_DURATION, + DOC_PERM_SYNC_DOCS_PROCESSED, + DOC_PERM_SYNC_DURATION, + DOC_PERM_SYNC_ERRORS, + GROUP_SYNC_DURATION, + GROUP_SYNC_ERRORS, + GROUP_SYNC_GROUPS_PROCESSED, + GROUP_SYNC_UPSERT_DURATION, + GROUP_SYNC_USERS_PROCESSED, + inc_doc_perm_sync_docs_processed, + inc_doc_perm_sync_errors, + inc_group_sync_errors, + inc_group_sync_groups_processed, + inc_group_sync_users_processed, observe_doc_perm_sync_db_update_duration, + observe_doc_perm_sync_duration, + observe_group_sync_duration, + observe_group_sync_upsert_duration, ) -from onyx.server.metrics.perm_sync_metrics import observe_doc_perm_sync_duration -from onyx.server.metrics.perm_sync_metrics import observe_group_sync_duration -from onyx.server.metrics.perm_sync_metrics import observe_group_sync_upsert_duration # --- Doc permission sync: overall duration --- diff --git a/backend/tests/unit/server/metrics/test_pruning_metrics.py b/backend/tests/unit/server/metrics/test_pruning_metrics.py index 208c506f0d..9f80cda1ee 100644 --- a/backend/tests/unit/server/metrics/test_pruning_metrics.py +++ b/backend/tests/unit/server/metrics/test_pruning_metrics.py @@ -2,12 +2,14 @@ import pytest -from onyx.server.metrics.pruning_metrics import inc_pruning_rate_limit_error -from onyx.server.metrics.pruning_metrics import observe_pruning_diff_duration -from onyx.server.metrics.pruning_metrics import observe_pruning_enumeration_duration -from onyx.server.metrics.pruning_metrics import PRUNING_DIFF_DURATION -from onyx.server.metrics.pruning_metrics import PRUNING_ENUMERATION_DURATION -from onyx.server.metrics.pruning_metrics import PRUNING_RATE_LIMIT_ERRORS +from onyx.server.metrics.pruning_metrics import ( + inc_pruning_rate_limit_error, + observe_pruning_diff_duration, + observe_pruning_enumeration_duration, + PRUNING_DIFF_DURATION, + PRUNING_ENUMERATION_DURATION, + PRUNING_RATE_LIMIT_ERRORS, +) class TestObservePruningEnumerationDuration: diff --git a/backend/tests/unit/server/metrics/test_worker_health.py b/backend/tests/unit/server/metrics/test_worker_health.py index 1046b88ee1..b4fa71c032 100644 --- a/backend/tests/unit/server/metrics/test_worker_health.py +++ b/backend/tests/unit/server/metrics/test_worker_health.py @@ -3,8 +3,10 @@ import time from unittest.mock import MagicMock -from onyx.server.metrics.indexing_pipeline import WorkerHealthCollector -from onyx.server.metrics.indexing_pipeline import WorkerHeartbeatMonitor +from onyx.server.metrics.indexing_pipeline import ( + WorkerHealthCollector, + WorkerHeartbeatMonitor, +) class TestWorkerHeartbeatMonitor: diff --git a/backend/tests/unit/shared_configs/test_cors_config.py b/backend/tests/unit/shared_configs/test_cors_config.py index 9b2df75a03..aa7635f7b6 100644 --- a/backend/tests/unit/shared_configs/test_cors_config.py +++ b/backend/tests/unit/shared_configs/test_cors_config.py @@ -1,7 +1,6 @@ import pytest -from shared_configs.configs import cors_allow_credentials -from shared_configs.configs import parse_cors_allowed_origins +from shared_configs.configs import cors_allow_credentials, parse_cors_allowed_origins def test_empty_env_allows_all_origins_without_credentials() -> None: diff --git a/backend/tests/unit/tools/test_memory_tool_packets.py b/backend/tests/unit/tools/test_memory_tool_packets.py index e4b10e4425..f32b47a3cf 100644 --- a/backend/tests/unit/tools/test_memory_tool_packets.py +++ b/backend/tests/unit/tools/test_memory_tool_packets.py @@ -1,19 +1,22 @@ """Tests for memory tool streaming packet emissions.""" import queue -from unittest.mock import MagicMock -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from onyx.chat.emitter import Emitter from onyx.server.query_and_chat.placement import Placement from onyx.server.query_and_chat.session_loading import create_memory_packets -from onyx.server.query_and_chat.streaming_models import MemoryToolDelta -from onyx.server.query_and_chat.streaming_models import MemoryToolStart -from onyx.server.query_and_chat.streaming_models import SectionEnd -from onyx.tools.tool_implementations.memory.memory_tool import MemoryTool -from onyx.tools.tool_implementations.memory.memory_tool import MemoryToolOverrideKwargs +from onyx.server.query_and_chat.streaming_models import ( + MemoryToolDelta, + MemoryToolStart, + SectionEnd, +) +from onyx.tools.tool_implementations.memory.memory_tool import ( + MemoryTool, + MemoryToolOverrideKwargs, +) from onyx.tools.tool_implementations.memory.models import MemoryToolResponse diff --git a/backend/tests/utils/aws_secrets.py b/backend/tests/utils/aws_secrets.py index 7b1b162187..e843bc1668 100644 --- a/backend/tests/utils/aws_secrets.py +++ b/backend/tests/utils/aws_secrets.py @@ -35,14 +35,11 @@ AWS SSO Authentication: import logging import os from collections.abc import Sequence -from typing import cast -from typing import overload +from typing import cast, overload from dotenv import dotenv_values -from tests.utils.secret_names import AnySecret -from tests.utils.secret_names import DeploySecret -from tests.utils.secret_names import TestSecret +from tests.utils.secret_names import AnySecret, DeploySecret, TestSecret logger = logging.getLogger(__name__) diff --git a/examples/assistants-api/topics_analyzer.py b/examples/assistants-api/topics_analyzer.py index 13a9298738..2dc12a31e8 100644 --- a/examples/assistants-api/topics_analyzer.py +++ b/examples/assistants-api/topics_analyzer.py @@ -1,9 +1,7 @@ import argparse import os import time -from datetime import datetime -from datetime import timedelta -from datetime import timezone +from datetime import datetime, timedelta, timezone from openai import OpenAI diff --git a/loadtest/locustfile.py b/loadtest/locustfile.py index 4735306023..6938ef8a5a 100644 --- a/loadtest/locustfile.py +++ b/loadtest/locustfile.py @@ -12,15 +12,17 @@ import os import prometheus_exporter # noqa: F401 (registers the /metrics exporter) from onyx_client.chat_user import BasicChatUser -from scenarios import ChatWithSearchUser -from scenarios import CompressionUser -from scenarios import DeepResearchUser -from scenarios import DisconnectUser -from scenarios import FileAttachmentUser -from scenarios import HealthProbeUser -from scenarios import LongConversationUser -from scenarios import MultiToolUser -from scenarios import ThreadHogUser +from scenarios import ( + ChatWithSearchUser, + CompressionUser, + DeepResearchUser, + DisconnectUser, + FileAttachmentUser, + HealthProbeUser, + LongConversationUser, + MultiToolUser, + ThreadHogUser, +) __all__ = [ "BasicChatUser", diff --git a/loadtest/mock_llm/app.py b/loadtest/mock_llm/app.py index 70df90a729..d5bcd3562e 100644 --- a/loadtest/mock_llm/app.py +++ b/loadtest/mock_llm/app.py @@ -64,26 +64,25 @@ import uuid from collections.abc import AsyncGenerator from typing import Any -from fastapi import FastAPI -from fastapi import HTTPException -from fastapi import Response -from fastapi.responses import JSONResponse -from fastapi.responses import StreamingResponse +from fastapi import FastAPI, HTTPException, Response +from fastapi.responses import JSONResponse, StreamingResponse -from mock_llm.models import AssistantMessage -from mock_llm.models import ChatCompletionChunk -from mock_llm.models import ChatCompletionRequest -from mock_llm.models import ChatCompletionResponse -from mock_llm.models import ChatMessage -from mock_llm.models import Choice -from mock_llm.models import ChunkChoice -from mock_llm.models import ChunkDelta -from mock_llm.models import StreamToolCall -from mock_llm.models import StreamToolCallFunction -from mock_llm.models import ToolCall -from mock_llm.models import ToolCallFunction -from mock_llm.models import ToolDefinition -from mock_llm.models import Usage +from mock_llm.models import ( + AssistantMessage, + ChatCompletionChunk, + ChatCompletionRequest, + ChatCompletionResponse, + ChatMessage, + Choice, + ChunkChoice, + ChunkDelta, + StreamToolCall, + StreamToolCallFunction, + ToolCall, + ToolCallFunction, + ToolDefinition, + Usage, +) app = FastAPI() diff --git a/loadtest/mock_llm/models.py b/loadtest/mock_llm/models.py index bc33b103ab..234d9330c0 100644 --- a/loadtest/mock_llm/models.py +++ b/loadtest/mock_llm/models.py @@ -7,11 +7,9 @@ never 422 a request the real OpenAI API would accept. from __future__ import annotations -from typing import Any -from typing import Literal +from typing import Any, Literal -from pydantic import BaseModel -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict ################################################ # Request diff --git a/loadtest/onyx_client/chat_user.py b/loadtest/onyx_client/chat_user.py index 2f86c9e007..651930109d 100644 --- a/loadtest/onyx_client/chat_user.py +++ b/loadtest/onyx_client/chat_user.py @@ -23,12 +23,9 @@ import time import uuid from typing import Any -from locust import constant -from locust import HttpUser -from locust import task +from locust import constant, HttpUser, task -from onyx_client.env import env_float -from onyx_client.env import env_int +from onyx_client.env import env_float, env_int from onyx_client.stream_parser import ChatStreamAnalyzer DEFAULT_MESSAGES = [ diff --git a/loadtest/onyx_client/stream_parser.py b/loadtest/onyx_client/stream_parser.py index 9016a445ec..75fa19d60f 100644 --- a/loadtest/onyx_client/stream_parser.py +++ b/loadtest/onyx_client/stream_parser.py @@ -12,8 +12,7 @@ monkey-patching, where importing onyx.* (grpc, psycopg, etc.) breaks. from __future__ import annotations import json -from dataclasses import dataclass -from dataclasses import field +from dataclasses import dataclass, field # Packet type strings (subset of StreamingType in # backend/onyx/server/query_and_chat/streaming_models.py — keep in sync). diff --git a/loadtest/prometheus_collector.py b/loadtest/prometheus_collector.py index 2289932a7c..aa05ff35e0 100644 --- a/loadtest/prometheus_collector.py +++ b/loadtest/prometheus_collector.py @@ -10,8 +10,7 @@ from __future__ import annotations from collections.abc import Iterator from typing import Any -from prometheus_client.core import CounterMetricFamily -from prometheus_client.core import GaugeMetricFamily +from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily _LABELS = ["name", "method"] diff --git a/loadtest/scenarios/health_probe.py b/loadtest/scenarios/health_probe.py index c73ed70329..49a25e2321 100644 --- a/loadtest/scenarios/health_probe.py +++ b/loadtest/scenarios/health_probe.py @@ -30,11 +30,8 @@ from __future__ import annotations import os import time -from locust import constant_pacing -from locust import HttpUser -from locust import task -from onyx_client.env import env_float -from onyx_client.env import env_int +from locust import constant_pacing, HttpUser, task +from onyx_client.env import env_float, env_int class HealthProbeUser(HttpUser): diff --git a/loadtest/tests/test_stream_parser.py b/loadtest/tests/test_stream_parser.py index 8e009cc154..320c8dc4fd 100644 --- a/loadtest/tests/test_stream_parser.py +++ b/loadtest/tests/test_stream_parser.py @@ -4,9 +4,11 @@ from __future__ import annotations import json -from onyx_client.stream_parser import ChatStreamAnalyzer -from onyx_client.stream_parser import FIRST_ANSWER_TOKEN -from onyx_client.stream_parser import FIRST_PACKET +from onyx_client.stream_parser import ( + ChatStreamAnalyzer, + FIRST_ANSWER_TOKEN, + FIRST_PACKET, +) def _feed_all(analyzer: ChatStreamAnalyzer, lines: list[str]) -> list[str]: diff --git a/pyproject.toml b/pyproject.toml index 3c21666d1b..b8be925892 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -288,7 +288,6 @@ ignore = [ select = ["ARG", "E", "F", "G004", "I", "S", "W"] [tool.ruff.lint.isort] -force-single-line = true order-by-type = false known-first-party = ["onyx", "ee", "tests", "shared_configs", "model_server"]