diff --git a/.devcontainer/claude-code/CLAUDE.md b/.devcontainer/claude-code/CLAUDE.md index ec440915af..f0b62d27aa 100644 --- a/.devcontainer/claude-code/CLAUDE.md +++ b/.devcontainer/claude-code/CLAUDE.md @@ -16,7 +16,6 @@ Each is also exported as an env var: - Postgres: `relational_db` (`POSTGRES_HOST`) - Redis: `cache` (`REDIS_HOST`) -- Vespa: `index` (`VESPA_HOST`) - Model server: `inference_model_server` (`MODEL_SERVER_HOST`) - OpenSearch: `opensearch` (`OPENSEARCH_HOST`) - MinIO / S3: `minio:9000` (`S3_ENDPOINT_URL=http://minio:9000`) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 7ae714e5c7..47bbb4445c 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -30,8 +30,7 @@ "REDIS_HOST": "cache", "S3_ENDPOINT_URL": "http://minio:9000", "SSH_AUTH_SOCK": "/tmp/ssh-agent.sock", - "UV_LINK_MODE": "copy", - "VESPA_HOST": "index" + "UV_LINK_MODE": "copy" }, "remoteUser": "${localEnv:DEVCONTAINER_REMOTE_USER:dev}", "updateRemoteUserUID": false, diff --git a/.github/workflows/helm-chart-releases.yml b/.github/workflows/helm-chart-releases.yml index 499dc17c59..e697493f26 100644 --- a/.github/workflows/helm-chart-releases.yml +++ b/.github/workflows/helm-chart-releases.yml @@ -69,7 +69,6 @@ jobs: - name: Add required Helm repositories run: | helm repo add ingress-nginx https://kubernetes.github.io/ingress-nginx - helm repo add onyx-vespa https://onyx-dot-app.github.io/vespa-helm-charts helm repo add opensearch https://opensearch-project.github.io/helm-charts helm repo add cloudnative-pg https://cloudnative-pg.github.io/charts helm repo add ot-container-kit https://ot-container-kit.github.io/helm-charts diff --git a/.github/workflows/pr-external-dependency-unit-tests.yml b/.github/workflows/pr-external-dependency-unit-tests.yml index f5126e87a0..bb5339c77e 100644 --- a/.github/workflows/pr-external-dependency-unit-tests.yml +++ b/.github/workflows/pr-external-dependency-unit-tests.yml @@ -80,7 +80,7 @@ jobs: external-dependency-unit-tests: needs: discover-test-dirs - # Use larger runner with more resources for Vespa + # Use larger runner with more resources for OpenSearch runs-on: - runs-on - runner=2cpu-linux-arm64 diff --git a/.github/workflows/pr-helm-chart-testing.yml b/.github/workflows/pr-helm-chart-testing.yml index e933d0734a..5a10d938e1 100644 --- a/.github/workflows/pr-helm-chart-testing.yml +++ b/.github/workflows/pr-helm-chart-testing.yml @@ -87,7 +87,6 @@ jobs: run: | echo "=== Adding Helm repositories ===" helm repo add ingress-nginx https://kubernetes.github.io/ingress-nginx - helm repo add vespa https://onyx-dot-app.github.io/vespa-helm-charts helm repo add opensearch https://opensearch-project.github.io/helm-charts helm repo add cloudnative-pg https://cloudnative-pg.github.io/charts helm repo add ot-container-kit https://ot-container-kit.github.io/helm-charts @@ -196,7 +195,6 @@ jobs: --helm-extra-set-args="\ --set=nginx.enabled=false \ --set=minio.enabled=false \ - --set=vespa.enabled=false \ --set=opensearch.enabled=true \ --set=auth.opensearch.enabled=true \ --set=auth.userauth.values.user_auth_secret=test-secret \ diff --git a/.github/workflows/pr-python-tests.yml b/.github/workflows/pr-python-tests.yml index 7a5b9de389..03f446eaf7 100644 --- a/.github/workflows/pr-python-tests.yml +++ b/.github/workflows/pr-python-tests.yml @@ -53,7 +53,7 @@ jobs: REDIS_CLOUD_PYTEST_PASSWORD: ${{ secrets.REDIS_CLOUD_PYTEST_PASSWORD }} DISABLE_TELEMETRY: "true" # Unit tests run without a live backend; EE code paths that depend on - # Redis/Vespa/etc are not available, so disable enforcement here. + # Redis/OpenSearch/etc are not available, so disable enforcement here. LICENSE_ENFORCEMENT_ENABLED: "false" steps: diff --git a/.vscode/.env.k8s.template b/.vscode/.env.k8s.template index eeda210aa1..fabbde2a30 100644 --- a/.vscode/.env.k8s.template +++ b/.vscode/.env.k8s.template @@ -96,9 +96,6 @@ INDEXING_MODEL_SERVER_HOST=onyx-indexing-model-service # Leave the placeholder as-is; it will be replaced before the api_server # launches. To fetch manually: # kubectl -n onyx get secret onyx-opensearch -o jsonpath='{.data.opensearch_admin_password}' | base64 -d -ONYX_DISABLE_VESPA=true -ENABLE_OPENSEARCH_INDEXING_FOR_ONYX=true -ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX=false OPENSEARCH_HOST=onyx-opensearch-master.onyx.svc.cluster.local OPENSEARCH_REST_API_PORT=9200 OPENSEARCH_ADMIN_USERNAME=admin diff --git a/.vscode/launch.json b/.vscode/launch.json index f721c5cc7d..bd947003ac 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -358,7 +358,7 @@ "--loglevel=INFO", "--hostname=light@%n", "-Q", - "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration" + "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim" ], "presentation": { "group": "2", @@ -646,7 +646,7 @@ "--loglevel=INFO", "--hostname=light@%n", "-Q", - "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration" + "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim" ], "presentation": { "group": "2", diff --git a/.zed/debug.json b/.zed/debug.json index ec5af0d0aa..4bb4fa81b6 100644 --- a/.zed/debug.json +++ b/.zed/debug.json @@ -156,7 +156,7 @@ "--loglevel=INFO", "--hostname=light@%n", "-Q", - "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration", + "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim", ], "env": { "LOG_LEVEL": "INFO", diff --git a/.zed/process-compose.k8s.yaml b/.zed/process-compose.k8s.yaml index cd5725b495..6f3f260beb 100644 --- a/.zed/process-compose.k8s.yaml +++ b/.zed/process-compose.k8s.yaml @@ -67,7 +67,7 @@ processes: -A onyx.background.celery.versioned_apps.light worker --pool=threads --concurrency=64 --prefetch-multiplier=8 --loglevel=INFO --hostname=light@%n - -Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration + -Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim working_dir: ${ONYX_ZED_ROOT}/backend env_file: ${ONYX_ZED_ROOT}/.vscode/.env.k8s environment: *celery_environment diff --git a/.zed/tasks.json b/.zed/tasks.json index 063da32c19..4f53fa6f9e 100644 --- a/.zed/tasks.json +++ b/.zed/tasks.json @@ -343,7 +343,7 @@ "--loglevel=INFO", "--hostname=light@%n", "-Q", - "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration", + "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim", ], "cwd": "$ZED_WORKTREE_ROOT/backend", "env": { diff --git a/backend/.gitignore b/backend/.gitignore index 80a9a5fd09..7223eb83bc 100644 --- a/backend/.gitignore +++ b/backend/.gitignore @@ -6,7 +6,6 @@ site_crawls/ api_keys.py *ipynb .env* -vespa-app.zip dynamic_config_storage/ celerybeat-schedule* onyx/connectors/salesforce/data/ diff --git a/backend/AGENTS.md b/backend/AGENTS.md index e8b877a584..a111a4ecee 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -13,9 +13,10 @@ tests. Additive to the root `AGENTS.md`. PostgreSQL implementation of `CacheBackend`, so its queries belong there. - When creating new FastAPI APIs, do NOT use the `response_model` field. Instead, just type the function. -- OpenSearch is the current document index backend for search and indexing. Some legacy modules, - Celery task names, and migration helpers still mention Vespa; treat those as compatibility or - migration artifacts unless the active `DocumentIndex` factory/config path explicitly uses them. +- OpenSearch is the only document index backend for search and indexing. Onyx no longer uses + Vespa. Some live Celery names still say "vespa" (the `vespa_metadata_sync` queue, the + `check_for_vespa_sync_task` task, the `onyx.background.celery.tasks.vespa` module); they sync + the document index and keep their names so running deployments are not disrupted. - Do not use `getattr`: it hides attribute access from the type checker. Use plain attribute access when the name is statically known. A genuinely dynamic lookup needs an `# ods: ignore[getattr]` comment with a brief justification diff --git a/backend/Dockerfile b/backend/Dockerfile index 553042a687..56fff08dc9 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -20,7 +20,6 @@ RUN apt-get update && \ apt-get install -y \ cmake \ curl \ - zip \ ca-certificates \ libgnutls30 \ libblkid1 \ @@ -77,14 +76,12 @@ RUN groupadd -g 1001 onyx && \ chown onyx:onyx /var/log/onyx # Install runtime system dependencies only (no build toolchain). -# zip for Vespa step further down # ca-certificates for HTTPS # libxmlsec1-openssl is the runtime counterpart to the build-time libxmlsec1-dev # Interactive debugging tools (curl, vim, nano, psql, ps) are deliberately absent to keep the # shipped image minimal; they live in the dev stage below, published with a -dev tag suffix. RUN apt-get update && \ apt-get install -y \ - zip \ ca-certificates \ libgnutls30 \ libblkid1 \ diff --git a/backend/alembic/versions/12635f6655b7_drive_canonical_ids.py b/backend/alembic/versions/12635f6655b7_drive_canonical_ids.py index 63b963bb48..41bd854126 100644 --- a/backend/alembic/versions/12635f6655b7_drive_canonical_ids.py +++ b/backend/alembic/versions/12635f6655b7_drive_canonical_ids.py @@ -12,16 +12,49 @@ from urllib.parse import urlparse, urlunparse from httpx import HTTPStatusError import httpx from onyx.db.search_settings import SearchSettings -from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client -from onyx.document_index.vespa.shared_utils.utils import ( - replace_invalid_doc_id_characters, -) -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT from onyx.utils.logger import setup_logger import os +from typing import cast logger = setup_logger() +# Onyx no longer uses Vespa. These helpers were inlined from the removed Vespa +# document index module so this migration stays importable. +_VESPA_APP_CONTAINER_URL = os.environ.get("VESPA_CLOUD_URL") or ( + f"http://{os.environ.get('VESPA_HOST') or 'localhost'}" + f":{os.environ.get('VESPA_PORT') or '8081'}" +) +DOCUMENT_ID_ENDPOINT = ( + f"{_VESPA_APP_CONTAINER_URL}/document/v1/default/{{index_name}}/docid" +) + + +# Managed Vespa (Vespa Cloud) used mutual TLS. Self-hosted Vespa served a +# self-signed certificate. +_MANAGED_VESPA = os.environ.get("MANAGED_VESPA", "").lower() == "true" +_VESPA_CLOUD_CERT_PATH = os.environ.get("VESPA_CLOUD_CERT_PATH") +_VESPA_CLOUD_KEY_PATH = os.environ.get("VESPA_CLOUD_KEY_PATH") +_VESPA_REQUEST_TIMEOUT = int(os.environ.get("VESPA_REQUEST_TIMEOUT") or "15") + + +def get_vespa_http_client() -> httpx.Client: + return httpx.Client( + cert=( + cast(tuple[str, str], (_VESPA_CLOUD_CERT_PATH, _VESPA_CLOUD_KEY_PATH)) + if _MANAGED_VESPA + else None + ), + verify=_MANAGED_VESPA, + timeout=_VESPA_REQUEST_TIMEOUT, + http2=True, + ) + + +def replace_invalid_doc_id_characters(text: str) -> str: + """Replaces the document ID characters that Vespa did not accept.""" + return text.replace("'", "_") + + # revision identifiers, used by Alembic. revision = "12635f6655b7" down_revision = "58c50ef19f08" diff --git a/backend/alembic/versions/3067343245d1_drop_opensearch_migration_tables.py b/backend/alembic/versions/3067343245d1_drop_opensearch_migration_tables.py new file mode 100644 index 0000000000..68c63ab6f5 --- /dev/null +++ b/backend/alembic/versions/3067343245d1_drop_opensearch_migration_tables.py @@ -0,0 +1,178 @@ +"""drop opensearch migration tables + +These tables tracked the Vespa to OpenSearch migration. Every instance has +finished that migration, and nothing reads or writes the tables anymore. + +Revision ID: 3067343245d1 +Revises: b7c2e4f1a9d3 +Create Date: 2026-09-30 13:47:48.331663 + +""" + +import logging +import time + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.exc import DBAPIError + +logger = logging.getLogger("alembic.runtime.migration") + +# revision identifiers, used by Alembic. +revision = "3067343245d1" +down_revision = "b7c2e4f1a9d3" +branch_labels = None +depends_on = None + +_LOCK_TIMEOUT = "5s" +_DROP_ATTEMPTS = 12 +_RETRY_DELAY_S = 5 +# SQLSTATE lock_not_available: the error a lock_timeout raises. +_LOCK_NOT_AVAILABLE = "55P03" + + +def _is_lock_timeout(error: DBAPIError) -> bool: + # psycopg2 and SQLAlchemy's asyncpg adapter both set `pgcode` on the + # driver error, but the DBAPI error type does not declare it. + pgcode = getattr(error.orig, "pgcode", None) # ods: ignore[getattr] + return pgcode == _LOCK_NOT_AVAILABLE + + +def _drop_table_with_bounded_lock_wait(table_name: str) -> None: + """Drops a table whose foreign key points at `document`. + + The drop needs an ACCESS EXCLUSIVE lock on `document`. While it waits for + that lock, every other query on `document` queues behind it. A short + lock_timeout keeps that stall short when a long transaction (for example + an indexing batch) holds a lock on `document`. Each attempt runs in a + savepoint, so a timeout does not abort the migration transaction. + """ + bind = op.get_bind() + for attempt in range(1, _DROP_ATTEMPTS + 1): + try: + with bind.begin_nested(): + bind.execute(sa.text(f"SET LOCAL lock_timeout = '{_LOCK_TIMEOUT}'")) + op.drop_table(table_name) + break + except DBAPIError as e: + # Only a lock timeout is worth a retry. Raise other errors at once. + if attempt == _DROP_ATTEMPTS or not _is_lock_timeout(e): + raise + logger.warning( + "Could not lock `document` to drop %s (attempt %s/%s). Retrying in %ss.", + table_name, + attempt, + _DROP_ATTEMPTS, + _RETRY_DELAY_S, + ) + time.sleep(_RETRY_DELAY_S) + bind.execute(sa.text("SET LOCAL lock_timeout = DEFAULT")) + + +def upgrade() -> None: + op.drop_table("opensearch_tenant_migration_record") + _drop_table_with_bounded_lock_wait("opensearch_document_migration_record") + + +def downgrade() -> None: + op.create_table( + "opensearch_document_migration_record", + sa.Column("document_id", sa.String(), nullable=False), + sa.Column("status", sa.String(), nullable=False, server_default="pending"), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("attempts_count", sa.Integer(), nullable=False, server_default="0"), + sa.Column("last_attempt_at", sa.DateTime(timezone=True), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.PrimaryKeyConstraint("document_id"), + sa.ForeignKeyConstraint( + ["document_id"], + ["document.id"], + ondelete="CASCADE", + ), + ) + op.create_index( + "ix_opensearch_document_migration_record_status", + "opensearch_document_migration_record", + ["status"], + ) + op.create_index( + "ix_opensearch_document_migration_record_attempts_count", + "opensearch_document_migration_record", + ["attempts_count"], + ) + op.create_index( + "ix_opensearch_document_migration_record_created_at", + "opensearch_document_migration_record", + ["created_at"], + ) + + op.create_table( + "opensearch_tenant_migration_record", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "document_migration_record_table_population_status", + sa.String(), + nullable=False, + server_default="pending", + ), + sa.Column( + "num_times_observed_no_additional_docs_to_populate_migration_table", + sa.Integer(), + nullable=False, + server_default="0", + ), + sa.Column( + "overall_document_migration_status", + sa.String(), + nullable=False, + server_default="pending", + ), + sa.Column( + "num_times_observed_no_additional_docs_to_migrate", + sa.Integer(), + nullable=False, + server_default="0", + ), + sa.Column( + "last_updated_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.Column("vespa_visit_continuation_token", sa.Text(), nullable=True), + sa.Column( + "total_chunks_migrated", sa.Integer(), nullable=False, server_default="0" + ), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + nullable=False, + server_default=sa.func.now(), + ), + sa.Column("migration_completed_at", sa.DateTime(timezone=True), nullable=True), + sa.Column( + "enable_opensearch_retrieval", + sa.Boolean(), + nullable=False, + server_default="false", + ), + sa.Column( + "total_chunks_errored", sa.Integer(), nullable=False, server_default="0" + ), + sa.Column( + "total_chunks_in_vespa", sa.Integer(), nullable=False, server_default="0" + ), + sa.Column("approx_chunk_count_in_vespa", sa.Integer(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.execute( + sa.text( + "CREATE UNIQUE INDEX idx_opensearch_tenant_migration_singleton " + "ON opensearch_tenant_migration_record ((true))" + ) + ) diff --git a/backend/alembic/versions/90e3b9af7da4_tag_fix.py b/backend/alembic/versions/90e3b9af7da4_tag_fix.py index eb2f81f8d3..0821b1e08b 100644 --- a/backend/alembic/versions/90e3b9af7da4_tag_fix.py +++ b/backend/alembic/versions/90e3b9af7da4_tag_fix.py @@ -14,15 +14,45 @@ from typing import cast from typing import Generator from alembic import op +import httpx import sqlalchemy as sa -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT from onyx.db.search_settings import SearchSettings from shared_configs.configs import MULTI_TENANT -from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client logger = logging.getLogger("alembic.runtime.migration") +# Onyx no longer uses Vespa. These helpers were inlined from the removed Vespa +# document index module so this migration stays importable. +_VESPA_APP_CONTAINER_URL = os.environ.get("VESPA_CLOUD_URL") or ( + f"http://{os.environ.get('VESPA_HOST') or 'localhost'}" + f":{os.environ.get('VESPA_PORT') or '8081'}" +) +DOCUMENT_ID_ENDPOINT = ( + f"{_VESPA_APP_CONTAINER_URL}/document/v1/default/{{index_name}}/docid" +) + + +# Managed Vespa (Vespa Cloud) used mutual TLS. Self-hosted Vespa served a +# self-signed certificate. +_MANAGED_VESPA = os.environ.get("MANAGED_VESPA", "").lower() == "true" +_VESPA_CLOUD_CERT_PATH = os.environ.get("VESPA_CLOUD_CERT_PATH") +_VESPA_CLOUD_KEY_PATH = os.environ.get("VESPA_CLOUD_KEY_PATH") +_VESPA_REQUEST_TIMEOUT = int(os.environ.get("VESPA_REQUEST_TIMEOUT") or "15") + + +def get_vespa_http_client() -> httpx.Client: + return httpx.Client( + cert=( + cast(tuple[str, str], (_VESPA_CLOUD_CERT_PATH, _VESPA_CLOUD_KEY_PATH)) + if _MANAGED_VESPA + else None + ), + verify=_MANAGED_VESPA, + timeout=_VESPA_REQUEST_TIMEOUT, + http2=True, + ) + # revision identifiers, used by Alembic. revision = "90e3b9af7da4" 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 46c5c14bfd..692b3ffc8b 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 @@ -114,7 +114,6 @@ DOCUMENT_PERMISSIONS_UPDATE_STOP_AFTER = 10 * 60 DOCUMENT_PERMISSIONS_UPDATE_MAX_WAIT = 60 -# 5 seconds more than RetryDocumentIndex STOP_AFTER+MAX_WAIT LIGHT_SOFT_TIME_LIMIT = 105 LIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15 diff --git a/backend/ee/onyx/db/user_group.py b/backend/ee/onyx/db/user_group.py index 1ef419c971..84837422c3 100644 --- a/backend/ee/onyx/db/user_group.py +++ b/backend/ee/onyx/db/user_group.py @@ -577,7 +577,7 @@ def _mark_user_group__cc_pair_relationships_outdated__no_commit( def _current_cc_pair_ids(db_user_group: UserGroup) -> list[int]: """The cc_pairs currently attached to the group — is_current junction rows only. - A removed cc_pair keeps a stale ``is_current=False`` row until the Vespa sync + A removed cc_pair keeps a stale ``is_current=False`` row until the document index sync deletes it, and the plain ``cc_pairs`` relationship has no is_current filter, so it would still surface the removed pair. Reading it as "current" lets a removed (possibly public / out-of-scope) pair be re-attached without re-clearing the @@ -803,8 +803,8 @@ def update_user_group( user_group_update: UserGroupUpdate, ) -> UserGroup: """If successful, this can set db_user_group.is_up_to_date = False. - That will be processed by check_for_vespa_user_groups_sync_task and trigger - a long running background sync to Vespa. + That will be processed by check_for_vespa_sync_task and trigger + a long running background sync to the document index. """ # Gate before any read so a non-manager can't confirm the group exists; the # cc_pair scope check below needs the group row and runs after. diff --git a/backend/ee/onyx/document_index/vespa/app_config/cloud-services.xml.jinja b/backend/ee/onyx/document_index/vespa/app_config/cloud-services.xml.jinja deleted file mode 100644 index c2cf45807e..0000000000 --- a/backend/ee/onyx/document_index/vespa/app_config/cloud-services.xml.jinja +++ /dev/null @@ -1,48 +0,0 @@ - - - - - - - - - - - - - - - - - -{{ document_elements }} - - - - - - - - - - 2 - - - - - - - - 3 - 750 - 350 - 300 - - - - 2 - - - diff --git a/backend/ee/onyx/search/process_search_query.py b/backend/ee/onyx/search/process_search_query.py index 1bd6bdcd06..66835151f0 100644 --- a/backend/ee/onyx/search/process_search_query.py +++ b/backend/ee/onyx/search/process_search_query.py @@ -20,7 +20,7 @@ from onyx.context.search.pipeline import merge_individual_chunks, search_pipelin 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 -from onyx.document_index.interfaces_new import DocumentIndex +from onyx.document_index.interfaces import DocumentIndex from onyx.llm.factory import get_default_llm from onyx.secondary_llm_flows.document_filter import select_sections_for_expansion from onyx.tools.tool_implementations.search.search_utils import ( @@ -81,8 +81,7 @@ def stream_search_query( """ # Get document index. search_settings = get_current_search_settings(db_session) - # This flow is for search so we do not get all indices. - document_index = get_default_document_index(search_settings, None, db_session) + document_index = get_default_document_index(search_settings, None) # Determine queries to execute original_query = request.search_query diff --git a/backend/onyx/access/models.py b/backend/onyx/access/models.py index 7f0d807c62..9d24c28026 100644 --- a/backend/onyx/access/models.py +++ b/backend/onyx/access/models.py @@ -70,7 +70,7 @@ class ExternalAccess: class DocExternalAccess: """ This is just a class to wrap the external access and the document ID - together. It's used for syncing document permissions to Vespa. + together. It's used for syncing document permissions to the document index. """ external_access: ExternalAccess diff --git a/backend/onyx/background/celery/apps/app_base.py b/backend/onyx/background/celery/apps/app_base.py index 465c00cbd7..07ebb788f1 100644 --- a/backend/onyx/background/celery/apps/app_base.py +++ b/backend/onyx/background/celery/apps/app_base.py @@ -34,15 +34,9 @@ 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.app_configs import DISABLE_VECTOR_DB 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.vespa.shared_utils.utils import wait_for_vespa_with_timeout -from onyx.httpx.httpx_pool import HttpxPool from onyx.redis.redis_connector import RedisConnector from onyx.redis.redis_connector_delete import RedisConnectorDelete from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync @@ -451,8 +445,6 @@ def on_worker_ready(sender: Any, **kwargs: Any) -> None: # noqa: ARG001 def on_worker_shutdown(sender: Any, **kwargs: Any) -> None: # noqa: ARG001 - HttpxPool.close_all() - hostname: str = cast(str, sender.hostname) path = make_probe_path("readiness", hostname) path.unlink(missing_ok=True) @@ -668,35 +660,23 @@ def reset_tenant_id( def wait_for_document_index_or_shutdown() -> None: """ - Waits for all configured document indices to become ready subject to a - timeout. + Waits for the document index to become ready subject to a timeout. Raises WorkerShutdown if the timeout is reached. """ if DISABLE_VECTOR_DB: - logger.info( - "DISABLE_VECTOR_DB is set — skipping Vespa/OpenSearch readiness check." - ) + logger.info("DISABLE_VECTOR_DB is set — skipping OpenSearch readiness check.") return - if not ONYX_DISABLE_VESPA: - if not wait_for_vespa_with_timeout(): - msg = ( - "[Vespa] Readiness probe did not succeed within the timeout. Exiting..." - ) - logger.error(msg) - raise WorkerShutdown(msg) + # Imported here: opensearchpy costs ~18 MB and not every worker needs it. + from onyx.document_index.opensearch.client import ( + wait_for_opensearch_with_timeout, + ) - if ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: - # Imported here: opensearchpy costs ~18 MB and not every worker needs it. - from onyx.document_index.opensearch.client import ( - wait_for_opensearch_with_timeout, - ) - - if not wait_for_opensearch_with_timeout(): - msg = "[OpenSearch] Readiness probe did not succeed within the timeout. Exiting..." - logger.error(msg) - raise WorkerShutdown(msg) + if not wait_for_opensearch_with_timeout(): + msg = "[OpenSearch] Readiness probe did not succeed within the timeout. Exiting..." + logger.error(msg) + raise WorkerShutdown(msg) # File for validating worker liveness @@ -730,7 +710,7 @@ def get_bootsteps() -> list[type]: return [LivenessProbe] -# Task modules that require a vector DB (Vespa/OpenSearch). +# Task modules that require a vector DB (OpenSearch). # When DISABLE_VECTOR_DB is True these are excluded from autodiscover lists. _VECTOR_DB_TASK_MODULES: set[str] = { "onyx.background.celery.tasks.connector_deletion", @@ -738,7 +718,6 @@ _VECTOR_DB_TASK_MODULES: set[str] = { "onyx.background.celery.tasks.docfetching", "onyx.background.celery.tasks.pruning", "onyx.background.celery.tasks.vespa", - "onyx.background.celery.tasks.opensearch_migration", "onyx.background.celery.tasks.doc_permission_syncing", "onyx.background.celery.tasks.hierarchyfetching", # EE modules that are vector-DB-dependent diff --git a/backend/onyx/background/celery/apps/light.py b/backend/onyx/background/celery/apps/light.py index 3f870bc7e7..0978f9fe8a 100644 --- a/backend/onyx/background/celery/apps/light.py +++ b/backend/onyx/background/celery/apps/light.py @@ -5,12 +5,6 @@ from celery.apps.worker import Worker 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, - 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 ( @@ -110,17 +104,6 @@ def on_worker_init(sender: Worker, **kwargs: Any) -> None: max_overflow=EXTRA_CONCURRENCY, ) - if MANAGED_VESPA: - httpx_init_vespa_pool( - sender.concurrency + EXTRA_CONCURRENCY, # ty: ignore[unresolved-attribute] - ssl_cert=VESPA_CLOUD_CERT_PATH, - ssl_key=VESPA_CLOUD_KEY_PATH, - ) - else: - httpx_init_vespa_pool( - sender.concurrency + EXTRA_CONCURRENCY # ty: ignore[unresolved-attribute] - ) - app_base.wait_for_redis(sender, **kwargs) app_base.wait_for_db(sender, **kwargs) app_base.wait_for_document_index_or_shutdown() @@ -163,7 +146,6 @@ celery_app.autodiscover_tasks( "onyx.background.celery.tasks.index_reclaim", "onyx.background.celery.tasks.doc_permission_syncing", "onyx.background.celery.tasks.docprocessing", - "onyx.background.celery.tasks.opensearch_migration", ] ) ) diff --git a/backend/onyx/background/celery/celery_utils.py b/backend/onyx/background/celery/celery_utils.py index 3416dfb9cb..4381c1ad4b 100644 --- a/backend/onyx/background/celery/celery_utils.py +++ b/backend/onyx/background/celery/celery_utils.py @@ -2,15 +2,11 @@ import time from collections.abc import Generator, Iterator, Sequence from datetime import datetime, timezone from pathlib import Path -from typing import Any, TypeVar, cast +from typing import Any, TypeVar -import httpx from pydantic import BaseModel -from onyx.configs.app_configs import ( - MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE, - VESPA_REQUEST_TIMEOUT, -) +from onyx.configs.app_configs import MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE 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 ( @@ -32,7 +28,6 @@ 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, @@ -278,28 +273,6 @@ def celery_is_worker_primary(worker: Any) -> bool: return False -def httpx_init_vespa_pool( - max_keepalive_connections: int, - timeout: int = VESPA_REQUEST_TIMEOUT, - ssl_cert: str | None = None, - ssl_key: str | None = None, -) -> None: - httpx_cert = None - httpx_verify = False - if ssl_cert and ssl_key: - httpx_cert = cast(tuple[str, str], (ssl_cert, ssl_key)) - httpx_verify = True - - HttpxPool.init_client( - name="vespa", - cert=httpx_cert, - verify=httpx_verify, - timeout=timeout, - http2=False, - limits=httpx.Limits(max_keepalive_connections=max_keepalive_connections), - ) - - def make_probe_path(probe: str, hostname: str) -> Path: """templates the path for a k8s probe file. diff --git a/backend/onyx/background/celery/tasks/beat_schedule.py b/backend/onyx/background/celery/tasks/beat_schedule.py index 2e70e689c2..c6816f3731 100644 --- a/backend/onyx/background/celery/tasks/beat_schedule.py +++ b/backend/onyx/background/celery/tasks/beat_schedule.py @@ -7,11 +7,8 @@ from celery.schedules import crontab 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 ( @@ -312,28 +309,6 @@ if SCHEDULED_EVAL_DATASET_NAMES: } ) -# Add OpenSearch migration task if enabled. -if ( - ENABLE_OPENSEARCH_INDEXING_FOR_ONYX - and not DISABLE_OPENSEARCH_MIGRATION_TASK - and not ONYX_DISABLE_VESPA -): - beat_task_templates.append( - { - "name": "migrate-chunks-from-vespa-to-opensearch", - "task": OnyxCeleryTask.MIGRATE_CHUNKS_FROM_VESPA_TO_OPENSEARCH_TASK, - # Try to enqueue an invocation of this task with this frequency. - "schedule": timedelta(seconds=120), # 2 minutes - "options": { - "priority": OnyxCeleryPriority.LOW, - # If the task was not dequeued in this time, revoke it. - "expires": BEAT_EXPIRES_DEFAULT, - "queue": OnyxCeleryQueues.OPENSEARCH_MIGRATION, - }, - } - ) - - # Beat task names that require a vector DB. Filtered out when DISABLE_VECTOR_DB. _VECTOR_DB_BEAT_TASK_NAMES: set[str] = { "check-for-indexing", @@ -347,7 +322,6 @@ _VECTOR_DB_BEAT_TASK_NAMES: set[str] = { "check-for-index-attempt-cleanup", "check-for-doc-permissions-sync", "check-for-external-group-sync", - "migrate-chunks-from-vespa-to-opensearch", } if DISABLE_VECTOR_DB: diff --git a/backend/onyx/background/celery/tasks/docfetching/tasks.py b/backend/onyx/background/celery/tasks/docfetching/tasks.py index 00643f48d9..2373c722b2 100644 --- a/backend/onyx/background/celery/tasks/docfetching/tasks.py +++ b/backend/onyx/background/celery/tasks/docfetching/tasks.py @@ -356,7 +356,7 @@ def docfetching_proxy_task( 2) upserts documents to postgres (index_doc_batch_prepare) 3) chunks each document (optionally adds context for contextual rag) 4) embeds chunks (embed_chunks_with_failure_handling) via a call to the model server - 5) write chunks to vespa (write_chunks_to_vector_db_with_backoff) + 5) write chunks to the document index (write_chunks_to_vector_db_with_backoff) 6) update document and indexing metadata in postgres 7) pulls all document IDs from the source and compares those IDs to locally stored documents and deletes all locally stored IDs missing from the most recently pulled document ID list diff --git a/backend/onyx/background/celery/tasks/docprocessing/tasks.py b/backend/onyx/background/celery/tasks/docprocessing/tasks.py index 86df930345..40df444ee0 100644 --- a/backend/onyx/background/celery/tasks/docprocessing/tasks.py +++ b/backend/onyx/background/celery/tasks/docprocessing/tasks.py @@ -20,7 +20,6 @@ from onyx.background.celery.celery_redis import ( 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 ( @@ -47,12 +46,7 @@ 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.app_configs import PERSISTENT_INDEXING from onyx.configs.constants import ( CELERY_GENERIC_BEAT_LOCK_TIMEOUT, CELERY_INDEXING_LOCK_TIMEOUT, @@ -115,7 +109,6 @@ from onyx.file_store.document_batch_storage import ( get_document_batch_storage, ) from onyx.file_store.staging import cleanup_staged_files_for_attempt -from onyx.httpx.httpx_pool import HttpxPool from onyx.redis.redis_connector import RedisConnector from onyx.redis.redis_docprocessing import RedisDocprocessing from onyx.redis.redis_pool import ( @@ -1689,14 +1682,6 @@ def _docprocessing_task( redis_connector = RedisConnector(tenant_id, cc_pair_id) r = get_redis_client(tenant_id=tenant_id) - # 20 is the documented default for httpx max_keepalive_connections - if MANAGED_VESPA: - httpx_init_vespa_pool( - 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH - ) - else: - httpx_init_vespa_pool(20) - # dummy lock to satisfy linter per_batch_lock: RedisLock | None = None @@ -1707,7 +1692,7 @@ def _docprocessing_task( try: # Inside the try so a failed first-use import still marks the attempt failed. - from onyx.document_index.factory import get_all_document_indices + from onyx.document_index.factory import get_default_document_index from onyx.indexing.adapters.document_indexing_adapter import ( DocumentIndexingBatchAdapter, ) @@ -1755,7 +1740,8 @@ def _docprocessing_task( ) # Phase 1: fast DB reads to set up the pipeline. Session closes before - # the slow embedding + Vespa work begins, returning the connection to the pool. + # the slow embedding + document index work begins, returning the connection + # to the pool. with get_session_with_current_tenant() as db_session: # matches parts of _run_indexing index_attempt = get_index_attempt( @@ -1798,10 +1784,9 @@ def _docprocessing_task( callback=callback, ) - document_indices = get_all_document_indices( + document_index = get_default_document_index( index_attempt.search_settings, None, - httpx_client=HttpxPool.get("vespa"), ) # Set up metadata for this batch @@ -1856,7 +1841,7 @@ def _docprocessing_task( # real work happens here! index_pipeline_result = run_indexing_pipeline( embedder=embedding_model, - document_indices=document_indices, + document_index=document_index, ignore_time_skip=True, # Documents are already filtered during extraction index_to_secondary=index_to_secondary, tenant_id=tenant_id, diff --git a/backend/onyx/background/celery/tasks/index_reclaim/tasks.py b/backend/onyx/background/celery/tasks/index_reclaim/tasks.py index 29648ba0ad..320e419973 100644 --- a/backend/onyx/background/celery/tasks/index_reclaim/tasks.py +++ b/backend/onyx/background/celery/tasks/index_reclaim/tasks.py @@ -60,7 +60,7 @@ from onyx.db.search_settings import ( get_search_settings_by_id, record_failure__no_commit, ) -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.client import OpenSearchIndexClient from onyx.document_index.opensearch.index_reclaim import ( ReclaimOutcome, diff --git a/backend/onyx/background/celery/tasks/monitoring/tasks.py b/backend/onyx/background/celery/tasks/monitoring/tasks.py index c68ff38f3e..e3f49fe3f0 100644 --- a/backend/onyx/background/celery/tasks/monitoring/tasks.py +++ b/backend/onyx/background/celery/tasks/monitoring/tasks.py @@ -185,7 +185,6 @@ def _collect_queue_metrics(redis_celery: Redis) -> list[Metric]: "user_file_port_queue_length": OnyxCeleryQueues.USER_FILE_PORT, "monitoring_queue_length": OnyxCeleryQueues.MONITORING, "sandbox_queue_length": OnyxCeleryQueues.SANDBOX, - "opensearch_migration_queue_length": OnyxCeleryQueues.OPENSEARCH_MIGRATION, } for name, queue in queue_mappings.items(): @@ -985,9 +984,6 @@ def monitor_celery_queues_helper( ) n_monitoring = celery_get_queue_length(OnyxCeleryQueues.MONITORING, r_celery) n_sandbox = celery_get_queue_length(OnyxCeleryQueues.SANDBOX, r_celery) - n_opensearch_migration = celery_get_queue_length( - OnyxCeleryQueues.OPENSEARCH_MIGRATION, r_celery - ) n_docfetching_prefetched = celery_get_unacked_task_ids( OnyxCeleryQueues.CONNECTOR_DOC_FETCHING, r_celery @@ -1022,7 +1018,6 @@ def monitor_celery_queues_helper( f"capability_checks={n_capability_checks} " f"monitoring={n_monitoring} " f"sandbox={n_sandbox} " - f"opensearch_migration={n_opensearch_migration} " ) diff --git a/backend/onyx/background/celery/tasks/opensearch_migration/__init__.py b/backend/onyx/background/celery/tasks/opensearch_migration/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/backend/onyx/background/celery/tasks/opensearch_migration/constants.py b/backend/onyx/background/celery/tasks/opensearch_migration/constants.py deleted file mode 100644 index 3fc6015bd2..0000000000 --- a/backend/onyx/background/celery/tasks/opensearch_migration/constants.py +++ /dev/null @@ -1,56 +0,0 @@ -# Tasks are expected to cease execution and do cleanup after the soft time -# limit. In principle they are also forceably terminated after the hard time -# limit, in practice this does not happen since we use threadpools for Celery -# task execution, and we simple hope that the total task time plus cleanup does -# not exceed this. Therefore tasks should regularly check their timeout and lock -# status. The lock timeout is the maximum time the lock manager (Redis in this -# case) will enforce the lock, independent of what is happening in the task. To -# reduce the chances that a task is still doing work while a lock has expired, -# make the lock timeout well above the task timeouts. In practice we should -# never see locks be held for this long anyway because a task should release the -# lock after its cleanup which happens at most after its soft timeout. - -# Constants corresponding to migrate_documents_from_vespa_to_opensearch_task. -from onyx.configs.app_configs import OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE - -MIGRATION_TASK_SOFT_TIME_LIMIT_S = 60 * 5 # 5 minutes. -MIGRATION_TASK_TIME_LIMIT_S = 60 * 6 # 6 minutes. -# The maximum time the lock can be held for. Will automatically be released -# after this time. -MIGRATION_TASK_LOCK_TIMEOUT_S = 60 * 7 # 7 minutes. -assert MIGRATION_TASK_SOFT_TIME_LIMIT_S < MIGRATION_TASK_TIME_LIMIT_S, ( - "The soft time limit must be less than the time limit." -) -assert MIGRATION_TASK_TIME_LIMIT_S < MIGRATION_TASK_LOCK_TIMEOUT_S, ( - "The time limit must be less than the lock timeout." -) -# Time to wait to acquire the lock. -MIGRATION_TASK_LOCK_BLOCKING_TIMEOUT_S = 60 * 2 # 2 minutes. - -# Constants corresponding to check_for_documents_for_opensearch_migration_task. -CHECK_FOR_DOCUMENTS_TASK_SOFT_TIME_LIMIT_S = 60 # 60 seconds / 1 minute. -CHECK_FOR_DOCUMENTS_TASK_TIME_LIMIT_S = 90 # 90 seconds. -# The maximum time the lock can be held for. Will automatically be released -# after this time. -CHECK_FOR_DOCUMENTS_TASK_LOCK_TIMEOUT_S = 120 # 120 seconds / 2 minutes. -assert ( - CHECK_FOR_DOCUMENTS_TASK_SOFT_TIME_LIMIT_S < CHECK_FOR_DOCUMENTS_TASK_TIME_LIMIT_S -), "The soft time limit must be less than the time limit." -assert ( - CHECK_FOR_DOCUMENTS_TASK_TIME_LIMIT_S < CHECK_FOR_DOCUMENTS_TASK_LOCK_TIMEOUT_S -), "The time limit must be less than the lock timeout." -# Time to wait to acquire the lock. -CHECK_FOR_DOCUMENTS_TASK_LOCK_BLOCKING_TIMEOUT_S = 30 # 30 seconds. - -TOTAL_ALLOWABLE_DOC_MIGRATION_ATTEMPTS_BEFORE_PERMANENT_FAILURE = 15 - -# WARNING: Do not change these values without knowing what changes also need to -# be made to OpenSearchTenantMigrationRecord. -GET_VESPA_CHUNKS_PAGE_SIZE = OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE -GET_VESPA_CHUNKS_SLICE_COUNT = 4 - -# String used to indicate in the vespa_visit_continuation_token mapping that the -# slice has finished and there is nothing left to visit. -FINISHED_VISITING_SLICE_CONTINUATION_TOKEN = ( - "FINISHED_VISITING_SLICE_CONTINUATION_TOKEN" -) diff --git a/backend/onyx/background/celery/tasks/opensearch_migration/tasks.py b/backend/onyx/background/celery/tasks/opensearch_migration/tasks.py deleted file mode 100644 index 89cf937d27..0000000000 --- a/backend/onyx/background/celery/tasks/opensearch_migration/tasks.py +++ /dev/null @@ -1,334 +0,0 @@ -"""Celery tasks for migrating documents from Vespa to OpenSearch.""" - -import time -import traceback - -from celery import Task, shared_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, - GET_VESPA_CHUNKS_PAGE_SIZE, - MIGRATION_TASK_LOCK_BLOCKING_TIMEOUT_S, - MIGRATION_TASK_LOCK_TIMEOUT_S, - MIGRATION_TASK_SOFT_TIME_LIMIT_S, - 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, - 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, - get_vespa_visit_state, - is_migration_completed, - mark_migration_completed_time_if_not_set_with_commit, - try_insert_opensearch_tenant_migration_record_with_commit, - 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 ( - OpenSearchDocumentIndex, -) -from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client -from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex -from onyx.indexing.models import IndexingSetting -from onyx.redis.redis_pool import get_redis_client -from shared_configs.configs import MULTI_TENANT -from shared_configs.contextvars import get_current_tenant_id - - -def is_continuation_token_done_for_all_slices( - continuation_token_map: dict[int, str | None], -) -> bool: - return all( - continuation_token == FINISHED_VISITING_SLICE_CONTINUATION_TOKEN - for continuation_token in continuation_token_map.values() - ) - - -# shared_task allows this task to be shared across celery app instances. -@shared_task( # ty: ignore[invalid-argument-type] - name=OnyxCeleryTask.MIGRATE_CHUNKS_FROM_VESPA_TO_OPENSEARCH_TASK, - # Does not store the task's return value in the result backend. - ignore_result=True, - # WARNING: This is here just for rigor but since we use threads for Celery - # this config is not respected and timeout logic must be implemented in the - # task. - soft_time_limit=MIGRATION_TASK_SOFT_TIME_LIMIT_S, - # WARNING: This is here just for rigor but since we use threads for Celery - # this config is not respected and timeout logic must be implemented in the - # task. - time_limit=MIGRATION_TASK_TIME_LIMIT_S, - # Passed in self to the task to get task metadata. - bind=True, -) -def migrate_chunks_from_vespa_to_opensearch_task( - self: Task, # noqa: ARG001 - *, - tenant_id: str, -) -> bool | None: - """ - Periodic task to migrate chunks from Vespa to OpenSearch via the Visit API. - - Uses Vespa's Visit API to iterate through ALL chunks in bulk (not - per-document), transform them, and index them into OpenSearch. Progress is - tracked via a continuation token map stored in the - OpenSearchTenantMigrationRecord. - - The first time we see no continuation token map and non-zero chunks - migrated, we consider the migration complete and all subsequent invocations - are no-ops. - - We divide the index into GET_VESPA_CHUNKS_SLICE_COUNT independent slices - where progress is tracked for each slice. - - Returns: - None if OpenSearch migration is not enabled, or if the lock could not be - acquired; effectively a no-op. True if the task completed - successfully. False if the task errored. - """ - # 1. Check if we should run the task. - # 1.a. If OpenSearch indexing is disabled, we don't run the task. - if not ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: - task_logger.warning( - "OpenSearch migration is not enabled, skipping chunk migration task." - ) - return None - task_logger.info("Starting chunk-level migration from Vespa to OpenSearch.") - task_start_time = time.monotonic() - - # 1.b. Only one instance per tenant of this task may run concurrently at - # once. If we fail to acquire a lock, we assume it is because another task - # has one and we exit. - r = get_redis_client() - lock: RedisLock = r.lock( - name=OnyxRedisLocks.OPENSEARCH_MIGRATION_BEAT_LOCK, - # The maximum time the lock can be held for. Will automatically be - # released after this time. - timeout=MIGRATION_TASK_LOCK_TIMEOUT_S, - # .acquire will block until the lock is acquired. - blocking=True, - # Time to wait to acquire the lock. - blocking_timeout=MIGRATION_TASK_LOCK_BLOCKING_TIMEOUT_S, - ) - if not lock.acquire(): - task_logger.warning( - "The OpenSearch migration task timed out waiting for the lock." - ) - return None - else: - task_logger.info( - f"Acquired the OpenSearch migration lock. Took {time.monotonic() - task_start_time:.3f} seconds. " - f"Token: {lock.local.token}" - ) - - # 2. Prepare to migrate. - total_chunks_migrated_this_task = 0 - total_chunks_errored_this_task = 0 - try: - # 2.a. Double-check that tenant info is correct. - if tenant_id != get_current_tenant_id(): - err_str = ( - f"Tenant ID mismatch in the OpenSearch migration task: " - f"{tenant_id} != {get_current_tenant_id()}. This should never happen." - ) - task_logger.error(err_str) - return False - - # Do as much as we can with a DB session in one spot to not hold a - # session during a migration batch. - with get_session_with_current_tenant() as db_session: - # 2.b. Immediately check to see if this tenant is done, to save - # having to do any other work. This function does not require a - # migration record to necessarily exist. - if is_migration_completed(db_session): - return True - - # 2.c. Try to insert the OpenSearchTenantMigrationRecord table if it - # does not exist. - try_insert_opensearch_tenant_migration_record_with_commit(db_session) - - # 2.d. Get search settings. - search_settings = get_current_search_settings(db_session) - indexing_setting = IndexingSetting.from_db_model(search_settings) - - task_logger.debug( - "Verified tenant info, migration record, and search settings." - ) - - # 2.e. Build sanitized to original doc ID mapping to check for - # conflicts in the event we sanitize a doc ID to an - # already-existing doc ID. - # We reconstruct this mapping for every task invocation because - # a document may have been added in the time between two tasks. - sanitized_doc_start_time = time.monotonic() - sanitized_to_original_doc_id_mapping = ( - build_sanitized_to_original_doc_id_mapping(db_session) - ) - task_logger.debug( - f"Built sanitized_to_original_doc_id_mapping with {len(sanitized_to_original_doc_id_mapping)} entries " - f"in {time.monotonic() - sanitized_doc_start_time:.3f} seconds." - ) - - # 2.f. Get the current migration state. - continuation_token_map, total_chunks_migrated = get_vespa_visit_state( - db_session - ) - # 2.f.1. Double-check that the migration state does not imply - # completion. Really we should never have to enter this block as we - # would expect is_migration_completed to return True, but in the - # strange event that the migration is complete but the migration - # completed time was never stamped, we do so here. - if is_continuation_token_done_for_all_slices(continuation_token_map): - task_logger.info( - f"OpenSearch migration COMPLETED for tenant {tenant_id}. Total chunks migrated: {total_chunks_migrated}." - ) - mark_migration_completed_time_if_not_set_with_commit(db_session) - return True - task_logger.debug( - f"Read the tenant migration record. Total chunks migrated: {total_chunks_migrated}. " - f"Continuation token map: {continuation_token_map}" - ) - - with get_vespa_http_client( - timeout=VESPA_MIGRATION_REQUEST_TIMEOUT_S - ) as vespa_client: - # 2.g. Create the OpenSearch and Vespa document indexes. - tenant_state = TenantState(tenant_id=tenant_id, multitenant=MULTI_TENANT) - opensearch_document_index = OpenSearchDocumentIndex( - tenant_state=tenant_state, - index_name=search_settings.index_name, - embedding_dim=indexing_setting.final_embedding_dim, - vector_quantization=indexing_setting.vector_quantization, - ) - vespa_document_index = VespaDocumentIndex( - index_name=search_settings.index_name, - tenant_state=tenant_state, - large_chunks_enabled=False, - httpx_client=vespa_client, - ) - - # 2.h. Get the approximate chunk count in Vespa as of this time to - # update the migration record. - approx_chunk_count_in_vespa: int | None = None - get_chunk_count_start_time = time.monotonic() - try: - approx_chunk_count_in_vespa = vespa_document_index.get_chunk_count() - except Exception: - # This failure should not be blocking. - task_logger.exception( - "Error getting approximate chunk count in Vespa. Moving on..." - ) - task_logger.debug( - f"Took {time.monotonic() - get_chunk_count_start_time:.3f} seconds to attempt to get " - f"approximate chunk count in Vespa. Got {approx_chunk_count_in_vespa}." - ) - - # 3. Do the actual migration in batches until we run out of time. - while ( - time.monotonic() - task_start_time < MIGRATION_TASK_SOFT_TIME_LIMIT_S - and lock.owned() - ): - # 3.a. Get the next batch of raw chunks from Vespa. - get_vespa_chunks_start_time = time.monotonic() - raw_vespa_chunks, next_continuation_token_map = ( - vespa_document_index.get_all_raw_document_chunks_paginated( - continuation_token_map=continuation_token_map, - page_size=GET_VESPA_CHUNKS_PAGE_SIZE, - ) - ) - task_logger.debug( - f"Read {len(raw_vespa_chunks)} chunks from Vespa in {time.monotonic() - get_vespa_chunks_start_time:.3f} " - f"seconds. Next continuation token map: {next_continuation_token_map}" - ) - - # 3.b. Transform the raw chunks to OpenSearch chunks in memory. - opensearch_document_chunks, errored_chunks = ( - transform_vespa_chunks_to_opensearch_chunks( - raw_vespa_chunks, - tenant_state, - sanitized_to_original_doc_id_mapping, - ) - ) - if len(opensearch_document_chunks) != len(raw_vespa_chunks): - task_logger.error( - f"Migration task error: Number of candidate chunks to migrate ({len(opensearch_document_chunks)}) does " - f"not match number of chunks in Vespa ({len(raw_vespa_chunks)}). {len(errored_chunks)} chunks " - "errored." - ) - - # 3.c. Index the OpenSearch chunks into OpenSearch. - index_opensearch_chunks_start_time = time.monotonic() - opensearch_document_index.index_raw_chunks( - chunks=opensearch_document_chunks - ) - task_logger.debug( - f"Indexed {len(opensearch_document_chunks)} chunks into OpenSearch in " - f"{time.monotonic() - index_opensearch_chunks_start_time:.3f} seconds." - ) - - total_chunks_migrated_this_task += len(opensearch_document_chunks) - total_chunks_errored_this_task += len(errored_chunks) - - # Do as much as we can with a DB session in one spot to not hold a - # session during a migration batch. - with get_session_with_current_tenant() as db_session: - # 3.d. Update the migration state. - update_vespa_visit_progress_with_commit( - db_session, - continuation_token_map=next_continuation_token_map, - chunks_processed=len(opensearch_document_chunks), - chunks_errored=len(errored_chunks), - approx_chunk_count_in_vespa=approx_chunk_count_in_vespa, - ) - - # 3.e. Get the current migration state. Even thought we - # technically have it in-memory since we just wrote it, we - # want to reference the DB as the source of truth at all - # times. - continuation_token_map, total_chunks_migrated = ( - get_vespa_visit_state(db_session) - ) - # 3.e.1. Check if the migration is done. - if is_continuation_token_done_for_all_slices( - continuation_token_map - ): - task_logger.info( - f"OpenSearch migration COMPLETED for tenant {tenant_id}. Total chunks migrated: {total_chunks_migrated}." - ) - mark_migration_completed_time_if_not_set_with_commit(db_session) - return True - task_logger.debug( - f"Read the tenant migration record. Total chunks migrated: {total_chunks_migrated}. " - f"Continuation token map: {continuation_token_map}" - ) - except Exception: - traceback.print_exc() - task_logger.exception("Error in the OpenSearch migration task.") - return False - finally: - if lock.owned(): - lock.release() - task_logger.debug("Released the OpenSearch migration lock.") - else: - task_logger.warning( - "The OpenSearch migration lock was not owned on completion of the migration task." - ) - - task_logger.info( - f"OpenSearch chunk migration task pausing (time limit reached). " - f"Total chunks migrated this task: {total_chunks_migrated_this_task}. " - f"Total chunks errored this task: {total_chunks_errored_this_task}. " - f"Elapsed: {time.monotonic() - task_start_time:.3f}s. " - "Will resume from continuation token on next invocation." - ) - - return True diff --git a/backend/onyx/background/celery/tasks/opensearch_migration/transformer.py b/backend/onyx/background/celery/tasks/opensearch_migration/transformer.py deleted file mode 100644 index fdfdd8dbb3..0000000000 --- a/backend/onyx/background/celery/tasks/opensearch_migration/transformer.py +++ /dev/null @@ -1,348 +0,0 @@ -import traceback -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, - 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 - -logger = setup_logger(__name__) - - -FIELDS_NEEDED_FOR_TRANSFORMATION: list[str] = [ - DOCUMENT_ID, - CHUNK_ID, - TITLE, - TITLE_EMBEDDING, - CONTENT, - EMBEDDINGS, - SOURCE_TYPE, - METADATA_LIST, - DOC_UPDATED_AT, - HIDDEN, - BOOST, - SEMANTIC_IDENTIFIER, - IMAGE_FILE_NAME, - SOURCE_LINKS, - BLURB, - DOC_SUMMARY, - CHUNK_CONTEXT, - METADATA_SUFFIX, - DOCUMENT_SETS, - USER_PROJECT, - PERSONAS, - PRIMARY_OWNERS, - SECONDARY_OWNERS, - ACCESS_CONTROL_LIST, -] -if MULTI_TENANT: - FIELDS_NEEDED_FOR_TRANSFORMATION.append(TENANT_ID) - - -def _extract_content_vector(embeddings: Any) -> list[float]: - """Extracts the full chunk embedding vector from Vespa's embeddings tensor. - - Vespa stores embeddings as a tensor(t{},x[dim]) where 't' maps - embedding names (like "full_chunk") to vectors. The API can return this in - different formats: - 1. Direct list: {"full_chunk": [...]} - 2. Blocks format: {"blocks": {"full_chunk": [0.1, 0.2, ...]}} - 3. Possibly other formats. - - We only support formats 1 and 2. Any other supplied format will raise an - error. - - Raises: - ValueError: If the embeddings format is not supported. - - Returns: - The full chunk content embedding vector as a list of floats. - """ - if isinstance(embeddings, dict): - # Handle format 1. - full_chunk_embedding = embeddings.get(FULL_CHUNK_EMBEDDING_KEY) - if isinstance(full_chunk_embedding, list): - # Double check that within the list we have floats and not another - # list or dict. - if not full_chunk_embedding: - raise ValueError("Full chunk embedding is empty.") - if isinstance(full_chunk_embedding[0], float): - return full_chunk_embedding - - # Handle format 2. - blocks = embeddings.get("blocks") - if isinstance(blocks, dict): - full_chunk_embedding = blocks.get(FULL_CHUNK_EMBEDDING_KEY) - if isinstance(full_chunk_embedding, list): - # Double check that within the list we have floats and not another - # list or dict. - if not full_chunk_embedding: - raise ValueError("Full chunk embedding is empty.") - if isinstance(full_chunk_embedding[0], float): - return full_chunk_embedding - - raise ValueError(f"Unknown embedding format: {type(embeddings)}") - - -def _extract_title_vector(title_embedding: Any | None) -> list[float] | None: - """Extract the title embedding vector. - - Returns None if no title embedding exists. - - Vespa returns title_embedding as tensor(x[dim]) which can be in - formats: - 1. Direct list: [0.1, 0.2, ...] - 2. Values format: {"values": [0.1, 0.2, ...]} - 3. Possibly other formats. - - Only formats 1 and 2 are supported. Any other supplied format will raise an - error. - - Raises: - ValueError: If the title embedding format is not supported. - - Returns: - The title embedding vector as a list of floats. - """ - if title_embedding is None: - return None - - # Handle format 1. - if isinstance(title_embedding, list): - # Double check that within the list we have floats and not another - # list or dict. - if not title_embedding: - return None - if isinstance(title_embedding[0], float): - return title_embedding - - # Handle format 2. - if isinstance(title_embedding, dict): - # Try values format. - values = title_embedding.get("values") - if values is not None and isinstance(values, list): - # Double check that within the list we have floats and not another - # list or dict. - if not values: - return None - if isinstance(values[0], float): - return values - - raise ValueError(f"Unknown title embedding format: {type(title_embedding)}") - - -def _transform_vespa_document_sets_to_opensearch_document_sets( - vespa_document_sets: dict[str, int] | None, -) -> list[str] | None: - if not vespa_document_sets: - return None - return list(vespa_document_sets.keys()) - - -def _transform_vespa_acl_to_opensearch_acl( - vespa_acl: dict[str, int] | None, -) -> tuple[bool, list[str]]: - if not vespa_acl: - return False, [] - acl_list = list(vespa_acl.keys()) - is_public = PUBLIC_DOC_PAT in acl_list - if is_public: - acl_list.remove(PUBLIC_DOC_PAT) - return is_public, acl_list - - -def transform_vespa_chunks_to_opensearch_chunks( - vespa_chunks: list[dict[str, Any]], - tenant_state: TenantState, - sanitized_to_original_doc_id_mapping: dict[str, str], -) -> tuple[list[DocumentChunk], list[dict[str, Any]]]: - result: list[DocumentChunk] = [] - errored_chunks: list[dict[str, Any]] = [] - for vespa_chunk in vespa_chunks: - try: - # This should exist; fail loudly if it does not. - vespa_document_id: str = vespa_chunk[DOCUMENT_ID] - if not vespa_document_id: - raise ValueError("Missing document_id in Vespa chunk.") - # Vespa doc IDs were sanitized using - # replace_invalid_doc_id_characters. This was a poor design choice - # and we don't want this in OpenSearch; whatever restrictions there - # may be on indexed chunk ID should have no bearing on the chunk's - # document ID field, even if document ID is an argument to the chunk - # ID. Deliberately choose to use the real doc ID supplied to this - # function. - if vespa_document_id in sanitized_to_original_doc_id_mapping: - logger.warning( - "Migration warning: Vespa document ID %s does not match the document ID supplied %s. The Vespa ID will be discarded.", - vespa_document_id, - sanitized_to_original_doc_id_mapping[vespa_document_id], - ) - document_id = sanitized_to_original_doc_id_mapping.get( - vespa_document_id, vespa_document_id - ) - - # This should exist; fail loudly if it does not. - chunk_index: int = vespa_chunk[CHUNK_ID] - - title: str | None = vespa_chunk.get(TITLE) - # WARNING: Should supply format.tensors=short-value to the Vespa - # client in order to get a supported format for the tensors. - title_vector: list[float] | None = _extract_title_vector( - vespa_chunk.get(TITLE_EMBEDDING) - ) - - # This should exist; fail loudly if it does not. - content: str = vespa_chunk[CONTENT] - if not content: - raise ValueError( - f"Missing content in Vespa chunk with document ID {vespa_document_id} and chunk index {chunk_index}." - ) - # This should exist; fail loudly if it does not. - # WARNING: Should supply format.tensors=short-value to the Vespa - # client in order to get a supported format for the tensors. - content_vector: list[float] = _extract_content_vector( - vespa_chunk[EMBEDDINGS] - ) - if not content_vector: - raise ValueError( - f"Missing content_vector in Vespa chunk with document ID {vespa_document_id} and chunk index {chunk_index}." - ) - - # This should exist; fail loudly if it does not. - source_type: str = vespa_chunk[SOURCE_TYPE] - if not source_type: - raise ValueError( - f"Missing source_type in Vespa chunk with document ID {vespa_document_id} and chunk index {chunk_index}." - ) - - metadata_list: list[str] | None = vespa_chunk.get(METADATA_LIST) - - _raw_doc_updated_at: int | None = vespa_chunk.get(DOC_UPDATED_AT) - last_updated: datetime | None = ( - datetime.fromtimestamp(_raw_doc_updated_at, tz=timezone.utc) - if _raw_doc_updated_at is not None - else None - ) - - hidden: bool = vespa_chunk.get(HIDDEN, False) - - # This should exist; fail loudly if it does not. - global_boost: int = vespa_chunk[BOOST] - - # This should exist; fail loudly if it does not. - semantic_identifier: str = vespa_chunk[SEMANTIC_IDENTIFIER] - if not semantic_identifier: - raise ValueError( - f"Missing semantic_identifier in Vespa chunk with document ID {vespa_document_id} and chunk " - f"index {chunk_index}." - ) - - image_file_id: str | None = vespa_chunk.get(IMAGE_FILE_NAME) - source_links: str | None = vespa_chunk.get(SOURCE_LINKS) - blurb: str = vespa_chunk.get(BLURB, "") - doc_summary: str = vespa_chunk.get(DOC_SUMMARY, "") - chunk_context: str = vespa_chunk.get(CHUNK_CONTEXT, "") - metadata_suffix: str | None = vespa_chunk.get(METADATA_SUFFIX) - document_sets: list[str] | None = ( - _transform_vespa_document_sets_to_opensearch_document_sets( - vespa_chunk.get(DOCUMENT_SETS) - ) - ) - user_projects: list[int] | None = vespa_chunk.get(USER_PROJECT) - personas: list[int] | None = vespa_chunk.get(PERSONAS) - primary_owners: list[str] | None = vespa_chunk.get(PRIMARY_OWNERS) - secondary_owners: list[str] | None = vespa_chunk.get(SECONDARY_OWNERS) - - is_public, acl_list = _transform_vespa_acl_to_opensearch_acl( - vespa_chunk.get(ACCESS_CONTROL_LIST) - ) - if not is_public and not acl_list: - logger.warning( - "Migration warning: Vespa chunk with document ID %s and chunk index %s has no public ACL and no access control list. This does not make sense as it implies the document is never searchable. Continuing with the migration...", - vespa_document_id, - chunk_index, - ) - - chunk_tenant_id: str | None = vespa_chunk.get(TENANT_ID) - if MULTI_TENANT: - if not chunk_tenant_id: - raise ValueError( - "Missing tenant_id in Vespa chunk in a multi-tenant environment." - ) - if chunk_tenant_id != tenant_state.tenant_id: - raise ValueError( - f"Chunk tenant_id {chunk_tenant_id} does not match expected tenant_id {tenant_state.tenant_id}" - ) - - opensearch_chunk = DocumentChunk( - # We deliberately choose to use the doc ID supplied to this function - # over the Vespa doc ID. - document_id=document_id, - chunk_index=chunk_index, - title=title, - title_vector=title_vector, - content=content, - content_vector=content_vector, - source_type=source_type, - metadata_list=metadata_list, - last_updated=last_updated, - public=is_public, - access_control_list=acl_list, - hidden=hidden, - global_boost=global_boost, - semantic_identifier=semantic_identifier, - image_file_id=image_file_id, - source_links=source_links, - blurb=blurb, - doc_summary=doc_summary, - chunk_context=chunk_context, - metadata_suffix=metadata_suffix, - document_sets=document_sets, - user_projects=user_projects, - personas=personas, - primary_owners=primary_owners, - secondary_owners=secondary_owners, - tenant_id=tenant_state, - ) - - result.append(opensearch_chunk) - except Exception: - traceback.print_exc() - logger.exception( - "Migration error: Error transforming Vespa chunk with document ID %s and chunk index %s into an OpenSearch chunk. Continuing with the migration...", - vespa_chunk.get(DOCUMENT_ID), - vespa_chunk.get(CHUNK_ID), - ) - errored_chunks.append(vespa_chunk) - - return result, errored_chunks diff --git a/backend/onyx/background/celery/tasks/port/tasks.py b/backend/onyx/background/celery/tasks/port/tasks.py index 7055c19492..e19dc7e738 100644 --- a/backend/onyx/background/celery/tasks/port/tasks.py +++ b/backend/onyx/background/celery/tasks/port/tasks.py @@ -613,7 +613,7 @@ def _resolve_port_target_settings(db_session: Session) -> SearchSettings | None: if port_backfill_has_pending_work(db_session, present.id): return present # Backfill drained: unpin the source so we stop re-checking a done job, the - # source index can be reclaimed, and the reindex/vespa guards read "not + # source index can be reclaimed, and the reindex/document-index-sync guards read "not # backfilling". Orphans were already swept per-attempt inside run_port_attempt. present.port_backfill_source_id = None db_session.commit() diff --git a/backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py b/backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py deleted file mode 100644 index 1c618af19c..0000000000 --- a/backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py +++ /dev/null @@ -1,48 +0,0 @@ -import httpx -from tenacity import ( - retry, - retry_if_exception_type, - stop_after_delay, - wait_random_exponential, -) - -from onyx.document_index.interfaces_new import DocumentIndex, MetadataUpdateRequest - - -class RetryDocumentIndex: - """A wrapper class to help with specific retries against Vespa involving - read timeouts. - - wait_random_exponential implements full jitter as per this article: - https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/""" - - MAX_WAIT = 30 - - # STOP_AFTER + MAX_WAIT should be slightly less (5?) than the celery soft_time_limit - STOP_AFTER = 70 - - def __init__(self, index: DocumentIndex): - self.index: DocumentIndex = index - - @retry( - retry=retry_if_exception_type(httpx.ReadTimeout), - wait=wait_random_exponential(multiplier=1, max=MAX_WAIT), - stop=stop_after_delay(STOP_AFTER), - ) - def delete( - self, - doc_id: str, - chunk_count: int | None = None, - ) -> int: - return self.index.delete(doc_id, chunk_count=chunk_count) - - @retry( - retry=retry_if_exception_type(httpx.ReadTimeout), - wait=wait_random_exponential(multiplier=1, max=MAX_WAIT), - stop=stop_after_delay(STOP_AFTER), - ) - def update( - self, - update_requests: list[MetadataUpdateRequest], - ) -> None: - self.index.update(update_requests) diff --git a/backend/onyx/background/celery/tasks/shared/tasks.py b/backend/onyx/background/celery/tasks/shared/tasks.py index 3fabdffbf6..dae0a2822b 100644 --- a/backend/onyx/background/celery/tasks/shared/tasks.py +++ b/backend/onyx/background/celery/tasks/shared/tasks.py @@ -1,17 +1,13 @@ import time from datetime import datetime from enum import Enum -from http import HTTPStatus -import httpx from celery import Task, shared_task from celery.exceptions import SoftTimeLimitExceeded from sqlalchemy.orm import Session -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, OnyxCeleryTask from onyx.db.connector_credential_pair import get_connector_credential_pair from onyx.db.document import ( @@ -35,16 +31,14 @@ from onyx.db.port_orphan_candidate import ( ) from onyx.db.relationships import delete_document_references_from_kg from onyx.db.search_settings import get_active_search_settings -from onyx.document_index.factory import get_all_document_indices -from onyx.document_index.interfaces_new import MetadataUpdateRequest -from onyx.httpx.httpx_pool import HttpxPool +from onyx.document_index.factory import get_default_document_index +from onyx.document_index.interfaces import MetadataUpdateRequest from onyx.redis.redis_pool import get_redis_client from onyx.server.documents.models import ConnectorCredentialPairIdentifier DOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES = 3 -# 5 seconds more than RetryDocumentIndex STOP_AFTER+MAX_WAIT LIGHT_SOFT_TIME_LIMIT = 105 LIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15 @@ -133,8 +127,7 @@ def document_by_cc_pair_cleanup_task( try: # Phase 1: read DB state, then release the connection before the - # document-index HTTP call (OpenSearch on cloud, Vespa on legacy - # deployments). Holding a pg transaction across the round-trip pins + # document-index HTTP call. Holding a pg transaction across the round-trip pins # a pgbouncer slot for the duration and saturates the pool under # bulk-deletion fan-out across the light worker fleet. chunk_count: int | None = None @@ -216,34 +209,27 @@ def document_by_cc_pair_cleanup_task( ) db_session.commit() - # Build document-index clients outside the DB session — construction + # Build the document-index client outside the DB session — construction # can take a few seconds to connect to the document index server, # and we don't want to pin a pgbouncer slot while that happens. - # This flow is for updates and deletion so we get all indices. - document_indices = get_all_document_indices( + document_index = get_default_document_index( primary_search_settings, secondary_search_settings, - httpx_client=HttpxPool.get("vespa"), ) - retry_document_indices: list[RetryDocumentIndex] = [ - RetryDocumentIndex(document_index) for document_index in document_indices - ] # Phase 2: document-index I/O — no DB connection held. if action == DocumentCleanupAction.DELETE: - for retry_document_index in retry_document_indices: - _ = retry_document_index.delete( - document_id, - chunk_count=chunk_count, - ) + _ = document_index.delete( + document_id, + chunk_count=chunk_count, + ) elif action == DocumentCleanupAction.UPDATE: assert update_request is not None - for retry_document_index in retry_document_indices: - # TODO(andrei): Previously there was a comment here saying - # it was ok if a doc did not exist in the document index. I - # don't agree with that claim, so keep an eye on this task - # to see if this raises. - retry_document_index.update([update_request]) + # TODO(andrei): Previously there was a comment here saying + # it was ok if a doc did not exist in the document index. I + # don't agree with that claim, so keep an eye on this task + # to see if this raises. + document_index.update([update_request]) # Phase 3: write back to PG in a fresh transaction. if action == DocumentCleanupAction.DELETE: @@ -282,77 +268,40 @@ def document_by_cc_pair_cleanup_task( except SoftTimeLimitExceeded: task_logger.info(f"SoftTimeLimitExceeded exception. doc={document_id}") completion_status = OnyxCeleryTaskCompletionStatus.SOFT_TIME_LIMIT - except Exception as ex: - e: Exception | None = None - while True: - if isinstance(ex, RetryError): - task_logger.warning( - f"Tenacity retry failed: num_attempts={ex.last_attempt.attempt_number}" - ) + except Exception as e: + task_logger.exception( + f"document_by_cc_pair_cleanup_task exceptioned: doc={document_id}" + ) - # only set the inner exception if it is of type Exception - e_temp = ex.last_attempt.exception() - if isinstance(e_temp, Exception): - e = e_temp - else: - e = ex - - if isinstance(e, httpx.HTTPStatusError): - if e.response.status_code == HTTPStatus.BAD_REQUEST: - task_logger.exception( - f"Non-retryable HTTPStatusError: doc={document_id} status={e.response.status_code}" - ) - # non-retryable failure removed nothing -> doc stays live; drop its candidate - with get_session_with_current_tenant() as db_session: - _clear_port_orphan_candidate_for_live_doc( - db_session, connector_id, credential_id, document_id - ) - db_session.commit() - completion_status = ( - OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION - ) - break - - task_logger.exception( - f"document_by_cc_pair_cleanup_task exceptioned: doc={document_id}" + completion_status = OnyxCeleryTaskCompletionStatus.RETRYABLE_EXCEPTION + if self.max_retries is not None and self.request.retries >= self.max_retries: + # This is the last attempt! mark the document as dirty in the db so that it + # eventually gets fixed out of band via stale document reconciliation + task_logger.warning( + f"Max celery task retries reached. Marking doc as dirty for reconciliation: doc={document_id}" ) - - completion_status = OnyxCeleryTaskCompletionStatus.RETRYABLE_EXCEPTION - if ( - self.max_retries is not None - and self.request.retries >= self.max_retries - ): - # This is the last attempt! mark the document as dirty in the db so that it - # eventually gets fixed out of band via stale document reconciliation - task_logger.warning( - f"Max celery task retries reached. Marking doc as dirty for reconciliation: doc={document_id}" + with get_session_with_current_tenant() as db_session: + # delete the cc pair relationship now and let reconciliation clean it up + # in the document index + delete_document_by_connector_credential_pair__no_commit( + db_session=db_session, + document_id=document_id, + connector_credential_pair_identifier=ConnectorCredentialPairIdentifier( + connector_id=connector_id, + credential_id=credential_id, + ), ) - with get_session_with_current_tenant() as db_session: - # delete the cc pair relationship now and let reconciliation clean it up - # in vespa - delete_document_by_connector_credential_pair__no_commit( - db_session=db_session, - document_id=document_id, - connector_credential_pair_identifier=ConnectorCredentialPairIdentifier( - connector_id=connector_id, - credential_id=credential_id, - ), - ) - mark_document_as_modified(document_id, db_session) - # helper keeps the candidate if this removed the doc's last link - _clear_port_orphan_candidate_for_live_doc( - db_session, connector_id, credential_id, document_id - ) - db_session.commit() - completion_status = ( - OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION + mark_document_as_modified(document_id, db_session) + # helper keeps the candidate if this removed the doc's last link + _clear_port_orphan_candidate_for_live_doc( + db_session, connector_id, credential_id, document_id ) - break - + db_session.commit() + completion_status = OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION + else: # Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64 countdown = 2 ** (self.request.retries + 4) self.retry(exc=e, countdown=countdown) # this will raise a celery exception - break # we won't hit this, but it looks weird not to have it finally: task_logger.info( f"document_by_cc_pair_cleanup_task completed: status={completion_status.value} doc={document_id}" 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 8c06d9e4a2..6cdd00cf7e 100644 --- a/backend/onyx/background/celery/tasks/user_file_processing/tasks.py +++ b/backend/onyx/background/celery/tasks/user_file_processing/tasks.py @@ -14,19 +14,12 @@ 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.background.task_utils import send_user_file_delete_task from onyx.chat.incognito import ( sweep_incognito_generated_files, sweep_stale_incognito_user_files, ) -from onyx.configs.app_configs import ( - DISABLE_VECTOR_DB, - MANAGED_VESPA, - VESPA_CLOUD_CERT_PATH, - VESPA_CLOUD_KEY_PATH, -) +from onyx.configs.app_configs import DISABLE_VECTOR_DB from onyx.configs.constants import ( CELERY_GENERIC_BEAT_LOCK_TIMEOUT, CELERY_USER_FILE_DELETE_TASK_EXPIRES, @@ -59,8 +52,9 @@ from onyx.db.user_file import ( fetch_user_files_with_access_relationships, mark_user_file_reconcile_pending, ) -from onyx.document_index.factory import get_all_document_indices -from onyx.document_index.interfaces_new import ( +from onyx.document_index.factory import get_default_document_index +from onyx.document_index.interfaces import ( + DocumentIndex, MetadataUpdateRequest, SecondaryIndexDocumentMissingError, ) @@ -73,7 +67,6 @@ 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 ( UserFileDeletingSkip, UserFileIndexingAdapter, @@ -393,14 +386,6 @@ def _process_user_file_with_indexing( Opens its own DB session for the indexing pipeline. The caller should not hold an open session when calling this function. """ - # 20 is the documented default for httpx max_keepalive_connections - if MANAGED_VESPA: - httpx_init_vespa_pool( - 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH - ) - else: - httpx_init_vespa_pool(20) - with get_session_with_current_tenant() as db_session: user_file = db_session.get(UserFile, _as_uuid(user_file_id)) if user_file is None or user_file.status == UserFileStatus.DELETING: @@ -421,11 +406,7 @@ def _process_user_file_with_indexing( embedding_model = DefaultIndexingEmbedder.from_db_search_settings( search_settings=current_search_settings, ) - document_indices = get_all_document_indices( - current_search_settings, - None, - httpx_client=HttpxPool.get("vespa"), - ) + document_index = get_default_document_index(current_search_settings, None) adapter = UserFileIndexingAdapter( tenant_id=tenant_id, db_session=db_session, @@ -433,7 +414,7 @@ def _process_user_file_with_indexing( try: index_pipeline_result = run_indexing_pipeline( embedder=embedding_model, - document_indices=document_indices, + document_index=document_index, ignore_time_skip=True, db_session=db_session, tenant_id=tenant_id, @@ -505,11 +486,7 @@ def _index_user_file_to_secondary( embedder = DefaultIndexingEmbedder.from_db_search_settings( search_settings=bound_secondary, ) - document_indices = get_all_document_indices( - bound_secondary, - None, - httpx_client=HttpxPool.get("vespa"), - ) + document_index = get_default_document_index(bound_secondary, None) adapter = UserFileIndexingAdapter( tenant_id=tenant_id, db_session=db_session, @@ -517,7 +494,7 @@ def _index_user_file_to_secondary( try: result = run_indexing_pipeline( embedder=embedder, - document_indices=document_indices, + document_index=document_index, ignore_time_skip=True, # skip the content_hash gate, else the PRESENT run's hash no-ops this write index_to_secondary=True, @@ -606,26 +583,22 @@ def _supply_user_file_to_secondary(user_file_id: str, tenant_id: str) -> bool: def _sync_metadata_and_reconcile_secondary( - retry_indices: list[RetryDocumentIndex], + document_index: DocumentIndex, update_request: MetadataUpdateRequest, user_file_id: str, tenant_id: str, ) -> bool: - """Apply the metadata update to every index; if the secondary is still porting and lacks + """Apply the metadata update to the index; if the secondary is still porting and lacks the doc, supply its content instead. Returns whether the secondary now matches PRESENT.""" - secondary_missing = False - for retry_index in retry_indices: - try: - retry_index.update([update_request]) - except SecondaryIndexDocumentMissingError: - task_logger.debug( - f"user_file={user_file_id} missing from a still-porting index; " - "supplying content." - ) - secondary_missing = True - if not secondary_missing: - return True - return _supply_user_file_to_secondary(user_file_id, tenant_id) + try: + document_index.update([update_request]) + except SecondaryIndexDocumentMissingError: + task_logger.debug( + f"user_file={user_file_id} missing from a still-porting index; " + "supplying content." + ) + return _supply_user_file_to_secondary(user_file_id, tenant_id) + return True def process_user_file_impl( @@ -883,19 +856,10 @@ def delete_user_file_impl( return try: - skip_vespa = DISABLE_VECTOR_DB - retry_document_indices: list[RetryDocumentIndex] = [] + document_index: DocumentIndex | None = None chunk_count_from_db: int | None = None file_id: str = "" - if not skip_vespa: - if MANAGED_VESPA: - httpx_init_vespa_pool( - 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH - ) - else: - httpx_init_vespa_pool(20) - # Phase 1: short read session — extract everything needed for slow I/O with get_session_with_current_tenant() as db_session: user_file = db_session.get(UserFile, _as_uuid(user_file_id)) @@ -908,17 +872,12 @@ def delete_user_file_impl( file_id = user_file.file_id chunk_count_from_db = user_file.chunk_count - if not skip_vespa: + if not DISABLE_VECTOR_DB: active_search_settings = get_active_search_settings(db_session) - document_indices = get_all_document_indices( + document_index = get_default_document_index( search_settings=active_search_settings.primary, secondary_search_settings=active_search_settings.secondary, - httpx_client=HttpxPool.get("vespa"), ) - retry_document_indices = [ - RetryDocumentIndex(document_index) - for document_index in document_indices - ] # Record the deletion before the index delete (below) so a racing port's # sweep removes any chunk its create-only copy resurrects. No-op when no @@ -935,20 +894,18 @@ def delete_user_file_impl( db_session.commit() # Phase 2: vector DB deletes + file store deletes (no DB session held). - # Pass the DB chunk count when known; otherwise None, which each document - # index resolves itself (Vespa fans out to find chunks, OpenSearch deletes - # by document id). This keeps the path backend-agnostic. - if not skip_vespa: + # Pass the DB chunk count when known; otherwise None, which the document + # index resolves itself. + if document_index is not None: chunk_count: int | None = ( chunk_count_from_db if chunk_count_from_db is not None and chunk_count_from_db > 0 else None ) - for retry_document_index in retry_document_indices: - retry_document_index.delete( - user_file_id, - chunk_count=chunk_count, - ) + document_index.delete( + user_file_id, + chunk_count=chunk_count, + ) file_store = get_default_file_store() blob_deleted = True @@ -1099,23 +1056,14 @@ def project_sync_user_file_impl( return try: - # Phase 1: short read session — extract all data needed for Vespa, then - # release the connection before the network-bound update calls. - retry_document_indices: list[RetryDocumentIndex] = [] + # Phase 1: short read session — extract all data needed for the document + # index, then release the connection before the network-bound update calls. + document_index: DocumentIndex | None = None project_ids: list[int] = [] persona_ids: list[int] = [] file_id_str: str = "" chunk_count: int | None = None access: DocumentAccess | None = None - skip_vespa = DISABLE_VECTOR_DB - - if not skip_vespa: - if MANAGED_VESPA: - httpx_init_vespa_pool( - 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH - ) - else: - httpx_init_vespa_pool(20) with get_session_with_current_tenant() as db_session: user_files = fetch_user_files_with_access_relationships( @@ -1130,7 +1078,7 @@ def project_sync_user_file_impl( ) return - if not skip_vespa: + if not DISABLE_VECTOR_DB: active_search_settings = get_active_search_settings(db_session) # INSTANT-promoted primary still backfilling: defer updates to # not-yet-ported files, else the create-only port reinstalls a stale ACL. @@ -1140,16 +1088,11 @@ def project_sync_user_file_impl( db_session, active_search_settings.primary.id ) ) - document_indices = get_all_document_indices( + document_index = get_default_document_index( search_settings=active_search_settings.primary, secondary_search_settings=active_search_settings.secondary, - httpx_client=HttpxPool.get("vespa"), primary_backfill_in_progress=primary_backfill_in_progress, ) - retry_document_indices = [ - RetryDocumentIndex(document_index) - for document_index in document_indices - ] project_ids = [project.id for project in user_file.projects] persona_ids = [p.id for p in user_file.assistants if not p.deleted] @@ -1161,7 +1104,7 @@ def project_sync_user_file_impl( # Phase 2: index update calls (no DB session held) secondary_consistent = True - if not skip_vespa: + if document_index is not None: update_request = MetadataUpdateRequest( document_ids=[file_id_str], doc_id_to_chunk_cnt={ @@ -1172,7 +1115,7 @@ def project_sync_user_file_impl( persona_ids=set(persona_ids), ) secondary_consistent = _sync_metadata_and_reconcile_secondary( - retry_document_indices, update_request, user_file_id, tenant_id + document_index, update_request, user_file_id, tenant_id ) task_logger.info(f"project_sync_user_file_impl - User file id={user_file_id}") diff --git a/backend/onyx/background/celery/tasks/vespa/tasks.py b/backend/onyx/background/celery/tasks/vespa/tasks.py index 0046a36c03..98ca883400 100644 --- a/backend/onyx/background/celery/tasks/vespa/tasks.py +++ b/backend/onyx/background/celery/tasks/vespa/tasks.py @@ -1,19 +1,15 @@ import time from collections.abc import Callable from datetime import datetime -from http import HTTPStatus from typing import Any, cast -import httpx from celery import Celery, Task, shared_task from celery.exceptions import SoftTimeLimitExceeded from redis.lock import Lock as RedisLock from sqlalchemy.orm import Session -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, LIGHT_TIME_LIMIT, @@ -26,7 +22,7 @@ from onyx.background.celery.tasks.vespa.document_sync import ( reset_document_sync, try_generate_stale_document_sync_tasks, ) -from onyx.configs.app_configs import JOB_TIMEOUT, VESPA_SYNC_MAX_TASKS +from onyx.configs.app_configs import DOCUMENT_INDEX_SYNC_MAX_TASKS, JOB_TIMEOUT from onyx.configs.constants import ( CELERY_DOCUMENT_SYNC_BEAT_LOCK_TIMEOUT, OnyxCeleryTask, @@ -57,12 +53,11 @@ from onyx.db.sync_record import ( insert_sync_record, update_sync_record_status, ) -from onyx.document_index.factory import get_all_document_indices -from onyx.document_index.interfaces_new import ( +from onyx.document_index.factory import get_default_document_index +from onyx.document_index.interfaces 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, @@ -119,7 +114,12 @@ def check_for_vespa_sync_task(self: Task, *, tenant_id: str) -> bool | None: # 1/3: KICKOFF with get_session_with_current_tenant() as db_session: try_generate_stale_document_sync_tasks( - self.app, VESPA_SYNC_MAX_TASKS, db_session, r, lock_beat, tenant_id + self.app, + DOCUMENT_INDEX_SYNC_MAX_TASKS, + db_session, + r, + lock_beat, + tenant_id, ) # region document set scan @@ -265,7 +265,7 @@ def try_generate_document_set_sync_tasks( # Add all documents that need to be updated into the queue result = rds.generate_tasks( - VESPA_SYNC_MAX_TASKS, celery_app, db_session, r, lock_beat, tenant_id + DOCUMENT_INDEX_SYNC_MAX_TASKS, celery_app, db_session, r, lock_beat, tenant_id ) if result is None: return None @@ -340,7 +340,7 @@ def try_generate_user_group_sync_tasks( f"RedisUserGroup.generate_tasks starting. usergroup_id={usergroup.id}" ) result = rug.generate_tasks( - VESPA_SYNC_MAX_TASKS, celery_app, db_session, r, lock_beat, tenant_id + DOCUMENT_INDEX_SYNC_MAX_TASKS, celery_app, db_session, r, lock_beat, tenant_id ) if result is None: return None @@ -540,36 +540,29 @@ def document_index_metadata_sync_task( completion_status = OnyxCeleryTaskCompletionStatus.SKIPPED else: # Client construction can be slow, so it also stays outside the - # session. This flow is for updates so we get all indices. - document_indices = get_all_document_indices( + # session. + document_index = get_default_document_index( search_settings=primary_search_settings, secondary_search_settings=secondary_search_settings, - httpx_client=HttpxPool.get("vespa"), primary_backfill_in_progress=primary_backfill_in_progress, ) - retry_document_indices: list[RetryDocumentIndex] = [ - RetryDocumentIndex(document_index) - for document_index in document_indices - ] - # Phase 2: document-index I/O — no DB connection held. # Reindex-port: doc missing from a still-populating index (FUTURE, or the # INSTANT-promoted primary) — defer rather than fail; the fully-populated # index's write already committed in the pair. port_index_missing = False - for retry_document_index in retry_document_indices: - try: - # TODO(andrei): Previously there was a comment here saying - # it was ok if a doc did not exist in the document index. I - # don't agree with that claim, so keep an eye on this task - # to see if this raises. - retry_document_index.update([update_request]) - except SecondaryIndexDocumentMissingError: - task_logger.debug( - f"doc={document_id} not in a still-porting index; deferring sync." - ) - port_index_missing = True + try: + # TODO(andrei): Previously there was a comment here saying + # it was ok if a doc did not exist in the document index. I + # don't agree with that claim, so keep an eye on this task + # to see if this raises. + document_index.update([update_request]) + except SecondaryIndexDocumentMissingError: + task_logger.debug( + f"doc={document_id} not in a still-porting index; deferring sync." + ) + port_index_missing = True # Phase 3: write back to PG in a fresh transaction. # update db last. Worst case = we crash right before this and @@ -595,48 +588,18 @@ def document_index_metadata_sync_task( except SoftTimeLimitExceeded: task_logger.info(f"SoftTimeLimitExceeded exception. doc={document_id}") completion_status = OnyxCeleryTaskCompletionStatus.SOFT_TIME_LIMIT - except Exception as ex: - e: Exception | None = None - while True: - if isinstance(ex, RetryError): - task_logger.warning( - f"Tenacity retry failed: num_attempts={ex.last_attempt.attempt_number}" - ) + except Exception as e: + task_logger.exception( + f"document_index_metadata_sync_task exceptioned: doc={document_id}" + ) - # only set the inner exception if it is of type Exception - e_temp = ex.last_attempt.exception() - if isinstance(e_temp, Exception): - e = e_temp - else: - e = ex + completion_status = OnyxCeleryTaskCompletionStatus.RETRYABLE_EXCEPTION + if self.max_retries is not None and self.request.retries >= self.max_retries: + completion_status = OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION - if isinstance(e, httpx.HTTPStatusError): - if e.response.status_code == HTTPStatus.BAD_REQUEST: - task_logger.exception( - f"Non-retryable HTTPStatusError: doc={document_id} status={e.response.status_code}" - ) - completion_status = ( - OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION - ) - break - - task_logger.exception( - f"document_index_metadata_sync_task exceptioned: doc={document_id}" - ) - - completion_status = OnyxCeleryTaskCompletionStatus.RETRYABLE_EXCEPTION - if ( - self.max_retries is not None - and self.request.retries >= self.max_retries - ): - completion_status = ( - OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION - ) - - # Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64 - countdown = 2 ** (self.request.retries + 4) - self.retry(exc=e, countdown=countdown) # this will raise a celery exception - break # we won't hit this, but it looks weird not to have it + # Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64 + countdown = 2 ** (self.request.retries + 4) + self.retry(exc=e, countdown=countdown) # this will raise a celery exception finally: task_logger.info( f"document_index_metadata_sync_task completed: status={completion_status.value} doc={document_id}" diff --git a/backend/onyx/background/indexing/run_targeted_reindex.py b/backend/onyx/background/indexing/run_targeted_reindex.py index 1955de3842..720a62f9f5 100644 --- a/backend/onyx/background/indexing/run_targeted_reindex.py +++ b/backend/onyx/background/indexing/run_targeted_reindex.py @@ -40,12 +40,11 @@ from onyx.db.models import ( TargetedReindexJobTarget, ) from onyx.db.targeted_reindex import targets_to_connector_failures -from onyx.document_index.factory import get_all_document_indices +from onyx.document_index.factory import get_default_document_index 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, ) @@ -131,11 +130,7 @@ def _flush_batch( search_settings=search_settings, callback=None, ) - document_indices = get_all_document_indices( - search_settings, - None, - httpx_client=HttpxPool.get("vespa"), - ) + document_index = get_default_document_index(search_settings, None) metadata = IndexAttemptMetadata( attempt_id=attempt.id, connector_id=attempt.connector_credential_pair.connector.id, @@ -154,7 +149,7 @@ def _flush_batch( result = run_indexing_pipeline( embedder=embedder, - document_indices=document_indices, + document_index=document_index, ignore_time_skip=True, # FUTURE/secondary build: skip the PRESENT-only content_hash dedup. index_to_secondary=search_settings.status.is_future(), @@ -311,7 +306,7 @@ def process_targets_for_cc_pair( ) # Per-attempt pipeline run. Each attempt commits to its own - # search_settings's document_indices. + # search_settings.s document index. for attempt in cc_pair_attempts: for batch_num, batch in enumerate(chunked(docs, INDEX_BATCH_SIZE)): landed, failed = _flush_batch( diff --git a/backend/onyx/chat/chat_utils.py b/backend/onyx/chat/chat_utils.py index fc8bd313a4..9dd639cdf2 100644 --- a/backend/onyx/chat/chat_utils.py +++ b/backend/onyx/chat/chat_utils.py @@ -23,7 +23,6 @@ from onyx.chat.models import ( from onyx.configs.app_configs import DISABLE_VECTOR_DB from onyx.configs.constants import ( DEFAULT_PERSONA_ID, - TMP_DRALPHA_PERSONA_NAME, FileOrigin, MessageType, ) @@ -40,10 +39,6 @@ from onyx.db.enums import ( record_mode_persists_content, ) from onyx.db.file_record import FileRecordNotFoundError -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.persona import user_can_access_persona @@ -55,10 +50,6 @@ 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, 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, TOOL_CALL_RESPONSE_CROSS_MESSAGE, @@ -435,26 +426,6 @@ def extract_headers( return extracted_headers -def process_kg_commands( - message: str, - persona_name: str, - tenant_id: str, # noqa: ARG001 - db_session: Session, -) -> None: - # Temporarily, until we have a draft UI for the KG Operations/Management - # TODO: move to api endpoint once we get frontend - if not persona_name.startswith(TMP_DRALPHA_PERSONA_NAME): - return - - kg_config_settings = get_kg_config_settings() - if not is_kg_config_settings_enabled_valid(kg_config_settings): - return - - if message == "kg_setup": - populate_missing_default_entity_types__commit(db_session=db_session) - raise KGException("KG setup done") - - def _get_or_extract_plaintext( file_id: str, extract_fn: Callable[[], str], diff --git a/backend/onyx/configs/app_configs.py b/backend/onyx/configs/app_configs.py index feacd07ea1..29117c44ca 100644 --- a/backend/onyx/configs/app_configs.py +++ b/backend/onyx/configs/app_configs.py @@ -75,7 +75,7 @@ GENERATIVE_MODEL_ACCESS_CHECK_FREQ = int( # Controls whether users can use User Knowledge (personal documents) in assistants DISABLE_USER_KNOWLEDGE = os.environ.get("DISABLE_USER_KNOWLEDGE", "").lower() == "true" -# Disables vector DB (Vespa/OpenSearch) entirely. When True, connectors and RAG search +# Disables vector DB (OpenSearch) entirely. When True, connectors and RAG search # are disabled but core chat, tools, user file uploads, and Projects still work. DISABLE_VECTOR_DB = os.environ.get("DISABLE_VECTOR_DB", "").lower() == "true" @@ -542,37 +542,12 @@ OPENSEARCH_EXPLAIN_ENABLED = ( # existing indices need reindexing after a change. OPENSEARCH_TEXT_ANALYZER = os.environ.get("OPENSEARCH_TEXT_ANALYZER") or "english" -# This is the "base" config for now, the idea is that at least for our dev -# environments we always want to be dual indexing into both OpenSearch and Vespa -# to stress test the new codepaths. Only enable this if there is some instance -# of OpenSearch running for the relevant Onyx instance. -# NOTE: Now enabled on by default, unless the env indicates otherwise. -ENABLE_OPENSEARCH_INDEXING_FOR_ONYX = ( - os.environ.get("ENABLE_OPENSEARCH_INDEXING_FOR_ONYX", "true").lower() == "true" -) -# NOTE: This effectively does nothing anymore, admins can now toggle whether -# retrieval is through OpenSearch. This value is only used as a final fallback -# in case that doesn't work for whatever reason. -# Given that the "base" config above is true, this enables whether we want to -# retrieve from OpenSearch or Vespa. We want to be able to quickly toggle this -# in the event we see issues with OpenSearch retrieval in our dev environments. -ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX = ( - ENABLE_OPENSEARCH_INDEXING_FOR_ONYX - and os.environ.get("ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX", "").lower() == "true" -) -DISABLE_OPENSEARCH_MIGRATION_TASK = ( - os.environ.get("DISABLE_OPENSEARCH_MIGRATION_TASK", "").lower() == "true" -) -ONYX_DISABLE_VESPA = os.environ.get("ONYX_DISABLE_VESPA", "true").lower() == "true" # Whether we should check for and create an index if necessary every time we # instantiate an OpenSearchDocumentIndex on multitenant cloud. Defaults to True. VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT = ( os.environ.get("VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT", "true").lower() == "true" ) -OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE = int( - os.environ.get("OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE") or 500 -) # Lifetime of a point-in-time used to scan an index consistently (reindex port). # Each search extends the lease; an idle PIT self-expires after this. PIT_KEEP_ALIVE: str = os.environ.get("PIT_KEEP_ALIVE") or "5m" @@ -592,19 +567,11 @@ ONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH = ( == "true" ) -VESPA_HOST = os.environ.get("VESPA_HOST") or "localhost" -# NOTE: this is used if and only if the vespa config server is accessible via a -# different host than the main vespa application -VESPA_CONFIG_SERVER_HOST = os.environ.get("VESPA_CONFIG_SERVER_HOST") or VESPA_HOST -VESPA_PORT = os.environ.get("VESPA_PORT") or "8081" -VESPA_TENANT_PORT = os.environ.get("VESPA_TENANT_PORT") or "19071" -# the number of times to try and connect to vespa on startup before giving up -VESPA_NUM_ATTEMPTS_ON_STARTUP = int(os.environ.get("NUM_RETRIES_ON_STARTUP") or 10) - -VESPA_CLOUD_URL = os.environ.get("VESPA_CLOUD_URL", "") - -VESPA_CLOUD_CERT_PATH = os.environ.get("VESPA_CLOUD_CERT_PATH") -VESPA_CLOUD_KEY_PATH = os.environ.get("VESPA_CLOUD_KEY_PATH") +# The number of times to try to connect to the document index on startup before +# giving up. +DOCUMENT_INDEX_NUM_ATTEMPTS_ON_STARTUP = int( + os.environ.get("NUM_RETRIES_ON_STARTUP") or 10 +) # Number of documents in a batch during indexing (further batching done by chunks before passing to bi-encoder) INDEX_BATCH_SIZE = int(os.environ.get("INDEX_BATCH_SIZE") or 16) @@ -1188,8 +1155,9 @@ CELERY_WORKER_SCHEDULED_TASKS_CONCURRENCY = int( os.environ.get("CELERY_WORKER_SCHEDULED_TASKS_CONCURRENCY") or 4 ) -# The maximum number of tasks that can be queued up to sync to Vespa in a single pass -VESPA_SYNC_MAX_TASKS = 8192 +# The maximum number of tasks that can be queued up to sync to the document index +# in a single pass +DOCUMENT_INDEX_SYNC_MAX_TASKS = 8192 DB_YIELD_PER_DEFAULT = 64 @@ -1675,15 +1643,6 @@ AVERAGE_SUMMARY_EMBEDDINGS = ( MAX_TOKENS_FOR_FULL_INCLUSION = 4096 -# The intent was to have this be configurable per query, but I don't think any -# codepath was actually configuring this, so for the migrated Vespa interface -# we'll just use the default value, but also have it be configurable by env var. -RECENCY_BIAS_MULTIPLIER = float(os.environ.get("RECENCY_BIAS_MULTIPLIER") or 1.0) - -# Should match the rerank-count value set in -# backend/onyx/document_index/vespa/app_config/schemas/danswer_chunk.sd.jinja. -RERANK_COUNT = int(os.environ.get("RERANK_COUNT") or 1000) - # Flat per-image cost (cents) when litellm has no price for an image model. # Clamped to >= 0 so a misconfigured negative can't credit usage. DEFAULT_IMAGE_COST_CENTS = max( @@ -1762,11 +1721,6 @@ PROMPT_CACHE_CHAT_HISTORY = ( # gateway. Off by default. ENABLE_AZURE_IMAGE_CAP = os.environ.get("ENABLE_AZURE_IMAGE_CAP", "").lower() == "true" -# If set to `true` will enable additional logs about Vespa query performance -# (time spent on finding the right docs + time spent fetching summaries from disk) -LOG_VESPA_TIMING_INFORMATION = ( - os.environ.get("LOG_VESPA_TIMING_INFORMATION", "").lower() == "true" -) LOG_ENDPOINT_LATENCY = os.environ.get("LOG_ENDPOINT_LATENCY", "").lower() == "true" LOG_POSTGRES_LATENCY = os.environ.get("LOG_POSTGRES_LATENCY", "").lower() == "true" LOG_POSTGRES_CONN_COUNTS = ( @@ -1836,22 +1790,6 @@ CUSTOM_ANSWER_VALIDITY_CONDITIONS = json.loads( os.environ.get("CUSTOM_ANSWER_VALIDITY_CONDITIONS", "[]") ) -VESPA_REQUEST_TIMEOUT = int(os.environ.get("VESPA_REQUEST_TIMEOUT") or "15") -# This is the timeout for the client side of the Vespa migration task. When -# exceeded, an exception is raised in our code. This value should be higher than -# VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT. -VESPA_MIGRATION_REQUEST_TIMEOUT_S = int( - os.environ.get("VESPA_MIGRATION_REQUEST_TIMEOUT_S") or "120" -) -# This is the timeout Vespa uses on the server side to know when to wrap up its -# traversal and try to report partial results. This differs from the client -# timeout above which raises an exception in our code when exceeded. This -# timeout allows Vespa to return gracefully. This value should be lower than -# VESPA_MIGRATION_REQUEST_TIMEOUT_S. Formatted as s. -VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT = os.environ.get( - "VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT", "110s" -) - SYSTEM_RECURSION_LIMIT = int(os.environ.get("SYSTEM_RECURSION_LIMIT") or "1000") # Size of the api-server anyio threadpool that runs sync endpoints, including the @@ -1913,9 +1851,6 @@ AZURE_IMAGE_DEPLOYMENT_NAME = os.environ.get("AZURE_IMAGE_DEPLOYMENT_NAME") IMAGE_MODEL_NAME = os.environ.get("IMAGE_MODEL_NAME", "gpt-image-1") IMAGE_MODEL_PROVIDER = os.environ.get("IMAGE_MODEL_PROVIDER", "openai") -# Use managed Vespa (Vespa Cloud). If set, must also set VESPA_CLOUD_URL, VESPA_CLOUD_CERT_PATH and VESPA_CLOUD_KEY_PATH -MANAGED_VESPA = os.environ.get("MANAGED_VESPA", "").lower() == "true" - ENABLE_EMAIL_INVITES = os.environ.get("ENABLE_EMAIL_INVITES", "").lower() == "true" # When true, GET /users is restricted to callers with READ_USERS so non-admins @@ -2177,10 +2112,6 @@ AZURE_STORAGE_CONNECTION_STRING = ( ) AZURE_STORAGE_ACCOUNT_KEY = os.environ.get("AZURE_STORAGE_ACCOUNT_KEY") or None -# Forcing Vespa Language -# English: en, German:de, etc. See: https://docs.vespa.ai/en/linguistics.html -VESPA_LANGUAGE_OVERRIDE = os.environ.get("VESPA_LANGUAGE_OVERRIDE") - ##### # Default LLM API Keys (for cloud deployments) diff --git a/backend/onyx/configs/chat_configs.py b/backend/onyx/configs/chat_configs.py index 4184556225..be859ff63d 100644 --- a/backend/onyx/configs/chat_configs.py +++ b/backend/onyx/configs/chat_configs.py @@ -12,7 +12,6 @@ MAX_CHUNKS_FED_TO_CHAT = int(os.environ.get("MAX_CHUNKS_FED_TO_CHAT") or 25) MAX_LLM_CYCLES: int = int(os.environ.get("MAX_LLM_CYCLES") or 6) # 1 / (1 + DOC_TIME_DECAY * doc-age-in-years), set to 0 to have no decay -# Capped in Vespa at 0.5 DOC_TIME_DECAY = float( os.environ.get("DOC_TIME_DECAY") or 0.5 # Hits limit at 2 years by default ) @@ -98,8 +97,6 @@ HARD_DELETE_CHATS = os.environ.get("HARD_DELETE_CHATS", "").lower() == "true" NUM_INTERNET_SEARCH_RESULTS = int(os.environ.get("NUM_INTERNET_SEARCH_RESULTS") or 10) NUM_INTERNET_SEARCH_CHUNKS = int(os.environ.get("NUM_INTERNET_SEARCH_CHUNKS") or 50) -VESPA_SEARCHER_THREADS = int(os.environ.get("VESPA_SEARCHER_THREADS") or 2) - # Whether or not to use the semantic & keyword search expansions for Basic Search USE_SEMANTIC_KEYWORD_EXPANSIONS_BASIC_SEARCH = ( os.environ.get("USE_SEMANTIC_KEYWORD_EXPANSIONS_BASIC_SEARCH", "false").lower() diff --git a/backend/onyx/configs/constants.py b/backend/onyx/configs/constants.py index fbaae6d2cb..dfe3ebd487 100644 --- a/backend/onyx/configs/constants.py +++ b/backend/onyx/configs/constants.py @@ -241,7 +241,7 @@ CELERY_USER_FILE_DELETE_TASK_EXPIRES = 60 # 1 minute (in seconds) # Per-doc metadata-sync task expiry: bounds queue growth if consumers stall. An # expired task's doc stays needs_sync / secondary_only_sync_pending and is -# re-enqueued on the next vespa-sync beat pass, so dropping it is safe. +# re-enqueued on the next document index sync beat pass, so dropping it is safe. CELERY_DOCUMENT_SYNC_TASK_EXPIRES = 60 * 60 # 1 hour (in seconds) # Max queue depth before the delete beat stops enqueuing more delete tasks. @@ -366,11 +366,6 @@ class BlobType(str, Enum): OCI_STORAGE = "oci_storage" -class DocumentIndexType(str, Enum): - COMBINED = "combined" # Vespa - SPLIT = "split" # Typesense + Qdrant - - class QueryHistoryType(str, Enum): DISABLED = "disabled" ANONYMIZED = "anonymized" @@ -518,8 +513,6 @@ class OnyxCeleryQueues: # Scheduled tasks queue (Craft scheduled-task executor) SCHEDULED_TASKS = "scheduled_tasks" - OPENSEARCH_MIGRATION = "opensearch_migration" - class OnyxRedisLocks: PRIMARY_WORKER = "da_lock:primary_worker" @@ -539,7 +532,6 @@ class OnyxRedisLocks: CHECK_CONNECTOR_EXTERNAL_GROUP_SYNC_BEAT_LOCK = ( "da_lock:check_connector_external_group_sync_beat" ) - OPENSEARCH_MIGRATION_BEAT_LOCK = "da_lock:opensearch_migration_beat" OPENSEARCH_VERIFY_INDEX_LOCK_PREFIX = "da_lock:opensearch_verify_index" SECURITY_SETTINGS = "da_lock:security_settings" @@ -744,16 +736,6 @@ class OnyxCeleryTask: SCHEDULED_TASKS_RUN = "scheduled_tasks_run" SCHEDULED_TASKS_CLEANUP_STUCK = "scheduled_tasks_cleanup_stuck" - CHECK_FOR_DOCUMENTS_FOR_OPENSEARCH_MIGRATION_TASK = ( - "check_for_documents_for_opensearch_migration_task" - ) - MIGRATE_DOCUMENTS_FROM_VESPA_TO_OPENSEARCH_TASK = ( - "migrate_documents_from_vespa_to_opensearch_task" - ) - MIGRATE_CHUNKS_FROM_VESPA_TO_OPENSEARCH_TASK = ( - "migrate_chunks_from_vespa_to_opensearch_task" - ) - # this needs to correspond to the matching entry in supervisord ONYX_CELERY_BEAT_HEARTBEAT_KEY = "onyx:celery:beat:heartbeat" @@ -773,11 +755,6 @@ else: REDIS_SOCKET_KEEPALIVE_OPTIONS[getattr(socket, "TCP_KEEPIDLE")] = 60 # noqa: B009 # ods: ignore[getattr] -class OnyxCallTypes(str, Enum): - FIREFLIES = "FIREFLIES" - GONG = "GONG" - - NUM_DAYS_TO_KEEP_CHECKPOINTS = 7 # checkpoints are queried based on index attempts, so we need to keep index attempts for one more day NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS = NUM_DAYS_TO_KEEP_CHECKPOINTS + 1 diff --git a/backend/onyx/configs/kg_configs.py b/backend/onyx/configs/kg_configs.py index dbac445970..3626383624 100644 --- a/backend/onyx/configs/kg_configs.py +++ b/backend/onyx/configs/kg_configs.py @@ -1,37 +1,5 @@ import os -KG_RESEARCH_NUM_RETRIEVED_DOCS: int = int( - os.environ.get("KG_RESEARCH_NUM_RETRIEVED_DOCS", "25") -) - - -KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES: int = int( - os.environ.get("KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES", "10") -) - - -KG_ENTITY_EXTRACTION_TIMEOUT: int = int( - os.environ.get("KG_ENTITY_EXTRACTION_TIMEOUT", "15") -) - -KG_RELATIONSHIP_EXTRACTION_TIMEOUT: int = int( - os.environ.get("KG_RELATIONSHIP_EXTRACTION_TIMEOUT", "15") -) - -KG_STRATEGY_GENERATION_TIMEOUT: int = int( - os.environ.get("KG_STRATEGY_GENERATION_TIMEOUT", "20") -) - -KG_SQL_GENERATION_TIMEOUT: int = int(os.environ.get("KG_SQL_GENERATION_TIMEOUT", "40")) - -KG_SQL_GENERATION_TIMEOUT_OVERRIDE: int = int( - os.environ.get("KG_SQL_GENERATION_TIMEOUT_OVERRIDE", "40") -) - -KG_SQL_GENERATION_MAX_TOKENS: int = int( - os.environ.get("KG_SQL_GENERATION_MAX_TOKENS", "1500") -) - KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX: str = os.environ.get( "KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX", "allowed_docs" ) @@ -45,101 +13,11 @@ KG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX: str = os.environ.get( ) -KG_FILTER_CONSTRUCTION_TIMEOUT: int = int( - os.environ.get("KG_FILTER_CONSTRUCTION_TIMEOUT", "15") -) - - -KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT: int = int( - os.environ.get("KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT", "100") -) - -KG_FILTERED_SEARCH_TIMEOUT: int = int( - os.environ.get("KG_FILTERED_SEARCH_TIMEOUT", "30") -) - - -KG_OBJECT_SOURCE_RESEARCH_TIMEOUT: int = int( - os.environ.get("KG_OBJECT_SOURCE_RESEARCH_TIMEOUT", "30") -) - -KG_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION: int = int( - os.environ.get("KG_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION", "45") -) - -KG_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION: int = int( - os.environ.get("KG_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION", "15") -) - -KG_MAX_TOKENS_ANSWER_GENERATION: int = int( - os.environ.get("KG_MAX_TOKENS_ANSWER_GENERATION", "1024") -) - -KG_MAX_DEEP_SEARCH_RESULTS: int = int( - os.environ.get("KG_MAX_DEEP_SEARCH_RESULTS", "30") -) - - -KG_METADATA_TRACKING_THRESHOLD: int = int( - os.environ.get("KG_METADATA_TRACKING_THRESHOLD", "10") -) - - KG_DEFAULT_MAX_PARENT_RECURSION_DEPTH: int = int( os.environ.get("KG_DEFAULT_MAX_PARENT_RECURSION_DEPTH", "2") ) -_KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT: float = max( - 1e-3, - min(1, float(os.environ.get("KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT", "0.25"))), -) -_KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT: float = max( - 1e-3, - min(1, float(os.environ.get("KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT", "0.25"))), -) -_KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT: float = max( - 1e-3, - min(1, float(os.environ.get("KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT", "0.5"))), -) -_KG_NORMALIZATION_RERANK_NGRAM_SUMS: float = ( - _KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT - + _KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT - + _KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT -) - -KG_NORMALIZATION_RERANK_NGRAM_WEIGHTS: tuple[float, float, float] = ( - _KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT / _KG_NORMALIZATION_RERANK_NGRAM_SUMS, - _KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT / _KG_NORMALIZATION_RERANK_NGRAM_SUMS, - _KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT / _KG_NORMALIZATION_RERANK_NGRAM_SUMS, -) - - -KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT: float = max( - 0, - min(1, float(os.environ.get("KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT", "0.25"))), -) - - -KG_NORMALIZATION_RERANK_THRESHOLD: float = float( - os.environ.get("KG_NORMALIZATION_RERANK_THRESHOLD", "0.3") -) - - -KG_CLUSTERING_RETRIEVE_THRESHOLD: float = float( - os.environ.get("KG_CLUSTERING_RETRIEVE_THRESHOLD", "0.6") -) - - -KG_CLUSTERING_THRESHOLD: float = float( - os.environ.get("KG_CLUSTERING_THRESHOLD", "0.96") -) - -KG_MAX_SEARCH_DOCUMENTS: int = int(os.environ.get("KG_MAX_SEARCH_DOCUMENTS", "15")) - -KG_MAX_DECOMPOSITION_SEGMENTS: int = int( - os.environ.get("KG_MAX_DECOMPOSITION_SEGMENTS", "10") -) KG_BETA_ASSISTANT_DESCRIPTION = ( "The KG Beta assistant uses the Onyx Knowledge Graph (beta) structure \ to answer questions" diff --git a/backend/onyx/configs/model_configs.py b/backend/onyx/configs/model_configs.py index 110f080dfc..53b256c0c7 100644 --- a/backend/onyx/configs/model_configs.py +++ b/backend/onyx/configs/model_configs.py @@ -17,7 +17,6 @@ from shared_configs.configs import DEFAULT_DOCUMENT_ENCODER_MODEL DOCUMENT_ENCODER_MODEL = ( os.environ.get("DOCUMENT_ENCODER_MODEL") or DEFAULT_DOCUMENT_ENCODER_MODEL ) -# If the below is changed, Vespa deployment must also be changed DOC_EMBEDDING_DIM = int(os.environ.get("DOC_EMBEDDING_DIM") or 768) NORMALIZE_EMBEDDINGS = ( os.environ.get("NORMALIZE_EMBEDDINGS") or "true" diff --git a/backend/onyx/context/search/models.py b/backend/onyx/context/search/models.py index d6873ea5d2..72f7c09097 100644 --- a/backend/onyx/context/search/models.py +++ b/backend/onyx/context/search/models.py @@ -125,7 +125,8 @@ class BaseFilters(BaseModel): class UserFileFilters(BaseModel): - # Scopes search to user files tagged with a given project/persona in Vespa. + # Scopes search to user files tagged with a given project/persona in the + # document index. # These are NOT simply the IDs of the current project or persona — they are # only set when the persona's/project's user files overflowed the LLM # context window and must be searched via vector DB instead of being loaded @@ -207,7 +208,7 @@ class InferenceChunk(BaseChunk): # TODO(andrei): Ideally we could improve this to where each value is just a # list of strings. metadata: dict[str, str | list[str]] - # Matched sections in the chunk. Uses Vespa syntax e.g. TEXT + # Matched sections in the chunk. Matches are wrapped in TEXT # to specify that a set of words should be highlighted. For example: # ["the answer is 42", "he couldn't find an answer"] match_highlights: list[str] @@ -315,7 +316,7 @@ class SearchDoc(BaseModel): score: float | None = None is_relevant: bool | None = None relevance_explanation: str | None = None - # Matched sections in the doc. Uses Vespa syntax e.g. TEXT + # Matched sections in the doc. Matches are wrapped in TEXT # to specify that a set of words should be highlighted. For example: # ["the answer is 42", "the answer is 42""] match_highlights: list[str] diff --git a/backend/onyx/context/search/pipeline.py b/backend/onyx/context/search/pipeline.py index c8f690b78c..f8d0c713bf 100644 --- a/backend/onyx/context/search/pipeline.py +++ b/backend/onyx/context/search/pipeline.py @@ -21,7 +21,7 @@ from onyx.context.search.retrieval.search_runner import search_chunks from onyx.context.search.utils import inference_section_from_chunks from onyx.db.document_set import filter_document_set_names_by_user_access from onyx.db.models import User -from onyx.document_index.interfaces_new import DocumentIndex +from onyx.document_index.interfaces 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 @@ -266,7 +266,7 @@ def search_pipeline( # Pre-extracted persona search configuration (None when no persona) persona_search_info: PersonaSearchInfo | None, db_session: Session | None = None, - # Vespa metadata filters for overflowing user files. NOT the raw IDs + # Document index metadata filters for overflowing user files. NOT the raw IDs # of the current project/persona — only set when user files couldn't fit # in the LLM context and need to be searched via vector DB. project_id_filter: int | None = None, diff --git a/backend/onyx/context/search/retrieval/search_runner.py b/backend/onyx/context/search/retrieval/search_runner.py index b49d9763f4..73e077f6ac 100644 --- a/backend/onyx/context/search/retrieval/search_runner.py +++ b/backend/onyx/context/search/retrieval/search_runner.py @@ -12,7 +12,7 @@ from onyx.context.search.models import ( 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.document_index.interfaces import DocumentIndex, DocumentSectionRequest from onyx.federated_connectors.federated_retrieval import ( FederatedRetrievalInfo, get_federated_retrieval_functions, @@ -133,9 +133,7 @@ def search_chunks( if normal_search_enabled: if query_request.hybrid_alpha is not None and query_request.hybrid_alpha == 0.0: # Hybrid alpha explicitly set to keyword only —> do pure keyword - # search without computing an embedding. This branch is currently - # only set by OpenSearch-aware producers; Vespa would raise - # NotImplementedError on `keyword_retrieval`. + # search without computing an embedding. run_queries.append( ( lambda: _keyword_search(query_request, document_index), diff --git a/backend/onyx/db/connector.py b/backend/onyx/db/connector.py index 471583275f..f74cfb516a 100644 --- a/backend/onyx/db/connector.py +++ b/backend/onyx/db/connector.py @@ -13,7 +13,6 @@ from onyx.db.models import ( FederatedConnector, IndexAttempt, ) -from onyx.kg.models import KGConnectorData from onyx.server.documents.models import ConnectorBase, ObjectCreationIdResponse from onyx.server.models import StatusResponse from onyx.utils.logger import setup_logger @@ -372,29 +371,3 @@ def mark_ccpair_with_indexing_trigger( except Exception: db_session.rollback() raise - - -def get_kg_enabled_connectors(db_session: Session) -> list[KGConnectorData]: - """ - Retrieves a list of connector IDs that have not been KG processed for a given tenant. - Args: - db_session (Session): The database session to use - Returns: - list[KGConnectorData]: List of connector IDs with KG extraction enabled but have unprocessed documents - """ - try: - stmt = select(Connector.id, Connector.source, Connector.kg_coverage_days).where( - Connector.kg_processing_enabled - ) - result = db_session.execute(stmt) - - connector_results = [ - KGConnectorData(id=row[0], source=row[1].lower(), kg_coverage_days=row[2]) - for row in result.fetchall() - ] - - return connector_results - - except Exception as e: - logger.error("Error fetching unprocessed connector IDs: %s", str(e)) - raise e diff --git a/backend/onyx/db/document.py b/backend/onyx/db/document.py index c9a05becf8..b98d4d6454 100644 --- a/backend/onyx/db/document.py +++ b/backend/onyx/db/document.py @@ -1,7 +1,7 @@ import contextlib import time from collections.abc import Generator, Iterable, Sequence -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from typing import Any, NamedTuple from uuid import UUID @@ -32,7 +32,6 @@ from sqlalchemy.orm import Session from sqlalchemy.sql.expression import null 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 @@ -49,8 +48,6 @@ from onyx.db.models import ( ConnectorCredentialPair, Credential, DocumentByConnectorCredentialPair, - KGEntity, - KGRelationship, User, ) from onyx.db.models import Document as DbDocument @@ -173,7 +170,7 @@ def count_secondary_only_sync_pending_documents_for_cc_pairs( def count_documents_by_needs_sync_or_secondary_pending(session: Session) -> int: """count_documents_by_needs_sync plus docs whose FUTURE sync was deferred. - The vespa sync producer gates on this so a deferred-only backlog still + The document index sync producer gates on this so a deferred-only backlog still generates drain tasks — a deferred doc has its needs_sync already cleared, so count_documents_by_needs_sync alone would miss it. """ @@ -1550,7 +1547,7 @@ def prepare_to_modify_documents( ) -> Generator[TransactionalContext, None, None]: """Try and acquire locks for the documents to prevent other jobs from modifying them at the same time (e.g. avoid race conditions). This should be - called ahead of any modification to Vespa. Locks should be released by the + called ahead of any modification to the document index. Locks should be released by the caller as soon as updates are complete by finishing the transaction. NOTE: only one commit is allowed within the context manager returned by this function. @@ -1740,156 +1737,6 @@ def fetch_chunk_count_for_document( return db_session.execute(stmt).scalar_one_or_none() -def get_unprocessed_kg_document_batch_for_connector( - db_session: Session, - connector_id: int, - kg_coverage_start: datetime, - kg_max_coverage_days: int, - batch_size: int = 100, -) -> list[DbDocument]: - """ - Retrieves a batch of documents that have not been processed for knowledge graph extraction. - Args: - db_session (Session): The database session to use - connector_id (int): The ID of the connector to get documents for - batch_size (int): The maximum number of documents to retrieve - Returns: - list[DbDocument]: List of documents that need KG processing - """ - - stmt = ( - select(DbDocument) - .join( - DocumentByConnectorCredentialPair, - DbDocument.id == DocumentByConnectorCredentialPair.id, - ) - .where( - and_( - DocumentByConnectorCredentialPair.connector_id == connector_id, - DbDocument.doc_updated_at - >= max( - kg_coverage_start, - datetime.now() - timedelta(days=kg_max_coverage_days), - ), - or_( - DbDocument.kg_stage.is_(None), - DbDocument.kg_stage == KGStage.NOT_STARTED, - DbDocument.doc_updated_at > DbDocument.kg_processing_time, - ), - ) - ) - .distinct() - .limit(batch_size) - ) - - documents = db_session.scalars(stmt).all() - db_session.flush() - - return list(documents) - - -def get_kg_extracted_document_ids(db_session: Session) -> list[str]: - """ - Retrieves all document IDs where kg_stage is EXTRACTED. - Args: - db_session (Session): The database session to use - Returns: - list[str]: List of document IDs that have been KG processed - """ - stmt = select(DbDocument.id).where(DbDocument.kg_stage == KGStage.EXTRACTED) - - return list(db_session.scalars(stmt).all()) - - -def update_document_kg_info( - db_session: Session, document_id: str, kg_stage: KGStage -) -> None: - """Updates the knowledge graph related information for a document. - Args: - db_session (Session): The database session to use - document_id (str): The ID of the document to update - kg_stage (KGStage): The stage of the knowledge graph processing for the document - Raises: - ValueError: If the document with the given ID is not found - """ - stmt = ( - update(DbDocument) - .where(DbDocument.id == document_id) - .values( - kg_stage=kg_stage, - kg_processing_time=datetime.now(timezone.utc), - ) - ) - db_session.execute(stmt) - - -def update_document_kg_stage( - db_session: Session, - document_id: str, - kg_stage: KGStage, -) -> None: - stmt = ( - update(DbDocument).where(DbDocument.id == document_id).values(kg_stage=kg_stage) - ) - db_session.execute(stmt) - db_session.flush() - - -def get_all_kg_extracted_documents_info( - db_session: Session, -) -> list[str]: - """Retrieves the knowledge graph data for all documents that have been processed. - Args: - db_session (Session): The database session to use - Returns: - List[Tuple[str, dict]]: A list of tuples containing: - - str: The document ID - - dict: The KG data containing 'entities', 'relationships', and 'terms' - Only returns documents where kg_stage is EXTRACTED - """ - stmt = ( - select(DbDocument.id) - .where(DbDocument.kg_stage == KGStage.EXTRACTED) - .order_by(DbDocument.id) - ) - - results = db_session.execute(stmt).all() - return [str(doc_id) for doc_id in results] - - -def get_base_llm_doc_information( - db_session: Session, document_ids: list[str] -) -> list[str]: - stmt = select(DbDocument).where(DbDocument.id.in_(document_ids)) - results = db_session.execute(stmt).all() - - documents = [] - - for _doc_nr, doc in enumerate(results): - bare_doc = doc[0] - documents.append( - f"""* [{bare_doc.semantic_id}]({bare_doc.link}) ({bare_doc.doc_updated_at})""" - ) - - return documents[:KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES] - - -def get_document_updated_at( - document_id: str, - db_session: Session, -) -> datetime | None: - """Retrieves the doc_updated_at timestamp for a given document ID. - Args: - document_id (str): The ID of the document to query - db_session (Session): The database session to use - Returns: - Optional[datetime]: The doc_updated_at timestamp if found, None if document doesn't exist - """ - - stmt = select(DbDocument.doc_updated_at).where(DbDocument.id == document_id) - return db_session.execute(stmt).scalar_one_or_none() - - def reset_all_document_kg_stages(db_session: Session) -> int: """Reset the KG stage of all documents that are not in NOT_STARTED state to NOT_STARTED. @@ -1915,211 +1762,6 @@ def reset_all_document_kg_stages(db_session: Session) -> int: ) -def update_document_kg_stages( - db_session: Session, source_stage: KGStage, target_stage: KGStage -) -> int: - """Reset the KG stage only of documents back to NOT_STARTED. - Part of reset flow for documents that have been extracted but not clustered. - - Args: - db_session (Session): The database session to use - - Returns: - int: Number of documents that were reset - """ - stmt = ( - update(DbDocument) - .where(DbDocument.kg_stage == source_stage) - .values(kg_stage=target_stage) - ) - result = db_session.execute(stmt) - # The hasattr check is needed for type checking, even though rowcount - # is guaranteed to exist at runtime for UPDATE operations - return ( - result.rowcount # ty: ignore[invalid-return-type] - if hasattr(result, "rowcount") - else 0 - ) - - -def get_skipped_kg_documents(db_session: Session) -> list[str]: - """ - Retrieves all document IDs where kg_stage is SKIPPED. - Args: - db_session (Session): The database session to use - Returns: - list[str]: List of document IDs that have been skipped in KG processing - """ - stmt = select(DbDocument.id).where(DbDocument.kg_stage == KGStage.SKIPPED) - - return list(db_session.scalars(stmt).all()) - - -# def get_kg_doc_info_for_entity_name( -# db_session: Session, document_id: str, entity_type: str -# ) -> KGEntityDocInfo: -# """ -# Get the semantic ID and the link for an entity name. -# """ - -# result = ( -# db_session.query(Document.semantic_id, Document.link) -# .filter(Document.id == document_id) -# .first() -# ) - -# if result is None: -# return KGEntityDocInfo( -# doc_id=None, -# doc_semantic_id=None, -# doc_link=None, -# semantic_entity_name=f"{entity_type}:{document_id}", -# semantic_linked_entity_name=f"{entity_type}:{document_id}", -# ) - -# return KGEntityDocInfo( -# doc_id=document_id, -# doc_semantic_id=result[0], -# doc_link=result[1], -# semantic_entity_name=f"{entity_type.upper()}:{result[0]}", -# semantic_linked_entity_name=f"[{entity_type.upper()}:{result[0]}]({result[1]})", -# ) - - -def check_for_documents_needing_kg_processing( - db_session: Session, kg_coverage_start: datetime, kg_max_coverage_days: int -) -> bool: - """Check if there are any documents that need KG processing. - - A document needs KG processing if: - 1. It is associated with a connector that has kg_processing_enabled = true - 2. AND either: - - Its kg_stage is NOT_STARTED or NULL - - OR its last_updated timestamp is greater than its kg_processing_time - - Args: - db_session (Session): The database session to use - - Returns: - bool: True if there are any documents needing KG processing, False otherwise - """ - - stmt = ( - select(1) - .select_from(DbDocument) - .join( - DocumentByConnectorCredentialPair, - DbDocument.id == DocumentByConnectorCredentialPair.id, - ) - .join( - Connector, - DocumentByConnectorCredentialPair.connector_id == Connector.id, - ) - .where( - and_( - Connector.kg_processing_enabled.is_(True), - DbDocument.doc_updated_at - >= max( - kg_coverage_start, - datetime.now() - timedelta(days=kg_max_coverage_days), - ), - or_( - DbDocument.kg_stage.is_(None), - DbDocument.kg_stage == KGStage.NOT_STARTED, - DbDocument.doc_updated_at > DbDocument.kg_processing_time, - ), - ) - ) - .exists() - ) - - return db_session.execute(select(stmt)).scalar() or False - - -def check_for_documents_needing_kg_clustering(db_session: Session) -> bool: - """Check if there are any documents that need KG clustering. - - A document needs KG clustering if: - 1. It is associated with a connector that has kg_processing_enabled = true - 2. AND either: - - Its kg_stage is EXTRACTED - - OR its last_updated timestamp is greater than its kg_processing_time - - Args: - db_session (Session): The database session to use - - Returns: - bool: True if there are any documents needing KG clustering, False otherwise - """ - stmt = ( - select(1) - .select_from(DbDocument) - .join( - DocumentByConnectorCredentialPair, - DbDocument.id == DocumentByConnectorCredentialPair.id, - ) - .join( - ConnectorCredentialPair, - and_( - DocumentByConnectorCredentialPair.connector_id - == ConnectorCredentialPair.connector_id, - DocumentByConnectorCredentialPair.credential_id - == ConnectorCredentialPair.credential_id, - ), - ) - .join( - Connector, - ConnectorCredentialPair.connector_id == Connector.id, - ) - .where( - and_( - Connector.kg_processing_enabled.is_(True), - ConnectorCredentialPair.status - != ConnectorCredentialPairStatus.DELETING, - or_( - DbDocument.kg_stage == KGStage.EXTRACTED, - DbDocument.last_modified > DbDocument.kg_processing_time, - ), - ) - ) - .exists() - ) - - return db_session.execute(select(stmt)).scalar() or False - - -def get_document_kg_entities_and_relationships( - db_session: Session, document_id: str -) -> tuple[list[KGEntity], list[KGRelationship]]: - """ - Get the KG entities and relationships that references the document. - """ - entities = ( - db_session.query(KGEntity).filter(KGEntity.document_id == document_id).all() - ) - if not entities: - return [], [] - entity_id_names = [entity.id_name for entity in entities] - - relationships = ( - db_session.query(KGRelationship) - .filter( - or_( - KGRelationship.source_node.in_(entity_id_names), - KGRelationship.target_node.in_(entity_id_names), - KGRelationship.source_document == document_id, - ) - ) - .all() - ) - return entities, relationships - - -def get_num_chunks_for_document(db_session: Session, document_id: str) -> int: - stmt = select(DbDocument.chunk_count).where(DbDocument.id == document_id) - return db_session.execute(stmt).scalar_one_or_none() or 0 - - def update_document_metadata__no_commit( db_session: Session, document_id: str, diff --git a/backend/onyx/db/document_set.py b/backend/onyx/db/document_set.py index 13042d8349..77cf77eba1 100644 --- a/backend/onyx/db/document_set.py +++ b/backend/onyx/db/document_set.py @@ -417,7 +417,7 @@ def update_document_set( ) -> tuple[DocumentSetDBModel, list[DocumentSet__ConnectorCredentialPair]]: """If successful, this sets document_set_row.is_up_to_date = False. That will be processed via Celery in check_for_vespa_sync_task - and trigger a long running background sync to Vespa. + and trigger a long running background sync to the document index. """ # Check if we have either CC pairs or federated connectors (or both) if ( @@ -545,7 +545,7 @@ def mark_document_set_as_to_be_deleted( ) -> None: """Cleans up all document_set -> cc_pair relationships and marks the document set as needing an update. The actual document set row will be deleted by the background - job which syncs these changes to Vespa.""" + job which syncs these changes to the document index.""" try: document_set_row = get_document_set_by_id_for_user( @@ -570,7 +570,7 @@ def mark_document_set_as_to_be_deleted( ) # delete all federated connector mappings so the cleanup task can fully - # remove the document set once the Vespa sync completes + # remove the document set once the document index sync completes delete_stmt = delete(FederatedConnector__DocumentSet).where( FederatedConnector__DocumentSet.document_set_id == document_set_id ) diff --git a/backend/onyx/db/entities.py b/backend/onyx/db/entities.py index 0b4b9ce929..8c285b62ce 100644 --- a/backend/onyx/db/entities.py +++ b/backend/onyx/db/entities.py @@ -1,280 +1,10 @@ -import uuid -from datetime import datetime, timezone -from typing import List +from datetime import datetime -from sqlalchemy import func, literal, select, update -from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy import func 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, KGEntity, KGEntityExtractionStaging, KGEntityType -from onyx.kg.models import KGGroundingType, KGStage -from onyx.kg.utils.formatting_utils import make_entity_id - - -def upsert_staging_entity( - db_session: Session, - name: str, - entity_type: str, - document_id: str | None = None, - occurrences: int = 1, - attributes: dict[str, str] | None = None, - event_time: datetime | None = None, -) -> KGEntityExtractionStaging: - """Add or update a new staging entity to the database. - - Args: - db_session: SQLAlchemy session - name: Name of the entity - entity_type: Type of the entity (must match an existing KGEntityType) - document_id: ID of the document the entity belongs to - occurrences: Number of times this entity has been found - attributes: Attributes of the entity - event_time: Time the entity was added to the database - - Returns: - KGEntityExtractionStaging: The created entity - """ - entity_type = entity_type.upper() - name = name.title() - id_name = make_entity_id(entity_type, name) - attributes = attributes or {} - - entity_key = attributes.get("key") - entity_parent = attributes.get("parent") - - keep_attributes = { - attr_key: attr_val - for attr_key, attr_val in attributes.items() - if attr_key not in ("key", "parent") - } - - # Create new entity - stmt = ( - pg_insert(KGEntityExtractionStaging) - .values( - id_name=id_name, - name=name, - entity_type_id_name=entity_type, - entity_key=entity_key, - parent_key=entity_parent, - document_id=document_id, - occurrences=occurrences, - attributes=keep_attributes, - event_time=event_time, - ) - .on_conflict_do_update( - index_elements=["id_name"], - set_={ - "occurrences": KGEntityExtractionStaging.occurrences + occurrences, - }, - ) - .returning(KGEntityExtractionStaging) - ) - - result = db_session.execute(stmt).scalar() - if result is None: - raise RuntimeError( - f"Failed to create or increment staging entity with id_name: {id_name}" - ) - - # Update the document's kg_stage if document_id is provided - if document_id is not None: - db_session.query(Document).filter(Document.id == document_id).update( - { - "kg_stage": KGStage.EXTRACTED, - "kg_processing_time": datetime.now(timezone.utc), - } - ) - db_session.flush() - - return result - - -def transfer_entity( - db_session: Session, - entity: KGEntityExtractionStaging, -) -> KGEntity: - """Transfer an entity from the extraction staging table to the normalized table. - - Args: - db_session: SQLAlchemy session - entity: Entity to transfer - - Returns: - KGEntity: The transferred entity - """ - # Create the transferred entity - stmt = ( - pg_insert(KGEntity) - .values( - id_name=make_entity_id(entity.entity_type_id_name, uuid.uuid4().hex[:20]), - name=entity.name.casefold(), - entity_key=entity.entity_key, - parent_key=entity.parent_key, - alternative_names=entity.alternative_names or [], - entity_type_id_name=entity.entity_type_id_name, - document_id=entity.document_id, - occurrences=entity.occurrences, - attributes=entity.attributes, - event_time=entity.event_time, - ) - .on_conflict_do_update( - index_elements=["name", "entity_type_id_name", "document_id"], - set_={ - "occurrences": KGEntity.occurrences + entity.occurrences, - "attributes": KGEntity.attributes.op("||")( - literal(entity.attributes, JSONB) - ), - "entity_key": func.coalesce(KGEntity.entity_key, entity.entity_key), - "parent_key": func.coalesce(KGEntity.parent_key, entity.parent_key), - "event_time": entity.event_time, - "time_updated": datetime.now(), - }, - ) - .returning(KGEntity) - ) - new_entity = db_session.execute(stmt).scalar() - if new_entity is None: - raise RuntimeError(f"Failed to transfer entity with id_name: {entity.id_name}") - - # Update the document's kg_stage if document_id is provided - if entity.document_id is not None: - dbdocument.update_document_kg_info( - db_session, - document_id=entity.document_id, - kg_stage=KGStage.NORMALIZED, - ) - - # Update transferred - db_session.query(KGEntityExtractionStaging).filter( - KGEntityExtractionStaging.id_name == entity.id_name - ).update({"transferred_id_name": new_entity.id_name}) - db_session.flush() - - return new_entity - - -def merge_entities( - db_session: Session, parent: KGEntity, child: KGEntityExtractionStaging -) -> KGEntity: - """Merge an entity from the extraction staging table into - an existing entity in the normalized table. - - Args: - db_session: SQLAlchemy session - parent: Parent entity to merge into - child: Child staging entity to merge - - Returns: - KGEntity: The merged entity - """ - # check we're not merging two entities with different document_ids - if ( - parent.document_id is not None - and child.document_id is not None - and parent.document_id != child.document_id - ): - raise ValueError( - "Overwriting the document_id of an entity with a document_id already is not allowed" - ) - - # update the parent entity (only document_id, alternative_names, occurrences) - setting_doc = parent.document_id is None and child.document_id is not None - document_id = child.document_id if setting_doc else parent.document_id - alternative_names = set(parent.alternative_names or []) - alternative_names.update(child.alternative_names or []) - alternative_names.add(child.name.lower()) - alternative_names.discard(parent.name) - - stmt = ( - update(KGEntity) - .where(KGEntity.id_name == parent.id_name) - .values( - document_id=document_id, - alternative_names=list(alternative_names), - occurrences=parent.occurrences + child.occurrences, - attributes=parent.attributes | child.attributes, - entity_key=parent.entity_key or child.entity_key, - parent_key=parent.parent_key or child.parent_key, - ) - .returning(KGEntity) - ) - - result = db_session.execute(stmt).scalar() - if result is None: - raise RuntimeError(f"Failed to merge entities with id_name: {parent.id_name}") - - # Update the document's kg_stage if document_id is set - if setting_doc and child.document_id is not None: - dbdocument.update_document_kg_info( - db_session, - document_id=child.document_id, - kg_stage=KGStage.NORMALIZED, - ) - - # Update transferred - db_session.query(KGEntityExtractionStaging).filter( - KGEntityExtractionStaging.id_name == child.id_name - ).update({"transferred_id_name": parent.id_name}) - db_session.flush() - - return result - - -def get_kg_entity_by_document(db: Session, document_id: str) -> KGEntity | None: - """ - Check if a document_id exists in the kg_entities table and return its id_name if found. - - Args: - db: SQLAlchemy database session - document_id: The document ID to search for - - Returns: - The id_name of the matching KGEntity if found, None otherwise - """ - query = select(KGEntity).where(KGEntity.document_id == document_id) - result = db.execute(query).scalar() - return result - - -def get_grounded_entities_by_types( - db_session: Session, entity_types: List[str], grounding: KGGroundingType -) -> List[KGEntity]: - """Get all entities matching an entity_type. - - Args: - db_session: SQLAlchemy session - entity_types: List of entity types to filter by - - Returns: - List of KGEntity objects belonging to the specified entity types - """ - return ( - db_session.query(KGEntity) - .join(KGEntityType, KGEntity.entity_type_id_name == KGEntityType.id_name) - .filter(KGEntity.entity_type_id_name.in_(entity_types)) - .filter(KGEntityType.grounding == grounding) - .all() - ) - - -def get_document_id_for_entity(db_session: Session, entity_id_name: str) -> str | None: - """Get the document ID associated with an entity. - - Args: - db_session: SQLAlchemy database session - entity_id_name: The entity id_name to look up - - Returns: - The document ID if found, None otherwise - """ - entity = ( - db_session.query(KGEntity).filter(KGEntity.id_name == entity_id_name).first() - ) - return entity.document_id if entity else None +from onyx.db.models import KGEntity, KGEntityExtractionStaging, KGEntityType def delete_from_kg_entities_extraction_staging__no_commit( @@ -295,14 +25,6 @@ def delete_from_kg_entities__no_commit( ) -def get_entity_name(db_session: Session, entity_id_name: str) -> str | None: - """Get the name of an entity.""" - entity = ( - db_session.query(KGEntity).filter(KGEntity.id_name == entity_id_name).first() - ) - return entity.name if entity else None - - def get_entity_stats_by_grounded_source_name( db_session: Session, ) -> dict[str, tuple[datetime, int]]: diff --git a/backend/onyx/db/entity_type.py b/backend/onyx/db/entity_type.py index d902c8e7d3..17a22933c8 100644 --- a/backend/onyx/db/entity_type.py +++ b/backend/onyx/db/entity_type.py @@ -12,42 +12,6 @@ from onyx.server.kg.models import EntityType UNGROUNDED_SOURCE_NAME = "Ungrounded" -def get_entity_types_with_grounded_source_name( - db_session: Session, -) -> list[KGEntityType]: - """Get all entity types that have non-null grounded_source_name. - - Args: - db_session: SQLAlchemy session - - Returns: - List of KGEntityType objects that have grounded_source_name defined - """ - return ( - db_session.query(KGEntityType) - .filter(KGEntityType.grounded_source_name.isnot(None)) - .all() - ) - - -def get_entity_types( - db_session: Session, - active: bool | None = True, -) -> list[KGEntityType]: - # Query the database for all distinct entity types - - if active is None: - return db_session.query(KGEntityType).order_by(KGEntityType.id_name).all() - - else: - return ( - db_session.query(KGEntityType) - .filter(KGEntityType.active == active) - .order_by(KGEntityType.id_name) - .all() - ) - - def get_configured_entity_types(db_session: Session) -> dict[str, list[KGEntityType]]: # get entity types from configured sources configured_connector_sources = { diff --git a/backend/onyx/db/enums.py b/backend/onyx/db/enums.py index ab9dcfa402..aa565c61ef 100644 --- a/backend/onyx/db/enums.py +++ b/backend/onyx/db/enums.py @@ -376,22 +376,6 @@ class SwitchoverType(str, PyEnum): INSTANT = "instant" -class OpenSearchDocumentMigrationStatus(str, PyEnum): - """Status for Vespa to OpenSearch migration per document.""" - - PENDING = "pending" - COMPLETED = "completed" - FAILED = "failed" - PERMANENTLY_FAILED = "permanently_failed" - - -class OpenSearchTenantMigrationStatus(str, PyEnum): - """Status for tenant-level OpenSearch migration.""" - - PENDING = "pending" - COMPLETED = "completed" - - # Onyx Build Mode Enums class BuildSessionStatus(str, PyEnum): """Lifecycle of a build session. diff --git a/backend/onyx/db/feedback.py b/backend/onyx/db/feedback.py index bbd70e7db0..59ae11c26a 100644 --- a/backend/onyx/db/feedback.py +++ b/backend/onyx/db/feedback.py @@ -166,7 +166,7 @@ def create_doc_retrieval_feedback( clicked: bool = False, feedback: SearchFeedbackType | None = None, ) -> None: - """Creates a new Document feedback row and updates the boost value in Postgres and Vespa""" + """Creates a new Document feedback row and updates the boost value in Postgres and the document index""" db_doc = _fetch_db_doc_by_id(document_id, db_session) retrieval_feedback = DocumentRetrievalFeedback( diff --git a/backend/onyx/db/kg_config.py b/backend/onyx/db/kg_config.py index 2d5266f034..6cac3d39f2 100644 --- a/backend/onyx/db/kg_config.py +++ b/backend/onyx/db/kg_config.py @@ -37,14 +37,6 @@ def validate_kg_settings(kg_config_settings: KGConfigSettings) -> None: raise ValueError("KG_VENDOR_DOMAINS is not set") -def is_kg_config_settings_enabled_valid(kg_config_settings: KGConfigSettings) -> bool: - try: - validate_kg_settings(kg_config_settings) - return True - except Exception: - return False - - def enable_kg(enable_req: EnableKGConfigRequest) -> None: kg_config_settings = get_kg_config_settings() kg_config_settings.KG_ENABLED = True diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index 9b181e25c0..09334ca9cb 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -96,8 +96,6 @@ from onyx.db.enums import ( MCPServerStatus, MCPTransport, NotificationSeverity, - OpenSearchDocumentMigrationStatus, - OpenSearchTenantMigrationStatus, PatType, Permission, PermissionSyncStatus, @@ -983,7 +981,7 @@ class ConnectorCredentialPair(Base): ) # Determines how documents are processed after fetching: - # REGULAR: Full pipeline (chunk → embed → Vespa) + # REGULAR: Full pipeline (chunk → embed → document index) # FILE_SYSTEM: Write to file system only (for CLI agent sandbox) processing_mode: Mapped[ProcessingMode] = mapped_column( Enum(ProcessingMode, native_enum=False), @@ -1158,7 +1156,7 @@ class Document(Base): DateTime(timezone=True), nullable=True ) - # Number of chunks in the document (in Vespa) + # Number of chunks in the document (in the document index) # Only null for documents indexed prior to this change chunk_count: Mapped[int | None] = mapped_column(Integer, nullable=True) @@ -1167,13 +1165,13 @@ class Document(Base): # Null for documents indexed before this column was added. content_hash: Mapped[str | None] = mapped_column(String, nullable=True) - # last time any vespa relevant row metadata or the doc changed. + # last time any document-index-relevant row metadata or the doc changed. # does not include last_synced last_modified: Mapped[datetime.datetime | None] = mapped_column( DateTime(timezone=True), nullable=False, index=True, default=func.now() ) - # last successful sync to vespa + # last successful sync to the document index last_synced: Mapped[datetime.datetime | None] = mapped_column( DateTime(timezone=True), nullable=True, index=True ) @@ -1273,121 +1271,6 @@ class Document(Base): ) -class OpenSearchDocumentMigrationRecord(Base): - """Tracks the migration status of documents from Vespa to OpenSearch. - - This table can be dropped when the migration is complete for all Onyx - instances. - """ - - __tablename__ = "opensearch_document_migration_record" - - document_id: Mapped[str] = mapped_column( - String, - ForeignKey("document.id", ondelete="CASCADE"), - primary_key=True, - nullable=False, - index=True, - ) - status: Mapped[OpenSearchDocumentMigrationStatus] = mapped_column( - Enum(OpenSearchDocumentMigrationStatus, native_enum=False), - default=OpenSearchDocumentMigrationStatus.PENDING, - nullable=False, - index=True, - ) - error_message: Mapped[str | None] = mapped_column(Text, nullable=True) - attempts_count: Mapped[int] = mapped_column( - Integer, default=0, nullable=False, index=True - ) - last_attempt_at: Mapped[datetime.datetime | None] = mapped_column( - DateTime(timezone=True), nullable=True - ) - created_at: Mapped[datetime.datetime] = mapped_column( - DateTime(timezone=True), - server_default=func.now(), - nullable=False, - index=True, - ) - - document: Mapped["Document"] = relationship("Document") - - -class OpenSearchTenantMigrationRecord(Base): - """Tracks the state of the OpenSearch migration for a tenant. - - Should only contain one row. - - This table can be dropped when the migration is complete for all Onyx - instances. - """ - - __tablename__ = "opensearch_tenant_migration_record" - __table_args__ = ( - # Singleton pattern - unique index on constant ensures only one row. - Index("idx_opensearch_tenant_migration_singleton", text("(true)"), unique=True), - ) - - id: Mapped[int] = mapped_column(primary_key=True, nullable=False) - document_migration_record_table_population_status: Mapped[ - OpenSearchTenantMigrationStatus - ] = mapped_column( - Enum(OpenSearchTenantMigrationStatus, native_enum=False), - default=OpenSearchTenantMigrationStatus.PENDING, - nullable=False, - ) - num_times_observed_no_additional_docs_to_populate_migration_table: Mapped[int] = ( - mapped_column(Integer, default=0, nullable=False) - ) - overall_document_migration_status: Mapped[OpenSearchTenantMigrationStatus] = ( - mapped_column( - Enum(OpenSearchTenantMigrationStatus, native_enum=False), - default=OpenSearchTenantMigrationStatus.PENDING, - nullable=False, - ) - ) - num_times_observed_no_additional_docs_to_migrate: Mapped[int] = mapped_column( - Integer, - default=0, - nullable=False, - ) - last_updated_at: Mapped[datetime.datetime] = mapped_column( - DateTime(timezone=True), - server_default=func.now(), - onupdate=func.now(), - nullable=False, - ) - # Opaque continuation token from Vespa's Visit API. - # NULL means "not started". - # Otherwise contains a serialized mapping between slice ID and continuation - # token for that slice. - vespa_visit_continuation_token: Mapped[str | None] = mapped_column( - Text, nullable=True - ) - total_chunks_migrated: Mapped[int] = mapped_column( - Integer, default=0, nullable=False - ) - total_chunks_errored: Mapped[int] = mapped_column( - Integer, default=0, nullable=False - ) - total_chunks_in_vespa: Mapped[int] = mapped_column( - Integer, default=0, nullable=False - ) - created_at: Mapped[datetime.datetime] = mapped_column( - DateTime(timezone=True), - server_default=func.now(), - nullable=False, - ) - migration_completed_at: Mapped[datetime.datetime | None] = mapped_column( - DateTime(timezone=True), nullable=True - ) - enable_opensearch_retrieval: Mapped[bool] = mapped_column( - Boolean, nullable=False, default=False - ) - approx_chunk_count_in_vespa: Mapped[int | None] = mapped_column( - Integer, nullable=True - ) - - class KGEntityType(Base): __tablename__ = "kg_entity_type" @@ -5233,7 +5116,8 @@ class UserGroup(Base): id: Mapped[int] = mapped_column(primary_key=True) name: Mapped[str] = mapped_column(String, unique=True) - # whether or not changes to the UserGroup have been propagated to Vespa + # whether or not changes to the UserGroup have been propagated to the + # document index is_up_to_date: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) # tell the sync job to clean up the group is_up_for_deletion: Mapped[bool] = mapped_column( diff --git a/backend/onyx/db/opensearch_migration.py b/backend/onyx/db/opensearch_migration.py deleted file mode 100644 index 86726e25f7..0000000000 --- a/backend/onyx/db/opensearch_migration.py +++ /dev/null @@ -1,440 +0,0 @@ -"""Database operations for OpenSearch migration tracking. - -This module provides functions to track the progress of migrating documents -from Vespa to OpenSearch. -""" - -import json -from datetime import datetime, timezone - -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, - TOTAL_ALLOWABLE_DOC_MIGRATION_ATTEMPTS_BEFORE_PERMANENT_FAILURE, -) -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, - OpenSearchDocumentMigrationRecord, - OpenSearchTenantMigrationRecord, -) -from onyx.document_index.vespa.shared_utils.utils import ( - replace_invalid_doc_id_characters, -) -from onyx.utils.logger import setup_logger - -logger = setup_logger() - - -def get_paginated_document_batch( - db_session: Session, - limit: int, - prev_ending_document_id: str | None = None, -) -> list[str]: - """Gets a paginated batch of document IDs from the Document table. - - We need some deterministic ordering to ensure that we don't miss any - documents when paginating. This function uses the document ID. It is - possible a document is inserted above a spot this function has already - passed. In that event we assume that the document will be indexed into - OpenSearch anyway and we don't need to migrate. - TODO(andrei): Consider ordering on last_modified in addition to ID to better - match get_opensearch_migration_records_needing_migration. - - Args: - db_session: SQLAlchemy session. - limit: Number of document IDs to fetch. - prev_ending_document_id: Document ID to start after (for pagination). If - None, returns the first batch of documents. If not None, this should - be the last ordered ID which was fetched in a previous batch. - Defaults to None. - - Returns: - List of document IDs. - """ - stmt = select(Document.id).order_by(Document.id.asc()).limit(limit) - if prev_ending_document_id is not None: - stmt = stmt.where(Document.id > prev_ending_document_id) - return list(db_session.scalars(stmt).all()) - - -def get_last_opensearch_migration_document_id( - db_session: Session, -) -> str | None: - """ - Gets the last document ID in the OpenSearchDocumentMigrationRecord table. - - Returns None if no records are found. - """ - stmt = ( - select(OpenSearchDocumentMigrationRecord.document_id) - .order_by(OpenSearchDocumentMigrationRecord.document_id.desc()) - .limit(1) - ) - return db_session.scalars(stmt).first() - - -def create_opensearch_migration_records_with_commit( - db_session: Session, - document_ids: list[str], -) -> None: - """Creates new OpenSearchDocumentMigrationRecord records. - - Silently skips any document IDs that already have records. - """ - if not document_ids: - return - - values = [ - { - "document_id": document_id, - "status": OpenSearchDocumentMigrationStatus.PENDING, - } - for document_id in document_ids - ] - - stmt = insert(OpenSearchDocumentMigrationRecord).values(values) - stmt = stmt.on_conflict_do_nothing(index_elements=["document_id"]) - - db_session.execute(stmt) - db_session.commit() - - -def get_opensearch_migration_records_needing_migration( - db_session: Session, - limit: int, -) -> list[OpenSearchDocumentMigrationRecord]: - """Gets records of documents that need to be migrated. - - Properties: - - First tries documents with status PENDING. - - Of these, orders documents with the oldest last_modified to prioritize - documents that were modified a long time ago, as they are presumed to be - stable. This column is modified in many flows so is not a guarantee of the - document having been indexed. - - Then if there's room in the result, tries documents with status FAILED. - - Of these, first orders documents on the least attempts_count so as to have - a backoff for recently-failed docs. Then orders on last_modified as - before. - """ - result: list[OpenSearchDocumentMigrationRecord] = [] - - # Step 1: Fetch as many PENDING status records as possible ordered by - # last_modified (oldest first). last_modified lives on Document, so we join. - stmt_pending = ( - select(OpenSearchDocumentMigrationRecord) - .join(Document, OpenSearchDocumentMigrationRecord.document_id == Document.id) - .where( - OpenSearchDocumentMigrationRecord.status - == OpenSearchDocumentMigrationStatus.PENDING - ) - .order_by(Document.last_modified.asc()) - .limit(limit) - ) - result.extend(list(db_session.scalars(stmt_pending).all())) - remaining = limit - len(result) - - # Step 2: If more are needed, fetch records with status FAILED, ordered by - # attempts_count (lowest first), then last_modified (oldest first). - if remaining > 0: - stmt_failed = ( - select(OpenSearchDocumentMigrationRecord) - .join( - Document, - OpenSearchDocumentMigrationRecord.document_id == Document.id, - ) - .where( - OpenSearchDocumentMigrationRecord.status - == OpenSearchDocumentMigrationStatus.FAILED - ) - .order_by( - OpenSearchDocumentMigrationRecord.attempts_count.asc(), - Document.last_modified.asc(), - ) - .limit(remaining) - ) - result.extend(list(db_session.scalars(stmt_failed).all())) - - return result - - -def get_total_opensearch_migration_record_count( - db_session: Session, -) -> int: - """Gets the total number of OpenSearch migration records. - - Used to check whether every document has been tracked for migration. - """ - return db_session.query(OpenSearchDocumentMigrationRecord).count() - - -def get_total_document_count(db_session: Session) -> int: - """Gets the total number of documents. - - Used to check whether every document has been tracked for migration. - """ - return db_session.query(Document).count() - - -def try_insert_opensearch_tenant_migration_record_with_commit( - db_session: Session, -) -> None: - """Tries to insert the singleton row on OpenSearchTenantMigrationRecord. - - Does nothing if the row already exists. - """ - stmt = insert(OpenSearchTenantMigrationRecord).on_conflict_do_nothing( - index_elements=[text("(true)")] - ) - db_session.execute(stmt) - db_session.commit() - - -def increment_num_times_observed_no_additional_docs_to_migrate_with_commit( - db_session: Session, -) -> None: - """Increments the number of times observed no additional docs to migrate. - - Requires the OpenSearchTenantMigrationRecord to exist. - - Used to track when to stop the migration task. - """ - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - raise RuntimeError("OpenSearchTenantMigrationRecord not found.") - record.num_times_observed_no_additional_docs_to_migrate += 1 - db_session.commit() - - -def increment_num_times_observed_no_additional_docs_to_populate_migration_table_with_commit( - db_session: Session, -) -> None: - """ - Increments the number of times observed no additional docs to populate the - migration table. - - Requires the OpenSearchTenantMigrationRecord to exist. - - Used to track when to stop the migration check task. - """ - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - raise RuntimeError("OpenSearchTenantMigrationRecord not found.") - record.num_times_observed_no_additional_docs_to_populate_migration_table += 1 - db_session.commit() - - -def should_document_migration_be_permanently_failed( - opensearch_document_migration_record: OpenSearchDocumentMigrationRecord, -) -> bool: - return ( - opensearch_document_migration_record.status - == OpenSearchDocumentMigrationStatus.PERMANENTLY_FAILED - or ( - opensearch_document_migration_record.status - == OpenSearchDocumentMigrationStatus.FAILED - and opensearch_document_migration_record.attempts_count - >= TOTAL_ALLOWABLE_DOC_MIGRATION_ATTEMPTS_BEFORE_PERMANENT_FAILURE - ) - ) - - -def get_vespa_visit_state( - db_session: Session, -) -> tuple[dict[int, str | None], int]: - """Gets the current Vespa migration state from the tenant migration record. - - Requires the OpenSearchTenantMigrationRecord to exist. - - Returns: - Tuple of (continuation_token_map, total_chunks_migrated). - """ - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - raise RuntimeError("OpenSearchTenantMigrationRecord not found.") - if record.vespa_visit_continuation_token is None: - continuation_token_map: dict[int, str | None] = dict.fromkeys( - range(GET_VESPA_CHUNKS_SLICE_COUNT) - ) - else: - json_loaded_continuation_token_map = json.loads( - record.vespa_visit_continuation_token - ) - continuation_token_map = { - int(key): value for key, value in json_loaded_continuation_token_map.items() - } - return continuation_token_map, record.total_chunks_migrated - - -def update_vespa_visit_progress_with_commit( - db_session: Session, - continuation_token_map: dict[int, str | None], - chunks_processed: int, - chunks_errored: int, - approx_chunk_count_in_vespa: int | None, -) -> None: - """Updates the Vespa migration progress and commits. - - Requires the OpenSearchTenantMigrationRecord to exist. - - Args: - db_session: SQLAlchemy session. - continuation_token_map: The new continuation token map. None entry means - the visit is complete for that slice. - chunks_processed: Number of chunks processed in this batch (added to - the running total). - chunks_errored: Number of chunks errored in this batch (added to the - running errored total). - approx_chunk_count_in_vespa: Approximate number of chunks in Vespa. If - None, the existing value is used. - """ - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - raise RuntimeError("OpenSearchTenantMigrationRecord not found.") - record.vespa_visit_continuation_token = json.dumps(continuation_token_map) - record.total_chunks_migrated += chunks_processed - record.total_chunks_errored += chunks_errored - record.approx_chunk_count_in_vespa = ( - approx_chunk_count_in_vespa - if approx_chunk_count_in_vespa is not None - else record.approx_chunk_count_in_vespa - ) - db_session.commit() - - -def mark_migration_completed_time_if_not_set_with_commit( - db_session: Session, -) -> None: - """Marks the migration completed time if not set. - - Requires the OpenSearchTenantMigrationRecord to exist. - """ - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - raise RuntimeError("OpenSearchTenantMigrationRecord not found.") - if record.migration_completed_at is not None: - return - record.migration_completed_at = datetime.now(timezone.utc) - db_session.commit() - - -def is_migration_completed(db_session: Session) -> bool: - """Returns True if the migration is completed. - - Can be run even if the migration record does not exist. - """ - record = db_session.query(OpenSearchTenantMigrationRecord).first() - return record is not None and record.migration_completed_at is not None - - -def build_sanitized_to_original_doc_id_mapping( - db_session: Session, -) -> dict[str, str]: - """Pre-computes a mapping of sanitized -> original document IDs. - - Only includes documents whose ID contains single quotes (the only character - that gets sanitized by replace_invalid_doc_id_characters). For all other - documents, sanitized == original and no mapping entry is needed. - - Scans over all documents. - - Checks if the sanitized ID already exists as a genuine separate document in - the Document table. If so, raises as there is no way of resolving the - conflict in the migration. The user will need to reindex. - - Args: - db_session: SQLAlchemy session. - - Returns: - Dict mapping sanitized_id -> original_id, only for documents where - the IDs differ. Empty dict means no documents have single quotes - in their IDs. - """ - # Find all documents with single quotes in their ID. - stmt = select(Document.id).where(Document.id.contains("'")) - ids_with_quotes = list(db_session.scalars(stmt).all()) - - result: dict[str, str] = {} - for original_id in ids_with_quotes: - sanitized_id = replace_invalid_doc_id_characters(original_id) - if sanitized_id != original_id: - result[sanitized_id] = original_id - - # See if there are any documents whose ID is a sanitized ID of another - # document. If there is even one match, we cannot proceed. - stmt = select(Document.id).where(Document.id.in_(result.keys())) - ids_with_matches = list(db_session.scalars(stmt).all()) - if ids_with_matches: - raise RuntimeError( - f"Documents with IDs {ids_with_matches} have sanitized IDs that match other documents. " - "This is not supported and the user will need to reindex." - ) - - return result - - -def get_opensearch_migration_state( - db_session: Session, -) -> tuple[int, datetime | None, datetime | None, int | None]: - """Returns the state of the Vespa to OpenSearch migration. - - If the tenant migration record is not found, returns defaults of 0, None, - None, None. - - Args: - db_session: SQLAlchemy session. - - Returns: - Tuple of (total_chunks_migrated, created_at, migration_completed_at, - approx_chunk_count_in_vespa). - """ - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - return 0, None, None, None - return ( - record.total_chunks_migrated, - record.created_at, - record.migration_completed_at, - record.approx_chunk_count_in_vespa, - ) - - -def get_opensearch_retrieval_state( - db_session: Session, -) -> bool: - """Returns the state of the OpenSearch retrieval. - - If the tenant migration record is not found, defaults to - ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX. - - If ONYX_DISABLE_VESPA is True, always returns True. - """ - if ONYX_DISABLE_VESPA: - return True - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - return ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX - return record.enable_opensearch_retrieval - - -def set_enable_opensearch_retrieval_with_commit( - db_session: Session, - enable: bool, -) -> None: - """Sets the enable_opensearch_retrieval flag on the singleton record. - - Creates the record if it doesn't exist yet. - """ - try_insert_opensearch_tenant_migration_record_with_commit(db_session) - record = db_session.query(OpenSearchTenantMigrationRecord).first() - if record is None: - raise RuntimeError("OpenSearchTenantMigrationRecord not found.") - record.enable_opensearch_retrieval = enable - db_session.commit() diff --git a/backend/onyx/db/port_attempt.py b/backend/onyx/db/port_attempt.py index fb6ca27a28..06ad86430f 100644 --- a/backend/onyx/db/port_attempt.py +++ b/backend/onyx/db/port_attempt.py @@ -578,7 +578,7 @@ def count_consecutive_failed_port_attempts_no_progress( def any_future_port_in_progress(db_session: Session) -> bool: """True if any PortAttempt against a FUTURE SearchSettings is active - (NOT_STARTED / IN_PROGRESS). The vespa sync producer drops deferred-doc syncs + (NOT_STARTED / IN_PROGRESS). The document index sync producer drops deferred-doc syncs to LOW priority while a port runs so they don't starve normal needs_sync work.""" stmt = select( exists() diff --git a/backend/onyx/db/relationships.py b/backend/onyx/db/relationships.py index 72c1de73c1..99f4458344 100644 --- a/backend/onyx/db/relationships.py +++ b/backend/onyx/db/relationships.py @@ -1,579 +1,16 @@ -from typing import List - -from sqlalchemy import or_ -from sqlalchemy.dialects import postgresql -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, 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() -def upsert_staging_relationship( - db_session: Session, - relationship_id_name: str, - source_document_id: str | None, - occurrences: int = 1, -) -> KGRelationshipExtractionStaging: - """ - Add or update a new staging relationship to the database. - - Args: - db_session: SQLAlchemy database session - relationship_id_name: The ID name of the relationship in format "source__relationship__target" - source_document_id: ID of the source document - occurrences: Number of times this relationship has been found - Returns: - The created or updated KGRelationshipExtractionStaging object - - Raises: - sqlalchemy.exc.IntegrityError: If there's an error with the database operation - """ - # Generate a unique ID for the relationship - relationship_id_name = format_relationship_id(relationship_id_name) - ( - source_entity_id_name, - relationship_string, - target_entity_id_name, - ) = split_relationship_id(relationship_id_name) - - source_entity_type = get_entity_type(source_entity_id_name) - target_entity_type = get_entity_type(target_entity_id_name) - relationship_type = extract_relationship_type_id(relationship_id_name) - - # Insert the new relationship - stmt = ( - postgresql.insert(KGRelationshipExtractionStaging) - .values( - { - "id_name": relationship_id_name, - "source_node": source_entity_id_name, - "target_node": target_entity_id_name, - "source_node_type": source_entity_type, - "target_node_type": target_entity_type, - "type": relationship_string.lower(), - "relationship_type_id_name": relationship_type, - "source_document": source_document_id, - "occurrences": occurrences, - } - ) - .on_conflict_do_update( - index_elements=["id_name", "source_document"], - set_={ - "occurrences": KGRelationshipExtractionStaging.occurrences - + occurrences, - }, - ) - .returning(KGRelationshipExtractionStaging) - ) - - result = db_session.execute(stmt).scalar() - if result is None: - raise RuntimeError( - f"Failed to create or increment staging relationship with id_name: {relationship_id_name}" - ) - - # Update the document's kg_stage if source_document is provided - if source_document_id is not None: - dbdocument.update_document_kg_info( - db_session, - document_id=source_document_id, - kg_stage=KGStage.EXTRACTED, - ) - db_session.flush() # Flush to get any DB errors early - - return result - - -def upsert_relationship( - db_session: Session, - relationship_id_name: str, - source_document_id: str | None, - occurrences: int = 1, -) -> KGRelationship: - """ - Upsert a new relationship directly to the database. - - Args: - db_session: SQLAlchemy database session - relationship_id_name: The ID name of the relationship in format "source__relationship__target" - source_document_id: ID of the source document - occurrences: Number of times this relationship has been found - Returns: - The created or updated KGRelationship object - - Raises: - sqlalchemy.exc.IntegrityError: If there's an error with the database operation - """ - # Generate a unique ID for the relationship - relationship_id_name = format_relationship_id(relationship_id_name) - ( - source_entity_id_name, - relationship_string, - target_entity_id_name, - ) = split_relationship_id(relationship_id_name) - - source_entity_type = get_entity_type(source_entity_id_name) - target_entity_type = get_entity_type(target_entity_id_name) - relationship_type = extract_relationship_type_id(relationship_id_name) - - # Insert the new relationship - stmt = ( - postgresql.insert(KGRelationship) - .values( - { - "id_name": relationship_id_name, - "source_node": source_entity_id_name, - "target_node": target_entity_id_name, - "source_node_type": source_entity_type, - "target_node_type": target_entity_type, - "type": relationship_string.lower(), - "relationship_type_id_name": relationship_type, - "source_document": source_document_id, - "occurrences": occurrences, - } - ) - .on_conflict_do_update( - index_elements=["id_name", "source_document"], - set_={ - "occurrences": KGRelationship.occurrences + occurrences, - }, - ) - .returning(KGRelationship) - ) - - new_relationship = db_session.execute(stmt).scalar() - if new_relationship is None: - raise RuntimeError( - f"Failed to upsert relationship with id_name: {relationship_id_name}" - ) - db_session.flush() - return new_relationship - - -def transfer_relationship( - db_session: Session, - relationship: KGRelationshipExtractionStaging, - entity_translations: dict[str, str], -) -> KGRelationship: - """ - Transfer a relationship from the staging table to the normalized table. - """ - # Translate the source and target nodes - source_node = entity_translations[relationship.source_node] - target_node = entity_translations[relationship.target_node] - relationship_id_name = make_relationship_id( - source_node, relationship.type, target_node - ) - - # Create the transferred relationship - stmt = ( - pg_insert(KGRelationship) - .values( - id_name=relationship_id_name, - source_node=source_node, - target_node=target_node, - source_node_type=relationship.source_node_type, - target_node_type=relationship.target_node_type, - type=relationship.type, - relationship_type_id_name=relationship.relationship_type_id_name, - source_document=relationship.source_document, - occurrences=relationship.occurrences, - ) - .on_conflict_do_update( - index_elements=["id_name", "source_document"], - set_={ - "occurrences": KGRelationship.occurrences + relationship.occurrences, - }, - ) - .returning(KGRelationship) - ) - - new_relationship = db_session.execute(stmt).scalar() - if new_relationship is None: - raise RuntimeError( - f"Failed to transfer relationship with id_name: {relationship.id_name}" - ) - - # Update transferred - db_session.query(KGRelationshipExtractionStaging).filter( - KGRelationshipExtractionStaging.id_name == relationship.id_name, - KGRelationshipExtractionStaging.source_document == relationship.source_document, - ).update({"transferred": True}) - db_session.flush() - - return new_relationship - - -def upsert_staging_relationship_type( - db_session: Session, - source_entity_type: str, - relationship_type: str, - target_entity_type: str, - definition: bool = False, - extraction_count: int = 1, -) -> KGRelationshipTypeExtractionStaging: - """ - Add a new relationship type to the database. - - Args: - db_session: SQLAlchemy session - source_entity_type: Type of the source entity - relationship_type: Type of relationship - target_entity_type: Type of the target entity - definition: Whether this relationship type represents a definition (default False) - - Returns: - The created KGRelationshipTypeExtractionStaging object - """ - - id_name = make_relationship_type_id( - source_entity_type, relationship_type, target_entity_type - ) - - # Create new relationship type - stmt = ( - postgresql.insert(KGRelationshipTypeExtractionStaging) - .values( - { - "id_name": id_name, - "name": relationship_type, - "source_entity_type_id_name": source_entity_type.upper(), - "target_entity_type_id_name": target_entity_type.upper(), - "definition": definition, - "occurrences": extraction_count, - "type": relationship_type, # Using the relationship_type as the type - "active": True, # Setting as active by default - } - ) - .on_conflict_do_update( - index_elements=["id_name"], - set_={ - "occurrences": KGRelationshipTypeExtractionStaging.occurrences - + extraction_count, - }, - ) - .returning(KGRelationshipTypeExtractionStaging) - ) - - result = db_session.execute(stmt).scalar() - if result is None: - raise RuntimeError( - f"Failed to create or increment staging relationship type with id_name: {id_name}" - ) - db_session.flush() # Flush to get any DB errors early - - return result - - -def upsert_relationship_type( - db_session: Session, - source_entity_type: str, - relationship_type: str, - target_entity_type: str, - definition: bool = False, - extraction_count: int = 1, -) -> KGRelationshipType: - """ - Upsert a new relationship type directly to the database. - - Args: - db_session: SQLAlchemy session - source_entity_type: Type of the source entity - relationship_type: Type of relationship - target_entity_type: Type of the target entity - definition: Whether this relationship type represents a definition (default False) - - Returns: - The created KGRelationshipType object - """ - - id_name = make_relationship_type_id( - source_entity_type, relationship_type, target_entity_type - ) - - # Create new relationship type - stmt = ( - postgresql.insert(KGRelationshipType) - .values( - { - "id_name": id_name, - "name": relationship_type, - "source_entity_type_id_name": source_entity_type.upper(), - "target_entity_type_id_name": target_entity_type.upper(), - "definition": definition, - "occurrences": extraction_count, - "type": relationship_type, # Using the relationship_type as the type - "active": True, # Setting as active by default - } - ) - .on_conflict_do_update( - index_elements=["id_name"], - set_={ - "occurrences": KGRelationshipType.occurrences + extraction_count, - }, - ) - .returning(KGRelationshipType) - ) - - new_relationship_type = db_session.execute(stmt).scalar() - if new_relationship_type is None: - raise RuntimeError( - f"Failed to upsert relationship type with id_name: {id_name}" - ) - db_session.flush() - return new_relationship_type - - -def transfer_relationship_type( - db_session: Session, - relationship_type: KGRelationshipTypeExtractionStaging, -) -> KGRelationshipType: - """ - Transfer a relationship type from the staging table to the normalized table. - """ - stmt = ( - pg_insert(KGRelationshipType) - .values( - id_name=relationship_type.id_name, - name=relationship_type.name, - source_entity_type_id_name=relationship_type.source_entity_type_id_name, - target_entity_type_id_name=relationship_type.target_entity_type_id_name, - definition=relationship_type.definition, - occurrences=relationship_type.occurrences, - type=relationship_type.type, - active=relationship_type.active, - ) - .on_conflict_do_update( - index_elements=["id_name"], - set_={ - "occurrences": KGRelationshipType.occurrences - + relationship_type.occurrences, - }, - ) - .returning(KGRelationshipType) - ) - - new_relationship_type = db_session.execute(stmt).scalar() - if new_relationship_type is None: - raise RuntimeError( - f"Failed to transfer relationship type with id_name: {relationship_type.id_name}" - ) - - # Update transferred - db_session.query(KGRelationshipTypeExtractionStaging).filter( - KGRelationshipTypeExtractionStaging.id_name == relationship_type.id_name - ).update({"transferred": True}) - db_session.flush() - - return new_relationship_type - - -def delete_relationships_by_id_names( - db_session: Session, id_names: list[str], kg_stage: KGStage -) -> int: - """ - Delete relationships from the database based on a list of id_names. - - Args: - db_session: SQLAlchemy database session - id_names: List of relationship id_names to delete - - Returns: - Number of relationships deleted - - Raises: - sqlalchemy.exc.SQLAlchemyError: If there's an error during deletion - """ - - deleted_count = 0 - - if kg_stage == KGStage.EXTRACTED: - deleted_count = ( - db_session.query(KGRelationshipExtractionStaging) - .filter(KGRelationshipExtractionStaging.id_name.in_(id_names)) - .delete(synchronize_session=False) - ) - elif kg_stage == KGStage.NORMALIZED: - deleted_count = ( - db_session.query(KGRelationship) - .filter(KGRelationship.id_name.in_(id_names)) - .delete(synchronize_session=False) - ) - - db_session.flush() # Flush to ensure deletion is processed - return deleted_count - - -def delete_relationship_types_by_id_names( - db_session: Session, id_names: list[str], kg_stage: KGStage -) -> int: - """ - Delete relationship types from the database based on a list of id_names. - - Args: - db_session: SQLAlchemy database session - id_names: List of relationship type id_names to delete - - Returns: - Number of relationship types deleted - - Raises: - sqlalchemy.exc.SQLAlchemyError: If there's an error during deletion - """ - deleted_count = 0 - - if kg_stage == KGStage.EXTRACTED: - deleted_count = ( - db_session.query(KGRelationshipTypeExtractionStaging) - .filter(KGRelationshipTypeExtractionStaging.id_name.in_(id_names)) - .delete(synchronize_session=False) - ) - elif kg_stage == KGStage.NORMALIZED: - deleted_count = ( - db_session.query(KGRelationshipType) - .filter(KGRelationshipType.id_name.in_(id_names)) - .delete(synchronize_session=False) - ) - - db_session.flush() # Flush to ensure deletion is processed - return deleted_count - - -def get_relationships_for_entity_type_pairs( - db_session: Session, entity_type_pairs: list[tuple[str, str]] -) -> list["KGRelationshipType"]: - """ - Get relationship types from the database based on a list of entity type pairs. - - Args: - db_session: SQLAlchemy database session - entity_type_pairs: List of tuples where each tuple contains (source_entity_type, target_entity_type) - - Returns: - List of KGRelationshipType objects where source and target types match the provided pairs - """ - - conditions = [ - ( - (KGRelationshipType.source_entity_type_id_name == source_type) - & (KGRelationshipType.target_entity_type_id_name == target_type) - ) - for source_type, target_type in entity_type_pairs - ] - - return db_session.query(KGRelationshipType).filter(or_(*conditions)).all() - - -def get_allowed_relationship_type_pairs( - db_session: Session, entities: list[str] -) -> list[str]: - """ - Get the allowed relationship pairs for the given entities. - - Args: - db_session: SQLAlchemy database session - entities: List of entity type ID names to filter by - - Returns: - List of id_names from KGRelationshipType where source or target entity types - are in the provided entities list. We also filter out for now the catch-all - relationship types 'VENDOR____' - """ - - entity_types = list({get_entity_type(entity) for entity in entities}) - - return [ - row[0] - for row in ( - db_session.query(KGRelationshipType.id_name) - .filter( - or_( - KGRelationshipType.source_entity_type_id_name.in_(entity_types), - KGRelationshipType.target_entity_type_id_name.in_(entity_types), - ) - ) - .filter(~KGRelationshipType.source_entity_type_id_name.like("VENDOR::%")) - .distinct() - .all() - ) - ] - - -def get_relationships_of_entity(db_session: Session, entity_id: str) -> List[str]: - """Get all relationship ID names where the given entity is either the source or target node. - - Args: - db_session: SQLAlchemy session - entity_id: ID of the entity to find relationships for - - Returns: - List of relationship ID names where the entity is either source or target - """ - return [ - row[0] - for row in ( - db_session.query(KGRelationship.id_name) - .filter( - or_( - KGRelationship.source_node == entity_id, - KGRelationship.target_node == entity_id, - ) - ) - .all() - ) - ] - - -def get_relationship_types_of_entity_types( - db_session: Session, entity_types_id: str -) -> List[str]: - """Get all relationship ID names where the given entity is either the source or target node. - - Args: - db_session: SQLAlchemy session - entity_types_id: ID of the entity to find relationships for - - Returns: - List of relationship ID names where the entity is either source or target - """ - - entity_types_id = entity_types_id.removesuffix(":*") - - return [ - row[0] - for row in ( - db_session.query(KGRelationshipType.id_name) - .filter( - or_( - KGRelationshipType.source_entity_type_id_name == entity_types_id, - KGRelationshipType.target_entity_type_id_name == entity_types_id, - ) - ) - .all() - ) - ] - - def delete_document_references_from_kg(db_session: Session, document_id: str) -> None: # Delete relationships from normalized stage db_session.query(KGRelationship).filter( diff --git a/backend/onyx/db/swap_index.py b/backend/onyx/db/swap_index.py index e80bc07f07..2bee9e76a5 100644 --- a/backend/onyx/db/swap_index.py +++ b/backend/onyx/db/swap_index.py @@ -6,9 +6,9 @@ from sqlalchemy.orm import Session from onyx.configs.app_configs import ( DISABLE_VECTOR_DB, + DOCUMENT_INDEX_NUM_ATTEMPTS_ON_STARTUP, PORT_SWAP_VERIFY_DOCS_PER_UNIT, PORT_SWAP_VERIFY_RETRY_DELAY_S, - VESPA_NUM_ATTEMPTS_ON_STARTUP, ) from onyx.configs.constants import KV_REINDEX_KEY from onyx.db.connector_credential_pair import ( @@ -58,7 +58,7 @@ from onyx.db.user_file import ( filter_existing_user_file_ids_any_owner, sample_ported_user_file_ids, ) -from onyx.document_index.factory import get_all_document_indices +from onyx.document_index.factory import get_default_document_index from onyx.document_index.opensearch.port_copy import ( find_documents_missing_from_index, find_documents_with_no_chunks, @@ -161,44 +161,42 @@ def _perform_index_swap( new_search_settings.contextual_rag_model_configuration_id = None db_session.commit() - # This flow is for checking and possibly creating an index so we get all - # indices. - document_indices = get_all_document_indices(new_search_settings, None, None) + # Check for and possibly create the index for the new search settings. + document_index = get_default_document_index(new_search_settings, None) WAIT_SECONDS = 5 - for document_index in document_indices: - success = False - for x in range(VESPA_NUM_ATTEMPTS_ON_STARTUP): - try: - logger.notice( - "Document index %s swap (attempt %s/%s)...", - document_index.__class__.__name__, - x + 1, - VESPA_NUM_ATTEMPTS_ON_STARTUP, - ) - document_index.verify_and_create_index_if_necessary( - embedding_dim=new_search_settings.final_embedding_dim, - ) - - logger.notice("Document index swap complete.") - success = True - break - except Exception: - logger.exception( - "Document index swap for %s did not succeed. The document index services may not be ready yet. Retrying in %s seconds.", - document_index.__class__.__name__, - WAIT_SECONDS, - ) - time.sleep(WAIT_SECONDS) - - if not success: - logger.error( - "Document index swap for %s did not succeed. Attempt limit reached. (%s)", + success = False + for x in range(DOCUMENT_INDEX_NUM_ATTEMPTS_ON_STARTUP): + try: + logger.notice( + "Document index %s swap (attempt %s/%s)...", document_index.__class__.__name__, - VESPA_NUM_ATTEMPTS_ON_STARTUP, + x + 1, + DOCUMENT_INDEX_NUM_ATTEMPTS_ON_STARTUP, ) - return None + document_index.verify_and_create_index_if_necessary( + embedding_dim=new_search_settings.final_embedding_dim, + ) + + logger.notice("Document index swap complete.") + success = True + break + except Exception: + logger.exception( + "Document index swap for %s did not succeed. The document index services may not be ready yet. Retrying in %s seconds.", + document_index.__class__.__name__, + WAIT_SECONDS, + ) + time.sleep(WAIT_SECONDS) + + if not success: + logger.error( + "Document index swap for %s did not succeed. Attempt limit reached. (%s)", + document_index.__class__.__name__, + DOCUMENT_INDEX_NUM_ATTEMPTS_ON_STARTUP, + ) + return None return current_search_settings diff --git a/backend/onyx/document_index/FILTER_SEMANTICS.md b/backend/onyx/document_index/FILTER_SEMANTICS.md index 67ce43a28e..5dd3a619a2 100644 --- a/backend/onyx/document_index/FILTER_SEMANTICS.md +++ b/backend/onyx/document_index/FILTER_SEMANTICS.md @@ -1,8 +1,7 @@ # Vector DB Filter Semantics -How `IndexFilters` fields combine into the final query filter. Describes the active -**OpenSearch** backend. The deprecated **Vespa** backend differs in one respect: -`project_id_filter` there remains *additive* (see "project_id_filter" notes below). +How `IndexFilters` fields combine into the final query filter in the +**OpenSearch** document index. ## Filter categories @@ -36,9 +35,7 @@ Two ways to constrain time, both AND-ed into the query: (either bound may be open) on the document's creation / last-update time, AND-ed together when both are set. Use them to express a query's created-vs-updated intent. The persona `search_start_date` floor is folded into -`updated_at_range.start` in the search pipeline. The deprecated Vespa backend -enforces only `updated_at_range` (it has no `created_at` field, so -`created_at_range` widens rather than narrows there). +`updated_at_range.start` in the search pipeline. For wire compatibility, `BaseFilters` still accepts the deprecated `time_cutoff` request field and folds it into `updated_at_range.start` at @@ -87,8 +84,7 @@ Each of these can start a knowledge scope on its own: user files overflowed the LLM context window. - **`project_id_filter`** is a **primary** trigger. A chat inside a project is scoped to that project, so `project_id_filter` alone restricts the search to the project's files — - project chats do not search team knowledge. (Deprecated Vespa backend: `project_id_filter` - is still *additive* there and only widens an existing scope.) + project chats do not search team knowledge. ### No explicit knowledge attached @@ -125,14 +121,14 @@ When an explicit knowledge restriction is in effect **and** `project_id_filter` -- Document sets + project files overflowed AND ( document_sets contains "Engineering" - OR user_project contains 7 + OR user_projects contains 7 ) -- Persona user files + project files (won't happen in practice; -- custom personas ignore project files per the precedence rule) AND ( personas contains 42 - OR user_project contains 7 + OR user_projects contains 7 ) ``` @@ -144,21 +140,21 @@ The search is restricted to the project's files. -- Restricted to project files NOT hidden AND (acl contains ...) -AND (user_project contains 7) +AND (user_projects contains 7) ``` ## Field reference -| Filter field | Vespa field | Vespa type | Purpose | +| Filter field | OpenSearch field | OpenSearch type | Purpose | |---|---|---|---| -| `document_set` | `document_sets` | `weightedset` | Connector doc sets attached to assistant | -| `attached_document_ids` | `document_id` | `string` | Documents explicitly attached (OpenSearch only) | -| `hierarchy_node_ids` | `ancestor_hierarchy_node_ids` | `array` | Folder/space nodes (OpenSearch only) | -| `persona_id_filter` | `personas` | `array` | Persona tag for overflowing user files (**primary** trigger) | -| `project_id_filter` | `user_project` | `array` | Project tag for overflowing project files (**primary** trigger; restricts to project files — OpenSearch. Vespa keeps it additive, see notes) | -| `access_control_list` | `access_control_list` | `weightedset` | ACL entries for the requesting user | -| `source_type` | `source_type` | `string` | Connector source type (e.g. `web`, `jira`) | -| `tags` | `metadata_list` | `array` | Document metadata tags | -| `created_at_range` | `created_at` | `long` | Window on document creation time; see [Time filtering](#time-filtering) (OpenSearch only) | -| `updated_at_range` | `doc_updated_at` | `long` | Window on document update time; see [Time filtering](#time-filtering) | -| `tenant_id` | `tenant_id` | `string` | Tenant isolation (multi-tenant) | +| `document_set` | `document_sets` | `keyword` | Connector doc sets attached to assistant | +| `attached_document_ids` | `document_id` | `keyword` | Documents explicitly attached | +| `hierarchy_node_ids` | `ancestor_hierarchy_node_ids` | `integer` | Folder/space nodes | +| `persona_id_filter` | `personas` | `integer` | Persona tag for overflowing user files (**primary** trigger) | +| `project_id_filter` | `user_projects` | `integer` | Project tag for overflowing project files (**primary** trigger; restricts to project files) | +| `access_control_list` | `access_control_list`, `public` | `keyword`, `boolean` | ACL entries for the requesting user | +| `source_type` | `source_type` | `keyword` | Connector source type (e.g. `web`, `jira`) | +| `tags` | `metadata_list` | `keyword` | Document metadata tags | +| `created_at_range` | `created_at` | `date` | Window on document creation time; see [Time filtering](#time-filtering) | +| `updated_at_range` | `last_updated` | `date` | Window on document update time; see [Time filtering](#time-filtering) | +| `tenant_id` | `tenant_id` | `keyword` | Tenant isolation (multi-tenant) | diff --git a/backend/onyx/document_index/chunk_content_enrichment.py b/backend/onyx/document_index/chunk_content_enrichment.py index b4f19859a8..3155d27e0c 100644 --- a/backend/onyx/document_index/chunk_content_enrichment.py +++ b/backend/onyx/document_index/chunk_content_enrichment.py @@ -50,9 +50,8 @@ def cleanup_content_for_chunks( """ def _remove_title(chunk: InferenceChunkUncleaned) -> str: - # TODO(andrei): This was ported over from - # backend/onyx/document_index/vespa/vespa_document_index.py but I don't - # think this logic is correct. In Vespa at least we set the title field + # TODO(andrei): This was ported over from the former Vespa document + # index but I don't think this logic is correct. We set the title field # from the output of get_title_for_document_index, which is not # necessarily the same data that is prepended to the content; that comes # from title_prefix. diff --git a/backend/onyx/document_index/disabled.py b/backend/onyx/document_index/disabled.py index 2678ef8b13..a92d4ae0b2 100644 --- a/backend/onyx/document_index/disabled.py +++ b/backend/onyx/document_index/disabled.py @@ -2,14 +2,14 @@ Used as a safety net when DISABLE_VECTOR_DB is True. Any code path that accidentally reaches the vector DB layer will fail loudly instead of timing -out against a nonexistent Vespa/OpenSearch instance. +out against a nonexistent OpenSearch instance. """ from collections.abc import Iterable from onyx.context.search.enums import QueryType from onyx.context.search.models import IndexFilters, InferenceChunk -from onyx.document_index.interfaces_new import ( +from onyx.document_index.interfaces import ( DocumentIndex, DocumentInsertionRecord, DocumentSectionRequest, diff --git a/backend/onyx/document_index/document_index_utils.py b/backend/onyx/document_index/document_index_utils.py index 917865d947..c192e11db2 100644 --- a/backend/onyx/document_index/document_index_utils.py +++ b/backend/onyx/document_index/document_index_utils.py @@ -1,21 +1,6 @@ -import math -import uuid -from uuid import UUID - -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, - get_secondary_search_settings, -) -from onyx.document_index.vespa.internal_types import EnrichedDocumentIndexingInfo -from onyx.indexing.models import DocMetadataAwareIndexChunk, MultipassConfig -from shared_configs.configs import MULTI_TENANT - -DEFAULT_BATCH_SIZE = 30 -DEFAULT_INDEX_NAME = "danswer_chunk" +from onyx.indexing.models import MultipassConfig def should_use_multipass(search_settings: SearchSettings | None) -> bool: @@ -40,166 +25,3 @@ def get_multipass_config(search_settings: SearchSettings) -> MultipassConfig: return MultipassConfig( multipass_indexing=multipass, enable_large_chunks=enable_large_chunks ) - - -def get_both_index_properties( - db_session: Session, -) -> tuple[str, str | None, bool, bool | None]: - search_settings = get_current_search_settings(db_session) - config_1 = get_multipass_config(search_settings) - - search_settings_new = get_secondary_search_settings(db_session) - if not search_settings_new: - return search_settings.index_name, None, config_1.enable_large_chunks, None - - config_2 = get_multipass_config(search_settings) - return ( - search_settings.index_name, - search_settings_new.index_name, - config_1.enable_large_chunks, - config_2.enable_large_chunks, - ) - - -def translate_boost_count_to_multiplier(boost: int) -> float: - """Mapping boost integer values to a multiplier according to a sigmoid curve - Piecewise such that at many downvotes, its 0.5x the score and with many upvotes - it is 2x the score. This should be in line with the Vespa calculation.""" - # 3 in the equation below stretches it out to hit asymptotes slower - if boost < 0: - # 0.5 + sigmoid -> range of 0.5 to 1 - return 0.5 + (1 / (1 + math.exp(-1 * boost / 3))) - - # 2 x sigmoid -> range of 1 to 2 - return 2 / (1 + math.exp(-1 * boost / 3)) - - -# Assembles a list of Vespa chunk IDs for a document -# given the required context. This can be used to directly query -# Vespa's Document API. -def get_document_chunk_ids( - enriched_document_info_list: list[EnrichedDocumentIndexingInfo], - tenant_id: str, - large_chunks_enabled: bool, -) -> list[UUID]: - doc_chunk_ids = [] - - for enriched_document_info in enriched_document_info_list: - for chunk_index in range( - enriched_document_info.chunk_start_index, - enriched_document_info.chunk_end_index, - ): - if not enriched_document_info.old_version: - doc_chunk_ids.append( - get_uuid_from_chunk_info( - document_id=enriched_document_info.doc_id, - chunk_id=chunk_index, - tenant_id=tenant_id, - ) - ) - else: - doc_chunk_ids.append( - get_uuid_from_chunk_info_old( - document_id=enriched_document_info.doc_id, - chunk_id=chunk_index, - ) - ) - - if large_chunks_enabled and chunk_index % 4 == 0: - large_chunk_id = int(chunk_index / 4) - large_chunk_reference_ids = [ - large_chunk_id + i - for i in range(4) - if large_chunk_id + i < enriched_document_info.chunk_end_index - ] - if enriched_document_info.old_version: - doc_chunk_ids.append( - get_uuid_from_chunk_info_old( - document_id=enriched_document_info.doc_id, - chunk_id=large_chunk_id, - large_chunk_reference_ids=large_chunk_reference_ids, - ) - ) - else: - doc_chunk_ids.append( - get_uuid_from_chunk_info( - document_id=enriched_document_info.doc_id, - chunk_id=large_chunk_id, - tenant_id=tenant_id, - large_chunk_id=large_chunk_id, - ) - ) - - return doc_chunk_ids - - -def get_uuid_from_chunk_info( - *, - document_id: str, - chunk_id: int, - tenant_id: str, - large_chunk_id: int | None = None, -) -> UUID: - """NOTE: be VERY carefuly about changing this function. If changed without a migration, - this can cause deletion/update/insertion to function incorrectly.""" - doc_str = document_id - - # Web parsing URL duplicate catching - if doc_str and doc_str[-1] == "/": - doc_str = doc_str[:-1] - - chunk_index = ( - "large_" + str(large_chunk_id) if large_chunk_id is not None else str(chunk_id) - ) - unique_identifier_string = "_".join([doc_str, chunk_index]) - if MULTI_TENANT: - unique_identifier_string += "_" + tenant_id - - uuid_value = uuid.uuid5(uuid.NAMESPACE_X500, unique_identifier_string) - return uuid_value - - -def get_uuid_from_chunk_info_old( - *, - document_id: str, - chunk_id: int, - large_chunk_reference_ids: list[int] | None = None, -) -> UUID: - if large_chunk_reference_ids is None: - large_chunk_reference_ids = [] - doc_str = document_id - - # Web parsing URL duplicate catching - if doc_str and doc_str[-1] == "/": - doc_str = doc_str[:-1] - unique_identifier_string = "_".join([doc_str, str(chunk_id), "0"]) - if large_chunk_reference_ids: - unique_identifier_string += "_large" + "_".join( - [ - str(referenced_chunk_id) - for referenced_chunk_id in large_chunk_reference_ids - ] - ) - return uuid.uuid5(uuid.NAMESPACE_X500, unique_identifier_string) - - -def get_uuid_from_chunk(chunk: DocMetadataAwareIndexChunk) -> uuid.UUID: - return get_uuid_from_chunk_info( - document_id=chunk.source_document.id, - chunk_id=chunk.chunk_id, - tenant_id=chunk.tenant_id, - large_chunk_id=chunk.large_chunk_id, - ) - - -def get_uuid_from_chunk_old( - chunk: DocMetadataAwareIndexChunk, - large_chunk_reference_ids: list[int] | None = None, -) -> UUID: - if large_chunk_reference_ids is None: - large_chunk_reference_ids = [] - return get_uuid_from_chunk_info_old( - document_id=chunk.source_document.id, - chunk_id=chunk.chunk_id, - large_chunk_reference_ids=large_chunk_reference_ids, - ) diff --git a/backend/onyx/document_index/document_metadata.py b/backend/onyx/document_index/document_metadata.py index ff1c3c31f5..ed506488d3 100644 --- a/backend/onyx/document_index/document_metadata.py +++ b/backend/onyx/document_index/document_metadata.py @@ -1,8 +1,7 @@ """Document-level metadata used during indexing. -Previously declared in the now-removed `onyx.document_index.interfaces` module. Used by the indexing pipeline / Postgres upsert layer; not part of the search -backend interface. +backend interface (`onyx.document_index.interfaces`). """ from dataclasses import dataclass diff --git a/backend/onyx/document_index/factory.py b/backend/onyx/document_index/factory.py index f032cc3be5..ee977a04c9 100644 --- a/backend/onyx/document_index/factory.py +++ b/backend/onyx/document_index/factory.py @@ -1,23 +1,11 @@ -import httpx -from sqlalchemy.orm import Session - -from onyx.configs.app_configs import ( - DISABLE_VECTOR_DB, - ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, - ONYX_DISABLE_VESPA, -) +from onyx.configs.app_configs import DISABLE_VECTOR_DB 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, TenantState +from onyx.document_index.interfaces 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.indexing.models import IndexingSetting from shared_configs.configs import MULTI_TENANT from shared_configs.contextvars import get_current_tenant_id @@ -34,7 +22,7 @@ def build_opensearch_document_index( The reindex port needs the lone index (to scan a PIT / call index_raw_chunks), not the primary+secondary pair `get_default_document_index` returns. Shared - with `_build_opensearch_pair` so the construction lives in one place. + with `get_default_document_index` so the construction lives in one place. """ indexing_setting = IndexingSetting.from_db_model(search_settings) return OpenSearchDocumentIndex( @@ -45,11 +33,23 @@ def build_opensearch_document_index( ) -def _build_opensearch_pair( +def get_default_document_index( search_settings: SearchSettings, secondary_search_settings: SearchSettings | None, + *, primary_backfill_in_progress: bool = False, -) -> OpenSearchIndexPair: +) -> DocumentIndex: + """Gets the document index for retrieval and indexing. + + Returns the primary+secondary pair, with secondary None when no second + search settings exist. + + ``primary_backfill_in_progress`` marks the primary as an INSTANT + reindex-port target still backfilling (see OpenSearchIndexPair.update). + """ + if DISABLE_VECTOR_DB: + return DisabledDocumentIndex() + primary = build_opensearch_document_index(search_settings) if secondary_search_settings is None: return OpenSearchIndexPair( @@ -67,105 +67,3 @@ def _build_opensearch_pair( secondary_embedding_dim=secondary_indexing_setting.final_embedding_dim, primary_backfill_in_progress=primary_backfill_in_progress, ) - - -def _build_vespa_pair( - search_settings: SearchSettings, - secondary_search_settings: SearchSettings | None, - httpx_client: httpx.Client | None, -) -> VespaIndexPair: - tenant_state = _build_tenant_state() - primary = VespaDocumentIndex( - index_name=search_settings.index_name, - tenant_state=tenant_state, - large_chunks_enabled=search_settings.large_chunks_enabled, - httpx_client=httpx_client, - ) - if secondary_search_settings is None: - return VespaIndexPair( - primary=primary, - secondary=None, - secondary_index_name=None, - secondary_embedding_dim=None, - ) - secondary_indexing_setting = IndexingSetting.from_db_model( - secondary_search_settings - ) - secondary = VespaDocumentIndex( - index_name=secondary_search_settings.index_name, - tenant_state=tenant_state, - large_chunks_enabled=secondary_search_settings.large_chunks_enabled, - httpx_client=httpx_client, - ) - return VespaIndexPair( - primary=primary, - secondary=secondary, - secondary_index_name=secondary_search_settings.index_name, - secondary_embedding_dim=secondary_indexing_setting.final_embedding_dim, - ) - - -def get_default_document_index( - search_settings: SearchSettings, - secondary_search_settings: SearchSettings | None, - db_session: Session, - httpx_client: httpx.Client | None = None, -) -> DocumentIndex: - """Gets the default document index for retrieval. - - Returns one DocumentIndex (the primary+secondary pair, with secondary None - when no second search settings exist). For indexing flows that need to write - to *all* configured backends, use `get_all_document_indices`. - """ - if DISABLE_VECTOR_DB: - return DisabledDocumentIndex() - - opensearch_retrieval_enabled = get_opensearch_retrieval_state(db_session) - if ONYX_DISABLE_VESPA and not opensearch_retrieval_enabled: - raise ValueError( - "Bug: ONYX_DISABLE_VESPA is set but opensearch_retrieval_enabled is not set." - ) - - if opensearch_retrieval_enabled: - return _build_opensearch_pair(search_settings, secondary_search_settings) - return _build_vespa_pair(search_settings, secondary_search_settings, httpx_client) - - -def get_all_document_indices( - search_settings: SearchSettings, - secondary_search_settings: SearchSettings | None, - httpx_client: httpx.Client | None = None, - primary_backfill_in_progress: bool = False, -) -> list[DocumentIndex]: - """Gets every document index that should be written to. - - NOTE: Make sure the Vespa index object is returned first. In the rare event - that there is some conflict between indexing and the migration task, it is - assumed that the state of Vespa is more up-to-date than the state of - OpenSearch. - - ``primary_backfill_in_progress`` marks the OpenSearch primary as an INSTANT - reindex-port target still backfilling (see OpenSearchIndexPair.update). - """ - if DISABLE_VECTOR_DB: - return [DisabledDocumentIndex()] - - if ONYX_DISABLE_VESPA and not ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: - raise ValueError( - "Bug: ONYX_DISABLE_VESPA is set but ENABLE_OPENSEARCH_INDEXING_FOR_ONYX is not set." - ) - - result: list[DocumentIndex] = [] - if not ONYX_DISABLE_VESPA: - result.append( - _build_vespa_pair(search_settings, secondary_search_settings, httpx_client) - ) - if ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: - result.append( - _build_opensearch_pair( - search_settings, - secondary_search_settings, - primary_backfill_in_progress=primary_backfill_in_progress, - ) - ) - return result diff --git a/backend/onyx/document_index/interfaces_new.py b/backend/onyx/document_index/interfaces.py similarity index 100% rename from backend/onyx/document_index/interfaces_new.py rename to backend/onyx/document_index/interfaces.py diff --git a/backend/onyx/document_index/opensearch/client.py b/backend/onyx/document_index/opensearch/client.py index 0797b995f5..7f36744f68 100644 --- a/backend/onyx/document_index/opensearch/client.py +++ b/backend/onyx/document_index/opensearch/client.py @@ -32,7 +32,7 @@ from onyx.configs.app_configs import ( OPENSEARCH_VERIFY_CERTS, PIT_KEEP_ALIVE, ) -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.constants import ( DEFAULT_MAX_CHUNK_SIZE, OpenSearchAuthMethod, diff --git a/backend/onyx/document_index/opensearch/index_reclaim.py b/backend/onyx/document_index/opensearch/index_reclaim.py index 3f5869c2c5..36cb10803e 100644 --- a/backend/onyx/document_index/opensearch/index_reclaim.py +++ b/backend/onyx/document_index/opensearch/index_reclaim.py @@ -6,7 +6,7 @@ COMPLETE. Driven by the reclaim beat task (background/celery/tasks/index_reclaim from enum import Enum from onyx.configs.app_configs import OLD_INDEX_RECLAIM_DELETE_BATCH_SIZE -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.client import OpenSearchIndexClient from onyx.document_index.opensearch.schema import TENANT_ID_FIELD_NAME from onyx.utils.logger import setup_logger diff --git a/backend/onyx/document_index/opensearch/opensearch_document_index.py b/backend/onyx/document_index/opensearch/opensearch_document_index.py index e4f5cae2a1..722569135c 100644 --- a/backend/onyx/document_index/opensearch/opensearch_document_index.py +++ b/backend/onyx/document_index/opensearch/opensearch_document_index.py @@ -26,7 +26,7 @@ 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 ( +from onyx.document_index.interfaces import ( DocumentIndex, DocumentInsertionRecord, DocumentSectionRequest, @@ -797,7 +797,7 @@ class OpenSearchDocumentIndex(DocumentIndex): self, chunk_requests: list[DocumentSectionRequest], filters: IndexFilters, - # TODO(andrei): Remove this from the new interface at some point; we + # TODO(andrei): Remove this from the interface at some point; we # should not be exposing this. batch_retrieval: bool = False, # noqa: ARG002 # TODO(andrei): Add a param for whether to retrieve hidden docs. @@ -820,7 +820,7 @@ class OpenSearchDocumentIndex(DocumentIndex): # NOTE: Index filters includes metadata tags which were filtered # for invalid unicode at indexing time. In theory it would be # ideal to do filtering here as well, in practice we never did - # that in the Vespa codepath and have not seen issues in + # that in the former Vespa codepath and have not seen issues in # production, so we deliberately conform to the existing logic # in order to not unknowningly introduce a possible bug. index_filters=filters, @@ -875,7 +875,7 @@ class OpenSearchDocumentIndex(DocumentIndex): # NOTE: Index filters includes metadata tags which were filtered # for invalid unicode at indexing time. In theory it would be # ideal to do filtering here as well, in practice we never did - # that in the Vespa codepath and have not seen issues in + # that in the former Vespa codepath and have not seen issues in # production, so we deliberately conform to the existing logic # in order to not unknowningly introduce a possible bug. index_filters=filters, @@ -924,7 +924,7 @@ class OpenSearchDocumentIndex(DocumentIndex): # NOTE: Index filters includes metadata tags which were filtered # for invalid unicode at indexing time. In theory it would be # ideal to do filtering here as well, in practice we never did - # that in the Vespa codepath and have not seen issues in + # that in the former Vespa codepath and have not seen issues in # production, so we deliberately conform to the existing logic # in order to not unknowningly introduce a possible bug. index_filters=filters, @@ -968,7 +968,7 @@ class OpenSearchDocumentIndex(DocumentIndex): # NOTE: Index filters includes metadata tags which were filtered # for invalid unicode at indexing time. In theory it would be # ideal to do filtering here as well, in practice we never did - # that in the Vespa codepath and have not seen issues in + # that in the former Vespa codepath and have not seen issues in # production, so we deliberately conform to the existing logic # in order to not unknowningly introduce a possible bug. index_filters=filters, @@ -1031,7 +1031,7 @@ class OpenSearchDocumentIndex(DocumentIndex): ) -> None: """Indexes raw document chunks into OpenSearch. - Used by the Vespa migration task and the reindex port. The reindex port + Used by the reindex port. The reindex port passes use_create_only=True so its stale backlog snapshot can never overwrite a chunk a live/forward writer already owns in FUTURE (an existing chunk is a benign 409). The port is pure gap-fill backfill of diff --git a/backend/onyx/document_index/opensearch/port_copy.py b/backend/onyx/document_index/opensearch/port_copy.py index 7dc8fc37b8..ec81f84158 100644 --- a/backend/onyx/document_index/opensearch/port_copy.py +++ b/backend/onyx/document_index/opensearch/port_copy.py @@ -18,7 +18,7 @@ from collections.abc import Callable, Iterable from onyx.db.models import SearchSettings from onyx.document_index.factory import build_opensearch_document_index -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.client import OpenSearchIndexClient from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, diff --git a/backend/onyx/document_index/opensearch/schema.py b/backend/onyx/document_index/opensearch/schema.py index 829e9383b1..6782369e1f 100644 --- a/backend/onyx/document_index/opensearch/schema.py +++ b/backend/onyx/document_index/opensearch/schema.py @@ -21,7 +21,7 @@ from onyx.configs.app_configs import ( USING_AWS_MANAGED_OPENSEARCH, ) from onyx.db.enums import VectorQuantization -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.constants import ( DEFAULT_MAX_CHUNK_SIZE, EF_CONSTRUCTION, @@ -505,8 +505,6 @@ class DocumentSchema: vector_quantization ), }, - # TODO(andrei): This is a tensor in Vespa. Also look at feature - # parity for these other method fields. CONTENT_VECTOR_FIELD_NAME: { "type": "knn_vector", "dimension": vector_dimension, diff --git a/backend/onyx/document_index/opensearch/search.py b/backend/onyx/document_index/opensearch/search.py index d3551e2f07..1566aebb5d 100644 --- a/backend/onyx/document_index/opensearch/search.py +++ b/backend/onyx/document_index/opensearch/search.py @@ -11,7 +11,7 @@ from onyx.configs.app_configs import ( from onyx.configs.constants import INDEX_SEPARATOR, DocumentSource from onyx.context.search.models import IndexFilters, Tag, TimeRange from onyx.db.enums import VectorQuantization -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.constants import ( ASSUMED_DOCUMENT_AGE_DAYS, DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES, @@ -1407,15 +1407,14 @@ class DocumentQuery: # See https://docs.opensearch.org/latest/search-plugins/searching-data/highlight/#highlighter-types "type": "unified", # The length in chars of a match snippet. Somewhat - # arbitrarily-chosen. The Vespa codepath limited total - # highlights length to 400 chars. fragment_size * - # number_of_fragments = 400 should be good enough. + # arbitrarily-chosen. fragment_size * number_of_fragments + # = 400 chars of highlights should be good enough. "fragment_size": 100, # The number of snippets to return per field per document # hit. "number_of_fragments": 4, - # These tags wrap matched keywords and they match what Vespa - # used to return. Use them to minimize changes to our code. + # These tags wrap matched keywords. Downstream code expects + # them. "pre_tags": [""], "post_tags": [""], } diff --git a/backend/onyx/document_index/vespa/__init__.py b/backend/onyx/document_index/vespa/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/backend/onyx/document_index/vespa/app_config/schemas/danswer_chunk.sd.jinja b/backend/onyx/document_index/vespa/app_config/schemas/danswer_chunk.sd.jinja deleted file mode 100644 index 4643d3c9b9..0000000000 --- a/backend/onyx/document_index/vespa/app_config/schemas/danswer_chunk.sd.jinja +++ /dev/null @@ -1,379 +0,0 @@ -schema {{ schema_name }} { - # source, type, target triplets for kg_relationships - struct kg_relationship { - field source type string {} - field rel_type type string {} - field target type string {} - } - - document {{ schema_name }} { - {% if multi_tenant %} - field tenant_id type string { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - {% endif %} - # Not to be confused with the UUID generated for this chunk which is called documentid by default - field document_id type string { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - field chunk_id type int { - indexing: summary | attribute - } - # Displayed in the UI as the main identifier for the doc - field semantic_identifier type string { - indexing: summary | attribute - } - # Must have an additional field for whether to skip title embeddings - # This information cannot be extracted from either the title field nor title embedding - field skip_title type bool { - indexing: attribute - } - # May not always match the `semantic_identifier` e.g. for Slack docs the - # `semantic_identifier` will be the channel name, but the `title` will be empty - field title type string { - indexing: summary | index | attribute - index: enable-bm25 - } - field content type string { - indexing: summary | index - index: enable-bm25 - } - # duplication of `content` is far from ideal, but is needed for - # non-gram based highlighting for now. If the capability to re-use a - # single field to do both is added, `content_summary` should be removed - field content_summary type string { - indexing: summary | index - summary: dynamic - } - # Title embedding (x1) - field title_embedding type tensor(x[{{ dim }}]) { - indexing: attribute | index - attribute { - distance-metric: angular - } - } - # Content embeddings (chunk + optional mini chunks embeddings) - # "t" and "x" are arbitrary names, not special keywords - field embeddings type tensor(t{},x[{{ dim }}]) { - indexing: attribute | index - attribute { - distance-metric: angular - } - } - # Starting section of the doc, currently unused as it has been replaced by match highlighting - field blurb type string { - indexing: summary | attribute - } - field image_file_name type string { - indexing: summary | attribute - } - # https://docs.vespa.ai/en/attributes.html potential enum store for speed, but probably not worth it - field source_type type string { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - # Can also index links https://docs.vespa.ai/en/reference/schema-reference.html#attribute - # URL type matching - field source_links type string { - indexing: summary | attribute - } - field section_continuation type bool { - indexing: summary | attribute - } - # Technically this one should be int, but can't change without causing breaks to existing index - field boost type float { - indexing: summary | attribute - } - field hidden type bool { - indexing: summary | attribute - rank: filter - } - # Field to indicate whether a short chunk is a low content chunk - field aggregated_chunk_boost_factor type float { - indexing: attribute - } - - # Separate array fields for knowledge graph data - field kg_entities type array { - indexing: summary | attribute - attribute: fast-search - } - - field kg_relationships type array { - indexing: summary - struct-field source { - indexing: attribute - attribute: fast-search - } - struct-field rel_type { - indexing: attribute - attribute: fast-search - } - struct-field target { - indexing: attribute - attribute: fast-search - } - } - - field kg_terms type array { - indexing: summary | attribute - attribute: fast-search - } - - # Needs to have a separate Attribute list for efficient filtering - field metadata_list type array { - indexing: summary | attribute - rank:filter - attribute: fast-search - } - # If chunk is a large chunk, this will contain the ids of the smaller chunks - field large_chunk_reference_ids type array { - indexing: summary | attribute - } - field metadata type string { - indexing: summary | attribute - } - field chunk_context type string { - indexing: summary | attribute - } - field doc_summary type string { - indexing: summary | attribute - } - field metadata_suffix type string { - indexing: summary | attribute - } - field doc_updated_at type int { - indexing: summary | attribute - } - field primary_owners type array { - indexing: summary | attribute - } - field secondary_owners type array { - indexing: summary | attribute - } - field access_control_list type weightedset { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - field document_sets type weightedset { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - field user_file type int { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - field user_folder type int { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - field user_project type array { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - field personas type array { - indexing: summary | attribute - rank: filter - attribute: fast-search - } - } - - # If using different tokenization settings, the fieldset has to be removed, and the field must - # be specified in the yql like: - # + 'or ({grammar: "weakAnd", defaultIndex:"title"}userInput(@query)) ' - # + 'or ({grammar: "weakAnd", defaultIndex:"content"}userInput(@query)) ' - # Note: for BM-25, the ngram size (and whether ngrams are used) changes the range of the scores - fieldset default { - fields: content, title - } - - rank-profile default_rank { - inputs { - query(decay_factor) double - } - - function inline document_boost() { - # 0.5 to 2x score: piecewise sigmoid function stretched out by factor of 3 - # meaning requires 3x the number of feedback votes to have default sigmoid effect - expression: if(attribute(boost) < 0, 0.5 + (1 / (1 + exp(-attribute(boost) / 3))), 2 / (1 + exp(-attribute(boost) / 3))) - } - - function inline document_age() { - # Time in years (91.3 days ~= 3 Months ~= 1 fiscal quarter if no age found) - expression: max(if(isNan(attribute(doc_updated_at)) == 1, 7890000, now() - attribute(doc_updated_at)) / 31536000, 0) - } - - function inline aggregated_chunk_boost() { - # Aggregated boost factor, currently only used for information content classification - expression: if(isNan(attribute(aggregated_chunk_boost_factor)) == 1, 1.0, attribute(aggregated_chunk_boost_factor)) - } - - # Document score decays from 1 to 0.75 as age of last updated time increases - function inline recency_bias() { - expression: max(1 / (1 + query(decay_factor) * document_age), 0.75) - } - - match-features: recency_bias - } - - rank-profile hybrid_search_semantic_base_{{ dim }} inherits default, default_rank { - inputs { - query(query_embedding) tensor(x[{{ dim }}]) - } - - function title_vector_score() { - expression { - # If no good matching titles, then it should use the context embeddings rather than having some - # irrelevant title have a vector score of 1. This way at least it will be the doc with the highest - # matching content score getting the full score - max(closeness(field, embeddings), closeness(field, title_embedding)) - } - } - - # First phase must be vector to allow hits that have no keyword matches - first-phase { - expression: query(title_content_ratio) * closeness(field, title_embedding) + (1 - query(title_content_ratio)) * closeness(field, embeddings) - } - - # Weighted average between Vector Search and BM-25 - global-phase { - expression { - ( - # Weighted Vector Similarity Score - ( - query(alpha) * ( - (query(title_content_ratio) * normalize_linear(title_vector_score)) - + - ((1 - query(title_content_ratio)) * normalize_linear(closeness(field, embeddings))) - ) - ) - - + - - # Weighted Keyword Similarity Score - # Note: for the BM25 Title score, it requires decent stopword removal in the query - # This needs to be the case so there aren't irrelevant titles being normalized to a score of 1 - ( - (1 - query(alpha)) * ( - (query(title_content_ratio) * normalize_linear(bm25(title))) - + - ((1 - query(title_content_ratio)) * normalize_linear(bm25(content))) - ) - ) - ) - # Boost based on user feedback - * document_boost - # Decay factor based on time document was last updated - * recency_bias - # Boost based on aggregated boost calculation - * aggregated_chunk_boost - } - # Target hits for hybrid retrieval should be at least this value. - rerank-count: 1000 - } - - match-features { - bm25(title) - bm25(content) - closeness(field, title_embedding) - closeness(field, embeddings) - document_boost - recency_bias - aggregated_chunk_boost - closest(embeddings) - } - } - - - rank-profile hybrid_search_keyword_base_{{ dim }} inherits default, default_rank { - inputs { - query(query_embedding) tensor(x[{{ dim }}]) - } - - function title_vector_score() { - expression { - # If no good matching titles, then it should use the context embeddings rather than having some - # irrelevant title have a vector score of 1. This way at least it will be the doc with the highest - # matching content score getting the full score - max(closeness(field, embeddings), closeness(field, title_embedding)) - } - } - - # First phase must be vector to allow hits that have no keyword matches - first-phase { - expression: query(title_content_ratio) * bm25(title) + (1 - query(title_content_ratio)) * bm25(content) - } - - # Weighted average between Vector Search and BM-25 - global-phase { - expression { - ( - # Weighted Vector Similarity Score - ( - query(alpha) * ( - (query(title_content_ratio) * normalize_linear(title_vector_score)) - + - ((1 - query(title_content_ratio)) * normalize_linear(closeness(field, embeddings))) - ) - ) - - + - - # Weighted Keyword Similarity Score - # Note: for the BM25 Title score, it requires decent stopword removal in the query - # This needs to be the case so there aren't irrelevant titles being normalized to a score of 1 - ( - (1 - query(alpha)) * ( - (query(title_content_ratio) * normalize_linear(bm25(title))) - + - ((1 - query(title_content_ratio)) * normalize_linear(bm25(content))) - ) - ) - ) - # Boost based on user feedback - * document_boost - # Decay factor based on time document was last updated - * recency_bias - # Boost based on aggregated boost calculation - * aggregated_chunk_boost - } - # Target hits for hybrid retrieval should be at least this value. - rerank-count: 1000 - } - - match-features { - bm25(title) - bm25(content) - closeness(field, title_embedding) - closeness(field, embeddings) - document_boost - recency_bias - aggregated_chunk_boost - closest(embeddings) - } - } - - # Used when searching from the admin UI for a specific doc to hide / boost - # Very heavily prioritize title - rank-profile admin_search inherits default, default_rank { - first-phase { - expression: bm25(content) + (5 * bm25(title)) - } - } - - rank-profile random_ inherits default { - first-phase { - expression: random - } - } -} diff --git a/backend/onyx/document_index/vespa/app_config/services.xml.jinja b/backend/onyx/document_index/vespa/app_config/services.xml.jinja deleted file mode 100644 index 67cebc3466..0000000000 --- a/backend/onyx/document_index/vespa/app_config/services.xml.jinja +++ /dev/null @@ -1,47 +0,0 @@ - - - - - - - - - - - - - - 1 - - - {{ document_elements }} - - - - - - - - - 0.85 - - - - - - - - {{ num_search_threads }} - - - - - - - 3 - 750 - 350 - 300 - - - \ No newline at end of file diff --git a/backend/onyx/document_index/vespa/app_config/validation-overrides.xml.jinja b/backend/onyx/document_index/vespa/app_config/validation-overrides.xml.jinja deleted file mode 100644 index a5f620c091..0000000000 --- a/backend/onyx/document_index/vespa/app_config/validation-overrides.xml.jinja +++ /dev/null @@ -1,11 +0,0 @@ - - schema-removal - indexing-change - field-type-change - diff --git a/backend/onyx/document_index/vespa/chunk_retrieval.py b/backend/onyx/document_index/vespa/chunk_retrieval.py deleted file mode 100644 index c6a3b022ee..0000000000 --- a/backend/onyx/document_index/vespa/chunk_retrieval.py +++ /dev/null @@ -1,687 +0,0 @@ -import json -import string -import time -from collections.abc import Callable, Mapping -from datetime import datetime, timezone -from typing import Any, cast - -import httpx - -from onyx.background.celery.tasks.opensearch_migration.constants import ( - FINISHED_VISITING_SLICE_CONTINUATION_TOKEN, -) -from onyx.background.celery.tasks.opensearch_migration.transformer import ( - FIELDS_NEEDED_FOR_TRANSFORMATION, -) -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, - build_vespa_id_based_retrieval_yql, -) -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 -from shared_configs.configs import MULTI_TENANT - -logger = setup_logger() - - -def _process_dynamic_summary( - dynamic_summary: str, max_summary_length: int = 400 -) -> list[str]: - if not dynamic_summary: - return [] - - current_length = 0 - processed_summary: list[str] = [] - for summary_section in dynamic_summary.split(""): - # if we're past the desired max length, break at the last word - if current_length + len(summary_section) >= max_summary_length: - summary_section = summary_section[: max_summary_length - current_length] - summary_section = summary_section.lstrip() # remove any leading whitespace - - # handle the case where the truncated section is either just a - # single (partial) word or if it's empty - first_space = summary_section.find(" ") - if first_space == -1: - # add ``...`` to previous section - if processed_summary: - processed_summary[-1] += "..." - break - - # handle the valid truncated section case - summary_section = summary_section.rsplit(" ", 1)[0] - if summary_section[-1] in string.punctuation: - summary_section = summary_section[:-1] - summary_section += "..." - processed_summary.append(summary_section) - break - - processed_summary.append(summary_section) - current_length += len(summary_section) - - return processed_summary - - -def _vespa_hit_to_inference_chunk( - hit: dict[str, Any], null_score: bool = False -) -> InferenceChunkUncleaned: - fields = cast(dict[str, Any], hit["fields"]) - - # parse fields that are stored as strings, but are really json / datetime - metadata = json.loads(fields[METADATA]) if METADATA in fields else {} - updated_at = ( - datetime.fromtimestamp(fields[DOC_UPDATED_AT], tz=timezone.utc) - if DOC_UPDATED_AT in fields - else None - ) - - match_highlights = _process_dynamic_summary( - # fallback to regular `content` if the `content_summary` field - # isn't present - dynamic_summary=hit["fields"].get(CONTENT_SUMMARY, hit["fields"][CONTENT]), - ) - semantic_identifier = fields.get(SEMANTIC_IDENTIFIER, "") - if not semantic_identifier: - logger.error( - "Chunk with blurb: %s... has no Semantic Identifier", - fields.get(BLURB, "Unknown")[:50], - ) - - source_links = fields.get(SOURCE_LINKS, {}) - source_links_dict_unprocessed = ( - json.loads(source_links) if isinstance(source_links, str) else source_links - ) - source_links_dict = { - int(k): v - for k, v in cast(dict[str, str], source_links_dict_unprocessed).items() - } - - return InferenceChunkUncleaned( - chunk_id=fields[CHUNK_ID], - blurb=fields.get(BLURB, ""), # Unused - content=fields[CONTENT], # Includes extra title prefix and metadata suffix; - # also sometimes context for contextual rag - source_links=source_links_dict or {0: ""}, - section_continuation=fields[SECTION_CONTINUATION], - document_id=fields[DOCUMENT_ID], - source_type=fields[SOURCE_TYPE], - # still called `image_file_name` in Vespa for backwards compatibility - image_file_id=fields.get(IMAGE_FILE_NAME), - title=fields.get(TITLE), - semantic_identifier=fields[SEMANTIC_IDENTIFIER], - boost=fields.get(BOOST, 1), - score=None if null_score else hit.get("relevance", 0), - hidden=fields.get(HIDDEN, False), - primary_owners=fields.get(PRIMARY_OWNERS), - secondary_owners=fields.get(SECONDARY_OWNERS), - large_chunk_reference_ids=fields.get(LARGE_CHUNK_REFERENCE_IDS, []), - metadata=metadata, - metadata_suffix=fields.get(METADATA_SUFFIX), - doc_summary=fields.get(DOC_SUMMARY, ""), - chunk_context=fields.get(CHUNK_CONTEXT, ""), - match_highlights=match_highlights, - updated_at=updated_at, - ) - - -def get_chunks_via_visit_api( - chunk_request: VespaChunkRequest, - index_name: str, - filters: IndexFilters, - field_names: list[str] | None = None, - get_large_chunks: bool = False, - short_tensor_format: bool = False, -) -> list[dict]: - # Constructing the URL for the Visit API - # NOTE: visit API uses the same URL as the document API, but with different params - url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name) - - # build the list of fields to retrieve - field_set_list = ( - [f"{field_name}" for field_name in field_names] if field_names else [] - ) - acl_fieldset_entry = f"{ACCESS_CONTROL_LIST}" - if ( - field_set_list - and filters.access_control_list - and acl_fieldset_entry not in field_set_list - ): - field_set_list.append(acl_fieldset_entry) - - if MULTI_TENANT: - tenant_id_fieldset_entry = f"{TENANT_ID}" - if field_set_list and tenant_id_fieldset_entry not in field_set_list: - field_set_list.append(tenant_id_fieldset_entry) - - if field_set_list: - field_set = f"{index_name}:" + ",".join(field_set_list) - else: - field_set = None - - # build filters - selection = f"{index_name}.document_id=='{chunk_request.document_id}'" - - if chunk_request.is_capped: - selection += f" and {index_name}.chunk_id>={chunk_request.min_chunk_ind or 0}" - selection += f" and {index_name}.chunk_id<={chunk_request.max_chunk_ind}" - if not get_large_chunks: - selection += f" and {index_name}.large_chunk_reference_ids == null" - - # enforcing tenant_id through a == condition - if MULTI_TENANT: - if filters.tenant_id: - selection += f" and {index_name}.tenant_id=='{filters.tenant_id}'" - else: - raise ValueError("Tenant ID is required for multi-tenant") - - # Setting up the selection criteria in the query parameters - params = { - # NOTE: Document Selector Language doesn't allow `contains`, so we can't check - # for the ACL in the selection. Instead, we have to check as a postfilter - "selection": selection, - "continuation": None, - "wantedDocumentCount": 1_000, - "fieldSet": field_set, - } - # Vespa can supply tensors in various different formats. This explicitly - # asks to retrieve tensor data in "short-value" format. - if short_tensor_format: - params["format.tensors"] = "short-value" - - document_chunks: list[dict] = [] - while True: - try: - filtered_params = {k: v for k, v in params.items() if v is not None} - with get_vespa_http_client() as http_client: - response = http_client.get(url, params=filtered_params) - response.raise_for_status() - except httpx.HTTPError as e: - error_base = "Failed to query Vespa" - logger.error( - "%s:\n" - "Request URL: %s\n" - "Request Headers: %s\n" - "Request Payload: %s\n" - "Exception: %s", - error_base, - e.request.url, - e.request.headers, - params, - str(e), - ) - raise httpx.HTTPError(error_base) from e - - # Check if the response contains any documents - response_data = response.json() - - if "documents" in response_data: - for document in response_data["documents"]: - if filters.access_control_list: - document_acl = document["fields"].get(ACCESS_CONTROL_LIST) - if not document_acl or not any( - user_acl_entry in document_acl - for user_acl_entry in filters.access_control_list - ): - continue - - if MULTI_TENANT: - if not filters.tenant_id: - raise ValueError("Tenant ID is required for multi-tenant") - document_tenant_id = document["fields"].get(TENANT_ID) - if document_tenant_id != filters.tenant_id: - logger.error( - "Skipping document %s because " - "it does not belong to tenant %s. " - "This should never happen.", - document["document_id"], - filters.tenant_id, - ) - continue - - document_chunks.append(document) - - # Check for continuation token to handle pagination - if "continuation" in response_data and response_data["continuation"]: - params["continuation"] = response_data["continuation"] - else: - break # Exit loop if no continuation token - - return document_chunks - - -def get_all_chunks_paginated( - index_name: str, - tenant_state: TenantState, - continuation_token_map: dict[int, str | None], - page_size: int, -) -> tuple[list[dict], dict[int, str | None]]: - """Gets all chunks in Vespa matching the filters, paginated. - - Uses the Visit API with slicing. Each continuation token map entry is for a - different slice. The number of entries determines the number of slices. - - Args: - index_name: The name of the Vespa index to visit. - tenant_state: The tenant state to filter by. - continuation_token_map: Map of slice ID to a token returned by Vespa - representing a page offset. None to start from the beginning of the - slice. - page_size: Best-effort batch size for the visit. Defaults to 1,000. - - Returns: - Tuple of (list of chunk dicts, next continuation token or None). The - continuation token is None when the visit is complete. - """ - - def _get_all_chunks_paginated_for_slice( - index_name: str, - tenant_state: TenantState, - slice_id: int, - total_slices: int, - continuation_token: str | None, - page_size: int, - ) -> tuple[list[dict], str | None]: - if continuation_token == FINISHED_VISITING_SLICE_CONTINUATION_TOKEN: - logger.debug( - "Slice %s has finished visiting. Returning empty list and %s.", - slice_id, - FINISHED_VISITING_SLICE_CONTINUATION_TOKEN, - ) - return [], FINISHED_VISITING_SLICE_CONTINUATION_TOKEN - - url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name) - - selection: str = f"{index_name}.large_chunk_reference_ids == null" - if MULTI_TENANT: - selection += f" and {index_name}.tenant_id=='{tenant_state.tenant_id}'" - - field_set = f"{index_name}:" + ",".join(FIELDS_NEEDED_FOR_TRANSFORMATION) - - params: dict[str, str | int | None] = { - "selection": selection, - "fieldSet": field_set, - "wantedDocumentCount": page_size, - "format.tensors": "short-value", - "slices": total_slices, - "sliceId": slice_id, - # When exceeded, Vespa should return gracefully with partial - # results. Even if no hits are returned, Vespa should still return a - # new continuation token representing a new spot in the linear - # traversal. - "timeout": VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT, - } - if continuation_token is not None: - params["continuation"] = continuation_token - - response: httpx.Response | None = None - start_time = time.monotonic() - try: - with get_vespa_http_client( - # When exceeded, an exception is raised in our code. No progress - # is saved, and the task will retry this spot in the traversal - # later. - timeout=VESPA_MIGRATION_REQUEST_TIMEOUT_S - ) as http_client: - response = http_client.get(url, params=params) - response.raise_for_status() - except httpx.HTTPError as e: - error_base = ( - f"Failed to get chunks from Vespa slice {slice_id} with continuation token " - f"{continuation_token} in {time.monotonic() - start_time:.3f} seconds." - ) - logger.exception( - "Request URL: %s\nRequest Headers: %s\nRequest Payload: %s", - e.request.url, - e.request.headers, - params, - ) - error_message = ( - response.json().get("message") if response else "No response" - ) - logger.error("Error message from response: %s", error_message) - raise httpx.HTTPError(error_base) from e - - response_data = response.json() - - # NOTE: If we see a falsey value for "continuation" in the response we - # assume we are done and return - # FINISHED_VISITING_SLICE_CONTINUATION_TOKEN instead. - next_continuation_token = ( - response_data.get("continuation") - or FINISHED_VISITING_SLICE_CONTINUATION_TOKEN - ) - chunks = [chunk["fields"] for chunk in response_data.get("documents", [])] - if next_continuation_token == FINISHED_VISITING_SLICE_CONTINUATION_TOKEN: - logger.debug( - "Slice %s has finished visiting. Returning %s chunks and %s.", - slice_id, - len(chunks), - next_continuation_token, - ) - return chunks, next_continuation_token - - total_slices = len(continuation_token_map) - if total_slices < 1: - raise ValueError("continuation_token_map must have at least one entry.") - # We want to guarantee that these invocations are ordered by slice_id, - # because we read in the same order below when parsing parallel_results. - functions_with_args: list[tuple[Callable, tuple]] = [ - ( - _get_all_chunks_paginated_for_slice, - ( - index_name, - tenant_state, - slice_id, - total_slices, - continuation_token, - page_size, - ), - ) - for slice_id, continuation_token in sorted(continuation_token_map.items()) - ] - - parallel_results = run_functions_tuples_in_parallel( - functions_with_args, allow_failures=True - ) - if len(parallel_results) != total_slices: - raise RuntimeError( - f"Expected {total_slices} parallel results, but got {len(parallel_results)}." - ) - - chunks: list[dict] = [] - next_continuation_token_map: dict[int, str | None] = dict( - continuation_token_map.items() - ) - for i, parallel_result in enumerate(parallel_results): - if i not in next_continuation_token_map: - raise RuntimeError(f"Slice {i} is not in the continuation token map.") - if parallel_result is None: - logger.error( - "Failed to get chunks for slice %s of %s. " - "The continuation token for this slice will not be updated.", - i, - total_slices, - ) - continue - chunks.extend(parallel_result[0]) - next_continuation_token_map[i] = parallel_result[1] - - return chunks, next_continuation_token_map - - -# TODO(rkuo): candidate for removal if not being used -# @retry(tries=10, delay=1, backoff=2) -# def get_all_vespa_ids_for_document_id( -# document_id: str, -# index_name: str, -# filters: IndexFilters | None = None, -# get_large_chunks: bool = False, -# ) -> list[str]: -# document_chunks = get_chunks_via_visit_api( -# chunk_request=VespaChunkRequest(document_id=document_id), -# index_name=index_name, -# filters=filters or IndexFilters(access_control_list=None), -# field_names=[DOCUMENT_ID], -# get_large_chunks=get_large_chunks, -# ) -# return [chunk["id"].split("::", 1)[-1] for chunk in document_chunks] - - -def parallel_visit_api_retrieval( - index_name: str, - chunk_requests: list[VespaChunkRequest], - filters: IndexFilters, - get_large_chunks: bool = False, -) -> list[InferenceChunkUncleaned]: - functions_with_args: list[tuple[Callable, tuple]] = [ - ( - get_chunks_via_visit_api, - (chunk_request, index_name, filters, get_large_chunks), - ) - for chunk_request in chunk_requests - ] - - parallel_results = run_functions_tuples_in_parallel( - functions_with_args, allow_failures=True - ) - - # Any failures to retrieve would give a None, drop the Nones and empty lists - vespa_chunk_sets = [res for res in parallel_results if res] - - flattened_vespa_chunks = [] - for chunk_set in vespa_chunk_sets: - flattened_vespa_chunks.extend(chunk_set) - - inference_chunks = [ - _vespa_hit_to_inference_chunk(chunk, null_score=True) - for chunk in flattened_vespa_chunks - ] - - return inference_chunks - - -@retry_builder(tries=3, delay=1, backoff=2) -def query_vespa( - query_params: Mapping[str, str | int | float], -) -> list[InferenceChunkUncleaned]: - if "query" in query_params and not cast(str, query_params["query"]).strip(): - raise ValueError("No/empty query received") - - params = dict( - **query_params, - **( - { - "presentation.timing": True, - } - if LOG_VESPA_TIMING_INFORMATION - else {} - ), - ) - - if VESPA_LANGUAGE_OVERRIDE: - params["language"] = VESPA_LANGUAGE_OVERRIDE - - try: - with get_vespa_http_client() as http_client: - response = http_client.post(SEARCH_ENDPOINT, json=params) - response.raise_for_status() - except httpx.HTTPError as e: - response_text = ( - e.response.text if isinstance(e, httpx.HTTPStatusError) else None - ) - status_code = ( - e.response.status_code if isinstance(e, httpx.HTTPStatusError) else None - ) - yql_value = params.get("yql", "") - yql_length = len(str(yql_value)) - - # Log each detail on its own line so log collectors capture them - # as separate entries rather than truncating a single multiline msg - logger.error( - "Failed to query Vespa | status=%s | yql_length=%s | exception=%s", - status_code, - yql_length, - str(e), - ) - if response_text: - logger.error("Vespa error response: %s", response_text[:1000]) - logger.error("Vespa request URL: %s", e.request.url) - - # Re-raise with diagnostics so callers see what actually went wrong - raise httpx.HTTPError( - f"Failed to query Vespa (status={status_code}, yql_length={yql_length})" - ) from e - - response_json: dict[str, Any] = response.json() - - if LOG_VESPA_TIMING_INFORMATION: - logger.debug("Vespa timing info: %s", response_json.get("timing")) - hits = response_json["root"].get("children", []) - - if not hits: - logger.warning( - "No hits found for YQL Query: %s", - query_params.get("yql", "No YQL Query"), - ) - logger.debug("Vespa Response: %s", response.text) - - for hit in hits: - if hit["fields"].get(CONTENT) is None: - identifier = hit["fields"].get("documentid") or hit["id"] - logger.error( - "Vespa Index with Vespa ID %s has no contents. " - "This is invalid because the vector is not meaningful and keywordsearch cannot " - "fetch this document", - identifier, - ) - - filtered_hits = [hit for hit in hits if hit["fields"].get(CONTENT) is not None] - - inference_chunks = [_vespa_hit_to_inference_chunk(hit) for hit in filtered_hits] - - try: - num_retrieved_inference_chunks = len(inference_chunks) - num_retrieved_document_ids = len( - {chunk.document_id for chunk in inference_chunks} - ) - logger.info( - "Retrieved %s inference chunks for %s documents", - num_retrieved_inference_chunks, - num_retrieved_document_ids, - ) - except Exception as e: - # Debug logging only, should not fail the retrieval - logger.error("Error logging retrieval statistics: %s", e) - - # Good Debugging Spot - return inference_chunks - - -def _get_chunks_via_batch_search( - index_name: str, - chunk_requests: list[VespaChunkRequest], - filters: IndexFilters, - get_large_chunks: bool = False, -) -> list[InferenceChunkUncleaned]: - if not chunk_requests: - return [] - - filters_str = build_vespa_filters(filters=filters, include_hidden=True) - - yql = ( - YQL_BASE.format(index_name=index_name) - + filters_str - + build_vespa_id_based_retrieval_yql(chunk_requests[0]) - ) - chunk_requests.pop(0) - - for request in chunk_requests: - yql += " or " + build_vespa_id_based_retrieval_yql(request) - params: dict[str, str | int | float] = { - "yql": yql, - "hits": MAX_ID_SEARCH_QUERY_SIZE, - } - - inference_chunks = query_vespa(params) - if not get_large_chunks: - inference_chunks = [ - chunk for chunk in inference_chunks if not chunk.large_chunk_reference_ids - ] - inference_chunks.sort(key=lambda chunk: chunk.chunk_id) - return inference_chunks - - -def batch_search_api_retrieval( - index_name: str, - chunk_requests: list[VespaChunkRequest], - filters: IndexFilters, - get_large_chunks: bool = False, -) -> list[InferenceChunkUncleaned]: - retrieved_chunks: list[InferenceChunkUncleaned] = [] - capped_requests: list[VespaChunkRequest] = [] - uncapped_requests: list[VespaChunkRequest] = [] - chunk_count = 0 - for req_ind, request in enumerate(chunk_requests, start=1): - # All requests without a chunk range are uncapped - # Uncapped requests are retrieved using the Visit API - range = request.range - if range is None: - uncapped_requests.append(request) - continue - - if ( - chunk_count + range > MAX_ID_SEARCH_QUERY_SIZE - or req_ind % MAX_OR_CONDITIONS == 0 - ): - retrieved_chunks.extend( - _get_chunks_via_batch_search( - index_name=index_name, - chunk_requests=capped_requests, - filters=filters, - get_large_chunks=get_large_chunks, - ) - ) - capped_requests = [] - chunk_count = 0 - capped_requests.append(request) - chunk_count += range - - if capped_requests: - retrieved_chunks.extend( - _get_chunks_via_batch_search( - index_name=index_name, - chunk_requests=capped_requests, - filters=filters, - get_large_chunks=get_large_chunks, - ) - ) - - if uncapped_requests: - logger.debug("Retrieving %s uncapped requests", len(uncapped_requests)) - retrieved_chunks.extend( - parallel_visit_api_retrieval( - index_name, uncapped_requests, filters, get_large_chunks - ) - ) - - return retrieved_chunks diff --git a/backend/onyx/document_index/vespa/deletion.py b/backend/onyx/document_index/vespa/deletion.py deleted file mode 100644 index 9215683950..0000000000 --- a/backend/onyx/document_index/vespa/deletion.py +++ /dev/null @@ -1,68 +0,0 @@ -import concurrent.futures -from uuid import UUID - -import httpx - -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 - -logger = setup_logger() - - -CONTENT_SUMMARY = "content_summary" - - -@retry_builder(tries=10, delay=1, backoff=2) -def _retryable_http_delete(http_client: httpx.Client, url: str) -> None: - res = http_client.delete(url) - res.raise_for_status() - - -def _delete_vespa_chunk( - doc_chunk_id: UUID, index_name: str, http_client: httpx.Client -) -> None: - try: - _retryable_http_delete( - http_client, - f"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{doc_chunk_id}", - ) - except httpx.HTTPStatusError as e: - logger.error("Failed to delete chunk, details: %s", e.response.text) - raise - - -def delete_vespa_chunks( - doc_chunk_ids: list[UUID], - index_name: str, - http_client: httpx.Client, - executor: concurrent.futures.ThreadPoolExecutor | None = None, -) -> None: - """Deletes a list of chunks from a Vespa index in parallel. - - Args: - doc_chunk_ids: List of chunk IDs to delete. - index_name: Name of the index to delete from. - http_client: HTTP client to use for the request. - executor: Executor to use for the request. - """ - external_executor = True - - if not executor: - external_executor = False - executor = concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) - - try: - chunk_deletion_future = { - executor.submit( - _delete_vespa_chunk, doc_chunk_id, index_name, http_client - ): doc_chunk_id - for doc_chunk_id in doc_chunk_ids - } - for future in concurrent.futures.as_completed(chunk_deletion_future): - # Will raise exception if the deletion raised an exception - future.result() - - finally: - if not external_executor: - executor.shutdown(wait=True) diff --git a/backend/onyx/document_index/vespa/indexing_utils.py b/backend/onyx/document_index/vespa/indexing_utils.py deleted file mode 100644 index 10276d0da4..0000000000 --- a/backend/onyx/document_index/vespa/indexing_utils.py +++ /dev/null @@ -1,439 +0,0 @@ -import concurrent.futures -import json -import random -import time -import uuid -from abc import ABC, abstractmethod -from collections.abc import Callable -from datetime import datetime, timezone -from http import HTTPStatus - -import httpx - -from onyx.connectors.cross_connector_utils.miscellaneous_utils import ( - get_experts_stores_representations, -) -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, - 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, - 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 -from onyx.utils.text_processing import remove_invalid_unicode_chars - -logger = setup_logger() - -# Retry configuration constants -INDEXING_MAX_RETRIES = 5 -INDEXING_BASE_DELAY = 1.0 -INDEXING_MAX_DELAY = 60.0 - - -@retry_builder(tries=3, delay=1, backoff=2) -def _does_doc_chunk_exist( - doc_chunk_id: uuid.UUID, index_name: str, http_client: httpx.Client -) -> bool: - doc_url = f"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{doc_chunk_id}" - doc_fetch_response = http_client.get(doc_url) - if doc_fetch_response.status_code == 404: - return False - - if doc_fetch_response.status_code != 200: - logger.debug("Failed to check for document with URL %s", doc_url) - raise RuntimeError( - f"Unexpected fetch document by ID value from Vespa: " - f"error={doc_fetch_response.status_code} " - f"index={index_name} " - f"doc_chunk_id={doc_chunk_id}" - ) - return True - - -def _vespa_get_updated_at_attribute(t: datetime | None) -> int | None: - if not t: - return None - - if t.tzinfo != timezone.utc: - raise ValueError("Connectors must provide document update time in UTC") - - return int(t.timestamp()) - - -def get_existing_documents_from_chunks( - chunks: list[DocMetadataAwareIndexChunk], - index_name: str, - http_client: httpx.Client, - executor: concurrent.futures.ThreadPoolExecutor | None = None, -) -> set[str]: - external_executor = True - - if not executor: - external_executor = False - executor = concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) - - document_ids: set[str] = set() - try: - chunk_existence_future = { - executor.submit( - _does_doc_chunk_exist, - get_uuid_from_chunk(chunk), - index_name, - http_client, - ): chunk - for chunk in chunks - } - for future in concurrent.futures.as_completed(chunk_existence_future): - chunk = chunk_existence_future[future] - chunk_already_existed = future.result() - if chunk_already_existed: - document_ids.add(chunk.source_document.id) - - finally: - if not external_executor: - executor.shutdown(wait=True) - - return document_ids - - -def _index_vespa_chunk( - chunk: DocMetadataAwareIndexChunk, - index_name: str, - http_client: httpx.Client, - multitenant: bool, -) -> None: - json_header = { - "Content-Type": "application/json", - } - document = chunk.source_document - - # No minichunk documents in vespa, minichunk vectors are stored in the chunk itself - - vespa_chunk_id = str(get_uuid_from_chunk(chunk)) - - embeddings = chunk.embeddings - - embeddings_name_vector_map = {FULL_CHUNK_EMBEDDING_KEY: embeddings.full_embedding} - - if embeddings.mini_chunk_embeddings: - for ind, m_c_embed in enumerate(embeddings.mini_chunk_embeddings): - embeddings_name_vector_map[f"mini_chunk_{ind}"] = m_c_embed - - title = document.get_title_for_document_index() - - metadata_json = document.metadata - cleaned_metadata_json: dict[str, str | list[str]] = {} - for key, value in metadata_json.items(): - cleaned_key = remove_invalid_unicode_chars(key) - if isinstance(value, list): - cleaned_metadata_json[cleaned_key] = [ - remove_invalid_unicode_chars(item) for item in value - ] - else: - cleaned_metadata_json[cleaned_key] = remove_invalid_unicode_chars(value) - - metadata_list = document.get_metadata_str_attributes() - if metadata_list: - metadata_list = [ - remove_invalid_unicode_chars(metadata) for metadata in metadata_list - ] - - vespa_document_fields = { - DOCUMENT_ID: document.id, - CHUNK_ID: chunk.chunk_id, - BLURB: remove_invalid_unicode_chars(chunk.blurb), - TITLE: remove_invalid_unicode_chars(title) if title else None, - SKIP_TITLE_EMBEDDING: not title, - # For the BM25 index, the keyword suffix is used, the vector is already generated with the more - # natural language representation of the metadata section - CONTENT: remove_invalid_unicode_chars( - generate_enriched_content_for_chunk_text(chunk) - ), - # This duplication of `content` is needed for keyword highlighting - # Note that it's not exactly the same as the actual content - # which contains the title prefix and metadata suffix - CONTENT_SUMMARY: remove_invalid_unicode_chars(chunk.content), - SOURCE_TYPE: str(document.source.value), - SOURCE_LINKS: json.dumps(chunk.source_links), - SEMANTIC_IDENTIFIER: remove_invalid_unicode_chars(document.semantic_identifier), - SECTION_CONTINUATION: chunk.section_continuation, - LARGE_CHUNK_REFERENCE_IDS: chunk.large_chunk_reference_ids, - METADATA: json.dumps(cleaned_metadata_json), - # Save as a list for efficient extraction as an Attribute - METADATA_LIST: metadata_list, - METADATA_SUFFIX: remove_invalid_unicode_chars(chunk.metadata_suffix_keyword), - CHUNK_CONTEXT: chunk.chunk_context, - DOC_SUMMARY: chunk.doc_summary, - EMBEDDINGS: embeddings_name_vector_map, - TITLE_EMBEDDING: chunk.title_embedding, - DOC_UPDATED_AT: _vespa_get_updated_at_attribute(document.doc_updated_at), - PRIMARY_OWNERS: get_experts_stores_representations(document.primary_owners), - SECONDARY_OWNERS: get_experts_stores_representations(document.secondary_owners), - # the only `set` vespa has is `weightedset`, so we have to give each - # element an arbitrary weight - # rkuo: acl, docset and boost metadata are also updated through the metadata sync queue - # which only calls VespaIndex.update - ACCESS_CONTROL_LIST: dict.fromkeys(chunk.access.to_acl(), 1), - DOCUMENT_SETS: dict.fromkeys(chunk.document_sets, 1), - # still called `image_file_name` in Vespa for backwards compatibility - IMAGE_FILE_NAME: chunk.image_file_id, - USER_PROJECT: chunk.user_project if chunk.user_project is not None else [], - PERSONAS: chunk.personas if chunk.personas is not None else [], - BOOST: chunk.boost, - AGGREGATED_CHUNK_BOOST_FACTOR: chunk.aggregated_chunk_boost_factor, - } - - if multitenant: - if chunk.tenant_id: - vespa_document_fields[TENANT_ID] = chunk.tenant_id - vespa_url = f"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{vespa_chunk_id}" - logger.debug('Indexing to URL "%s"', vespa_url) - - # Retry logic with exponential backoff for rate limiting - for attempt in range(INDEXING_MAX_RETRIES): - try: - res = http_client.post( - vespa_url, headers=json_header, json={"fields": vespa_document_fields} - ) - res.raise_for_status() - return # Success, exit the function - except httpx.HTTPStatusError as e: - # Handle 429 rate limiting specifically - if e.response.status_code == HTTPStatus.TOO_MANY_REQUESTS: - if attempt < INDEXING_MAX_RETRIES - 1: - # Calculate exponential backoff with jitter - delay = min( - INDEXING_BASE_DELAY * (2**attempt), INDEXING_MAX_DELAY - ) * random.uniform(0.5, 1.0) - logger.warning( - "Rate limited while indexing document '%s' (attempt %s/%s). Vespa response: '%s'. Backing off for %s seconds.", - document.id, - attempt + 1, - INDEXING_MAX_RETRIES, - e.response.text, - format(delay, ".2f"), - ) - time.sleep(delay) - continue - else: - raise RuntimeError( - f"Failed to index document '{document.id}' after {INDEXING_MAX_RETRIES} attempts due to rate limiting" - ) from e - elif e.response.status_code == HTTPStatus.INSUFFICIENT_STORAGE: - logger.error( - "Failed to index document: '%s'. Got response: '%s'", - document.id, - e.response.text, - ) - logger.error( - "NOTE: HTTP Status 507 Insufficient Storage usually means " - "you need to allocate more memory or disk space to the " - "Vespa/index container." - ) - raise - else: - # For other HTTP errors, check if retryable - if e.response.status_code in ( - HTTPStatus.BAD_REQUEST, - HTTPStatus.UNAUTHORIZED, - HTTPStatus.FORBIDDEN, - HTTPStatus.NOT_FOUND, - ): - # Non-retryable errors - fail immediately - logger.error( - "Non-retryable HTTP %s error for document '%s'", - e.response.status_code, - document.id, - ) - raise - # Retry other errors with shorter backoff - if attempt < INDEXING_MAX_RETRIES - 1: - delay = INDEXING_BASE_DELAY * (1.5**attempt) - logger.warning( - "HTTP error %s while indexing document '%s' (attempt %s/%s). Retrying in %s seconds.", - e.response.status_code, - document.id, - attempt + 1, - INDEXING_MAX_RETRIES, - format(delay, ".2f"), - ) - time.sleep(delay) - continue - else: - logger.exception( - "Failed to index document: '%s'. Got response: '%s'", - document.id, - e.response.text, - ) - raise - except Exception as e: - # For non-HTTP errors, use simple retry logic - if attempt < INDEXING_MAX_RETRIES - 1: - delay = INDEXING_BASE_DELAY * (1.5**attempt) - logger.warning( - "Error while indexing document '%s' (attempt %s/%s): %s. Retrying in %s seconds.", - document.id, - attempt + 1, - INDEXING_MAX_RETRIES, - str(e), - format(delay, ".2f"), - ) - time.sleep(delay) - continue - else: - logger.exception("Failed to index document: '%s'", document.id) - raise - - -def batch_index_vespa_chunks( - chunks: list[DocMetadataAwareIndexChunk], - index_name: str, - http_client: httpx.Client, - multitenant: bool, - executor: concurrent.futures.ThreadPoolExecutor | None = None, -) -> None: - """Indexes a list of chunks in a Vespa index in parallel. - - Args: - chunks: List of chunks to index. - index_name: Name of the index to index into. - http_client: HTTP client to use for the request. - multitenant: Whether the index is multitenant. - executor: Executor to use for the request. - """ - external_executor = True - - if not executor: - external_executor = False - executor = concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) - - try: - chunk_index_future = { - executor.submit( - _index_vespa_chunk, chunk, index_name, http_client, multitenant - ): chunk - for chunk in chunks - } - for future in concurrent.futures.as_completed(chunk_index_future): - # Will raise exception if any indexing raised an exception - future.result() - - finally: - if not external_executor: - executor.shutdown(wait=True) - - -def clean_chunk_id_copy( - chunk: DocMetadataAwareIndexChunk, -) -> DocMetadataAwareIndexChunk: - clean_chunk = chunk.model_copy( - update={ - "source_document": chunk.source_document.model_copy( - update={ - "id": replace_invalid_doc_id_characters(chunk.source_document.id) - } - ) - } - ) - return clean_chunk - - -def check_for_final_chunk_existence( - minimal_doc_info: MinimalDocumentIndexingInfo, - start_index: int, - index_name: str, - http_client: httpx.Client, -) -> int: - index = start_index - while True: - doc_chunk_id = get_uuid_from_chunk_info_old( - document_id=minimal_doc_info.doc_id, - chunk_id=index, - large_chunk_reference_ids=[], - ) - if not _does_doc_chunk_exist(doc_chunk_id, index_name, http_client): - return index - index += 1 - - -class BaseHTTPXClientContext(ABC): - """Abstract base class for an HTTPX client context manager.""" - - @abstractmethod - def __enter__(self) -> httpx.Client: - pass - - @abstractmethod - def __exit__(self, exc_type, exc_value, traceback): - pass - - -class GlobalHTTPXClientContext(BaseHTTPXClientContext): - """Context manager for a global HTTPX client that does not close it.""" - - def __init__(self, client: httpx.Client): - self._client = client - - def __enter__(self) -> httpx.Client: - return self._client # Reuse the global client - - def __exit__(self, exc_type, exc_value, traceback): - pass # Do nothing; don't close the global client - - -class TemporaryHTTPXClientContext(BaseHTTPXClientContext): - """Context manager for a temporary HTTPX client that closes it after use.""" - - def __init__(self, client_factory: Callable[[], httpx.Client]): - self._client_factory = client_factory - self._client: httpx.Client | None = None # Client will be created in __enter__ - - def __enter__(self) -> httpx.Client: - self._client = self._client_factory() # Create a new client - return self._client - - def __exit__(self, exc_type, exc_value, traceback): - if self._client: - self._client.close() diff --git a/backend/onyx/document_index/vespa/internal_types.py b/backend/onyx/document_index/vespa/internal_types.py deleted file mode 100644 index c165e8c53f..0000000000 --- a/backend/onyx/document_index/vespa/internal_types.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Vespa-internal data types. - -These were previously declared in the now-removed -`onyx.document_index.interfaces` module. They survive only to support Vespa's -chunk-id deletion / range-retrieval plumbing, which has not been ported to the -new generic DocumentIndex interface. -""" - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class VespaChunkRequest: - document_id: str - min_chunk_ind: int | None = None - max_chunk_ind: int | None = None - - @property - def is_capped(self) -> bool: - # If the max chunk index is not None, then the chunk request is capped. - # If the min chunk index is None, we can assume the min is 0. - return self.max_chunk_ind is not None - - @property - def range(self) -> int | None: - if self.max_chunk_ind is not None: - return (self.max_chunk_ind - (self.min_chunk_ind or 0)) + 1 - return None - - -@dataclass -class MinimalDocumentIndexingInfo: - """Minimal information necessary for indexing a document.""" - - doc_id: str - chunk_start_index: int - - -@dataclass -class EnrichedDocumentIndexingInfo(MinimalDocumentIndexingInfo): - """Enriched information necessary for indexing a document, including - version and chunk range.""" - - old_version: bool - chunk_end_index: int diff --git a/backend/onyx/document_index/vespa/kg_interactions.py b/backend/onyx/document_index/vespa/kg_interactions.py deleted file mode 100644 index f6c8f78dd2..0000000000 --- a/backend/onyx/document_index/vespa/kg_interactions.py +++ /dev/null @@ -1,63 +0,0 @@ -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, - VespaDocumentIndex, -) -from onyx.utils.logger import setup_logger -from shared_configs.configs import MULTI_TENANT - -logger = setup_logger() - - -def update_kg_chunks_vespa_info( - kg_update_requests: list[KGUChunkUpdateRequest], - index_name: str, - tenant_id: str, -) -> None: - """ """ - vespa_index = VespaDocumentIndex( - index_name=index_name, - tenant_state=TenantState(tenant_id=tenant_id, multitenant=MULTI_TENANT), - large_chunks_enabled=False, - httpx_client=None, - ) - - vespa_index.kg_chunk_updates( - kg_update_requests=kg_update_requests, tenant_id=tenant_id - ) - - -def get_kg_vespa_info_update_requests_for_document( - document_id: str, -) -> list[KGUChunkUpdateRequest]: - """Get the kg_info update requests for a document.""" - # get all entities and relationships tied to the document - with get_session_with_current_tenant() as db_session: - entities, relationships = get_document_kg_entities_and_relationships( - db_session, document_id - ) - - # create the kg vespa info - kg_entities = {entity.id_name for entity in entities} - kg_relationships = {relationship.id_name for relationship in relationships} - - # get chunks in the document - with get_session_with_current_tenant() as db_session: - num_chunks = get_num_chunks_for_document(db_session, document_id) - - # get vespa update requests - return [ - KGUChunkUpdateRequest( - document_id=document_id, - chunk_id=chunk_id, - core_entity="unused", - entities=kg_entities, - relationships=kg_relationships or None, - ) - for chunk_id in range(num_chunks) - ] diff --git a/backend/onyx/document_index/vespa/shared_utils/utils.py b/backend/onyx/document_index/vespa/shared_utils/utils.py deleted file mode 100644 index a99c64d21e..0000000000 --- a/backend/onyx/document_index/vespa/shared_utils/utils.py +++ /dev/null @@ -1,114 +0,0 @@ -import time -from typing import cast - -import httpx - -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 - -logger = setup_logger() - -# NOTE: This does not seem to be used in reality despite the Vespa Docs pointing to this code -# See here for reference: https://docs.vespa.ai/en/documents.html -# https://github.com/vespa-engine/vespa/blob/master/vespajlib/src/main/java/com/yahoo/text/Text.java - -# Define allowed ASCII characters -ALLOWED_ASCII_CHARS: list[bool] = [False] * 0x80 -ALLOWED_ASCII_CHARS[0x9] = True # tab -ALLOWED_ASCII_CHARS[0xA] = True # newline -ALLOWED_ASCII_CHARS[0xD] = True # carriage return -for i in range(0x20, 0x7F): - ALLOWED_ASCII_CHARS[i] = True # printable ASCII chars -ALLOWED_ASCII_CHARS[0x7F] = True # del - discouraged, but allowed - - -def is_text_character(codepoint: int) -> bool: - """Returns whether the given codepoint is a valid text character.""" - if codepoint < 0x80: - return ALLOWED_ASCII_CHARS[codepoint] - if codepoint < 0xD800: - return True - if codepoint <= 0xDFFF: - return False - if codepoint < 0xFDD0: - return True - if codepoint <= 0xFDEF: - return False - if codepoint >= 0x10FFFE: - return False - return (codepoint & 0xFFFF) < 0xFFFE - - -def replace_invalid_doc_id_characters(text: str) -> str: - """Replaces invalid document ID characters in text. - NOTE: this must be called at the start of every vespa-related operation or else we - risk discrepancies -> silent failures on deletion/update/insertion.""" - # There may be a more complete set of replacements that need to be made but Vespa docs are unclear - # and users only seem to be running into this error with single quotes - return text.replace("'", "_") - - -def get_vespa_http_client( - no_timeout: bool = False, http2: bool = True, timeout: int | None = None -) -> httpx.Client: - """ - Configures and returns an HTTP client for communicating with Vespa, - including authentication if needed. - """ - return httpx.Client( - cert=( - cast(tuple[str, str], (VESPA_CLOUD_CERT_PATH, VESPA_CLOUD_KEY_PATH)) - if MANAGED_VESPA - else None - ), - verify=False if not MANAGED_VESPA else True, - timeout=None if no_timeout else (timeout or VESPA_REQUEST_TIMEOUT), - http2=http2, - ) - - -def wait_for_vespa_with_timeout(wait_interval: int = 5, wait_limit: int = 60) -> bool: - """Waits for Vespa to become ready subject to a timeout. - Returns True if Vespa is ready, False otherwise.""" - - time_start = time.monotonic() - logger.info("Vespa: Readiness probe starting.") - while True: - url = f"{VESPA_APP_CONTAINER_URL}/state/v1/health" - try: - client = get_vespa_http_client() - response = client.get(url) - response.raise_for_status() - - response_dict = response.json() - if response_dict["status"]["code"] == "up": - logger.info("Vespa: Readiness probe succeeded. Continuing...") - return True - except Exception as e: - logger.warning( - "Vespa: Readiness probe failed trying to connect to %s. Exception: %s", - url, - e, - ) - - time_elapsed = time.monotonic() - time_start - if time_elapsed > wait_limit: - logger.info( - "Vespa: Readiness probe did not succeed within the timeout (%s seconds).", - wait_limit, - ) - return False - - logger.info( - "Vespa: Readiness probe ongoing. elapsed=%s timeout=%s", - format(time_elapsed, ".1f"), - format(wait_limit, ".1f"), - ) - - time.sleep(wait_interval) 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 deleted file mode 100644 index 1365ef2783..0000000000 --- a/backend/onyx/document_index/vespa/shared_utils/vespa_request_builders.py +++ /dev/null @@ -1,294 +0,0 @@ -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, - 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 - -logger = setup_logger() - - -def build_tenant_id_filter(tenant_id: str) -> str: - return f'({TENANT_ID} contains "{tenant_id}")' - - -def build_vespa_filters( - filters: IndexFilters, - *, - include_hidden: bool = False, - remove_trailing_and: bool = False, # Set to True when using as a complete Vespa query -) -> str: - def _build_or_filters(key: str, vals: list[str] | None) -> str: - """For string-based 'contains' filters, e.g. WSET fields or array fields. - Returns a bare clause like '(key contains "v1" or key contains "v2")' or "".""" - if not key or not vals: - return "" - eq_elems = [f'{key} contains "{val}"' for val in vals if val] - if not eq_elems: - return "" - return f"({' or '.join(eq_elems)})" - - def _build_weighted_set_filter(key: str, vals: list[str] | None) -> str: - """Build a Vespa weightedSet filter for large value lists. - - Uses Vespa's native weightedSet() operator instead of OR-chained - 'contains' clauses. This is critical for fields like - access_control_list where a single user may have tens of thousands - of ACL entries — OR clauses at that scale cause Vespa to reject - the query with HTTP 400.""" - if not key or not vals: - return "" - filtered = [val for val in vals if val] - if not filtered: - return "" - items = ", ".join(f'"{val}":1' for val in filtered) - return f"weightedSet({key}, {{{items}}})" - - def _build_int_or_filters(key: str, vals: list[int] | None) -> str: - """For an integer field filter. - Returns a bare clause or "".""" - if vals is None or not vals: - return "" - eq_elems = [f"{key} = {val}" for val in vals] - return f"({' or '.join(eq_elems)})" - - def _build_kg_filter( - kg_entities: list[str] | None, - kg_relationships: list[str] | None, - kg_terms: list[str] | None, - ) -> str: - if not kg_entities and not kg_relationships and not kg_terms: - return "" - - combined_filter_parts = [] - - def _build_kge(entity: str) -> str: - GENERAL = "::*" - if entity.endswith(GENERAL): - return f'({{prefix: true}}"{entity.split(GENERAL, 1)[0]}")' - else: - return f'"{entity}"' - - if kg_entities: - filter_parts = [ - f"(kg_entities contains {_build_kge(kg_entity)})" - for kg_entity in kg_entities - ] - combined_filter_parts.append(f"({' or '.join(filter_parts)})") - - # TODO: handle complex nested relationship logic (e.g., A participated, and B or C participated) - if kg_relationships: - filter_parts = [] - for kg_relationship in kg_relationships: - source, rel_type, target = split_relationship_id(kg_relationship) - filter_parts.append( - "(kg_relationships contains sameElement(" - f"source contains {_build_kge(source)}," - f'rel_type contains "{rel_type}",' - f"target contains {_build_kge(target)}))" - ) - combined_filter_parts.append(f"{' and '.join(filter_parts)}") - - # TODO: remove kg terms entirely from prompts and codebase - - return f"({' and '.join(combined_filter_parts)})" - - def _build_kg_source_filters( - kg_sources: list[str] | None, - ) -> str: - if not kg_sources: - return "" - - source_phrases = [f'{DOCUMENT_ID} contains "{source}"' for source in kg_sources] - return f"({' or '.join(source_phrases)})" - - def _build_kg_chunk_id_zero_only_filter( - kg_chunk_id_zero_only: bool, - ) -> str: - if not kg_chunk_id_zero_only: - return "" - return "(chunk_id = 0)" - - def _build_time_filter( - cutoff: datetime | None, - cutoff_upper: datetime | None = None, - untimed_doc_cutoff: timedelta = timedelta(days=92), - ) -> str: - if not cutoff and not cutoff_upper: - return "" - - clauses: list[str] = [] - if cutoff: - # Untimed docs (no doc_updated_at) are only included for an old, open- - # ended lower bound. A bounded range excludes them — an undated doc - # cannot be shown to fall within [cutoff, cutoff_upper]. - include_untimed = ( - cutoff_upper is None - and datetime.now(timezone.utc) - untimed_doc_cutoff > cutoff - ) - cutoff_secs = int(cutoff.timestamp()) - if include_untimed: - clauses.append(f"!({DOC_UPDATED_AT} < {cutoff_secs})") - else: - clauses.append(f"({DOC_UPDATED_AT} >= {cutoff_secs})") - if cutoff_upper: - clauses.append(f"({DOC_UPDATED_AT} <= {int(cutoff_upper.timestamp())})") - - return " and ".join(clauses) - - def _build_user_project_filter( - project_id: int | None, - ) -> str: - if project_id is None: - return "" - try: - pid = int(project_id) - except Exception: - return "" - return f'({USER_PROJECT} contains "{pid}")' - - def _build_persona_filter( - persona_id: int | None, - ) -> str: - if persona_id is None: - return "" - try: - pid = int(persona_id) - except Exception: - logger.warning("Invalid persona ID: %s", persona_id) - return "" - return f'({PERSONAS} contains "{pid}")' - - def _append(parts: list[str], clause: str) -> None: - if clause: - parts.append(clause) - - # Collect all top-level filter clauses, then join with " and " at the end. - filter_parts: list[str] = [] - - if not include_hidden: - filter_parts.append(f"!({HIDDEN}=true)") - - # TODO: add error condition if MULTI_TENANT and no tenant_id filter is set - if filters.tenant_id and MULTI_TENANT: - filter_parts.append(build_tenant_id_filter(filters.tenant_id)) - - # ACL filters — use weightedSet for efficient matching against the - # access_control_list weightedset field. OR-chaining thousands - # of 'contains' clauses causes Vespa to reject the query (HTTP 400) - # for users with large numbers of external permission groups. - if filters.access_control_list is not None: - _append( - filter_parts, - _build_weighted_set_filter( - ACCESS_CONTROL_LIST, filters.access_control_list - ), - ) - - # Source type filters - source_strs = ( - [s.value for s in filters.source_type] if filters.source_type else None - ) - _append(filter_parts, _build_or_filters(SOURCE_TYPE, source_strs)) - - # Tag filters - tag_attributes = None - if filters.tags: - tag_attributes = [ - f"{tag.tag_key}{INDEX_SEPARATOR}{tag.tag_value}" for tag in filters.tags - ] - _append(filter_parts, _build_or_filters(METADATA_LIST, tag_attributes)) - - # Knowledge scope: explicit knowledge attachments restrict what an - # assistant can see. When none are set, the assistant can see - # everything. - # - # persona_id_filter is a primary trigger — a persona with user files IS - # explicit knowledge, so it can start a knowledge scope on its own. - # - # project_id_filter is additive — it widens the scope to also cover - # overflowing project files but never restricts on its own (a chat - # inside a project should still search team knowledge). - knowledge_scope_parts: list[str] = [] - - _append( - knowledge_scope_parts, _build_or_filters(DOCUMENT_SETS, filters.document_set) - ) - _append(knowledge_scope_parts, _build_persona_filter(filters.persona_id_filter)) - - # project_id_filter only widens an existing scope. - if knowledge_scope_parts: - _append( - knowledge_scope_parts, - _build_user_project_filter(filters.project_id_filter), - ) - - if len(knowledge_scope_parts) > 1: - filter_parts.append("(" + " or ".join(knowledge_scope_parts) + ")") - elif len(knowledge_scope_parts) == 1: - filter_parts.append(knowledge_scope_parts[0]) - - # Vespa only indexes doc_updated_at: created_at_range is dropped (widens - # rather than narrows). - updated_at_range = filters.updated_at_range - _append( - filter_parts, - _build_time_filter( - updated_at_range.start if updated_at_range else None, - updated_at_range.end if updated_at_range else None, - ), - ) - - # # Knowledge Graph Filters - # _append(filter_parts, _build_kg_filter( - # kg_entities=filters.kg_entities, - # kg_relationships=filters.kg_relationships, - # kg_terms=filters.kg_terms, - # )) - - # _append(filter_parts, _build_kg_source_filters(filters.kg_sources)) - - # _append(filter_parts, _build_kg_chunk_id_zero_only_filter( - # filters.kg_chunk_id_zero_only or False - # )) - - filter_str = " and ".join(filter_parts) - - if filter_str and not remove_trailing_and: - filter_str += " and " - - return filter_str - - -def build_vespa_id_based_retrieval_yql( - chunk_request: VespaChunkRequest, -) -> str: - id_based_retrieval_yql_section = ( - f'({DOCUMENT_ID} contains "{chunk_request.document_id}"' - ) - - if chunk_request.is_capped: - id_based_retrieval_yql_section += ( - f" and {CHUNK_ID} >= {chunk_request.min_chunk_ind or 0}" - ) - id_based_retrieval_yql_section += ( - f" and {CHUNK_ID} <= {chunk_request.max_chunk_ind}" - ) - - id_based_retrieval_yql_section += ")" - return id_based_retrieval_yql_section diff --git a/backend/onyx/document_index/vespa/vespa_document_index.py b/backend/onyx/document_index/vespa/vespa_document_index.py deleted file mode 100644 index 152409f25d..0000000000 --- a/backend/onyx/document_index/vespa/vespa_document_index.py +++ /dev/null @@ -1,1341 +0,0 @@ -import concurrent.futures -import io -import logging -import os -import random -import re -import time -import urllib.parse -import zipfile -from collections.abc import Generator, Iterable -from datetime import datetime, timedelta -from typing import Any, BinaryIO, cast -from uuid import UUID - -import httpx -import jinja2 -import requests -from pydantic import BaseModel - -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, InferenceChunk -from onyx.document_index.chunk_content_enrichment import cleanup_content_for_chunks -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, - GlobalHTTPXClientContext, - TemporaryHTTPXClientContext, - batch_index_vespa_chunks, - check_for_final_chunk_existence, - clean_chunk_id_copy, -) -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, - 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 -from onyx.tools.tool_implementations.search.constants import KEYWORD_QUERY_HYBRID_ALPHA -from onyx.utils.batching import batch_generator -from onyx.utils.logger import setup_logger -from onyx.utils.retry_wrapper import retry_builder -from shared_configs.configs import MULTI_TENANT -from shared_configs.model_server_models import Embedding - -logger = setup_logger(__name__) -# Set the logging level to WARNING to ignore INFO and DEBUG logs from httpx. By -# default it emits INFO-level logs for every request. -httpx_logger = logging.getLogger("httpx") -httpx_logger.setLevel(logging.WARNING) - - -VESPA_SCHEMA_JINJA_FILENAME = "danswer_chunk.sd.jinja" - - -class KGVespaChunkUpdateRequest(BaseModel): - document_id: str - chunk_id: int - url: str - update_request: dict[str, dict] - - -class KGUChunkUpdateRequest(BaseModel): - """Update KG fields for a document.""" - - document_id: str - chunk_id: int - core_entity: str - entities: set[str] | None = None - relationships: set[str] | None = None - terms: set[str] | None = None - - -class KGUDocumentUpdateRequest(BaseModel): - """Update KG fields for a document.""" - - document_id: str - entities: set[str] - relationships: set[str] - terms: set[str] - - -def _generate_kg_update_request( - kg_update_request: KGUChunkUpdateRequest, -) -> dict[str, dict]: - kg_update_dict: dict[str, dict] = {} - - if kg_update_request.entities is not None: - kg_update_dict["kg_entities"] = {"assign": list(kg_update_request.entities)} - - if kg_update_request.relationships is not None: - kg_update_dict["kg_relationships"] = {"assign": []} - for relationship in kg_update_request.relationships: - source, rel_type, target = split_relationship_id(relationship) - kg_update_dict["kg_relationships"]["assign"].append( - { - "source": source, - "rel_type": rel_type, - "target": target, - } - ) - - return kg_update_dict - - -def _in_memory_zip_from_file_bytes(file_contents: dict[str, bytes]) -> BinaryIO: - zip_buffer = io.BytesIO() - with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zipf: - for filename, content in file_contents.items(): - zipf.writestr(filename, content) - zip_buffer.seek(0) - return zip_buffer - - -def _create_document_xml_lines(doc_names: list[str | None] | list[str]) -> str: - doc_lines = [ - f'' - for doc_name in doc_names - if doc_name - ] - return "\n".join(doc_lines) - - -def _add_ngrams_to_schema(schema_content: str) -> str: - schema_content = re.sub( - r"(field title type string \{[^}]*indexing: summary \| index \| attribute)", - r"\1\n match {\n gram\n gram-size: 3\n }", - schema_content, - ) - schema_content = re.sub( - r"(field content type string \{[^}]*indexing: summary \| index)", - r"\1\n match {\n gram\n gram-size: 3\n }", - schema_content, - ) - return schema_content - - -def deploy_vespa_schemas( - primary_index_name: str, - primary_embedding_dim: int, - secondary_index_name: str | None, - secondary_embedding_dim: int | None, -) -> None: - """Deploys (or redeploys) the Vespa application package containing primary - (and optionally secondary) document schemas. - - A Vespa application deploy is a single atomic operation, so primary and - secondary schemas must be packaged together and shipped in one call — - unlike OpenSearch, which has independent per-index creation calls. - """ - if MULTI_TENANT: - logger.info( - "Skipping Vespa index setup for multitenant (would wipe all indices)" - ) - return - - # TODO(security): if schema generation ever takes user-controlled input, - # swap to `jinja2.Environment(autoescape=select_autoescape(["xml"]))`. - jinja_env = jinja2.Environment() # noqa: S701 — renders Vespa schema files, not HTML - - deploy_url = f"{VESPA_APPLICATION_ENDPOINT}/tenant/default/prepareandactivate" - logger.notice("Deploying Vespa application package to %s", deploy_url) - - vespa_schema_path = os.path.join( - os.getcwd(), "onyx", "document_index", "vespa", "app_config" - ) - schema_jinja_file = os.path.join( - vespa_schema_path, "schemas", VESPA_SCHEMA_JINJA_FILENAME - ) - services_jinja_file = os.path.join(vespa_schema_path, "services.xml.jinja") - overrides_jinja_file = os.path.join( - vespa_schema_path, "validation-overrides.xml.jinja" - ) - - with open(services_jinja_file, "r") as services_f: - schema_names = [primary_index_name, secondary_index_name] - doc_lines = _create_document_xml_lines(schema_names) - services_template_str = services_f.read() - services_template = jinja_env.from_string(services_template_str) - services = services_template.render( - document_elements=doc_lines, - num_search_threads=str(VESPA_SEARCHER_THREADS), - ) - - kv_store = get_shared_kv_store() - - needs_reindexing = False - try: - needs_reindexing = cast(bool, kv_store.load(KV_REINDEX_KEY)) - except Exception: - logger.debug("Could not load the reindexing flag. Using ngrams") - - with open(overrides_jinja_file, "r") as overrides_f: - overrides_template_str = overrides_f.read() - overrides_template = jinja_env.from_string(overrides_template_str) - now = datetime.now() - date_in_7_days = now + timedelta(days=7) - formatted_date = date_in_7_days.strftime("%Y-%m-%d") - overrides = overrides_template.render(until_date=formatted_date) - - zip_dict = { - "services.xml": services.encode("utf-8"), - "validation-overrides.xml": overrides.encode("utf-8"), - } - - with open(schema_jinja_file, "r") as schema_f: - template_str = schema_f.read() - - template = jinja_env.from_string(template_str) - schema = template.render( - multi_tenant=MULTI_TENANT, - schema_name=primary_index_name, - dim=primary_embedding_dim, - ) - schema = _add_ngrams_to_schema(schema) if needs_reindexing else schema - zip_dict[f"schemas/{primary_index_name}.sd"] = schema.encode("utf-8") - - if secondary_index_name: - if secondary_embedding_dim is None: - raise ValueError("Secondary index embedding dimension is required") - upcoming_schema = template.render( - multi_tenant=MULTI_TENANT, - schema_name=secondary_index_name, - dim=secondary_embedding_dim, - ) - zip_dict[f"schemas/{secondary_index_name}.sd"] = upcoming_schema.encode("utf-8") - - zip_file = _in_memory_zip_from_file_bytes(zip_dict) - headers = {"Content-Type": "application/zip"} - response = requests.post(deploy_url, headers=headers, data=zip_file) - if response.status_code != 200: - logger.error("Failed to prepare Vespa Onyx Index. Response: %s", response.text) - raise RuntimeError( - f"Failed to prepare Vespa Onyx Index. Response: {response.text}" - ) - - -def register_multitenant_vespa_indices( - indices: list[str], - embedding_dims: list[int], -) -> None: - """Registers a set of multi-tenant Vespa schemas in one application deploy.""" - if not MULTI_TENANT: - raise ValueError("Multi-tenant is not enabled") - - deploy_url = f"{VESPA_APPLICATION_ENDPOINT}/tenant/default/prepareandactivate" - logger.info("Deploying Vespa application package to %s", deploy_url) - - vespa_schema_path = os.path.join( - os.getcwd(), "onyx", "document_index", "vespa", "app_config" - ) - schema_jinja_file = os.path.join( - vespa_schema_path, "schemas", VESPA_SCHEMA_JINJA_FILENAME - ) - services_jinja_file = os.path.join(vespa_schema_path, "services.xml.jinja") - overrides_jinja_file = os.path.join( - vespa_schema_path, "validation-overrides.xml.jinja" - ) - - # TODO(security): if schema generation ever takes user-controlled input, - # swap to `jinja2.Environment(autoescape=select_autoescape(["xml"]))`. - jinja_env = jinja2.Environment() # noqa: S701 — renders Vespa schema files, not HTML - - with open(services_jinja_file, "r") as services_f: - schema_names = list(indices) - doc_lines = _create_document_xml_lines(schema_names) - services_template_str = services_f.read() - services_template = jinja_env.from_string(services_template_str) - services = services_template.render( - document_elements=doc_lines, - num_search_threads=str(VESPA_SEARCHER_THREADS), - ) - - kv_store = get_shared_kv_store() - needs_reindexing = False - try: - needs_reindexing = cast(bool, kv_store.load(KV_REINDEX_KEY)) - except Exception: - logger.debug("Could not load the reindexing flag. Using ngrams") - - with open(overrides_jinja_file, "r") as overrides_f: - overrides_template_str = overrides_f.read() - overrides_template = jinja_env.from_string(overrides_template_str) - now = datetime.now() - date_in_7_days = now + timedelta(days=7) - formatted_date = date_in_7_days.strftime("%Y-%m-%d") - overrides = overrides_template.render(until_date=formatted_date) - - zip_dict = { - "services.xml": services.encode("utf-8"), - "validation-overrides.xml": overrides.encode("utf-8"), - } - - with open(schema_jinja_file, "r") as schema_f: - schema_template_str = schema_f.read() - - schema_template = jinja_env.from_string(schema_template_str) - - for i, index_name in enumerate(indices): - embedding_dim = embedding_dims[i] - logger.info( - "Creating index: %s with embedding dimension: %s", - index_name, - embedding_dim, - ) - schema = schema_template.render( - multi_tenant=MULTI_TENANT, - schema_name=index_name, - dim=embedding_dim, - ) - schema = _add_ngrams_to_schema(schema) if needs_reindexing else schema - zip_dict[f"schemas/{index_name}.sd"] = schema.encode("utf-8") - - zip_file = _in_memory_zip_from_file_bytes(zip_dict) - headers = {"Content-Type": "application/zip"} - response = requests.post(deploy_url, headers=headers, data=zip_file) - if response.status_code != 200: - raise RuntimeError( - f"Failed to prepare Vespa Onyx Indexes. Response: {response.text}" - ) - - -class _VespaDeleteRequest: - def __init__(self, document_id: str, index_name: str) -> None: - self.document_id = document_id - encoded_doc_id = urllib.parse.quote_plus(self.document_id) - self.url = f"{VESPA_APPLICATION_ENDPOINT}/document/v1/{index_name}/{index_name}/docid/{encoded_doc_id}" - - -def _enrich_basic_chunk_info( - index_name: str, - http_client: httpx.Client, - document_id: str, - previous_chunk_count: int | None, - new_chunk_count: int, -) -> EnrichedDocumentIndexingInfo: - """Determines which chunks need to be deleted during document reindexing. - - When a document is reindexed, it may have fewer chunks than before. This - function identifies the range of old chunks that need to be deleted by - comparing the new chunk count with the previous chunk count. - - Example: - If a document previously had 10 chunks (0-9) and now has 7 chunks (0-6), - this function identifies that chunks 7-9 need to be deleted. - - Args: - index_name: The Vespa index/schema name. - http_client: HTTP client for making requests to Vespa. - document_id: The Vespa-sanitized ID of the document being reindexed. - previous_chunk_count: The total number of chunks the document had before - reindexing. None for documents using the legacy chunk ID system. - new_chunk_count: The total number of chunks the document has after - reindexing. This becomes the starting index for deletion since - chunks are 0-indexed. - - Returns: - EnrichedDocumentIndexingInfo with chunk_start_index set to - new_chunk_count (where deletion begins) and chunk_end_index set to - previous_chunk_count (where deletion ends). - """ - # Technically last indexed chunk index +1. - last_indexed_chunk = previous_chunk_count - # If the document has no `chunk_count` in the database, we know that it - # has the old chunk ID system and we must check for the final chunk index. - is_old_version = False - if last_indexed_chunk is None: - is_old_version = True - minimal_doc_info = MinimalDocumentIndexingInfo( - doc_id=document_id, chunk_start_index=new_chunk_count - ) - last_indexed_chunk = check_for_final_chunk_existence( - minimal_doc_info=minimal_doc_info, - start_index=new_chunk_count, - index_name=index_name, - http_client=http_client, - ) - - assert last_indexed_chunk is not None and last_indexed_chunk >= 0, ( - f"Bug: Last indexed chunk index is None or less than 0 for document: {document_id}." - ) - - enriched_doc_info = EnrichedDocumentIndexingInfo( - doc_id=document_id, - chunk_start_index=new_chunk_count, - chunk_end_index=last_indexed_chunk, - old_version=is_old_version, - ) - return enriched_doc_info - - -@retry_builder( - tries=3, - delay=1, - backoff=2, - exceptions=httpx.HTTPError, -) -def _update_single_chunk( - doc_chunk_id: UUID, - index_name: str, - doc_id: str, - http_client: httpx.Client, - update_request: MetadataUpdateRequest, -) -> None: - """Updates a single document chunk in Vespa. - - TODO(andrei): Couldn't this be batched? - - Args: - doc_chunk_id: The ID of the chunk to update. - index_name: The index the chunk belongs to. - doc_id: The ID of the document the chunk belongs to. Used only for - logging. - http_client: The HTTP client to use to make the request. - update_request: Metadata update request object received in the bulk - update method containing fields to update. - """ - - class _Boost(BaseModel): - model_config = {"frozen": True} - assign: float - - class _DocumentSets(BaseModel): - model_config = {"frozen": True} - assign: dict[str, int] - - class _AccessControl(BaseModel): - model_config = {"frozen": True} - assign: dict[str, int] - - class _Hidden(BaseModel): - model_config = {"frozen": True} - assign: bool - - class _UserProjects(BaseModel): - model_config = {"frozen": True} - assign: list[int] - - class _Personas(BaseModel): - model_config = {"frozen": True} - assign: list[int] - - class _VespaPutFields(BaseModel): - model_config = {"frozen": True} - # The names of these fields are based the Vespa schema. Changes to the - # schema require changes here. These names were originally found in - # backend/onyx/document_index/vespa_constants.py. - boost: _Boost | None = None - document_sets: _DocumentSets | None = None - access_control_list: _AccessControl | None = None - hidden: _Hidden | None = None - user_project: _UserProjects | None = None - personas: _Personas | None = None - - class _VespaPutRequest(BaseModel): - model_config = {"frozen": True} - fields: _VespaPutFields - - boost_update: _Boost | None = ( - _Boost(assign=update_request.boost) - if update_request.boost is not None - else None - ) - document_sets_update: _DocumentSets | None = ( - _DocumentSets(assign=dict.fromkeys(update_request.document_sets, 1)) - if update_request.document_sets is not None - else None - ) - access_update: _AccessControl | None = ( - _AccessControl(assign=dict.fromkeys(update_request.access.to_acl(), 1)) - if update_request.access is not None - else None - ) - hidden_update: _Hidden | None = ( - _Hidden(assign=update_request.hidden) - if update_request.hidden is not None - else None - ) - user_projects_update: _UserProjects | None = ( - _UserProjects(assign=list(update_request.project_ids)) - if update_request.project_ids is not None - else None - ) - personas_update: _Personas | None = ( - _Personas(assign=list(update_request.persona_ids)) - if update_request.persona_ids is not None - else None - ) - - vespa_put_fields = _VespaPutFields( - boost=boost_update, - document_sets=document_sets_update, - access_control_list=access_update, - hidden=hidden_update, - user_project=user_projects_update, - personas=personas_update, - ) - - vespa_put_request = _VespaPutRequest( - fields=vespa_put_fields, - ) - - vespa_url = f"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{doc_chunk_id}?create=true" - - try: - resp = http_client.put( - vespa_url, - headers={"Content-Type": "application/json"}, - json=vespa_put_request.model_dump( - exclude_none=True - ), # NOTE: Important to not produce null fields in the json. - ) - resp.raise_for_status() - except httpx.HTTPStatusError as e: - logger.error( - "Failed to update doc chunk %s (doc_id=%s). Code: %s. Details: %s", - doc_chunk_id, - doc_id, - e.response.status_code, - e.response.text, - ) - # Re-raise so the @retry decorator will catch and retry, unless the - # status code is < 5xx, in which case wrap the exception in something - # other than an HTTPError to skip retries. - if e.response.status_code >= 500: - raise - raise RuntimeError( - f"Non-retryable error updating chunk {doc_chunk_id}: {e}" - ) from e - - -class VespaDocumentIndex(DocumentIndex): - """Vespa-specific implementation of the DocumentIndex interface. - - This class provides document indexing, retrieval, and management operations - for a Vespa search engine instance. It handles the complete lifecycle of - document chunks within a specific Vespa index/schema. - """ - - def __init__( - self, - index_name: str, - tenant_state: TenantState, - large_chunks_enabled: bool, - httpx_client: httpx.Client | None = None, - ) -> None: - self._index_name = index_name - self._tenant_id = tenant_state.tenant_id - self._large_chunks_enabled = large_chunks_enabled - # NOTE: using `httpx` here since `requests` doesn't support HTTP2. This - # is beneficial for indexing / updates / deletes since we have to make a - # large volume of requests. - self._httpx_client_context: BaseHTTPXClientContext - if httpx_client: - # Use the provided client. Because this client is presumed global, - # it does not close after exiting a context manager. - self._httpx_client_context = GlobalHTTPXClientContext(httpx_client) - else: - # We did not receive a client, so create one what will close after - # exiting a context manager. - self._httpx_client_context = TemporaryHTTPXClientContext( - get_vespa_http_client - ) - self._multitenant = tenant_state.multitenant - - def verify_and_create_index_if_necessary(self, embedding_dim: int) -> None: - """Deploys (or redeploys) the Vespa application package containing this - index's schema. - - NOTE: For paired primary+secondary deployments use - :class:`VespaIndexPair`, which packages both schemas into a single - deploy call. Calling this on a single VespaDocumentIndex deploys an - application with only this one schema, which would wipe a sibling - secondary schema if one exists. - """ - deploy_vespa_schemas( - primary_index_name=self._index_name, - primary_embedding_dim=embedding_dim, - secondary_index_name=None, - secondary_embedding_dim=None, - ) - - def index( - self, - chunks: Iterable[DocMetadataAwareIndexChunk], - indexing_metadata: IndexingMetadata, - ) -> list[DocumentInsertionRecord]: - doc_id_to_chunk_cnt_diff = indexing_metadata.doc_id_to_chunk_cnt_diff - doc_id_to_previous_chunk_cnt = { - doc_id: chunk_cnt_diff.old_chunk_cnt - for doc_id, chunk_cnt_diff in doc_id_to_chunk_cnt_diff.items() - } - doc_id_to_new_chunk_cnt = { - doc_id: chunk_cnt_diff.new_chunk_cnt - for doc_id, chunk_cnt_diff in doc_id_to_chunk_cnt_diff.items() - } - assert ( - len(doc_id_to_chunk_cnt_diff) - == len(doc_id_to_previous_chunk_cnt) - == len(doc_id_to_new_chunk_cnt) - ), "Bug: Doc ID to chunk maps have different lengths." - - # Vespa has restrictions on valid characters, yet document IDs come from - # external w.r.t. this class. We need to sanitize them. - # - # Instead of materializing all cleaned chunks upfront, we stream them - # through a generator that cleans IDs and builds the original-ID mapping - # incrementally as chunks flow into Vespa. - def _clean_and_track( - chunks_iter: Iterable[DocMetadataAwareIndexChunk], - id_map: dict[str, str], - seen_ids: set[str], - ) -> Generator[DocMetadataAwareIndexChunk, None, None]: - """Cleans chunk IDs and builds the original-ID mapping - incrementally as chunks flow through, avoiding a separate - materialization pass.""" - for chunk in chunks_iter: - original_id = chunk.source_document.id - cleaned = clean_chunk_id_copy(chunk) - cleaned_id = cleaned.source_document.id - # Needed so the final DocumentInsertionRecord returned can have - # the original document ID. cleaned_chunks might not contain IDs - # exactly as callers supplied them. - id_map[cleaned_id] = original_id - seen_ids.add(cleaned_id) - yield cleaned - - new_document_id_to_original_document_id: dict[str, str] = {} - all_cleaned_doc_ids: set[str] = set() - - existing_docs: set[str] = set() - - with ( - concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor, - self._httpx_client_context as http_client, - ): - # We require the start and end index for each document in order to - # know precisely which chunks to delete. This information exists for - # documents that have `chunk_count` in the database, but not for - # `old_version` documents. - enriched_doc_infos: list[EnrichedDocumentIndexingInfo] = [ - _enrich_basic_chunk_info( - index_name=self._index_name, - http_client=http_client, - document_id=doc_id, - previous_chunk_count=doc_id_to_previous_chunk_cnt[doc_id], - new_chunk_count=doc_id_to_new_chunk_cnt[doc_id], - ) - for doc_id in doc_id_to_chunk_cnt_diff.keys() - # TODO(andrei), WARNING: Don't we need to sanitize these doc IDs? - ] - - for enriched_doc_info in enriched_doc_infos: - # If the document has previously indexed chunks, we know it - # previously existed and this is a reindex. - if enriched_doc_info.chunk_end_index: - existing_docs.add(enriched_doc_info.doc_id) - - # Now, for each doc, we know exactly where to start and end our - # deletion. So let's generate the chunk IDs for each chunk to - # delete. - # WARNING: This code seems to use - # indexing_metadata.doc_id_to_chunk_cnt_diff as the source of truth - # for which chunks to delete. This implies that the onus is on the - # caller to ensure doc_id_to_chunk_cnt_diff only contains docs - # relevant to the chunks argument to this method. This should not be - # the contract of DocumentIndex; and this code is only a refactor - # from old code. It would seem we should use all_cleaned_doc_ids as - # the source of truth. - chunks_to_delete = get_document_chunk_ids( - enriched_document_info_list=enriched_doc_infos, - tenant_id=self._tenant_id, - large_chunks_enabled=self._large_chunks_enabled, - ) - - # Delete old Vespa documents. - for doc_chunk_ids_batch in batch_generator(chunks_to_delete, BATCH_SIZE): - delete_vespa_chunks( - doc_chunk_ids=doc_chunk_ids_batch, - index_name=self._index_name, - http_client=http_client, - executor=executor, - ) - - # Insert new Vespa documents, streaming through the cleaning - # pipeline so chunks are never fully materialized. - cleaned_chunks = _clean_and_track( - chunks, - new_document_id_to_original_document_id, - all_cleaned_doc_ids, - ) - for chunk_batch in batch_generator( - cleaned_chunks, min(BATCH_SIZE, MAX_CHUNKS_PER_DOC_BATCH) - ): - batch_index_vespa_chunks( - chunks=chunk_batch, - index_name=self._index_name, - http_client=http_client, - multitenant=self._multitenant, - executor=executor, - ) - - return [ - DocumentInsertionRecord( - document_id=new_document_id_to_original_document_id[cleaned_doc_id], - already_existed=cleaned_doc_id in existing_docs, - ) - for cleaned_doc_id in all_cleaned_doc_ids - ] - - def delete(self, document_id: str, chunk_count: int | None = None) -> int: - total_chunks_deleted = 0 - - sanitized_doc_id = replace_invalid_doc_id_characters(document_id) - - with ( - concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor, - self._httpx_client_context as http_client, - ): - enriched_doc_info = _enrich_basic_chunk_info( - index_name=self._index_name, - http_client=http_client, - document_id=sanitized_doc_id, - previous_chunk_count=chunk_count, - new_chunk_count=0, - ) - chunks_to_delete = get_document_chunk_ids( - enriched_document_info_list=[enriched_doc_info], - tenant_id=self._tenant_id, - large_chunks_enabled=self._large_chunks_enabled, - ) - - for doc_chunk_ids_batch in batch_generator(chunks_to_delete, BATCH_SIZE): - total_chunks_deleted += len(doc_chunk_ids_batch) - delete_vespa_chunks( - doc_chunk_ids=doc_chunk_ids_batch, - index_name=self._index_name, - http_client=http_client, - executor=executor, - ) - - return total_chunks_deleted - - def update( - self, - update_requests: list[MetadataUpdateRequest], - ) -> None: - # WARNING: This method can be called by document_index_metadata_sync_task, - # which is kicked off by check_for_vespa_sync_task, notably before a document - # has finished indexing. In this way, chunk_count below could be unknown - # even for chunks not on the "old" chunk ID system; i.e. there could be - # a race condition. Passing in None to _enrich_basic_chunk_info should - # handle this, but a higher level TODO might be to not run update at all - # on connectors that are still indexing, and therefore do not yet have a - # chunk count because update_docs_chunk_count__no_commit has not been - # run yet. - with self._httpx_client_context as httpx_client: - # Each invocation of this method can contain multiple update requests. - for update_request in update_requests: - # Each update request can correspond to multiple documents. - for doc_id in update_request.document_ids: - # NOTE: -1 represents an unknown chunk count. - chunk_count = update_request.doc_id_to_chunk_cnt[doc_id] - sanitized_doc_id = replace_invalid_doc_id_characters(doc_id) - enriched_doc_info = _enrich_basic_chunk_info( - index_name=self._index_name, - http_client=httpx_client, - document_id=sanitized_doc_id, - previous_chunk_count=chunk_count if chunk_count >= 0 else None, - new_chunk_count=0, # WARNING: This semantically makes no sense and is misusing this function. - ) - - doc_chunk_ids = get_document_chunk_ids( - enriched_document_info_list=[enriched_doc_info], - tenant_id=self._tenant_id, - large_chunks_enabled=self._large_chunks_enabled, - ) - - for doc_chunk_id in doc_chunk_ids: - _update_single_chunk( - doc_chunk_id, - self._index_name, - # NOTE: Used only for logging, raw ID is ok here. - doc_id, - httpx_client, - update_request, - ) - - logger.info( - "Updated %s chunks for document %s.", len(doc_chunk_ids), doc_id - ) - - def id_based_retrieval( - self, - chunk_requests: list[DocumentSectionRequest], - filters: IndexFilters, - batch_retrieval: bool = False, - ) -> list[InferenceChunk]: - sanitized_chunk_requests = [ - VespaChunkRequest( - document_id=replace_invalid_doc_id_characters( - chunk_request.document_id - ), - min_chunk_ind=chunk_request.min_chunk_ind, - max_chunk_ind=chunk_request.max_chunk_ind, - ) - for chunk_request in chunk_requests - ] - - if self._multitenant: - filters = filters.model_copy(update={"tenant_id": self._tenant_id}) - - if batch_retrieval: - return cleanup_content_for_chunks( - batch_search_api_retrieval( - index_name=self._index_name, - chunk_requests=sanitized_chunk_requests, - filters=filters, - # No one was passing in this parameter in the legacy - # interface, it always defaulted to False. - get_large_chunks=False, - ) - ) - return cleanup_content_for_chunks( - parallel_visit_api_retrieval( - index_name=self._index_name, - chunk_requests=sanitized_chunk_requests, - filters=filters, - # No one was passing in this parameter in the legacy interface, - # it always defaulted to False. - get_large_chunks=False, - ) - ) - - def hybrid_retrieval( - self, - query: str, - query_embedding: Embedding, - final_keywords: list[str] | None, - query_type: QueryType, - filters: IndexFilters, - num_to_retrieve: int, - ) -> list[InferenceChunk]: - vespa_where_clauses = build_vespa_filters(filters) - # Avoid over-fetching a very large candidate set for global-phase reranking. - # Keep enough headroom for quality while capping cost on larger indices. - target_hits = min(max(4 * num_to_retrieve, 100), RERANK_COUNT) - - yql = ( - YQL_BASE.format(index_name=self._index_name) - + vespa_where_clauses - + f"(({{targetHits: {target_hits}}}nearestNeighbor(embeddings, query_embedding)) " - + f"or ({{targetHits: {target_hits}}}nearestNeighbor(title_embedding, query_embedding)) " - + 'or ({grammar: "weakAnd"}userInput(@query)) ' - + f'or ({{defaultIndex: "{CONTENT_SUMMARY}"}}userInput(@query)))' - ) - - final_query = " ".join(final_keywords) if final_keywords else query - - ranking_profile = ( - f"hybrid_search_{query_type.value}_base_{len(query_embedding)}" - ) - - logger.info("Selected ranking profile: %s", ranking_profile) - - logger.debug("Query YQL: %s", yql) - - # In this interface we do not pass in hybrid alpha. Tracing the codepath - # of the legacy Vespa interface, it so happens that KEYWORD always - # corresponds to an alpha of 0.2 (from KEYWORD_QUERY_HYBRID_ALPHA), and - # SEMANTIC to 0.5 (from HYBRID_ALPHA). HYBRID_ALPHA_KEYWORD was only - # used in dead code so we do not use it here. - hybrid_alpha = ( - KEYWORD_QUERY_HYBRID_ALPHA - if query_type == QueryType.KEYWORD - else HYBRID_ALPHA - ) - - params: dict[str, str | int | float] = { - "yql": yql, - "query": final_query, - "input.query(query_embedding)": str(query_embedding), - "input.query(decay_factor)": str(DOC_TIME_DECAY * RECENCY_BIAS_MULTIPLIER), - "input.query(alpha)": hybrid_alpha, - "input.query(title_content_ratio)": TITLE_CONTENT_RATIO, - "hits": num_to_retrieve, - "ranking.profile": ranking_profile, - "timeout": VESPA_TIMEOUT, - } - - return cleanup_content_for_chunks(query_vespa(params)) - - def keyword_retrieval( - self, - query: str, - filters: IndexFilters, - num_to_retrieve: int, - include_hidden: bool = False, - ) -> list[InferenceChunk]: - # Ported from the legacy Vespa admin_retrieval: pure-keyword search over - # weakAnd(userInput) plus a content_summary userInput pass for - # highlighting (n-gram highlighting is broken / not behaving as - # desired). The caller decides whether hidden documents are surfaced; - # admin search passes include_hidden=True. - vespa_where_clauses = build_vespa_filters( - filters, include_hidden=include_hidden - ) - yql = ( - YQL_BASE.format(index_name=self._index_name) - + vespa_where_clauses - + '({grammar: "weakAnd"}userInput(@query) ' - + f'or ({{defaultIndex: "{CONTENT_SUMMARY}"}}userInput(@query)))' - ) - - params: dict[str, str | int] = { - "yql": yql, - "query": query, - "hits": num_to_retrieve, - "ranking.profile": "admin_search", - "timeout": VESPA_TIMEOUT, - } - - return cleanup_content_for_chunks(query_vespa(params)) - - def semantic_retrieval( - self, - query_embedding: Embedding, - filters: IndexFilters, - num_to_retrieve: int, - ) -> list[InferenceChunk]: - raise NotImplementedError - - def random_retrieval( - self, - filters: IndexFilters, - num_to_retrieve: int = 100, - dirty: bool | None = None, # noqa: ARG002 - ) -> list[InferenceChunk]: - vespa_where_clauses = build_vespa_filters(filters, remove_trailing_and=True) - - yql = YQL_BASE.format(index_name=self._index_name) + vespa_where_clauses - - random_seed = random.randint(0, 1_000_000) - - params: dict[str, str | int | float] = { - "yql": yql, - "hits": num_to_retrieve, - "timeout": VESPA_TIMEOUT, - "ranking.profile": "random_", - "ranking.properties.random.seed": random_seed, - } - - return cleanup_content_for_chunks(query_vespa(params)) - - def get_raw_document_chunks(self, document_id: str) -> list[dict[str, Any]]: - """Gets all raw document chunks for a document as returned by Vespa. - - Used in the Vespa migration task. - - Args: - document_id: The ID of the document to get chunks for. - - Returns: - List of raw document chunks. - """ - # Vespa doc IDs are sanitized using replace_invalid_doc_id_characters. - sanitized_document_id = replace_invalid_doc_id_characters(document_id) - chunk_request = VespaChunkRequest(document_id=sanitized_document_id) - raw_chunks = get_chunks_via_visit_api( - chunk_request=chunk_request, - index_name=self._index_name, - filters=IndexFilters(access_control_list=None, tenant_id=self._tenant_id), - get_large_chunks=False, - short_tensor_format=True, - ) - # Vespa returns other metadata around the actual document chunk. The raw - # chunk we're interested in is in the "fields" field. - raw_document_chunks = [chunk["fields"] for chunk in raw_chunks] - return raw_document_chunks - - def get_all_raw_document_chunks_paginated( - self, - continuation_token_map: dict[int, str | None], - page_size: int, - ) -> tuple[list[dict[str, Any]], dict[int, str | None]]: - """Gets all the chunks in Vespa, paginated. - - Used in the chunk-level Vespa-to-OpenSearch migration task. - - Args: - continuation_token: Token returned by Vespa representing a page - offset. None to start from the beginning. Defaults to None. - page_size: Best-effort batch size for the visit. - - Returns: - Tuple of (list of chunk dicts, next continuation token or None). The - continuation token is None when the visit is complete. - """ - raw_chunks, next_continuation_token_map = get_all_chunks_paginated( - index_name=self._index_name, - tenant_state=TenantState( - tenant_id=self._tenant_id, multitenant=MULTI_TENANT - ), - continuation_token_map=continuation_token_map, - page_size=page_size, - ) - return raw_chunks, next_continuation_token_map - - def index_raw_chunks(self, chunks: list[dict[str, Any]]) -> None: - """Indexes raw document chunks into Vespa. - - To only be used in tests. Not for production. - """ - json_header = { - "Content-Type": "application/json", - } - with self._httpx_client_context as http_client: - for chunk in chunks: - chunk_id = str( - get_uuid_from_chunk_info( - document_id=chunk[DOCUMENT_ID], - chunk_id=chunk[CHUNK_ID], - tenant_id=self._tenant_id, - ) - ) - vespa_url = f"{DOCUMENT_ID_ENDPOINT.format(index_name=self._index_name)}/{chunk_id}" - response = http_client.post( - vespa_url, - headers=json_header, - json={"fields": chunk}, - ) - response.raise_for_status() - - def get_chunk_count(self) -> int: - """Returns the exact number of document chunks in Vespa for this tenant. - - Uses the Vespa Search API with `limit 0` and `ranking.profile=unranked` - to get an exact count without fetching any document data. - - Includes large chunks. There is no way to filter these out using the - Search API. - """ - where_clause = ( - f'tenant_id contains "{self._tenant_id}"' if self._multitenant else "true" - ) - yql = f"select documentid from {self._index_name} where {where_clause} limit 0" # noqa: S608 - Vespa YQL with internal index_name/tenant_id - params: dict[str, str | int] = { - "yql": yql, - "ranking.profile": "unranked", - "timeout": VESPA_TIMEOUT, - } - - with get_vespa_http_client() as http_client: - response = http_client.post(SEARCH_ENDPOINT, json=params) - response.raise_for_status() - response_data = response.json() - return response_data["root"]["fields"]["totalCount"] - - @property - def index_name(self) -> str: - return self._index_name - - @property - def httpx_client_context(self) -> BaseHTTPXClientContext: - return self._httpx_client_context - - @classmethod - def _apply_kg_chunk_updates_batched( - cls, - updates: list[KGVespaChunkUpdateRequest], - httpx_client: httpx.Client, - batch_size: int = BATCH_SIZE, - ) -> None: - """Runs a batch of KG chunk updates in parallel via the - ThreadPoolExecutor.""" - - @retry_builder(tries=3, delay=1, backoff=2, jitter=(0.0, 1.0)) - def _kg_update_chunk( - update: KGVespaChunkUpdateRequest, http_client: httpx.Client - ) -> httpx.Response: - return http_client.put( - update.url, - headers={"Content-Type": "application/json"}, - json=update.update_request, - ) - - with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor: - for update_batch in batch_generator(updates, batch_size): - future_to_document_id = { - executor.submit( - _kg_update_chunk, - update, - httpx_client, - ): update.document_id - for update in update_batch - } - for future in concurrent.futures.as_completed(future_to_document_id): - res = future.result() - try: - res.raise_for_status() - except httpx.HTTPStatusError: - # http_client is httpx.Client, so raise_for_status raises - # httpx.HTTPStatusError; logging here is the only way to - # attach the doc id, since the canonical exception - # propagates unchanged for callers. - logger.error( - "Failed to update document %s. Response: %s", - future_to_document_id[future], - res.text, - ) - raise - - def kg_chunk_updates( - self, - kg_update_requests: list[KGUChunkUpdateRequest], - tenant_id: str, - ) -> None: - """Applies a batch of knowledge-graph metadata updates to chunks in - this Vespa index.""" - - processed_updates_requests: list[KGVespaChunkUpdateRequest] = [] - update_start = time.monotonic() - - for kg_update_request in kg_update_requests: - kg_update_dict: dict[str, dict] = { - "fields": _generate_kg_update_request(kg_update_request) - } - if not kg_update_dict["fields"]: - logger.error("Update request received but nothing to update") - continue - - doc_chunk_id = get_uuid_from_chunk_info( - document_id=kg_update_request.document_id, - chunk_id=kg_update_request.chunk_id, - tenant_id=tenant_id, - large_chunk_id=None, - ) - - processed_updates_requests.append( - KGVespaChunkUpdateRequest( - document_id=kg_update_request.document_id, - chunk_id=kg_update_request.chunk_id, - url=f"{DOCUMENT_ID_ENDPOINT.format(index_name=self._index_name)}/{doc_chunk_id}", - update_request=kg_update_dict, - ) - ) - - with self._httpx_client_context as httpx_client: - self._apply_kg_chunk_updates_batched( - processed_updates_requests, httpx_client - ) - logger.debug( - "Updated %d vespa documents in %.2f seconds", - len(processed_updates_requests), - time.monotonic() - update_start, - ) - - -class VespaIndexPair(DocumentIndex): - """Pair wrapper that fans operations out to a primary Vespa index and an - optional secondary one. - - Mirrors today's `VespaIndex` semantics minus the OLD-interface translation: - `index` writes only to primary (a separate pipeline handles secondary - backfill); `delete`, `update`, and `verify_and_create_index_if_necessary` - fan out to both. - - Vespa's schema deploy is a single atomic operation, so - `verify_and_create_index_if_necessary` packages both schemas into one call - rather than calling each individually. - """ - - def __init__( - self, - primary: VespaDocumentIndex, - secondary: VespaDocumentIndex | None, - # Embedding info is needed at deploy time only, so we accept it on the - # pair rather than threading it through every retrieval call. - secondary_index_name: str | None, - secondary_embedding_dim: int | None, - ) -> None: - # All three secondary fields must be set together or all None — checked - # independently so a partially-set state surfaces here rather than - # deferring to a less informative ValueError inside deploy_vespa_schemas. - secondary_set = secondary is not None - name_set = secondary_index_name is not None - dim_set = secondary_embedding_dim is not None - if not (secondary_set == name_set == dim_set): - raise ValueError( - "Bug: secondary VespaDocumentIndex, secondary_index_name, and " - "secondary_embedding_dim must all be set together or all be None. " - f"Got: secondary={secondary_set}, index_name={name_set}, " - f"embedding_dim={dim_set}." - ) - self._primary = primary - self._secondary = secondary - self._secondary_index_name = secondary_index_name - self._secondary_embedding_dim = secondary_embedding_dim - - def verify_and_create_index_if_necessary(self, embedding_dim: int) -> None: - deploy_vespa_schemas( - primary_index_name=self._primary._index_name, - primary_embedding_dim=embedding_dim, - secondary_index_name=self._secondary_index_name, - secondary_embedding_dim=self._secondary_embedding_dim, - ) - - def index( - self, - chunks: Iterable[DocMetadataAwareIndexChunk], - indexing_metadata: IndexingMetadata, - ) -> list[DocumentInsertionRecord]: - # Secondary is filled by a separate pipeline; primary only here. - return self._primary.index(chunks, indexing_metadata) - - def delete(self, document_id: str, chunk_count: int | None = None) -> int: - total = self._primary.delete(document_id, chunk_count) - if self._secondary is not None: - total += self._secondary.delete(document_id, chunk_count) - return total - - def update(self, update_requests: list[MetadataUpdateRequest]) -> None: - self._primary.update(update_requests) - if self._secondary is not None: - self._secondary.update(update_requests) - - def id_based_retrieval( - self, - chunk_requests: list[DocumentSectionRequest], - filters: IndexFilters, - batch_retrieval: bool = False, - ) -> list[InferenceChunk]: - return self._primary.id_based_retrieval( - chunk_requests, filters, batch_retrieval - ) - - def hybrid_retrieval( - self, - query: str, - query_embedding: Embedding, - final_keywords: list[str] | None, - query_type: QueryType, - filters: IndexFilters, - num_to_retrieve: int, - ) -> list[InferenceChunk]: - return self._primary.hybrid_retrieval( - query, - query_embedding, - final_keywords, - query_type, - filters, - num_to_retrieve, - ) - - def keyword_retrieval( - self, - query: str, - filters: IndexFilters, - num_to_retrieve: int, - include_hidden: bool = False, - ) -> list[InferenceChunk]: - return self._primary.keyword_retrieval( - query, filters, num_to_retrieve, include_hidden=include_hidden - ) - - def semantic_retrieval( - self, - query_embedding: Embedding, - filters: IndexFilters, - num_to_retrieve: int, - ) -> list[InferenceChunk]: - return self._primary.semantic_retrieval( - query_embedding, filters, num_to_retrieve - ) - - def random_retrieval( - self, - filters: IndexFilters, - num_to_retrieve: int = 10, - dirty: bool | None = None, - ) -> list[InferenceChunk]: - return self._primary.random_retrieval(filters, num_to_retrieve, dirty) - - @property - def primary(self) -> VespaDocumentIndex: - return self._primary - - @property - def secondary(self) -> VespaDocumentIndex | None: - return self._secondary diff --git a/backend/onyx/document_index/vespa_constants.py b/backend/onyx/document_index/vespa_constants.py deleted file mode 100644 index 3caf5945ed..0000000000 --- a/backend/onyx/document_index/vespa_constants.py +++ /dev/null @@ -1,112 +0,0 @@ -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 - - -VESPA_CONFIG_SERVER_URL = ( - VESPA_CLOUD_URL or f"http://{VESPA_CONFIG_SERVER_HOST}:{VESPA_TENANT_PORT}" -) -VESPA_APPLICATION_ENDPOINT = f"{VESPA_CONFIG_SERVER_URL}/application/v2" - -# main search application -VESPA_APP_CONTAINER_URL = VESPA_CLOUD_URL or f"http://{VESPA_HOST}:{VESPA_PORT}" - - -# danswer_chunk below is defined in vespa/app_configs/schemas/danswer_chunk.sd.jinja -DOCUMENT_ID_ENDPOINT = ( - f"{VESPA_APP_CONTAINER_URL}/document/v1/default/{{index_name}}/docid" -) - -# the default document id endpoint is http://localhost:8080/document/v1/default/danswer_chunk/docid - -SEARCH_ENDPOINT = f"{VESPA_APP_CONTAINER_URL}/search/" - -# Since Vespa doesn't allow batching of inserts / updates, we use threads to -# parallelize the operations. -NUM_THREADS = 32 -MAX_ID_SEARCH_QUERY_SIZE = 400 -# Suspect that adding too many "or" conditions will cause Vespa to timeout and return -# an empty list of hits (with no error status and coverage: 0 and degraded) -MAX_OR_CONDITIONS = 10 -# up from 500ms for now, since we've seen quite a few timeouts -# in the long term, we are looking to improve the performance of Vespa -# so that we can bring this back to default -VESPA_TIMEOUT = "10s" -# The size of the batch to use for batched operations like inserts / updates. -# The batch will likely be sent to a threadpool of size NUM_THREADS. -BATCH_SIZE = 128 - -TENANT_ID = "tenant_id" -DOCUMENT_ID = "document_id" -CHUNK_ID = "chunk_id" -BLURB = "blurb" -CONTENT = "content" -SOURCE_LINKS = "source_links" -SEMANTIC_IDENTIFIER = "semantic_identifier" -TITLE = "title" -SKIP_TITLE_EMBEDDING = "skip_title" -SECTION_CONTINUATION = "section_continuation" -EMBEDDINGS = "embeddings" -TITLE_EMBEDDING = "title_embedding" -ACCESS_CONTROL_LIST = "access_control_list" -DOCUMENT_SETS = "document_sets" -USER_FILE = "user_file" -USER_FOLDER = "user_folder" -USER_PROJECT = "user_project" -PERSONAS = "personas" -LARGE_CHUNK_REFERENCE_IDS = "large_chunk_reference_ids" -METADATA = "metadata" -METADATA_LIST = "metadata_list" -METADATA_SUFFIX = "metadata_suffix" -DOC_SUMMARY = "doc_summary" -CHUNK_CONTEXT = "chunk_context" -BOOST = "boost" -AGGREGATED_CHUNK_BOOST_FACTOR = "aggregated_chunk_boost_factor" -DOC_UPDATED_AT = "doc_updated_at" # Indexed as seconds since epoch -PRIMARY_OWNERS = "primary_owners" -SECONDARY_OWNERS = "secondary_owners" -RECENCY_BIAS = "recency_bias" -HIDDEN = "hidden" -# for legacy reasons, called `name` in Vespa despite it really being an ID -IMAGE_FILE_NAME = "image_file_name" - -# Specific to Vespa, needed for highlighting matching keywords / section -CONTENT_SUMMARY = "content_summary" - -FULL_CHUNK_EMBEDDING_KEY = "full_chunk" - - -YQL_BASE = ( - f"select " # noqa: S608 - Vespa YQL with static column-name constants, not SQL - f"documentid, " - f"{DOCUMENT_ID}, " - f"{CHUNK_ID}, " - f"{BLURB}, " - f"{CONTENT}, " - f"{SOURCE_TYPE}, " - f"{SOURCE_LINKS}, " - f"{SEMANTIC_IDENTIFIER}, " - f"{TITLE}, " - f"{SECTION_CONTINUATION}, " - f"{IMAGE_FILE_NAME}, " - f"{BOOST}, " - f"{AGGREGATED_CHUNK_BOOST_FACTOR}, " - f"{HIDDEN}, " - f"{DOC_UPDATED_AT}, " - f"{PRIMARY_OWNERS}, " - f"{SECONDARY_OWNERS}, " - f"{LARGE_CHUNK_REFERENCE_IDS}, " - f"{METADATA}, " - f"{METADATA_SUFFIX}, " - f"{DOC_SUMMARY}, " - f"{CHUNK_CONTEXT}, " - f"{CONTENT_SUMMARY} " - f"from {{index_name}} where " -) diff --git a/backend/onyx/httpx/httpx_pool.py b/backend/onyx/httpx/httpx_pool.py deleted file mode 100644 index bbd3e9fe4f..0000000000 --- a/backend/onyx/httpx/httpx_pool.py +++ /dev/null @@ -1,60 +0,0 @@ -import threading -from typing import Any - -import httpx - - -def make_default_kwargs() -> dict[str, Any]: - return { - "http2": True, - "limits": httpx.Limits(), - } - - -class HttpxPool: - """Class to manage a global httpx Client instance""" - - _clients: dict[str, httpx.Client] = {} - _lock: threading.Lock = threading.Lock() - - # Default parameters for creation - - def __init__(self) -> None: - pass - - @classmethod - def _init_client(cls, **kwargs: Any) -> httpx.Client: - """Private helper method to create and return an httpx.Client.""" - merged_kwargs = {**(make_default_kwargs()), **kwargs} - return httpx.Client(**merged_kwargs) - - @classmethod - def init_client(cls, name: str, **kwargs: Any) -> None: - """Allow the caller to init the client with extra params.""" - with cls._lock: - if name not in cls._clients: - cls._clients[name] = cls._init_client(**kwargs) - - @classmethod - def close_client(cls, name: str) -> None: - """Allow the caller to close the client.""" - with cls._lock: - client = cls._clients.pop(name, None) - if client: - client.close() - - @classmethod - def close_all(cls) -> None: - """Close all registered clients.""" - with cls._lock: - for client in cls._clients.values(): - client.close() - cls._clients.clear() - - @classmethod - def get(cls, name: str) -> httpx.Client: - """Gets the httpx.Client. Will init to default settings if not init'd.""" - with cls._lock: - if name not in cls._clients: - cls._clients[name] = cls._init_client() - return cls._clients[name] diff --git a/backend/onyx/indexing/adapters/document_indexing_adapter.py b/backend/onyx/indexing/adapters/document_indexing_adapter.py index 54e468b862..a05e864d4d 100644 --- a/backend/onyx/indexing/adapters/document_indexing_adapter.py +++ b/backend/onyx/indexing/adapters/document_indexing_adapter.py @@ -47,7 +47,7 @@ class DocumentIndexingBatchAdapter(IndexingBatchAdapter): Keeps orchestration logic in the pipeline and side-effects in the adapter. Each phase opens its own short-lived session so no connection is held idle - during the long embedding and Vespa-write phases. + during the long embedding and document-index-write phases. """ def __init__( diff --git a/backend/onyx/indexing/indexing_pipeline.py b/backend/onyx/indexing/indexing_pipeline.py index 76977db3e0..ae8f22ccde 100644 --- a/backend/onyx/indexing/indexing_pipeline.py +++ b/backend/onyx/indexing/indexing_pipeline.py @@ -54,7 +54,7 @@ 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 ( +from onyx.document_index.interfaces import ( DocumentIndex, DocumentInsertionRecord, IndexingMetadata, @@ -174,7 +174,7 @@ class IndexingPipelineResult(BaseModel): # NOTE: need total_docs, since the pipeline can skip some docs # (e.g. not even insert them into Postgres) total_docs: int - # number of chunks that were inserted into Vespa + # number of chunks that were inserted into the document index total_chunks: int failures: list[ConnectorFailure] @@ -451,7 +451,7 @@ def index_doc_batch_with_handler( *, chunker: Chunker, embedder: IndexingEmbedder, - document_indices: list[DocumentIndex], + document_index: DocumentIndex, document_batch: list[Document], request_id: str | None, tenant_id: str, @@ -468,7 +468,7 @@ def index_doc_batch_with_handler( index_pipeline_result = index_doc_batch( chunker=chunker, embedder=embedder, - document_indices=document_indices, + document_index=document_index, document_batch=document_batch, request_id=request_id, tenant_id=tenant_id, @@ -1381,7 +1381,7 @@ def index_doc_batch( document_batch: list[Document], chunker: Chunker, embedder: IndexingEmbedder, - document_indices: list[DocumentIndex], + document_index: DocumentIndex, request_id: str | None, tenant_id: str, adapter: IndexingBatchAdapter, @@ -1531,7 +1531,8 @@ def index_doc_batch( ) # Filter to only successfully embedded chunks so - # doc_id_to_new_chunk_cnt reflects what's actually written to Vespa. + # doc_id_to_new_chunk_cnt reflects what's actually written to the + # document index. embedded_chunks = [ c for c in chunks if c.source_document.id not in embedding_failed_doc_ids ] @@ -1560,51 +1561,28 @@ def index_doc_batch( doc_id_to_chunk_cnt_diff=doc_id_to_chunk_cnt_diff, ) - primary_doc_idx_insertion_records: list[DocumentInsertionRecord] | None = ( - None + def _enriched_stream() -> Iterator[DocMetadataAwareIndexChunk]: + for chunk in chunk_store.stream(): + yield enricher.enrich_chunk(chunk, 1.0) + + vector_db_write_start = time.monotonic() + insertion_records, write_failures = write_chunks_to_vector_db_with_backoff( + document_index=document_index, + make_chunks=_enriched_stream, + indexing_metadata=indexing_metadata, + tenant_id=tenant_id, ) - primary_doc_idx_vector_db_write_failures: list[ConnectorFailure] | None = ( - None + vector_db_write_ms = max( + 0, int((time.monotonic() - vector_db_write_start) * 1000) ) - # Sum vector-db write time across all configured indices and - # record once per batch. Most deployments have a single index - # (primary), but during a switchover both primary and secondary - # are written; we want a single combined number per batch rather - # than one event per index (avoids inflating event_count). - vector_db_write_ms = 0 - for document_index in document_indices: - - def _enriched_stream() -> Iterator[DocMetadataAwareIndexChunk]: - for chunk in chunk_store.stream(): - yield enricher.enrich_chunk(chunk, 1.0) - - vector_db_write_start = time.monotonic() - insertion_records, write_failures = ( - write_chunks_to_vector_db_with_backoff( - document_index=document_index, - make_chunks=_enriched_stream, - indexing_metadata=indexing_metadata, - tenant_id=tenant_id, - ) - ) - vector_db_write_ms += max( - 0, int((time.monotonic() - vector_db_write_start) * 1000) - ) - - _verify_indexing_completeness( - insertion_records=insertion_records, - write_failures=write_failures, - embedding_failed_doc_ids=embedding_failed_doc_ids, - updatable_ids=updatable_ids, - document_index_name=document_index.__class__.__name__, - ) - # We treat the first document index we got as the primary one used - # for reporting the state of indexing. - if primary_doc_idx_insertion_records is None: - primary_doc_idx_insertion_records = insertion_records - if primary_doc_idx_vector_db_write_failures is None: - primary_doc_idx_vector_db_write_failures = write_failures + _verify_indexing_completeness( + insertion_records=insertion_records, + write_failures=write_failures, + embedding_failed_doc_ids=embedding_failed_doc_ids, + updatable_ids=updatable_ids, + document_index_name=document_index.__class__.__name__, + ) safe_record_single_event_if_set( IndexAttemptStage.VECTOR_DB_WRITE, attempt_id, vector_db_write_ms @@ -1630,10 +1608,8 @@ def index_doc_batch( # on the next sync. Hashes were pre-computed in get_docs_to_update. # Skipped for FUTURE writes: stamping the PRESENT-only hash would make # the PRESENT poll skip the doc (cross-index suppression). - if primary_doc_idx_insertion_records is not None and not index_to_secondary: - successfully_indexed_ids = { - r.document_id for r in primary_doc_idx_insertion_records - } + if not index_to_secondary: + successfully_indexed_ids = {r.document_id for r in insertion_records} update_docs_content_hash__no_commit( ids_to_new_hash={ doc_id: context.doc_id_to_content_hash[doc_id] @@ -1643,23 +1619,18 @@ def index_doc_batch( db_session=db_session, ) - assert primary_doc_idx_insertion_records is not None - assert primary_doc_idx_vector_db_write_failures is not None - _maybe_push_documents( adapter=adapter, filtered_documents=filtered_documents, - insertion_records=primary_doc_idx_insertion_records, + insertion_records=insertion_records, from_beginning=from_beginning, ) return IndexingPipelineResult( - new_docs=sum( - 1 for r in primary_doc_idx_insertion_records if not r.already_existed - ), + new_docs=sum(1 for r in insertion_records if not r.already_existed), total_docs=len(filtered_documents), total_chunks=len(embedding_result.successful_chunk_ids), - failures=primary_doc_idx_vector_db_write_failures + failures=write_failures + embedding_result.connector_failures + enrichment_partition.failures + filter_failures, @@ -1671,7 +1642,7 @@ def run_indexing_pipeline( document_batch: list[Document], request_id: str | None, embedder: IndexingEmbedder, - document_indices: list[DocumentIndex], + document_index: DocumentIndex, db_session: Session | None = None, tenant_id: str, adapter: IndexingBatchAdapter, @@ -1737,7 +1708,7 @@ def run_indexing_pipeline( return index_doc_batch_with_handler( chunker=chunker, embedder=embedder, - document_indices=document_indices, + document_index=document_index, document_batch=document_batch, request_id=request_id, tenant_id=tenant_id, diff --git a/backend/onyx/indexing/vector_db_insertion.py b/backend/onyx/indexing/vector_db_insertion.py index af532cac6a..3799a50c17 100644 --- a/backend/onyx/indexing/vector_db_insertion.py +++ b/backend/onyx/indexing/vector_db_insertion.py @@ -1,13 +1,11 @@ import time from collections.abc import Callable, Iterable -from http import HTTPStatus from itertools import chain, groupby -import httpx import sentry_sdk from onyx.connectors.models import ConnectorFailure, DocumentFailure -from onyx.document_index.interfaces_new import ( +from onyx.document_index.interfaces import ( DocumentIndex, DocumentInsertionRecord, IndexingMetadata, @@ -18,16 +16,6 @@ from onyx.utils.logger import setup_logger logger = setup_logger() -def _log_insufficient_storage_error(e: Exception) -> None: - if isinstance(e, httpx.HTTPStatusError): - if e.response.status_code == HTTPStatus.INSUFFICIENT_STORAGE: - logger.error( - "NOTE: HTTP Status 507 Insufficient Storage indicates " - "you need to allocate more memory or disk space to the " - "Vespa/index container." - ) - - def write_chunks_to_vector_db_with_backoff( document_index: DocumentIndex, make_chunks: Callable[[], Iterable[DocMetadataAwareIndexChunk]], @@ -67,9 +55,6 @@ def write_chunks_to_vector_db_with_backoff( e, ) - # give some specific logging on this common failure case. - _log_insufficient_storage_error(e) - # wait a couple seconds just to give the vector db a chance to recover time.sleep(2) @@ -108,9 +93,6 @@ def write_chunks_to_vector_db_with_backoff( "Failed to write document chunks for '%s' to vector db", doc_id ) - # give some specific logging on this common failure case. - _log_insufficient_storage_error(e) - failures.append( ConnectorFailure( failed_document=DocumentFailure( diff --git a/backend/onyx/kg/clustering/clustering.py b/backend/onyx/kg/clustering/clustering.py deleted file mode 100644 index 2a4cdc5dce..0000000000 --- a/backend/onyx/kg/clustering/clustering.py +++ /dev/null @@ -1,470 +0,0 @@ -import time -from collections.abc import Generator -from typing import cast - -from rapidfuzz.fuzz import ratio -from redis.lock import Lock as RedisLock -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, - KG_CLUSTERING_THRESHOLD, -) -from onyx.db.engine.sql_engine import get_session_with_current_tenant -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.kg.models import KGGroundingType -from onyx.kg.utils.formatting_utils import make_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 -from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA - -logger = setup_logger() - - -def _get_batch_untransferred_grounded_entities( - batch_size: int, -) -> Generator[list[KGEntityExtractionStaging], None, None]: - while True: - with get_session_with_current_tenant() as db_session: - batch = ( - db_session.query(KGEntityExtractionStaging) - .join( - KGEntityType, - KGEntityExtractionStaging.entity_type_id_name - == KGEntityType.id_name, - ) - .filter( - KGEntityType.grounding == KGGroundingType.GROUNDED, - KGEntityExtractionStaging.transferred_id_name.is_(None), - ) - .limit(batch_size) - .all() - ) - if not batch: - break - yield batch - - -def _get_batch_untransferred_relationship_types( - batch_size: int, -) -> Generator[list[KGRelationshipTypeExtractionStaging], None, None]: - while True: - with get_session_with_current_tenant() as db_session: - batch = ( - db_session.query(KGRelationshipTypeExtractionStaging) - .filter(KGRelationshipTypeExtractionStaging.transferred.is_(False)) - .limit(batch_size) - .all() - ) - if not batch: - break - yield batch - - -def _get_batch_untransferred_relationships( - batch_size: int, -) -> Generator[list[KGRelationshipExtractionStaging], None, None]: - while True: - with get_session_with_current_tenant() as db_session: - batch = ( - db_session.query(KGRelationshipExtractionStaging) - .filter(KGRelationshipExtractionStaging.transferred.is_(False)) - .limit(batch_size) - .all() - ) - if not batch: - break - yield batch - - -def _get_batch_entities_with_parent( - batch_size: int, -) -> Generator[list[KGEntityExtractionStaging], None, None]: - offset = 0 - - while True: - with get_session_with_current_tenant() as db_session: - batch = ( - db_session.query(KGEntityExtractionStaging) - .filter(KGEntityExtractionStaging.parent_key.isnot(None)) - .order_by(KGEntityExtractionStaging.id_name) - .offset(offset) - .limit(batch_size) - .all() - ) - if not batch: - break - # we can't filter out ""s earlier as it will mess up the pagination - yield [entity for entity in batch if entity.parent_key != ""] - offset += batch_size - - -def _get_batch_kg_processed_documents( - batch_size: int, -) -> Generator[list[Document], None, None]: - offset = 0 - - while True: - with get_session_with_current_tenant() as db_session: - batch = ( - db_session.query(Document) - .join( - KGEntityExtractionStaging, - Document.id == KGEntityExtractionStaging.document_id, - ) - .filter( - KGEntityExtractionStaging.transferred_id_name.is_not(None), - ) - .order_by(Document.id) - .offset(offset) - .limit(batch_size) - .all() - ) - if not batch: - break - yield batch - offset += batch_size - - -def _cluster_one_grounded_entity( - entity: KGEntityExtractionStaging, -) -> tuple[KGEntity, bool]: - """ - Cluster a single grounded entity. - """ - with get_session_with_current_tenant() as db_session: - # get entity name and filtering conditions - if entity.document_id is not None: - entity_name = cast( - str, - db_session.query(Document.semantic_id) - .filter(Document.id == entity.document_id) - .scalar(), - ).lower() - filtering = [KGEntity.document_id.is_(None)] - else: - entity_name = entity.name.lower() - filtering = [] - - # skip those with numbers so we don't cluster version1 and version2, etc. - similar_entities: list[KGEntity] = [] - if not any(char.isdigit() for char in entity_name): - # find similar entities, uses GIN index, very efficient - db_session.execute( - text( - "SET pg_trgm.similarity_threshold = " - + str(KG_CLUSTERING_RETRIEVE_THRESHOLD) - ) - ) - similar_entities = ( - db_session.query(KGEntity) - .filter( - # find entities of the same type with a similar name - *filtering, - KGEntity.entity_type_id_name == entity.entity_type_id_name, - getattr( # ods: ignore[getattr] - func, POSTGRES_DEFAULT_SCHEMA - ).similarity_op(KGEntity.name, entity_name), - ) - .all() - ) - - # find best match - best_score = -1.0 - best_entity = None - for similar in similar_entities: - # skip those with numbers so we don't cluster version1 and version2, etc. - if any(char.isdigit() for char in similar.name): - continue - score = ratio(similar.name, entity_name) - if score >= KG_CLUSTERING_THRESHOLD * 100 and score > best_score: - best_score = score - best_entity = similar - - # if there is a match, update the entity, otherwise create a new one - with get_session_with_current_tenant() as db_session: - if best_entity: - logger.debug("Merged %s with %s", entity.name, best_entity.name) - update_vespa = ( - best_entity.document_id is None and entity.document_id is not None - ) - transferred_entity = merge_entities( - db_session=db_session, parent=best_entity, child=entity - ) - else: - update_vespa = entity.document_id is not None - transferred_entity = transfer_entity(db_session=db_session, entity=entity) - - db_session.commit() - - return transferred_entity, update_vespa - - -def _create_one_parent_child_relationship(entity: KGEntityExtractionStaging) -> None: - """ - Creates a relationship between the entity and its parent, if it exists. - Then, updates the entity's parent to the next ancestor. - """ - with get_session_with_current_tenant() as db_session: - # find the next ancestor - parent = ( - db_session.query(KGEntity) - .filter(KGEntity.entity_key == entity.parent_key) - .first() - ) - - if parent is not None: - # create parent child relationship and relationship type - upsert_relationship_type( - db_session=db_session, - source_entity_type=parent.entity_type_id_name, - relationship_type="has_subcomponent", - target_entity_type=entity.entity_type_id_name, - ) - relationship_id_name = make_relationship_id( - parent.id_name, - "has_subcomponent", - cast(str, entity.transferred_id_name), - ) - upsert_relationship( - db_session=db_session, - relationship_id_name=relationship_id_name, - source_document_id=entity.document_id, - ) - - next_ancestor = parent.parent_key or "" - else: - next_ancestor = "" - - # set the staging entity's parent to the next ancestor - # if there is no parent or next ancestor, set to "" to differentiate from None - # None will mess up the pagination in _get_batch_entities_with_parent - db_session.query(KGEntityExtractionStaging).filter( - KGEntityExtractionStaging.id_name == entity.id_name - ).update({"parent_key": next_ancestor}) - db_session.commit() - - -def _transfer_one_relationship( - relationship: KGRelationshipExtractionStaging, -) -> None: - with get_session_with_current_tenant() as db_session: - # get the translations - staging_entity_id_names = { - relationship.source_node, - relationship.target_node, - } - entity_translations: dict[str, str] = { - entity.id_name: entity.transferred_id_name - for entity in db_session.query(KGEntityExtractionStaging) - .filter(KGEntityExtractionStaging.id_name.in_(staging_entity_id_names)) - .all() - if entity.transferred_id_name is not None - } - if len(entity_translations) != len(staging_entity_id_names): - logger.error( - "Missing entity translations for %s", - staging_entity_id_names - entity_translations.keys(), - ) - return - - # transfer the relationship - transfer_relationship( - db_session=db_session, - relationship=relationship, - entity_translations=entity_translations, - ) - db_session.commit() - - -def kg_clustering( - tenant_id: str, - index_name: str, - lock: RedisLock, - processing_chunk_batch_size: int = 16, -) -> None: - """ - Here we will cluster the extractions based on their cluster frameworks. - Initially, this will only focus on grounded entities with pre-determined - relationships, so 'clustering' is actually not yet required. - However, we may need to reconcile entities coming from different sources. - - The primary purpose of this function is to populate the actual KG tables - from the temp_extraction tables. - - This will change with deep extraction, where grounded-sourceless entities - can be extracted and then need to be clustered. - """ - logger.info("Starting kg clustering for tenant %s", tenant_id) - - kg_config_settings = get_kg_config_settings() - validate_kg_settings(kg_config_settings) - - last_lock_time = time.monotonic() - - # Cluster and transfer grounded entities sequentially - start_time = time.monotonic() - i_batch = 0 - for i_batch, untransferred_grounded_entities in enumerate( # noqa: B007 - _get_batch_untransferred_grounded_entities( - batch_size=processing_chunk_batch_size - ) - ): - for entity in untransferred_grounded_entities: - _cluster_one_grounded_entity(entity) - last_lock_time = extend_lock( - lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time - ) - # logger.debug(f"Transferred entities batch {i}") - # NOTE: we assume every entity is transferred, as we currently only have grounded entities - time_delta = time.monotonic() - start_time - logger.info( - "Finished transferring %s entity batches in %ss", - i_batch + 1, - format(time_delta, ".2f"), - ) - - # Create parent-child relationships in parallel - for _ in range(kg_config_settings.KG_MAX_PARENT_RECURSION_DEPTH): - for root_entities in _get_batch_entities_with_parent( - batch_size=processing_chunk_batch_size - ): - run_functions_tuples_in_parallel( - [ - (_create_one_parent_child_relationship, (root_entity,)) - for root_entity in root_entities - ] - ) - last_lock_time = extend_lock( - lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time - ) - logger.info("Finished creating all parent-child relationships") - - # Transfer the relationship types (no need to do in parallel as there's only a few) - start_time = time.monotonic() - i_batch = 0 - for i_batch, relationship_types in enumerate( # noqa: B007 - _get_batch_untransferred_relationship_types( - batch_size=processing_chunk_batch_size - ) - ): - with get_session_with_current_tenant() as db_session: - for relationship_type in relationship_types: - transfer_relationship_type(db_session, relationship_type) - db_session.commit() - last_lock_time = extend_lock( - lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time - ) - # logger.debug(f"Transferred relationship types batch {i}") - time_delta = time.monotonic() - start_time - logger.info( - "Finished transferring %s relationship type batches in %ss", - i_batch + 1, - format(time_delta, ".2f"), - ) - - # Transfer the relationships in parallel - start_time = time.monotonic() - i_batch = 0 - for i_batch, relationships in enumerate( # noqa: B007 - _get_batch_untransferred_relationships(batch_size=processing_chunk_batch_size) - ): - run_functions_tuples_in_parallel( - [ - (_transfer_one_relationship, (relationship,)) - for relationship in relationships - ] - ) - last_lock_time = extend_lock( - lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time - ) - # logger.debug(f"Transferred relationships batch {i}") - time_delta = time.monotonic() - start_time - logger.info( - "Finished transferring %s relationship batches in %ss", - i_batch + 1, - format(time_delta, ".2f"), - ) - - # Update vespa for each document - start_time = time.monotonic() - i_batch = 0 - for i_batch, documents in enumerate( # noqa: B007 - _get_batch_kg_processed_documents(batch_size=processing_chunk_batch_size) - ): - batch_update_requests = run_functions_tuples_in_parallel( - [ - (get_kg_vespa_info_update_requests_for_document, (document.id,)) - for document in documents - ] - ) - for update_requests, document in zip( - batch_update_requests, documents, strict=True - ): - try: - update_kg_chunks_vespa_info(update_requests, index_name, tenant_id) - except Exception as e: - logger.error("Error updating vespa for document %s: %s", document.id, e) - last_lock_time = extend_lock( - lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time - ) - # logger.debug(f"Updated vespa for documents batch {i}") - time_delta = time.monotonic() - start_time - logger.info( - "Finished updating %s document batches in %ss", - i_batch + 1, - format(time_delta, ".2f"), - ) - - # Delete the transferred objects from the staging tables - try: - with get_session_with_current_tenant() as db_session: - db_session.query(KGRelationshipExtractionStaging).filter( - KGRelationshipExtractionStaging.transferred.is_(True) - ).delete(synchronize_session=False) - db_session.commit() - except Exception as e: - logger.error("Error deleting relationships: %s", e) - - try: - with get_session_with_current_tenant() as db_session: - db_session.query(KGRelationshipTypeExtractionStaging).filter( - KGRelationshipTypeExtractionStaging.transferred.is_(True) - ).delete(synchronize_session=False) - db_session.commit() - except Exception as e: - logger.error("Error deleting relationship types: %s", e) - - try: - with get_session_with_current_tenant() as db_session: - db_session.query(KGEntityExtractionStaging).filter( - KGEntityExtractionStaging.transferred_id_name.is_not(None) - ).delete(synchronize_session=False) - db_session.commit() - except Exception as e: - logger.error("Error deleting entities: %s", e) - logger.info("Finished deleting all transferred staging entries") diff --git a/backend/onyx/kg/clustering/normalizations.py b/backend/onyx/kg/clustering/normalizations.py deleted file mode 100644 index b2581a76e3..0000000000 --- a/backend/onyx/kg/clustering/normalizations.py +++ /dev/null @@ -1,347 +0,0 @@ -import re -from collections import defaultdict -from typing import cast - -import numpy as np -from rapidfuzz.distance.DamerauLevenshtein import normalized_similarity -from sqlalchemy import Float, MetaData, String, Table, desc, func, select -from sqlalchemy.dialects.postgresql import ARRAY - -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, NormalizedRelationships -from onyx.kg.utils.embeddings import encode_string_batch -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 - -logger = setup_logger() - - -alphanum_regex = re.compile(r"[^a-z0-9]+") -rem_email_regex = re.compile(r"(?<=\S)@([a-z0-9-]+)\.([a-z]{2,6})$") - - -def _ngrams(sequence: str, n: int) -> list[tuple[str, ...]]: - """Generate n-grams from a sequence.""" - return [tuple(sequence[i : i + n]) for i in range(len(sequence) - n + 1)] - - -def _clean_name(entity_name: str) -> str: - """ - Clean an entity string by removing non-alphanumeric characters and email addresses. - If the name after cleaning is empty, return the original name in lowercase. - """ - cleaned_entity = entity_name.casefold() - return ( - alphanum_regex.sub("", rem_email_regex.sub("", cleaned_entity)) - or cleaned_entity - ) - - -def _normalize_one_entity( - entity: str, - attributes: dict[str, str], - allowed_docs_temp_view_name: str | None = None, -) -> str | None: - """ - Matches a single entity to the best matching entity of the same type. - """ - entity_type, entity_name = split_entity_id(entity) - if entity_name == "*": - return entity - - cleaned_entity = _clean_name(entity_name) - - # narrow filter to subtype if requested - type_filters = [KGEntity.entity_type_id_name == entity_type] - if "subtype" in attributes: - type_filters.append( - KGEntity.attributes.op("@>")({"subtype": attributes["subtype"]}) - ) - - # step 1: find entities containing the entity_name or something similar - with get_session_with_current_tenant() as db_session: - # get allowed documents - metadata = MetaData() - if allowed_docs_temp_view_name is None: - raise ValueError("allowed_docs_temp_view_name is not available") - - effective_schema_allowed_docs_temp_view_name = ( - allowed_docs_temp_view_name.split(".")[-1] - ) - - allowed_docs_temp_view = Table( - effective_schema_allowed_docs_temp_view_name, - metadata, - autoload_with=db_session.get_bind(), - ) - - # generate trigrams of the queried entity Q - query_trigrams = db_session.query( - getattr(func, POSTGRES_DEFAULT_SCHEMA) # ods: ignore[getattr] - .show_trgm(cleaned_entity) - .cast(ARRAY(String(3))) - .label("trigrams") - ).cte("query") - - candidates = cast( - list[tuple[str, str, float]], - db_session.query( - KGEntity.id_name, - KGEntity.name, - ( - # for each entity E, compute score = | Q ∩ E | / min(|Q|, |E|) - func.cardinality( - func.array( - select(func.unnest(KGEntity.name_trigrams)) - .correlate(KGEntity) - .intersect( - select( - func.unnest(query_trigrams.c.trigrams) - ).correlate(query_trigrams) - ) - .scalar_subquery() - ) - ).cast(Float) - / func.least( - func.cardinality(query_trigrams.c.trigrams), - func.cardinality(KGEntity.name_trigrams), - ) - ).label("score"), - ) - .select_from(KGEntity, query_trigrams) - .outerjoin( - allowed_docs_temp_view, - KGEntity.document_id == allowed_docs_temp_view.c.allowed_doc_id, - ) - .filter( - *type_filters, - KGEntity.name_trigrams.overlap(query_trigrams.c.trigrams), - # Add filter for allowed docs - either document_id is NULL or it's in allowed_docs - ( - KGEntity.document_id.is_(None) - | allowed_docs_temp_view.c.allowed_doc_id.isnot(None) - ), - ) - .order_by(desc("score")) - .limit(KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT) - .all(), - ) - if not candidates: - return None - - # step 2: do a weighted ngram analysis and damerau levenshtein distance to rerank - n1, n2, n3 = ( - set(_ngrams(cleaned_entity, 1)), - set(_ngrams(cleaned_entity, 2)), - set(_ngrams(cleaned_entity, 3)), - ) - for i, (candidate_id_name, candidate_name, _) in enumerate(candidates): - cleaned_candidate = _clean_name(candidate_name) - h_n1, h_n2, h_n3 = ( - set(_ngrams(cleaned_candidate, 1)), - set(_ngrams(cleaned_candidate, 2)), - set(_ngrams(cleaned_candidate, 3)), - ) - - # compute ngram overlap, renormalize scores if the names are too short for larger ngrams - grams_used = min(2, len(cleaned_entity) - 1, len(cleaned_candidate) - 1) - W_n1, W_n2, W_n3 = KG_NORMALIZATION_RERANK_NGRAM_WEIGHTS - ngram_score = ( - # compute | Q ∩ E | / min(|Q|, |E|) for unigrams and bigrams (trigrams already computed) - W_n1 * len(n1 & h_n1) / max(1, min(len(n1), len(h_n1))) - + W_n2 * len(n2 & h_n2) / max(1, min(len(n2), len(h_n2))) - + W_n3 * len(n3 & h_n3) / max(1, min(len(n3), len(h_n3))) - ) / (W_n1, W_n1 + W_n2, 1.0)[grams_used] - - # compute damerau levenshtein distance to fuzzy match against typos - W_leven = KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT - leven_score = normalized_similarity(cleaned_entity, cleaned_candidate) - - # combine scores - score = (1.0 - W_leven) * ngram_score + W_leven * leven_score - candidates[i] = (candidate_id_name, candidate_name, score) - candidates = sorted( - filter(lambda x: x[2] > KG_NORMALIZATION_RERANK_THRESHOLD, candidates), - key=lambda x: x[2], - reverse=True, - ) - if not candidates: - return None - - return candidates[0][0] - - -def _get_existing_normalized_relationships( - raw_relationships: list[str], -) -> dict[str, dict[str, list[str]]]: - """ - Get existing normalized relationships from the database. - """ - - relationship_type_map: dict[str, dict[str, list[str]]] = defaultdict( - lambda: defaultdict(list) - ) - relationship_pairs = list( - { - ( - get_entity_type(split_relationship_id(relationship)[0]), - get_entity_type(split_relationship_id(relationship)[2]), - ) - for relationship in raw_relationships - } - ) - - with get_session_with_current_tenant() as db_session: - relationships = get_relationships_for_entity_type_pairs( - db_session, relationship_pairs - ) - - for relationship in relationships: - relationship_type_map[relationship.source_entity_type_id_name][ - relationship.target_entity_type_id_name - ].append(relationship.id_name) - - return relationship_type_map - - -def normalize_entities( - raw_entities: list[str], - raw_entities_w_attributes: list[str], - allowed_docs_temp_view_name: str | None = None, -) -> NormalizedEntities: - """ - Match each entity against a list of normalized entities using fuzzy matching. - Returns the best matching normalized entity for each input entity. - - Args: - raw_entities: list of entity strings to normalize, w/o attributes - raw_entities_w_attributes: list of entity strings to normalize, w/ attributes - - Returns: - list of normalized entity strings - """ - normalized_entities: list[str] = [] - normalized_entities_w_attributes: list[str] = [] - normalized_map: dict[str, str] = {} - - entity_attributes = [ - get_attributes(attr_entity) for attr_entity in raw_entities_w_attributes - ] - - mapping: list[str | None] = run_functions_tuples_in_parallel( - [ - (_normalize_one_entity, (entity, attributes, allowed_docs_temp_view_name)) - for entity, attributes in zip(raw_entities, entity_attributes, strict=True) - ] - ) - for entity, attributes, normalized_entity in zip( - raw_entities, entity_attributes, mapping, strict=True - ): - if normalized_entity is not None: - normalized_entities.append(normalized_entity) - normalized_entities_w_attributes.append( - make_entity_w_attributes(normalized_entity, attributes) - ) - normalized_map[entity] = format_entity_id_for_models(normalized_entity) - else: - logger.warning("No normalized entity found for %s", entity) - normalized_map[entity] = format_entity_id_for_models(entity) - - return NormalizedEntities( - entities=normalized_entities, - entities_w_attributes=normalized_entities_w_attributes, - entity_normalization_map=normalized_map, - ) - - -def normalize_relationships( - raw_relationships: list[str], entity_normalization_map: dict[str, str] -) -> NormalizedRelationships: - """ - Normalize relationships using entity mappings and relationship string matching. - - Args: - relationships: list of relationships in format "source__relation__target" - entity_normalization_map: Mapping of raw entities to normalized ones (or None) - - Returns: - NormalizedRelationships containing normalized relationships and mapping - """ - # Placeholder for normalized relationship structure - nor_relationships = _get_existing_normalized_relationships(raw_relationships) - - normalized_rels: list[str] = [] - normalization_map: dict[str, str] = {} - - for raw_rel in raw_relationships: - # 1. Split and normalize entities - try: - source, rel_string, target = split_relationship_id(raw_rel) - except ValueError: - raise ValueError(f"Invalid relationship format: {raw_rel}") - - # Check if entities are in normalization map and not None - norm_source = entity_normalization_map.get(source) - norm_target = entity_normalization_map.get(target) - - if norm_source is None or norm_target is None: - logger.warning("No normalized entities found for %s", raw_rel) - continue - - # 2. Find candidate normalized relationships - candidate_rels = [] - norm_source_type = get_entity_type(format_entity_id_for_models(norm_source)) - norm_target_type = get_entity_type(format_entity_id_for_models(norm_target)) - if ( - norm_source_type in nor_relationships - and norm_target_type in nor_relationships[norm_source_type] - ): - candidate_rels = [ - split_relationship_id(rel)[1] - for rel in nor_relationships[norm_source_type][norm_target_type] - ] - - if not candidate_rels: - logger.warning("No candidate relationships found for %s", raw_rel) - continue - - # 3. Encode and find best match - strings_to_encode = [rel_string] + candidate_rels - vectors = encode_string_batch(strings_to_encode) - - # Get raw relation vector and candidate vectors - raw_vector = vectors[0] - candidate_vectors = vectors[1:] - - # Calculate dot products - dot_products = np.dot(candidate_vectors, raw_vector) - best_match_idx = np.argmax(dot_products) - - # Create normalized relationship - norm_rel = make_relationship_id( - norm_source, candidate_rels[best_match_idx], norm_target - ) - normalized_rels.append(norm_rel) - normalization_map[raw_rel] = norm_rel - - return NormalizedRelationships( - relationships=normalized_rels, relationship_normalization_map=normalization_map - ) diff --git a/backend/onyx/kg/extractions/extraction_processing.py b/backend/onyx/kg/extractions/extraction_processing.py deleted file mode 100644 index e64039bcd5..0000000000 --- a/backend/onyx/kg/extractions/extraction_processing.py +++ /dev/null @@ -1,602 +0,0 @@ -import time -from typing import Any - -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, - 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, upsert_staging_entity -from onyx.db.entity_type import get_entity_types -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 - -logger = setup_logger() - - -def _get_classification_extraction_instructions() -> dict[ - str | None, dict[str, KGEntityTypeInstructions] -]: - """ - Prepare the classification instructions for the given source. - """ - - classification_instructions_dict: dict[ - str | None, dict[str, KGEntityTypeInstructions] - ] = {} - - with get_session_with_current_tenant() as db_session: - entity_types = get_entity_types(db_session, active=True) - - for entity_type in entity_types: - grounded_source_name = entity_type.grounded_source_name - - if grounded_source_name not in classification_instructions_dict: - classification_instructions_dict[grounded_source_name] = {} - - if grounded_source_name is None: - continue - - attributes = entity_type.parsed_attributes - classification_attributes = { - option: info - for option, info in attributes.classification_attributes.items() - if info.extraction - } - classification_options = ", ".join(classification_attributes.keys()) - classification_enabled = ( - len(classification_options) > 0 and len(classification_attributes) > 0 - ) - - classification_instructions_dict[grounded_source_name][entity_type.id_name] = ( - KGEntityTypeInstructions( - metadata_attribute_conversion=attributes.metadata_attribute_conversion, - classification_instructions=KGClassificationInstructions( - classification_enabled=classification_enabled, - classification_options=classification_options, - classification_class_definitions=classification_attributes, - ), - extraction_instructions=KGExtractionInstructions( - deep_extraction=entity_type.deep_extraction, - active=entity_type.active, - ), - entity_filter_attributes=attributes.entity_filter_attributes, - ) - ) - - return classification_instructions_dict - - -def _get_batch_documents_enhanced_metadata( - unprocessed_document_batch: list[Document], - source_type_classification_extraction_instructions: dict[ - str, KGEntityTypeInstructions - ], - connector_source: str, -) -> dict[str, KGEnhancedDocumentMetadata]: - """ - Get the entity types for the given unprocessed documents. - """ - - kg_document_meta_data_dict: dict[str, KGEnhancedDocumentMetadata] = { - document.id: KGEnhancedDocumentMetadata( - entity_type=None, - metadata_attribute_conversion=None, - document_metadata=None, - deep_extraction=False, - classification_enabled=False, - classification_instructions=None, - skip=True, - ) - for document in unprocessed_document_batch - } - - batch_entity = None - if len(source_type_classification_extraction_instructions) == 1: - # if source only has one entity type, the document must be of that type - batch_entity = list(source_type_classification_extraction_instructions.keys())[ - 0 - ] - - # the documents can be of multiple entity types. We need to identify the entity type for each document - batch_metadata = get_batch_documents_metadata( - [ - unprocessed_document.id - for unprocessed_document in unprocessed_document_batch - ], - connector_source, - ) - - for metadata in batch_metadata: - document_id = metadata.document_id - doc_entity = None - - if not isinstance(document_id, str): - continue - - chunk_metadata = metadata.source_metadata - - if batch_entity: - doc_entity = batch_entity - else: - # TODO: make this a helper function - if not chunk_metadata: - continue - - for ( - potential_entity_type - ) in source_type_classification_extraction_instructions.keys(): - potential_entity_type_attribute_filters = ( - source_type_classification_extraction_instructions[ - potential_entity_type - ].entity_filter_attributes - or {} - ) - - if not potential_entity_type_attribute_filters: - continue - - if all( - chunk_metadata.get(attribute) - == potential_entity_type_attribute_filters.get(attribute) - for attribute in potential_entity_type_attribute_filters - ): - doc_entity = potential_entity_type - break - - if doc_entity is None: - continue - - entity_instructions = source_type_classification_extraction_instructions[ - doc_entity - ] - - kg_document_meta_data_dict[document_id] = KGEnhancedDocumentMetadata( - entity_type=doc_entity, - metadata_attribute_conversion=( - source_type_classification_extraction_instructions[ - doc_entity - ].metadata_attribute_conversion - ), - document_metadata=chunk_metadata, - deep_extraction=entity_instructions.extraction_instructions.deep_extraction, - classification_enabled=entity_instructions.classification_instructions.classification_enabled, - classification_instructions=entity_instructions.classification_instructions, - skip=False, - ) - - return kg_document_meta_data_dict - - -def kg_extraction( - tenant_id: str, - index_name: str, - lock: RedisLock, - processing_chunk_batch_size: int = 8, -) -> None: - """ - This extraction will try to extract from all chunks that have not been kg-processed yet. - - Approach: - - Get all connectors that are enabled for KG extraction - - For each enabled connector: - - Get unprocessed documents (using a generator) - - For each batch of unprocessed documents: - - Classify each document to select proper ones - - Get and extract from chunks - - Update chunks in Vespa - - Update temporary KG extraction tables - - Update document table to set kg_extracted = True - """ - - logger.info("Starting kg extraction for tenant %s", tenant_id) - - kg_config_settings = get_kg_config_settings() - validate_kg_settings(kg_config_settings) - - # get connector ids that are enabled for KG extraction - with get_session_with_current_tenant() as db_session: - kg_enabled_connectors = get_kg_enabled_connectors(db_session) - - document_classification_extraction_instructions = ( - _get_classification_extraction_instructions() - ) - - # get entity type info - with get_session_with_current_tenant() as db_session: - all_entity_types = get_entity_types(db_session) - active_entity_types = { - entity_type.id_name - for entity_type in get_entity_types(db_session, active=True) - } - - # entity_type: (metadata: conversion property) - entity_metadata_conversion_instructions = { - entity_type.id_name: entity_type.parsed_attributes.metadata_attribute_conversion - for entity_type in all_entity_types - } - - # Track which metadata attributes are possible for each entity type - metadata_tracker = EntityTypeMetadataTracker() - metadata_tracker.import_typeinfo() - - last_lock_time = time.monotonic() - - # Iterate over connectors that are enabled for KG extraction - for kg_enabled_connector in kg_enabled_connectors: - connector_id = kg_enabled_connector.id - connector_coverage_days = kg_enabled_connector.kg_coverage_days - connector_source = kg_enabled_connector.source - - document_batch_counter = 0 - - # iterate over un-kg-processed documents in connector - while True: - # get a batch of unprocessed documents - with get_session_with_current_tenant() as db_session: - unprocessed_document_batch = ( - get_unprocessed_kg_document_batch_for_connector( - db_session, - connector_id, - kg_coverage_start=kg_config_settings.KG_COVERAGE_START_DATE, - kg_max_coverage_days=connector_coverage_days - or kg_config_settings.KG_MAX_COVERAGE_DAYS, - batch_size=processing_chunk_batch_size, - ) - ) - - if len(unprocessed_document_batch) == 0: - logger.info( - "No unprocessed documents found for connector %s. Processed %s batches.", - connector_id, - document_batch_counter, - ) - break - - document_batch_counter += 1 - last_lock_time = extend_lock( - lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time - ) - logger.info("Processing document batch %s", document_batch_counter) - - # Get the document attributes and entity types - batch_metadata = _get_batch_documents_enhanced_metadata( - unprocessed_document_batch, - document_classification_extraction_instructions.get( - connector_source, {} - ), - connector_source, - ) - - # mark docs in unprocessed_document_batch as EXTRACTING - for unprocessed_document in unprocessed_document_batch: - if batch_metadata[unprocessed_document.id].entity_type is None: - # info for after the connector has been processed - kg_stage = KGStage.SKIPPED - logger.debug( - "Document %s is not of any entity type", unprocessed_document.id - ) - elif batch_metadata[unprocessed_document.id].skip: - # info for after the connector has been processed. But no message as there may be many - # purposefully skipped documents - kg_stage = KGStage.SKIPPED - else: - kg_stage = KGStage.EXTRACTING - - with get_session_with_current_tenant() as db_session: - update_document_kg_stage( - db_session, - unprocessed_document.id, - kg_stage, - ) - - if kg_stage == KGStage.EXTRACTING: - delete_from_kg_relationships__no_commit( - db_session, [unprocessed_document.id] - ) - delete_from_kg_entities__no_commit( - db_session, [unprocessed_document.id] - ) - db_session.commit() - - # Iterate over batches of unprocessed documents - # For each document: - # - extract implied entities and relationships - # - if deep extraction is enabled, extract entities and relationships with LLM - # - if deep extraction and classification are enabled, classify document - # - update postgres with - # - extracted entities (with classification) and relationships - # - kg_stage of the processed document - - documents_to_process = [x.id for x in unprocessed_document_batch] - batch_implied_extraction: dict[str, KGImpliedExtractionResults] = {} - batch_deep_extraction_args: list[ - tuple[str, KGEnhancedDocumentMetadata, KGImpliedExtractionResults] - ] = [] - - for unprocessed_document in unprocessed_document_batch: - if ( - unprocessed_document.id not in documents_to_process - or batch_metadata[unprocessed_document.id].entity_type is None - or batch_metadata[unprocessed_document.id].skip - ): - with get_session_with_current_tenant() as db_session: - update_document_kg_stage( - db_session, - unprocessed_document.id, - KGStage.SKIPPED, - ) - db_session.commit() - continue - - # 1. perform (implicit) KG 'extractions' on the documents that should be processed - # This is really about assigning document meta-data to KG entities/relationships or KG entity attributes - # General approach: - # - vendor emails to Employee-type entities + relationship to current primary grounded entity - # - external account emails to Account-type entities + relationship to current primary grounded entity - # - non-email owners to KG current entity's attributes, no relationships - # We also collect email addresses of vendors and external accounts to inform chunk processing - batch_implied_extraction[unprocessed_document.id] = ( - kg_implied_extraction( - unprocessed_document, - batch_metadata[unprocessed_document.id], - active_entity_types, - kg_config_settings, - ) - ) - - # 2. prepare inputs for deep extraction and classification - if batch_metadata[unprocessed_document.id].deep_extraction: - batch_deep_extraction_args.append( - ( - unprocessed_document.id, - batch_metadata[unprocessed_document.id], - batch_implied_extraction[unprocessed_document.id], - ) - ) - - # 2. perform deep extraction and classification in parallel - batch_deep_extraction_func_calls = [ - ( - kg_deep_extraction, - ( - *arg, - tenant_id, - index_name, - kg_config_settings, - ), - ) - for arg in batch_deep_extraction_args - ] - batch_deep_extractions: dict[str, KGDocumentDeepExtractionResults] = dict( - zip( - # Only the deep-extraction documents have a result. Skipped - # documents are not in `batch_deep_extraction_args`. - [arg[0] for arg in batch_deep_extraction_args], - run_functions_tuples_in_parallel(batch_deep_extraction_func_calls), - strict=True, - ) - ) - - # Collect entities and relationships to upsert - batch_entities: list[tuple[str | None, str]] = [] - batch_relationships: list[tuple[str, str]] = [] - entity_classification: dict[str, str] = {} - - for document_id, implied_metadata in batch_implied_extraction.items(): - batch_entities += [ - (None, entity) for entity in implied_metadata.implied_entities - ] - batch_entities.append((document_id, implied_metadata.document_entity)) - batch_relationships += [ - (document_id, relationship) - for relationship in implied_metadata.implied_relationships - ] - - for document_id, deep_extraction_result in batch_deep_extractions.items(): - batch_entities += [ - (None, entity) - for entity in deep_extraction_result.deep_extracted_entities - ] - for relationship in deep_extraction_result.deep_extracted_relationships: - source_entity, _, target_entity = split_relationship_id( - relationship - ) - if ( - source_entity in active_entity_types - and target_entity in active_entity_types - ): - batch_relationships += [(document_id, relationship)] - - classification_result = deep_extraction_result.classification_result - if not classification_result: - continue - entity_classification[classification_result.document_entity] = ( - classification_result.classification_class - ) - - # Populate the KG database with the extracted entities, relationships, and terms - for potential_document_id, entity in batch_entities: - # verify the entity is valid - parts = split_entity_id(entity) - if len(parts) != 2: - logger.error( - "Invalid entity %s in aggregated_kg_extractions.entities", - entity, - ) - continue - - entity_type, entity_name = parts - entity_type = entity_type.upper() - entity_name = entity_name.capitalize() - - if entity_type not in active_entity_types: - continue - - try: - with get_session_with_current_tenant() as db_session: - entity_attributes: dict[str, Any] = {} - - if potential_document_id: - entity_attributes = ( - batch_metadata[potential_document_id].document_metadata - or {} - ) - - # only keep selected attributes (and translate the attribute names) - metadata_attributes = entity_metadata_conversion_instructions[ - entity_type - ] - keep_attributes = { - metadata_attributes[attr_name].name: attr_val - for attr_name, attr_val in entity_attributes.items() - if ( - attr_name in metadata_attributes - and metadata_attributes[attr_name].keep - ) - } - - # add the classification result to the attributes - if entity in entity_classification: - keep_attributes["classification"] = entity_classification[ - entity - ] - - event_time = None - if potential_document_id: - event_time = get_document_updated_at( - potential_document_id, db_session - ) - - upserted_entity = upsert_staging_entity( - db_session=db_session, - name=entity_name, - entity_type=entity_type, - document_id=potential_document_id, - occurrences=1, - attributes=keep_attributes, - event_time=event_time, - ) - metadata_tracker.track_metadata( - entity_type, upserted_entity.attributes - ) - - db_session.commit() - except Exception as e: - logger.error("Error adding entity %s. Error message: %s", entity, e) - - for document_id, relationship in batch_relationships: - relationship_split = split_relationship_id(relationship) - - if len(relationship_split) != 3: - logger.error( - "Invalid relationship %s in aggregated_kg_extractions.relationships", - relationship, - ) - continue - - source_entity, relationship_type, target_entity = relationship_split - - source_entity_type = get_entity_type(source_entity) - target_entity_type = get_entity_type(target_entity) - - if ( - source_entity_type not in active_entity_types - or target_entity_type not in active_entity_types - ): - continue - - relationship_type_id_name = extract_relationship_type_id(relationship) - - with get_session_with_current_tenant() as db_session: - try: - upsert_staging_relationship_type( - db_session=db_session, - source_entity_type=source_entity_type.upper(), - relationship_type=relationship_type, - target_entity_type=target_entity_type.upper(), - definition=False, - extraction_count=1, - ) - db_session.commit() - except Exception as e: - logger.error( - "Error adding relationship type %s to the database: %s", - relationship_type_id_name, - e, - ) - - with get_session_with_current_tenant() as db_session: - try: - upsert_staging_relationship( - db_session=db_session, - relationship_id_name=relationship, - source_document_id=document_id, - occurrences=1, - ) - db_session.commit() - except Exception as e: - logger.error( - "Error adding relationship %s to the database: %s", - relationship, - e, - ) - - # Populate the Documents table with the kg information for the documents - - for processed_document in documents_to_process: - with get_session_with_current_tenant() as db_session: - update_document_kg_info( - db_session, - processed_document, - KGStage.EXTRACTED, - ) - db_session.commit() - - # Update the the Skipped Docs back to Not Started - with get_session_with_current_tenant() as db_session: - skipped_documents = get_skipped_kg_documents(db_session) - for document_id in skipped_documents: - update_document_kg_stage( - db_session, - document_id, - KGStage.NOT_STARTED, - ) - db_session.commit() - - metadata_tracker.export_typeinfo() diff --git a/backend/onyx/kg/models.py b/backend/onyx/kg/models.py index e5c317dd37..4131fbbfd8 100644 --- a/backend/onyx/kg/models.py +++ b/backend/onyx/kg/models.py @@ -88,75 +88,6 @@ class KGEntityTypeDefinition(BaseModel): entity_values: list[str] = [] -class KGChunkFormat(BaseModel): - connector_id: int | None = None - document_id: str - chunk_id: int - title: str - content: str - primary_owners: list[str] - secondary_owners: list[str] - source_type: str - metadata: dict[str, str | list[str]] | None = None - - -class KGPerson(BaseModel): - name: str - company: str - employee: bool - - -class NormalizedEntities(BaseModel): - entities: list[str] - entities_w_attributes: list[str] - entity_normalization_map: dict[str, str] - - -class NormalizedRelationships(BaseModel): - relationships: list[str] - relationship_normalization_map: dict[str, str] - - -class KGMetadataContent(BaseModel): - document_id: str - source_type: str - source_metadata: dict[str, Any] | None = None - - -class KGClassificationInstructions(BaseModel): - classification_enabled: bool - classification_options: str - classification_class_definitions: dict[str, KGEntityTypeClassificationInfo] - - -class KGExtractionInstructions(BaseModel): - deep_extraction: bool - active: bool - - -class KGEntityTypeInstructions(BaseModel): - metadata_attribute_conversion: dict[str, KGAttributeProperty] - classification_instructions: KGClassificationInstructions - extraction_instructions: KGExtractionInstructions - entity_filter_attributes: dict[str, Any] | None = None - - -class KGEnhancedDocumentMetadata(BaseModel): - entity_type: str | None - metadata_attribute_conversion: dict[str, KGAttributeProperty] | None - document_metadata: dict[str, Any] | None - deep_extraction: bool - classification_enabled: bool - classification_instructions: KGClassificationInstructions | None - skip: bool - - -class KGConnectorData(BaseModel): - id: int - source: str - kg_coverage_days: int | None - - class KGStage(str, Enum): EXTRACTED = "extracted" NORMALIZED = "normalized" @@ -165,26 +96,3 @@ class KGStage(str, Enum): NOT_STARTED = "not_started" EXTRACTING = "extracting" DO_NOT_EXTRACT = "do_not_extract" - - -class KGClassificationResult(BaseModel): - document_entity: str - classification_class: str - - -class KGImpliedExtractionResults(BaseModel): - document_entity: str - implied_entities: set[str] - implied_relationships: set[str] - company_participant_emails: set[str] - account_participant_emails: set[str] - - -class KGDocumentDeepExtractionResults(BaseModel): - classification_result: KGClassificationResult | None - deep_extracted_entities: set[str] - deep_extracted_relationships: set[str] - - -class KGException(Exception): - pass diff --git a/backend/onyx/kg/resets/reset_source.py b/backend/onyx/kg/resets/reset_source.py deleted file mode 100644 index 9eae3c97c3..0000000000 --- a/backend/onyx/kg/resets/reset_source.py +++ /dev/null @@ -1,102 +0,0 @@ -from redis.lock import Lock as RedisLock -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, - 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 - - -def reset_source_kg_index( - source_name: str | None, tenant_id: str, index_name: str, lock: RedisLock -) -> None: - """ - Resets the knowledge graph index and vespa for a source. - """ - # reset vespa for the source - reset_vespa_kg_index(tenant_id, index_name, lock, source_name) - - with get_session_with_current_tenant() as db_session: - if source_name is None: - reset_full_kg_index__commit(db_session) - return - - # get all the entity types for the given source - entity_types = [ - et.id_name - for et in db_session.query(KGEntityType) - .filter(KGEntityType.grounded_source_name == source_name) - .all() - ] - if not entity_types: - raise ValueError(f"There are no entity types for the source {source_name}") - - # delete the entity type from the knowledge graph - for entity_type in entity_types: - db_session.query(KGRelationship).filter( - or_( - KGRelationship.source_node_type == entity_type, - KGRelationship.target_node_type == entity_type, - ) - ).delete() - db_session.query(KGRelationshipType).filter( - or_( - KGRelationshipType.source_entity_type_id_name == entity_type, - KGRelationshipType.target_entity_type_id_name == entity_type, - ) - ).delete() - db_session.query(KGEntity).filter( - KGEntity.entity_type_id_name == entity_type - ).delete() - db_session.query(KGRelationshipExtractionStaging).filter( - or_( - KGRelationshipExtractionStaging.source_node_type == entity_type, - KGRelationshipExtractionStaging.target_node_type == entity_type, - ) - ).delete() - db_session.query(KGEntityExtractionStaging).filter( - KGEntityExtractionStaging.entity_type_id_name == entity_type - ).delete() - db_session.query(KGRelationshipTypeExtractionStaging).filter( - or_( - KGRelationshipTypeExtractionStaging.source_entity_type_id_name - == entity_type, - KGRelationshipTypeExtractionStaging.target_entity_type_id_name - == entity_type, - ) - ).delete() - db_session.commit() - - with get_session_with_current_tenant() as db_session: - # get all the documents for the given source - kg_connectors = [ - connector.id - for connector in db_session.query(Connector) - .filter(Connector.source == DocumentSource(source_name)) - .all() - ] - document_ids = [ - cc_pair.id - for cc_pair in db_session.query(DocumentByConnectorCredentialPair) - .filter(DocumentByConnectorCredentialPair.connector_id.in_(kg_connectors)) - .all() - ] - - # reset the kg stage for the documents - db_session.query(Document).filter(Document.id.in_(document_ids)).update( - {"kg_stage": KGStage.NOT_STARTED} - ) - db_session.commit() diff --git a/backend/onyx/kg/resets/reset_vespa.py b/backend/onyx/kg/resets/reset_vespa.py deleted file mode 100644 index d4584f5650..0000000000 --- a/backend/onyx/kg/resets/reset_vespa.py +++ /dev/null @@ -1,129 +0,0 @@ -import time -from typing import Any - -from redis.lock import Lock as RedisLock - -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, 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, - 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 -from shared_configs.configs import MULTI_TENANT - -logger = setup_logger() - - -def _reset_vespa_for_doc(document_id: str, tenant_id: str, index_name: str) -> None: - vespa_index = VespaDocumentIndex( - index_name=index_name, - tenant_state=TenantState(tenant_id=tenant_id, multitenant=MULTI_TENANT), - large_chunks_enabled=False, - httpx_client=None, - ) - - reset_update_dict: dict[str, Any] = { - "fields": { - "kg_entities": {"assign": []}, - "kg_relationships": {"assign": []}, - "kg_terms": {"assign": []}, - } - } - - with get_session_with_current_tenant() as db_session: - num_chunks = get_num_chunks_for_document(db_session, document_id) - - vespa_requests: list[KGVespaChunkUpdateRequest] = [] - for chunk_num in range(num_chunks): - doc_chunk_id = get_uuid_from_chunk_info( - document_id=document_id, - chunk_id=chunk_num, - tenant_id=tenant_id, - large_chunk_id=None, - ) - vespa_requests.append( - KGVespaChunkUpdateRequest( - document_id=document_id, - chunk_id=chunk_num, - url=f"{DOCUMENT_ID_ENDPOINT.format(index_name=vespa_index.index_name)}/{doc_chunk_id}", - update_request=reset_update_dict, - ) - ) - - with vespa_index.httpx_client_context as httpx_client: - vespa_index._apply_kg_chunk_updates_batched(vespa_requests, httpx_client) - - -def reset_vespa_kg_index( - tenant_id: str, index_name: str, lock: RedisLock, source_name: str | None = None -) -> None: - """ - Reset the kg info in vespa for all documents of a given source name, - or all documents from kg grounded sources if source_name is None. - """ - logger.info( - "Resetting kg vespa index %s for tenant %s, source: %s", - index_name, - tenant_id, - source_name or "all", - ) - - last_lock_time = time.monotonic() - - # Get all documents that need a vespa reset - with get_session_with_current_tenant() as db_session: - if source_name: - # get all connectors of the given source name - kg_connectors = [ - connector.id - for connector in db_session.query(Connector) - .filter(Connector.source == DocumentSource(source_name)) - .all() - ] - else: - # get all connectors that have kg enabled - kg_sources = [ - DocumentSource(et.grounded_source_name) - for et in db_session.query(KGEntityType) - .filter( - KGEntityType.grounded_source_name.is_not(None), - KGEntityType.active.is_(True), - ) - .distinct() - .all() - ] - kg_connectors = [ - connector.id - for connector in db_session.query(Connector) - .filter(Connector.source.in_(kg_sources)) - .all() - ] - - # Get all the documents for the given connectors - document_ids = [ - cc_pair.id - for cc_pair in db_session.query(DocumentByConnectorCredentialPair) - .filter(DocumentByConnectorCredentialPair.connector_id.in_(kg_connectors)) - .all() - ] - - # Reset the kg fields - for document_id in document_ids: - _reset_vespa_for_doc(document_id, tenant_id, index_name) - last_lock_time = extend_lock( - lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time - ) - - logger.info( - "Finished resetting kg vespa index %s for tenant %s, source: %s", - index_name, - tenant_id, - source_name or "all", - ) diff --git a/backend/onyx/kg/utils/embeddings.py b/backend/onyx/kg/utils/embeddings.py deleted file mode 100644 index 40256fc5e4..0000000000 --- a/backend/onyx/kg/utils/embeddings.py +++ /dev/null @@ -1,24 +0,0 @@ -from typing import List - -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, - EmbedTextType, -) -from shared_configs.configs import MODEL_SERVER_HOST, MODEL_SERVER_PORT - - -def encode_string_batch(strings: List[str]) -> np.ndarray: - with get_session_with_current_tenant() as db_session: - current_search_settings = get_current_search_settings(db_session) - model = EmbeddingModel.from_db_model( - search_settings=current_search_settings, - server_host=MODEL_SERVER_HOST, - server_port=MODEL_SERVER_PORT, - ) - # Get embeddings while session is still open - embedding = model.encode(strings, text_type=EmbedTextType.QUERY) - return np.array(embedding) diff --git a/backend/onyx/kg/utils/extraction_utils.py b/backend/onyx/kg/utils/extraction_utils.py deleted file mode 100644 index e4a06d702a..0000000000 --- a/backend/onyx/kg/utils/extraction_utils.py +++ /dev/null @@ -1,748 +0,0 @@ -import json - -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, KGEntityType, KGRelationshipType -from onyx.db.tag import get_structured_tags_for_document -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.interfaces import GenerationContext -from onyx.llm.models import GenerationRequest, UserMessage -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.framework.create import ensure_trace -from onyx.tracing.framework.traces import TraceContentMode -from onyx.utils.logger import setup_logger - -logger = setup_logger() - -# Deep extraction returns a full JSON entity/relationship graph per chunk batch, -# so it needs more than the default invoke timeout. -_DEEP_EXTRACTION_TIMEOUT_S = 300 - -KG_DOCUMENT_PROCESSING_TRACE_NAME = "kg_document_processing" - - -def get_entity_types_str(active: bool | None = None) -> str: - """ - Format the entity types into a string for the LLM. - """ - with get_session_with_current_tenant() as db_session: - entity_types = get_entity_types(db_session, active) - - entity_types_list: list[str] = [] - for entity_type in entity_types: - if entity_type.description: - entity_description = "\n - Description: " + entity_type.description - else: - entity_description = "" - - if entity_type.entity_values: - allowed_values = "\n - Allowed Values: " + ", ".join( - entity_type.entity_values - ) - else: - allowed_values = "" - - attributes = entity_type.parsed_attributes - - entity_type_attribute_list: list[str] = [] - for attribute, values in attributes.attribute_values.items(): - entity_type_attribute_list.append( - f"{attribute}: {trackinfo_to_str(values)}" - ) - - if attributes.classification_attributes: - entity_type_attribute_list.append( - # TODO: restructure classification attribute to be a dict of attribute name to classification info - # e.g., {scope: {internal: prompt, external: prompt}, sentiment: {positive: prompt, negative: prompt}} - "classification: one of: " - + ", ".join(attributes.classification_attributes.keys()) - ) - if entity_type_attribute_list: - entity_attributes = "\n - Attributes:\n - " + "\n - ".join( - entity_type_attribute_list - ) - else: - entity_attributes = "" - - entity_types_list.append( - entity_type.id_name - + entity_description - + allowed_values - + entity_attributes - ) - - return "\n".join(entity_types_list) - - -def get_relationship_types_str(active: bool | None = None) -> str: - """ - Format the relationship types into a string for the LLM. - """ - with get_session_with_current_tenant() as db_session: - active_filters = [] - if active is not None: - active_filters.append(KGRelationshipType.active == active) - - relationship_types = ( - db_session.query(KGRelationshipType).filter(*active_filters).all() - ) - - relationship_types_list = [] - for rel_type in relationship_types: - # Format as "source_type__relationship_type__target_type" - formatted_type = make_relationship_type_id( - rel_type.source_entity_type_id_name, - rel_type.type, - rel_type.target_entity_type_id_name, - ) - relationship_types_list.append(formatted_type) - - return "\n".join(relationship_types_list) - - -def kg_process_owners( - owner_emails: list[str], - document_entity_id: str, - relationship_type: str, - kg_config_settings: KGConfigSettings, - active_entity_types: set[str], -) -> tuple[set[str], set[str], set[str], set[str]]: - owner_entities: set[str] = set() - owner_relationships: set[str] = set() - company_participant_emails: set[str] = set() - account_participant_emails: set[str] = set() - - for owner_email in owner_emails: - if extract_email(owner_email) is None: - continue - - process_results = kg_process_person( - owner_email, - document_entity_id, - relationship_type, - kg_config_settings, - active_entity_types, - ) - if process_results is None: - continue - - ( - owner_entity, - owner_relationship, - company_participant_email, - account_participant_email, - ) = process_results - - owner_entities.add(owner_entity) - owner_relationships.add(owner_relationship) - if company_participant_email: - company_participant_emails.add(company_participant_email) - if account_participant_email: - account_participant_emails.add(account_participant_email) - - return ( - owner_entities, - owner_relationships, - company_participant_emails, - account_participant_emails, - ) - - -def kg_implied_extraction( - document: Document, - doc_metadata: KGEnhancedDocumentMetadata, - active_entity_types: set[str], - kg_config_settings: KGConfigSettings, -) -> KGImpliedExtractionResults: - """ - Generate entities, relationships, and attributes for a document. - """ - - # Get document entity and metadata stuff from the KGEnhancedDocumentMetadata - document_entity_type = doc_metadata.entity_type - document_metadata = doc_metadata.document_metadata or {} - metadata_attribute_conversion = doc_metadata.metadata_attribute_conversion - if document_entity_type is None or metadata_attribute_conversion is None: - raise ValueError("Entity type and metadata attributes are required") - - implied_entities: set[str] = set() - implied_relationships: set[str] = set() - - # Quantity needed for call processing - participants from vendor - company_participant_emails: set[str] = set() - # Quantity needed for call processing - external participants - account_participant_emails: set[str] = set() - - # Chunk treatment variables - - document_is_from_call = document_entity_type.lower() in ( - call_type.value.lower() for call_type in OnyxCallTypes - ) - - # Get core entity - - document_id = document.id - primary_owners = document.primary_owners - secondary_owners = document.secondary_owners - - with get_session_with_current_tenant() as db_session: - document_entity = get_kg_entity_by_document(db_session, document_id) - - if document_entity: - document_entity_id = document_entity.id_name - else: - document_entity_id = make_entity_id(document_entity_type, document_id) - - # Get implied entities and relationships from primary/secondary owners - - if document_is_from_call: - ( - implied_entities, - implied_relationships, - company_participant_emails, - account_participant_emails, - ) = kg_process_owners( - owner_emails=(primary_owners or []) + (secondary_owners or []), - document_entity_id=document_entity_id, - relationship_type="participates_in", - kg_config_settings=kg_config_settings, - active_entity_types=active_entity_types, - ) - else: - ( - implied_entities, - implied_relationships, - company_participant_emails, - account_participant_emails, - ) = kg_process_owners( - owner_emails=primary_owners or [], - document_entity_id=document_entity_id, - relationship_type="leads", - kg_config_settings=kg_config_settings, - active_entity_types=active_entity_types, - ) - - ( - participant_entities, - participant_relationships, - company_emails, - account_emails, - ) = kg_process_owners( - owner_emails=secondary_owners or [], - document_entity_id=document_entity_id, - relationship_type="participates_in", - kg_config_settings=kg_config_settings, - active_entity_types=active_entity_types, - ) - implied_entities.update(participant_entities) - implied_relationships.update(participant_relationships) - company_participant_emails.update(company_emails) - account_participant_emails.update(account_emails) - - # Get implied entities and relationships from document metadata - for metadata, value in document_metadata.items(): - # get implication property for this metadata - if metadata not in metadata_attribute_conversion: - continue - if ( - implication_property := metadata_attribute_conversion[ - metadata - ].implication_property - ) is None: - continue - - if not isinstance(value, str) and not isinstance(value, list): - continue - values: list[str] = [value] if isinstance(value, str) else value - - # create implied entities and relationships - for item in values: - if ( - implication_property.implied_entity_type - == KGAttributeEntityOption.FROM_EMAIL - ): - # determine entity type from email - email = extract_email(item) - if email is None: - continue - process_results = kg_process_person( - email=email, - document_entity_id=document_entity_id, - relationship_type=implication_property.implied_relationship_name, - kg_config_settings=kg_config_settings, - active_entity_types=active_entity_types, - ) - if process_results is None: - continue - - implied_entity, implied_relationship, _, _ = process_results - implied_entities.add(implied_entity) - implied_relationships.add(implied_relationship) - else: - # use the given entity type - entity_type = implication_property.implied_entity_type - if entity_type not in active_entity_types: - continue - - implied_entity = make_entity_id(entity_type, item) - implied_entities.add(implied_entity) - implied_relationships.add( - make_relationship_id( - implied_entity, - implication_property.implied_relationship_name, - document_entity_id, - ) - ) - - return KGImpliedExtractionResults( - document_entity=document_entity_id, - implied_entities=implied_entities, - implied_relationships=implied_relationships, - company_participant_emails=company_participant_emails, - account_participant_emails=account_participant_emails, - ) - - -def kg_deep_extraction( - document_id: str, - metadata: KGEnhancedDocumentMetadata, - implied_extraction: KGImpliedExtractionResults, - tenant_id: str, - index_name: str, - kg_config_settings: KGConfigSettings, -) -> KGDocumentDeepExtractionResults: - """Perform one document's deep extraction and classification workflow.""" - with ensure_trace( - KG_DOCUMENT_PROCESSING_TRACE_NAME, - content_mode=TraceContentMode.METADATA_ONLY, - ): - return _kg_deep_extraction( - document_id=document_id, - metadata=metadata, - implied_extraction=implied_extraction, - tenant_id=tenant_id, - index_name=index_name, - kg_config_settings=kg_config_settings, - ) - - -def _kg_deep_extraction( - document_id: str, - metadata: KGEnhancedDocumentMetadata, - implied_extraction: KGImpliedExtractionResults, - tenant_id: str, - index_name: str, - kg_config_settings: KGConfigSettings, -) -> KGDocumentDeepExtractionResults: - result = KGDocumentDeepExtractionResults( - classification_result=None, - deep_extracted_entities=set(), - deep_extracted_relationships=set(), - ) - - entity_types_str = get_entity_types_str(active=True) - relationship_types_str = get_relationship_types_str(active=True) - - for i, chunk_batch in enumerate( - get_document_vespa_contents(document_id, index_name, tenant_id) - ): - # use first batch for classification - if i == 0 and metadata.classification_enabled: - if not metadata.classification_instructions: - raise ValueError( - "Classification is enabled but no instructions are provided" - ) - result.classification_result = kg_classify_document( - document_entity=implied_extraction.document_entity, - chunk_batch=chunk_batch, - implied_extraction=implied_extraction, - classification_instructions=metadata.classification_instructions, - kg_config_settings=kg_config_settings, - ) - - # deep extract from this chunk batch - chunk_batch_results = kg_deep_extract_chunks( - document_entity=implied_extraction.document_entity, - chunk_batch=chunk_batch, - implied_extraction=implied_extraction, - kg_config_settings=kg_config_settings, - entity_types_str=entity_types_str, - relationship_types_str=relationship_types_str, - ) - if chunk_batch_results is not None: - result.deep_extracted_entities.update( - chunk_batch_results.deep_extracted_entities - ) - result.deep_extracted_relationships.update( - chunk_batch_results.deep_extracted_relationships - ) - - return result - - -def kg_classify_document( - document_entity: str, - chunk_batch: list[KGChunkFormat], - implied_extraction: KGImpliedExtractionResults, - classification_instructions: KGClassificationInstructions, - kg_config_settings: KGConfigSettings, -) -> KGClassificationResult | None: - # currently, classification is only done for calls - # TODO: add support (or use same prompt and format) for non-call documents - entity_type = get_entity_type(document_entity) - if entity_type not in (call_type.value for call_type in OnyxCallTypes): - return None - - # prepare prompt - company_participants = implied_extraction.company_participant_emails - account_participants = implied_extraction.account_participant_emails - content = ( - f"Title: {chunk_batch[0].title}:\nVendor Participants:\n" - + "".join(f" - {participant}\n" for participant in company_participants) - + "Other Participants:\n" - + "".join(f" - {participant}\n" for participant in account_participants) - + "Call Content:\n" - + "\n".join(chunk.content for chunk in chunk_batch) - ) - category_list = { - cls: definition.description - for cls, definition in classification_instructions.classification_class_definitions.items() - } - prompt = CALL_DOCUMENT_CLASSIFICATION_PROMPT.format( - beginning_of_call_content=content, - category_list=category_list, - category_options=classification_instructions.classification_options, - vendor=kg_config_settings.KG_VENDOR, - ) - - # classify with LLM with Braintrust tracing - llm = get_default_llm() - try: - prompt_msg = UserMessage(content=prompt) - response = llm.invoke( - GenerationRequest(messages=[prompt_msg]), - context=GenerationContext( - flow=LLMFlow.KG_DOCUMENT_CLASSIFICATION, - content_mode=TraceContentMode.METADATA_ONLY, - ), - ) - raw_classification_result = response.text - - classification_result = ( - raw_classification_result.replace("```json", "").replace("```", "").strip() - ) - # no json parsing here because of reasoning output - classification_class = classification_result.split("CATEGORY:")[1].strip() - - if ( - classification_class - in classification_instructions.classification_class_definitions - ): - return KGClassificationResult( - document_entity=document_entity, - classification_class=classification_class, - ) - except Exception as e: - logger.error( - "Failed to classify document %s. Error: %s", document_entity, str(e) - ) - return None - - -def kg_deep_extract_chunks( - document_entity: str, - chunk_batch: list[KGChunkFormat], - implied_extraction: KGImpliedExtractionResults, - kg_config_settings: KGConfigSettings, - entity_types_str: str, - relationship_types_str: str, -) -> KGDocumentDeepExtractionResults | None: - # currently, calls are treated differently - # TODO: either treat some other documents differently too, or ideally all the same way - entity_type = get_entity_type(document_entity) - is_call = entity_type in (call_type.value for call_type in OnyxCallTypes) - - content = "\n".join(chunk.content for chunk in chunk_batch) - - # prepare prompt - if is_call: - company_participants_str = "".join( - f" - {participant}\n" - for participant in implied_extraction.company_participant_emails - ) - account_participants_str = "".join( - f" - {participant}\n" - for participant in implied_extraction.account_participant_emails - ) - llm_context = CALL_CHUNK_PREPROCESSING_PROMPT.format( - participant_string=company_participants_str, - account_participant_string=account_participants_str, - vendor=kg_config_settings.KG_VENDOR, - content=content, - ) - else: - llm_context = GENERAL_CHUNK_PREPROCESSING_PROMPT.format( - vendor=kg_config_settings.KG_VENDOR, - content=content, - ) - prompt = MASTER_EXTRACTION_PROMPT.format( - entity_types=entity_types_str, - relationship_types=relationship_types_str, - ).replace("---content---", llm_context) - - # extract with LLM with Braintrust tracing - llm = get_default_llm() - try: - prompt_msg = UserMessage(content=prompt) - response = llm.invoke( - GenerationRequest(messages=[prompt_msg]), - context=GenerationContext( - flow=LLMFlow.KG_DEEP_EXTRACTION, - total_timeout_s=_DEEP_EXTRACTION_TIMEOUT_S, - content_mode=TraceContentMode.METADATA_ONLY, - ), - ) - raw_extraction_result = response.text - - cleaned_response = ( - raw_extraction_result.replace("{{", "{") - .replace("}}", "}") - .replace("```json\n", "") - .replace("\n```", "") - .replace("\n", "") - ) - first_bracket = cleaned_response.find("{") - last_bracket = cleaned_response.rfind("}") - cleaned_response = cleaned_response[first_bracket : last_bracket + 1] - parsed_result = json.loads(cleaned_response) - return KGDocumentDeepExtractionResults( - classification_result=None, - deep_extracted_entities=set(parsed_result.get("entities", [])), - deep_extracted_relationships={ - rel.replace(" ", "_") for rel in parsed_result.get("relationships", []) - }, - ) - except Exception as e: - failed_chunks = [chunk.chunk_id for chunk in chunk_batch] - logger.error( - "Failed to process chunks %s from document %s. Error: %s", - failed_chunks, - document_entity, - str(e), - ) - return None - - -def kg_process_person( - email: str, - document_entity_id: str, - relationship_type: str, - kg_config_settings: KGConfigSettings, - active_entity_types: set[str], -) -> tuple[str, str, str, str] | None: - """ - Create an employee or account entity from an email address, and a relationship to - the entity from the document that the email is from. - - Returns: - tuple containing (person_entity, person_relationship, company_participant_email, - and account_participant_email), or None if the created entity is not of an - active entity type or is from an ignored email domain. - """ - kg_person = kg_email_processing(email, kg_config_settings) - if any( - domain.lower() in kg_person.company.lower() - for domain in kg_config_settings.KG_IGNORE_EMAIL_DOMAINS - ): - return None - - person_entity = None - if kg_person.employee and "EMPLOYEE" in active_entity_types: - person_entity = make_entity_id("EMPLOYEE", kg_person.name) - elif not kg_person.employee and "ACCOUNT" in active_entity_types: - person_entity = make_entity_id("ACCOUNT", kg_person.company) - - if person_entity: - is_account = person_entity.startswith("ACCOUNT") - participant_email = f"{kg_person.name} -- ({kg_person.company})" - return ( - person_entity, - make_relationship_id(person_entity, relationship_type, document_entity_id), - participant_email if not is_account else "", - participant_email if is_account else "", - ) - - return None - - -def get_batch_documents_metadata( - document_ids: list[str], connector_source: str -) -> list[KGMetadataContent]: - """ - Gets the metadata for a batch of documents. - """ - batch_metadata: list[KGMetadataContent] = [] - source_type = DocumentSource(connector_source).value - - with get_session_with_current_tenant() as db_session: - for document_id in document_ids: - # get document metadata - metadata = get_structured_tags_for_document(document_id, db_session) - - batch_metadata.append( - KGMetadataContent( - document_id=document_id, - source_type=source_type, - source_metadata=metadata, - ) - ) - return batch_metadata - - -def trackinfo_to_str( - trackinfo: KGAttributeTrackInfo | None, -) -> str: - """Convert trackinfo to an LLM friendly string""" - if trackinfo is None: - return "" - - if trackinfo.type == KGAttributeTrackType.LIST: - if trackinfo.values is None: - return "a list of any suitable values" - return "a list with possible values: " + ", ".join(trackinfo.values) - elif trackinfo.type == KGAttributeTrackType.VALUE: - if trackinfo.values is None: - return "any suitable value" - return "one of: " + ", ".join(trackinfo.values) - - -def trackinfo_to_dict(trackinfo: KGAttributeTrackInfo | None) -> dict | None: - if trackinfo is None: - return None - return { - "type": trackinfo.type, - "values": (list(trackinfo.values) if trackinfo.values else None), - } - - -class EntityTypeMetadataTracker: - def __init__(self) -> None: - """ - Tracks the possible values the metadata attributes can take for each entity type. - """ - # entity type -> attribute -> trackinfo - self.entity_attr_info: dict[str, dict[str, KGAttributeTrackInfo | None]] = {} - self.entity_allowed_attrs: dict[str, set[str]] = {} - - def import_typeinfo(self) -> None: - """ - Loads the metadata tracking information from the database. - """ - with get_session_with_current_tenant() as db_session: - entity_types = db_session.query(KGEntityType).all() - - for entity_type in entity_types: - self.entity_attr_info[entity_type.id_name] = ( - entity_type.parsed_attributes.attribute_values - ) - self.entity_allowed_attrs[entity_type.id_name] = { - attr.name - for attr in entity_type.parsed_attributes.metadata_attribute_conversion.values() - } - - def export_typeinfo(self) -> None: - """ - Exports the metadata tracking information to the database. - """ - with get_session_with_current_tenant() as db_session: - for entity_type_id_name, attribute_values in self.entity_attr_info.items(): - db_session.query(KGEntityType).filter( - KGEntityType.id_name == entity_type_id_name - ).update( - { - KGEntityType.attributes: KGEntityType.attributes.op("||")( - { - "attribute_values": { - attr: trackinfo_to_dict(info) - for attr, info in attribute_values.items() - } - } - ) - }, - synchronize_session=False, - ) - db_session.commit() - - def track_metadata( - self, entity_type: str, attributes: dict[str, str | list[str]] - ) -> None: - """ - Tracks which values are possible for the given attributes. - If the attribute value is a list, we track the values in the list rather than the list itself. - If we see to many different values, we stop tracking the attribute. - """ - for attribute, value in attributes.items(): - # ignore types/metadata we are not tracking - if entity_type not in self.entity_attr_info: - continue - if attribute not in self.entity_allowed_attrs[entity_type]: - continue - - # determine if the attribute is a list or a value - trackinfo = self.entity_attr_info[entity_type].get(attribute, None) - if trackinfo is None: - trackinfo = KGAttributeTrackInfo( - type=( - KGAttributeTrackType.VALUE - if isinstance(value, str) - else KGAttributeTrackType.LIST - ), - values=set(), - ) - self.entity_attr_info[entity_type][attribute] = trackinfo - - # None means marked as don't track - if trackinfo.values is None: - continue - - # track the value - if isinstance(value, str): - trackinfo.values.add(value) - else: - trackinfo.type = KGAttributeTrackType.LIST - trackinfo.values.update(value) - - # if we see to many different values, we stop tracking - if len(trackinfo.values) > KG_METADATA_TRACKING_THRESHOLD: - trackinfo.values = None diff --git a/backend/onyx/kg/utils/formatting_utils.py b/backend/onyx/kg/utils/formatting_utils.py deleted file mode 100644 index 284c9d68fc..0000000000 --- a/backend/onyx/kg/utils/formatting_utils.py +++ /dev/null @@ -1,132 +0,0 @@ -import re - -from onyx.db.kg_config import KGConfigSettings -from onyx.kg.models import KGPerson - - -def format_entity_id(entity_id_name: str) -> str: - return make_entity_id(*split_entity_id(entity_id_name)) - - -def make_entity_id(entity_type: str, entity_name: str) -> str: - return f"{entity_type.upper()}::{entity_name.lower()}" - - -def split_entity_id(entity_id_name: str) -> list[str]: - return entity_id_name.split("::") - - -def get_entity_type(entity_id_name: str) -> str: - return entity_id_name.split("::", 1)[0].upper() - - -def format_entity_id_for_models(entity_id_name: str) -> str: - entity_split = entity_id_name.split("::") - if len(entity_split) == 2: - entity_type, entity_name = entity_split - separator = "::" - elif len(entity_split) > 2: - raise ValueError(f"Entity {entity_id_name} is not in the correct format") - else: - entity_name = entity_id_name - separator = entity_type = "" - - formatted_entity_type = entity_type.strip().upper() - formatted_entity_name = entity_name.strip().replace('"', "").replace("'", "") - - return f"{formatted_entity_type}{separator}{formatted_entity_name}" - - -def get_attributes(entity_w_attributes: str) -> dict[str, str]: - """ - Extract attributes from an entity string. - E.g., "TYPE::Entity--[attr1: value1, attr2: value2]" -> {"attr1": "value1", "attr2": "value2"} - """ - attr_split = entity_w_attributes.split("--") - if len(attr_split) != 2: - raise ValueError(f"Invalid entity with attributes: {entity_w_attributes}") - - match = re.search(r"\[(.*)\]", attr_split[1]) - if not match: - return {} - - attr_list_str = match.group(1) - return { - attr_split[0].strip(): attr_split[1].strip() - for attr in attr_list_str.split(",") - if len(attr_split := attr.split(":", 1)) == 2 - } - - -def make_entity_w_attributes(entity: str, attributes: dict[str, str]) -> str: - return f"{entity}--[{', '.join(f'{k}: {v}' for k, v in attributes.items())}]" - - -def format_relationship_id(relationship_id_name: str) -> str: - return make_relationship_id(*split_relationship_id(relationship_id_name)) - - -def make_relationship_id( - source_node: str, relationship_type: str, target_node: str -) -> str: - return f"{format_entity_id(source_node)}__{relationship_type.lower()}__{format_entity_id(target_node)}" - - -def split_relationship_id(relationship_id_name: str) -> list[str]: - return relationship_id_name.split("__") - - -def format_relationship_type_id(relationship_type_id_name: str) -> str: - return make_relationship_type_id( - *split_relationship_type_id(relationship_type_id_name) - ) - - -def make_relationship_type_id( - source_node_type: str, relationship_type: str, target_node_type: str -) -> str: - return f"{source_node_type.upper()}__{relationship_type.lower()}__{target_node_type.upper()}" - - -def split_relationship_type_id(relationship_type_id_name: str) -> list[str]: - return relationship_type_id_name.split("__") - - -def extract_relationship_type_id(relationship_id_name: str) -> str: - source_node, relationship_type, target_node = split_relationship_id( - relationship_id_name - ) - return make_relationship_type_id( - get_entity_type(source_node), relationship_type, get_entity_type(target_node) - ) - - -def extract_email(email: str) -> str | None: - """ - Extract an email from an arbitrary string (if any). - Only the first email is returned. - """ - match = re.search(r"([A-Za-z0-9._+-]+@[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)+)", email) - return match.group(0) if match else None - - -def kg_email_processing(email: str, kg_config_settings: KGConfigSettings) -> KGPerson: - """ - Process the email. - """ - name, company_domain = email.split("@") - assert isinstance(company_domain, str) - assert isinstance(kg_config_settings.KG_VENDOR_DOMAINS, list) - assert isinstance(kg_config_settings.KG_VENDOR, str) - - employee = any( - domain in company_domain for domain in kg_config_settings.KG_VENDOR_DOMAINS - ) - if employee: - company = kg_config_settings.KG_VENDOR - else: - # TODO: maybe store a list of domains for each account and use that to match - # right now, gmail and other random domains are being converted into accounts - company = company_domain.title() - - return KGPerson(name=name, company=company, employee=employee) diff --git a/backend/onyx/kg/utils/lock_utils.py b/backend/onyx/kg/utils/lock_utils.py deleted file mode 100644 index 5b3d6ecd58..0000000000 --- a/backend/onyx/kg/utils/lock_utils.py +++ /dev/null @@ -1,12 +0,0 @@ -import time - -from redis.lock import Lock as RedisLock - - -def extend_lock(lock: RedisLock, timeout: int, last_lock_time: float) -> float: - current_time = time.monotonic() - if current_time - last_lock_time >= (timeout / 4): - lock.reacquire() - last_lock_time = current_time - - return last_lock_time diff --git a/backend/onyx/kg/vespa/vespa_interactions.py b/backend/onyx/kg/vespa/vespa_interactions.py deleted file mode 100644 index 7e8c6f0dab..0000000000 --- a/backend/onyx/kg/vespa/vespa_interactions.py +++ /dev/null @@ -1,80 +0,0 @@ -import json -from collections.abc import Generator - -from onyx.context.search.models import IndexFilters -from onyx.document_index.vespa.chunk_retrieval import get_chunks_via_visit_api -from onyx.document_index.vespa.internal_types import VespaChunkRequest -from onyx.kg.models import KGChunkFormat -from onyx.utils.logger import setup_logger - -logger = setup_logger() - - -def get_document_vespa_contents( - document_id: str, - index_name: str, - tenant_id: str, - batch_size: int = 8, -) -> Generator[list[KGChunkFormat], None, None]: - """ - Retrieves chunks from Vespa for the given document IDs and converts them to KGChunks. - - Args: - document_id (str): ID of the document to fetch chunks for - index_name (str): Name of the Vespa index - tenant_id (str): ID of the tenant - batch_size (int): Number of chunks to fetch per batch - - Yields: - list[KGChunk]: Batches of chunks ready for KG processing - """ - - current_batch: list[KGChunkFormat] = [] - - # get all chunks for the document - # TODO: revisit the visit function - chunks = get_chunks_via_visit_api( - chunk_request=VespaChunkRequest(document_id=document_id), - index_name=index_name, - filters=IndexFilters(access_control_list=None, tenant_id=tenant_id), - field_names=[ - "document_id", - "chunk_id", - "title", - "content", - "metadata", - "primary_owners", - "secondary_owners", - "source_type", - ], - get_large_chunks=False, - ) - - # Convert Vespa chunks to KGChunks - # kg_chunks: list[KGChunkFormat] = [] - - for _i, chunk in enumerate(chunks): - fields = chunk["fields"] - if isinstance(fields.get("metadata", {}), str): - fields["metadata"] = json.loads(fields["metadata"]) - current_batch.append( - KGChunkFormat( - connector_id=None, # We may need to adjust this - document_id=fields.get("document_id"), - chunk_id=fields.get("chunk_id"), - primary_owners=fields.get("primary_owners", []), - secondary_owners=fields.get("secondary_owners", []), - source_type=fields.get("source_type", ""), - title=fields.get("title", ""), - content=fields.get("content", ""), - metadata=fields.get("metadata", {}), - ) - ) - - if len(current_batch) >= batch_size: - yield current_batch - current_batch = [] - - # Yield any remaining chunks - if current_batch: - yield current_batch diff --git a/backend/onyx/main.py b/backend/onyx/main.py index abe6ed6a05..fb95b69aa6 100644 --- a/backend/onyx/main.py +++ b/backend/onyx/main.py @@ -129,9 +129,6 @@ from onyx.server.manage.image_generation.api import ( from onyx.server.manage.llm.api import admin_router as llm_admin_router from onyx.server.manage.llm.api import basic_router as llm_router from onyx.server.manage.oauth_test import router as oauth_test_admin_router -from onyx.server.manage.opensearch_migration.api import ( - admin_router as opensearch_migration_admin_router, -) from onyx.server.manage.search_settings import router as search_settings_router from onyx.server.manage.slack_bot import router as slack_bot_management_router from onyx.server.manage.sso.api import admin_router as sso_admin_router @@ -632,9 +629,6 @@ def get_application(lifespan_override: Lifespan | None = None) -> FastAPI: include_router_with_global_prefix_prepended(application, voice_admin_router) include_router_with_global_prefix_prepended(application, voice_router) include_router_with_global_prefix_prepended(application, voice_websocket_router) - include_router_with_global_prefix_prepended( - application, opensearch_migration_admin_router - ) include_router_with_global_prefix_prepended(application, cost_override_router) include_router_with_global_prefix_prepended(application, user_usage_router) include_router_with_global_prefix_prepended(application, admin_usage_router) diff --git a/backend/onyx/prompts/kg_prompts.py b/backend/onyx/prompts/kg_prompts.py index 3dada247ae..ad96ff724c 100644 --- a/backend/onyx/prompts/kg_prompts.py +++ b/backend/onyx/prompts/kg_prompts.py @@ -1,1458 +1,6 @@ -# Standards -SEPARATOR_LINE = "-------" -SEPARATOR_LINE_LONG = "---------------" -NO_EXTRACTION = "No extraction of knowledge graph objects was feasible." -YES = "yes" -NO = "no" - -# Framing/Support/Template Prompts -ENTITY_TYPE_SETTING_PROMPT = f""" -{SEPARATOR_LINE} -{{entity_types}} -{SEPARATOR_LINE} -""".strip() - -RELATIONSHIP_TYPE_SETTING_PROMPT = f""" -Here are the types of relationships: -{SEPARATOR_LINE} -{{relationship_types}} -{SEPARATOR_LINE} -""".strip() - -EXTRACTION_FORMATTING_PROMPT = r""" -{{"entities": [::' (please use that capitalization). If allowed options \ -are provided above, you can only extract those types of entities! Again, there should be an 'Other' \ -option. Pick this if none of the others apply.>], -"relationships": [::__\ -__::'>], -"terms": ['>] -}} -""".strip() - -QUERY_ENTITY_EXTRACTION_FORMATTING_PROMPT = r""" -{{"entities": [::' (please use that capitalization)>. Each entity \ -also should be followed by a list of comma-separated attribute filters for the entity, if referred to in the \ -question for that entity. CRITICAL: you can only use attributes that are mentioned above for the \ -entity type in question. Example: 'ACCOUNT::* -- [account_type: customer, status: active]' if the question is \ -'list all customer accounts', and ACCOUNT was an entity type with these attribute key/values allowed.] \ -"time_filter": -}} -""".strip() - -QUERY_RELATIONSHIP_EXTRACTION_FORMATTING_PROMPT = r""" -{{"relationships": [::__\ -__::'>] -}} -""".strip() - -EXAMPLE_1 = r""" -{{"entities": ["ACCOUNT::Nike", "CONCERN::*"], - "relationships": ["ACCOUNT::Nike__had__CONCERN::*"], "terms": []}} -""".strip() - -EXAMPLE_2 = r""" -{{"entities": ["ACCOUNT::Nike", "CONCERN::performance"], - "relationships": ["ACCOUNT::*__had_issues__CONCERN::performance"], "terms": ["performance issue"]}} -""".strip() - -EXAMPLE_3 = r""" -{{"entities": ["ACCOUNT::Nike", "CONCERN::performance", "CONCERN::user_experience"], - "relationships": ["ACCOUNT::Nike__had__CONCERN::performance", - "ACCOUNT::Nike__solved__CONCERN::user_experience"], - "terms": ["performance", "user experience"]}} -""".strip() - -EXAMPLE_4 = r""" -{{"entities": ["ACCOUNT::Nike", "FEATURE::dashboard", "CONCERN::performance"], - "relationships": ["ACCOUNT::Nike__had__CONCERN::performance", - "ACCOUNT::Nike__had_issues__FEATURE::dashboard", - "ACCOUNT::NIKE__gets_value_from__FEATURE::dashboard"], - "terms": ["value", "performance"]}} -""".strip() - -RELATIONSHIP_EXAMPLE_1 = r""" -'Which issues did Nike report?' and the extracted entities were found to be: - - "ACCOUNT::Nike", "CONCERN::*" - -then a valid relationship extraction could be: - -{{"relationships": ["ACCOUNT::Nike__had__CONCERN::*"]}} -""".strip() - -RELATIONSHIP_EXAMPLE_2 = r""" -'Did Nike say anything about performance issues?' and the extracted entities were found to be: - -"ACCOUNT::Nike", "CONCERN::performance" - -then a much more suitable relationship extraction could be: -{{"relationships": ["ACCOUNT::*__had_issues__CONCERN::performance"]}} -""".strip() - -RELATIONSHIP_EXAMPLE_3 = r""" -'Did Nike report some performance issues with our solution? And were they happy that the user experience issue got solved?', \ -and the extracted entities were found to be: - -"ACCOUNT::Nike", "CONCERN::performance", "CONCERN::user_experience" - -then a valid relationship extraction could be: - -{{"relationships": ["ACCOUNT::Nike__had__CONCERN::performance", - "ACCOUNT::Nike__solved__CONCERN::user_experience"]}} -""".strip() - -RELATIONSHIP_EXAMPLE_4 = r""" -'Nike reported some performance issues with our dashboard solution, but do they think it delivers great value nevertheless?' \ -and the extracted entities were found to be: - -"ACCOUNT::Nike", "FEATURE::dashboard", "CONCERN::performance" - -then a valid relationship extraction could be: -Example 4: - -{{"relationships": ["ACCOUNT::Nike__had__CONCERN::performance", - "ACCOUNT::Nike__had_issues__FEATURE::dashboard", - "ACCOUNT::NIKE__gets_value_from__FEATURE::dashboard"]}} - -Explanation: - - Nike did report performance concerns - - Nike had problems with the dashboard, which is a feature - - We are interested in the value relationship between Nike and the dashboard feature - -""".strip() - -RELATIONSHIP_EXAMPLE_5 = r""" -'In which emails did Nike discuss their issues with the dashboard?' \ -and the extracted entities were found to be: - -"ACCOUNT::Nike", "FEATURE::dashboard", "EMAIL::*" - -then a valid relationship extraction could be: - -{{"relationships": ["ACCOUNT::Nike__had__CONCERN::*", - "ACCOUNT::Nike__had_issues__FEATURE::dashboard", - "ACCOUNT::NIKE__in__EMAIL::*", - "EMAIL::*__discusses__FEATURE::dashboard", - "EMAIL::*Nike__had__CONCERN::* "]}} -Explanation: - - Nike did report unspecified concerns - - Nike had problems with the dashboard, which is a feature - - We are interested in emails that Nike exchanged with us -""".strip() - -RELATIONSHIP_EXAMPLE_6 = r""" -'List the last 5 emails that Lisa exchanged with Nike:' \ -and the extracted entities were found to be: - -"ACCOUNT::Nike", "EMAIL::*", "EMPLOYEE::Lisa" - -then a valid relationship extraction could be: - -{{"relationships": ["ACCOUNT::Nike__had__CONCERN::*", - "ACCOUNT::Nike__had_issues__FEATURE::dashboard", - "ACCOUNT::NIKE__in__EMAIL::*"]}} -Explanation: - - Nike did report unspecified concerns - - Nike had problems with the dashboard, which is a feature - - We are interested in emails that Nike exchanged with us -""".strip() - - -ENTITY_EXAMPLE_1 = r""" -{{"entities": ["ACCOUNT::Nike--[]", "CONCERN::*--[]"]}} -""".strip() - -ENTITY_EXAMPLE_2 = r""" -{{"entities": ["ACCOUNT::Nike--[]", "CONCERN::performance--[]"]}} -""".strip() - -ENTITY_EXAMPLE_3 = r""" -{{"entities": ["ACCOUNT::*--[]", "CONCERN::performance--[]", "CONCERN::user_experience--[]"]}} -""".strip() - -ENTITY_EXAMPLE_4 = r""" -{{"entities": ["ACCOUNT::*--[]", "CONCERN::performance--[degree: severe]"]}} -""".strip() - -MASTER_EXTRACTION_PROMPT = f""" -You are an expert in the area of knowledge extraction in order to construct a knowledge graph. You are given a text \ -and asked to extract entities, relationships, and terms from it that you can reliably identify. - -Here are the entity types that are available for extraction. Some of them may have a description, others \ -should be obvious. Also, for a given entity allowed options may be provided. If allowed options are provided, \ -you can only extract those types of entities! If no allowed options are provided, take your best guess. - -You can ONLY extract entities of these types and relationships between objects of these types: -{SEPARATOR_LINE} -{ENTITY_TYPE_SETTING_PROMPT} -{SEPARATOR_LINE} -Please format your answer in this format: -{SEPARATOR_LINE} -{EXTRACTION_FORMATTING_PROMPT} -{SEPARATOR_LINE} - -The list above here is the exclusive, only list of entities you can choose from! - -Here are some important additional instructions. (For the purpose of illustration, assume that ] - "ACCOUNT", "CONCERN", and "FEATURE" are all in the list of entity types above, and shown actual \ -entities fall into allowed options. Note that this \ -is just assumed for these examples, but you MUST use only the entities above for the actual extraction!) - -- You can either extract specific entities if a specific entity is referred to, or you can refer to the entity type. -* if the entity type is referred to in general, you would use '*' as the entity name in the extraction. -As an example, if the text would say: - 'Nike reported that they had issues' -then a valid extraction could be: -Example 1: -{EXAMPLE_1} - -* If on the other hand the text would say: -'Nike reported that they had performance issues' -then a much more suitable extraction could be: -Example 2: -{EXAMPLE_2} - -- You can extract multiple relationships between the same two entity types. -As an example, if the text would say: -'Nike reported some performance issues with our solution, but they are very happy that the user experience issue got solved.' -then a valid extraction could be: -Example 3: -{EXAMPLE_3} - -- You can extract multiple relationships between the same two actual entities if you think that \ -there are multiple relationships between them based on the text. -As an example, if the text would say: -'Nike reported some performance issues with our dashboard solution, but they think it delivers great value.' -then a valid extraction could be: -Example 4: -{EXAMPLE_4} - -Note that effectively a three-way relationship (Nike - performance issues - dashboard) extracted as two individual \ -relationships. - -- Again, - - you should only extract entities belonging to the entity types above - but do extract all that you \ -can reliably identify in the text - - use refer to 'all' entities in an entity type listed above by using '*' as the entity name - - only extract important relationships that signify something non-trivial, expressing things like \ -needs, wants, likes, dislikes, plans, interests, lack of interests, problems the account is having, etc. - - you MUST only use the initial list of entities provided! Ignore the entities in the examples unless \ -they are also part of the initial list of entities! This is essential! - - only extract relationships between the entities extracted first! - - -{SEPARATOR_LINE} - -Here is the text you are asked to extract knowledge from, if needed with additional information about any participants: -{SEPARATOR_LINE} ----content--- -{SEPARATOR_LINE} -""".strip() - - -QUERY_ENTITY_EXTRACTION_PROMPT = f""" -You are an expert in the area of knowledge extraction and using knowledge graphs. You are given a question \ -and asked to extract entities (with attributes if applicable) that you can reliably identify, which will then -be matched with a known entity in the knowledge graph. You are also asked to extract time constraints information \ -from the QUESTION. Some time constraints will be captured by entity attributes if \ -the entity type has a fitting attribute (example: 'created_at' could be a candidate for that), other times -we will extract an explicit time filter if no attribute fits. (Note regarding 'last', 'first', etc.: DO NOT \ -imply the need for a time filter just because the question asks for something that is not the current date. \ -They will relate to ordering that we will handle separately later). - -In case useful, today is ---today_date--- and the user asking is ---user_name---, which may or may not be relevant. -Here are the entity types that are available for extraction. Some of them may have \ -a description, others should be obvious. Also, notice that some may have attributes associated with them, which will \ -be important later. -You can ONLY extract entities of these types: -{SEPARATOR_LINE} -{ENTITY_TYPE_SETTING_PROMPT} -{SEPARATOR_LINE} - -The list above here is the exclusive, only list of entities you can choose from! - -Also, note that there are fixed relationship types between these entities. Please consider those \ -as well so to make sure that you are not missing implicit entities! Implicit entities are often \ -in verbs ('emailed to', 'talked to', ...). Also, they may be used to connect entities that are \ -clearly in the question. - -{SEPARATOR_LINE} -{RELATIONSHIP_TYPE_SETTING_PROMPT} -{SEPARATOR_LINE} - -Here are some important additional instructions. (For the purpose of illustration, assume that \ - "ACCOUNT", "CONCERN", "EMAIL", and "FEATURE" are all in the list of entity types above, and the \ -attribute options for "CONCERN" include 'degree' with possible values that include 'severe'. Note that this \ -is just assumed for these examples, but you MUST use only the entities above for the actual extraction!) - -- You can either extract specific entities if a specific entity is referred to, or you can refer to the entity type. -* if the entity type is referred to in general, you would use '*' as the entity name in the extraction. -As an example, if the question would say: - 'Which issues did Nike report?' -then a valid entity and term extraction could be: -Example 1: -{ENTITY_EXAMPLE_1} - -* If on the other hand the question would say: -'Did Nike say anything about performance issues?' -then a much more suitable entity and term extraction could be: -Example 2: -{ENTITY_EXAMPLE_2} - -* Then, if the question is: -'Who reported performance issues?' -then a suitable entity and term extraction could be: -Example 3: -{ENTITY_EXAMPLE_3} - -* Then, if we inquire about an entity with a specific attribute : -'Who reported severe performance issues?' -then a suitable entity and term extraction could be: -Example 3: -{ENTITY_EXAMPLE_4} - -- Again, - - you should only extract entities belonging to the entity types above - but do extract all that you \ -can reliably identify in the text - - if you refer to all/any/an unspecified entity of an entity type listed above, use '*' as the entity name - - similarly, if a specific entity type is referred to in general, you should use '*' as the entity name - - you MUST only use the initial list of entities provided! Ignore the entities in the examples unless \ -they are also part of the initial list of entities! This is essential! - - don't forget to provide answers also to the event filtering and whether documents need to be inspected! - - 'who' often refers to individuals or accounts. - - see whether any of the entities are supposed to be narrowed down by an attribute value. The precise attribute \ -and the value would need to be taken from the specification, as the question may use different words and the \ -actual attribute may be implied. - - don't just look at the entities that are mentioned in the question but also those that the question \ -may be about. - - be very careful that you only extract attributes that are listed above for the entity type in question! Do \ -not make up attributes even if they are implied! Particularly if there is a relationship type that would \ -actually represent that information, you MUST not extract the information as an attribute. We \ -will extract the relationship type later. - - For the values of attributes, look at the possible values above! For example 'open' may refer to \ -'backlog', 'todo', 'in progress', etc. In cases like that construct a ';'-separated list of values that you think may fit \ -what is implied in the question (in the exanple: 'open; backlog; todo; in progress'). - -Also, if you think the name or the title of an entity is given but name or title are not mentioned \ -explicitly as an attribute, then you should indeed extract the name/title as the entity name. - -{SEPARATOR_LINE} - -Here is the question you are asked to extract desired entities and time filters from: -{SEPARATOR_LINE} ----content--- -{SEPARATOR_LINE} - -Please format your answer in this format: -{SEPARATOR_LINE} -{QUERY_ENTITY_EXTRACTION_FORMATTING_PROMPT} -{SEPARATOR_LINE} - -""".strip() - - -QUERY_RELATIONSHIP_EXTRACTION_PROMPT = f""" -You are an expert in the area of knowledge extraction and using knowledge graphs. You are given a question \ -and previously you were asked to identify known entities in the question. Now you are asked to extract \ -the relationships between the entities you have identified earlier. - -First off as background, here are the entity types that are known to the system: -{SEPARATOR_LINE} ----entity_types--- -{SEPARATOR_LINE} - - -Here are the entities you have identified earlier: -{SEPARATOR_LINE} ----identified_entities--- -{SEPARATOR_LINE} - -Note that the notation for the entities is ::. - -Here are the options for the relationship types(!) between the entities you have identified earlier \ -as well as relationship types between the identified entities and other entities \ -not explicitly mentioned: -{SEPARATOR_LINE} ----relationship_type_options--- -{SEPARATOR_LINE} - -These types are, if any were identified, formatted as \ -____, and they \ -limit the allowed relationships that you can extract. You would then though use the actual full entities as in: - -::____::. - -Note: should be a word or two that captures the nature \ -of the relationship. Common relationships may be: 'likes', 'dislikes', 'uses', 'is interested in', 'mentions', \ -'addresses', 'participates in', etc., but look at the text to find the most appropriate relationship. \ -Use spaces here for word separation. - -Please format your answer in this format: -{SEPARATOR_LINE} -{QUERY_RELATIONSHIP_EXTRACTION_FORMATTING_PROMPT} -{SEPARATOR_LINE} - -The list above here is the exclusive, only list of entities and relationship types you can choose from! - -Here are some important additional instructions. (For the purpose of illustration, assume that ] - "ACCOUNT", "CONCERN", and "FEATURE" are all in the list of entity types above. Note that this \ -is just assumed for these examples, but you MUST use only the entities above for the actual extraction!) - -- You can either extract specific entities if a specific entity is referred to, or you can refer to the entity type. -* if the entity type is referred to in general, you would use '*' as the entity name in the extraction. - -As an example, if the question would say: - -{RELATIONSHIP_EXAMPLE_1} - -* If on the other hand the question would say: - -{RELATIONSHIP_EXAMPLE_2} - -- You can extract multiple relationships between the same two entity types. -For example 3, if the question would say: - -{RELATIONSHIP_EXAMPLE_3} - -- You can extract multiple relationships between the same two actual entities if you think that \ -there are multiple relationships between them based on the question. -As an example, if the question would say: - -{RELATIONSHIP_EXAMPLE_4} - -Note that effectively a three-way relationship (Nike - performance issues - dashboard) extracted as two individual \ -relationships. - -- Again, - - you can only extract relationships between the entities extracted earlier - - you can only extract the relationships that match the listed relationship types - - if in doubt and there are multiple relationships between the same two entities, you can extract \ -all of those that may fit with the question. - - be really thinking through the question which type of relationships should be extracted and which should not. - -Other important notes: - - For questions that really try to explore in general what a certain entity was involved in like 'what did Paul Smith do \ -in the last 3 months?', and Paul Smith has been extracted i.e. as an entity of type 'EMPLOYEE', then you need to extract \ -all of the possible relationships an empoyee Paul Smith could have. - - You are not forced to use all or any of the relationship types listed above. Really look at the question to \ - determine which relationships are explicitly or implicitly referred to in the question. - -{SEPARATOR_LINE} - -Here is the question you are asked to extract desired entities, relationships, and terms from: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} -""".strip() - - -GENERAL_CHUNK_PREPROCESSING_PROMPT = """ -This is a part of a document that you need to extract information (entities, relationships) from. - -Note: when you extract relationships, please make sure that: - - if you see a relationship for one of our employees, you should extract the relationship both for the employee AND \ - VENDOR::{vendor}. - - if you see a relationship for one of the representatives of other accounts, you should extract the relationship \ -only for the account ACCOUNT::! - --- -And here is the content: -{content} -""".strip() - - -### Source-specific prompts - -CALL_CHUNK_PREPROCESSING_PROMPT = """ -This is a call between employees of the VENDOR's company and representatives of one or more ACCOUNTs (usually one). \ -When you extract information based on the instructions, please make sure that you properly attribute the information \ -to the correct employee and account. \ - -Here are the participants (name component of email) from us ({vendor}): -{participant_string} - -Here are the participants (name component of email) from the other account(s): -{account_participant_string} - -In the text it should be easy to associate a name with the email, and then with the account ('us' vs 'them'). If in doubt, \ -look at the context and try to identify whether the statement comes from the other account. If you are not sure, ignore. - -Note: when you extract relationships, please make sure that: - - if you see a relationship for one of our employees, you should extract the relationship both for the employee AND \ - VENDOR::{vendor}. - - if you see a relationship for one of the representatives of other accounts, you should extract the relationship \ -only for the account ACCOUNT::! - --- -And here is the content: -{content} -""".strip() - - -CALL_DOCUMENT_CLASSIFICATION_PROMPT = """ -This is the beginning of a call between employees of the VENDOR's company ({vendor}) and other participants. - -Your task is to classify the call into one of the following categories: -{category_options} - -Please also consider the participants when you perform your classification task - they can be important indicators \ -for the category. - -Please format your answer as a string in the format: - -REASONING: - CATEGORY: - --- -And here is the beginning of the call, including title and participants: - -{beginning_of_call_content} -""".strip() - - -STRATEGY_GENERATION_PROMPT = f""" -Now you need to decide what type of strategy to use to answer a given question, how ultimately \ -the answer should be formatted to match the user's expectation, and what an appropriate question \ -to/about 'one object or one set of objects' may be, should the answer logically benefit from a divide \ -and conquer strategy, or it naturally relates to one or few individual objects. Also, you are \ -supposed to determine whether a divide and conquer strategy would be appropriate. - - -Here are the entity types that are available in the knowledge graph: -{SEPARATOR_LINE} ----possible_entities--- -{SEPARATOR_LINE} - -Here are the relationship types that are available in the knowledge graph: -{SEPARATOR_LINE} ----possible_relationships--- -{SEPARATOR_LINE} - -Here is the question whose answer is ultimately sought: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -And here are the entities and relationships that have been extracted earlier from this question: -{SEPARATOR_LINE} ----entities--- ----relationships--- -{SEPARATOR_LINE} - -Here are more instructions: - -a) Regarding the strategy, there are three aspects to it: - -a1) "Search Type": -Should the question be answered as a SEARCH ('filtered search'), or as a SQL ('SQL query search')? - -The options are: -1. SEARCH: A filtered search simply uses the entities and relationships that you extracted earlier and \ -applies them as filters to search the underlying documents, which are properly indexed. Examples are \ -'what did Nike say about the Analyzer product?', or 'what did I say in my calls with Nike about pricing?'. So this \ -is used really when there is *no implicit or explicit constraint or requirements* on underlying source documents \ -outside of filters, and there is no ordering, no limiting their number, etc. So use this for a question that \ -tries to get information *across* documents which may be filtered by their related relationships and entities, but without \ -other constraints. - -2. SQL: Choose this option if the question either requires counting of entities (e.g. 'how many calls...'), or \ -if the query refers to specific entities that first need to be identified and then analyzed/searched/listed. \ -Examples here are 'what did I say about pricing in my call with Nike last week?' (the specific call needs to \ -be identified first and then analyzed), \ -'what are the next steps of our two largest opportunities?', or 'summarize my 3 most recent customer calls'. So \ -this is used if there *are implicit constraints* on the underlying source documents beyond filtering, including \ -ordering, limiting, etc. Use this also if the answer expects to analyze each source independently as part \ -of the overall answer. - -Note: - - here, you should look at the extracted entities and relationships and judge whether using them as filters \ -(using an *and*) would be appropriate to identify the range of underlying sources, or whether more \ -calculations would be needed to find the underlying sources ('last 2...', etc.) . - - It is also *critical* to look at the attributes of the entities! You only can use the given attributes (and their - values, if given) as where conditions etc in a SQL statement. So if you think you would 'want - to' have a where condition but there is not appropriate attribute, then you should not use the SQL strategy - but the SEARCH strategy. (A Search can always look through data and see what is the best fit, SQL needs to - be more specific.). On the other hand, if the question maps well to the entities and attributes, then - SQL may be a good choice. - - Likely, if there are questions 'about something', then this only is used in a SQL statement or a filter \ - if it shows up as an entity or relationship in the extracted entities and relationships. Otherwise, it will \ - be part of the analysis/search. not the document identification. - - again, note that we can only FILTER (SEARCH) or COMPUTE (SQL) using the extracted entities (and their attributes) - and relationships. \ - So do not think that if there is another term in the question, it should be included in the SQL statement. \ - It cannot. - - -a2) "Search Strategy": -If a SQL search is chosen, i.e., documents have to be identified first, there are two approaches: -1. SIMPLE: You think you can answer the question using a database that is aware of the entities, relationships \ -above, and is generally suitable if it is enough to either list or count entities, return dates, etc. Usually, \ -'SIMPLE' is chosen for questions of the form 'how many...' (always), or 'list the...' (often), 'when was...', \ -'what did (someone) work on...'etc. Often it is also used in cases like 'what did John work on since April?'. Here, \ -the user would expect to just see the list. So chose 'SIMPLE' here unless there are REALLY CLEAR \ -follow-up instructions for each item (like 'summarize...' , 'analyze...', 'what are the main points of...'.) If \ -it is a 'what did...'-type question, choose 'SIMPLE'! - -2. DEEP: You think you really should ALSO leverage the actual text of sources to answer the question, which sits \ -in a vector database. Examples are 'what is discussed in...', 'summarize', 'what is the discussion about...',\ -'how does... relate to...', 'are there any mentions of... in..', 'what are the main points in...', \ -'what are the next steps...', etc. Those are usually questions 'about' \ -the entities retrieved from the knowledge graph, or questions about the underlying sources. - -Your task is to decide which of the two strategies to use. - -a3) "Relationship Detection": -You need to evaluate whether the question involves any relationships between entities (of the same type) \ -or between entities and relationships. Respond with 'RELATIONSHIPS' or 'NO_RELATIONSHIPS'. - -b) Regarding the format of the answer: there are also two types of formats available to you: - -1. LIST: The user would expect an answer as a bullet point list of objects, likely with text associated with each \ -bullet point (or sub-bullet). This will be clearer once the data is available. -2. TEXT: The user would expect the questions to be answered in text form. - -Your task is to decide which of the two formats to use. - - -c) Regarding the broken down question for one object: - -Always generate a broken_down_question if the question pertains ultimately to a specific objects, even if it seems to be \ -a singular object. - -- If the question is of type 'how many...', or similar, then imagine that the individual objects have been \ -found and you want to ask each object something that illustrates why/in what what that object relates to the \ -question. (question: 'How many cars are fast?' -> broken_down_question: 'How fast is this car?') - -- Assume the answer would either i) best be generated by first analyzing one object at a time, then aggregating \ -the results, or ii) directly relates to one or few objects found through matching suitable criteria. - -- The key is to drop any filtering/criteria matching as the objects are already filtered by the criteria. Also, do not \ -try to verify here whether the object in question actually satisfies a filter criteria, but rather see \ -what it says/does etc. In other words, use this to identify more details about the object, as it relates \ -to the original question. -(Example: question: 'What did our oil & gas customers say about the new product?' -> broken_down_question: \ -'What did this customer say about the new product?', -or: -question: 'What was in the email from Frank?' -> broken_down_question: 'What is in this email?') - - -d) Regarding the divide and conquer strategy: - -You are supposed to decide whether a divide and conquer strategy would be appropriate. That means, do you think \ -that in order to answer the question, it would be good to first analyze one object at a time, and then aggregate the \ -results? Or should the information rather be analyzed as a whole? This would be 'yes' or 'no'. - -Please answer in json format in this form: - -{{ - "search_type": , - "search_strategy": , - "relationship_detection": , - "format": , - "broken_down_question": , - "divide_and_conquer": -}} - -Do not include any other text or explanations. -""" - -SOURCE_DETECTION_PROMPT = f""" -You are an expert in generating, understanding and analyzing SQL statements. - -You are given an original SQL statement that returns a list of entities from a table or \ -an aggregation of entities from a table. Your task will be to \ -identify the source documents that are relevant to what the SQL statement is returning. - -The task is actually quite simple. There are two tables involved - relationship_table and entity_table. \ -relationship_table was used to generate the original SQL statement. Again, returning entities \ -or aggregations of entities. The second table, entity_table contains the entities and \ -the corresponding source_documents. All you need to do is to appropriately join the \ -entity_table table on the entities that would be retrieved from the original SQL statement, \ -and then return the source_documents from the entity_table table. - -For your orientation, the relationship_table table has this structure: - - Table name: relationship_table - - Columns: - - relationship (str): The name of the RELATIONSHIP, combining the nature of the relationship and the names of the entities. \ -It is of the form \ -::____:: \ -[example: ACCOUNT::Nike__has__CONCERN::performance]. Note that this is NOT UNIQUE! - - source_entity (str): the id of the source ENTITY/NODE in the relationship [example: ACCOUNT::Nike] - - source_entity_attributes (json): the attributes of the source entity/node [example: {{"account_type": "customer"}}] - - target_entity (str): the id of the target ENTITY/NODE in the relationship [example: CONCERN::performance] - - target_entity_attributes (json): the attributes of the target entity/node [example: {{"degree": "severe"}}] - - source_entity_type (str): the type of the source entity/node [example: ACCOUNT]. Only the entity types provided \ - below are valid. - - target_entity_type (str): the type of the target entity/node [example: CONCERN]. Only the entity types provided \ - below are valid. - - relationship_type (str): the type of the relationship, formatted as \ -____. So the explicit entity_names have \ -been removed. [example: ACCOUNT__has__CONCERN] - - source_date (str): the 'event' date of the source document [example: 2021-01-01] - -The second table, entity_table, has this structure: - - Table name: entity_table - - Columns: - - entity (str): The name of the ENTITY, which is unique in this table. source_entity and target_entity \ -in the relationship_table table are the same as entity in this table. - - source_document (str): the id of the document that contains the entity. - -Again, ultimately, your task is to join the entity_table table on the entities that would be retrieved from the \ -original SQL statement, and then return the source_documents from the entity_table table. - -The way to do that is to create a common table expression for the original SQL statement and join the \ -entity_table table suitably on the entities. - -Here is the *original* SQL statement: -{SEPARATOR_LINE} ----original_sql_statement--- -{SEPARATOR_LINE} - -Please structure your answer using , ,, start and end tags as in: - -[think very briefly through the problem step by step, not more than 2-3 sentences] \ -[the new SQL statement that returns the source documents involved in the original SQL statement] -""".strip() - -ENTITY_SOURCE_DETECTION_PROMPT = f""" -You are an expert in generating, understanding and analyzing SQL statements. - -You are given a SQL statement that returned an aggregation of entities in a table. \ -Your task will be to identify the source documents for the entities involved in \ -the answer. For example, should the original SQL statement be \ -'SELECT COUNT(entity) FROM entity_table where entity_type = "ACCOUNT"' \ -then you should return the source documents that contain the entities of type 'ACCOUNT'. - -The table has this structure: - - Table name: entity_table - - Columns: - - entity (str): The name of the ENTITY, combining the nature of the entity and the id of the entity. \ -It is of the form :: [example: ACCOUNT::625482894]. - - entity_type (str): the type of the entity [example: ACCOUNT]. - - entity_attributes (json): the attributes of the entity [example: {{"priority": "high", "status": "active"}}] - - source_document (str): the id of the document that contains the entity. Note that the combination of \ -id_name and source_document IS UNIQUE! - - source_date (timestamp): the 'event' date of the source document [example: 2025-04-25 21:43:31.054741+00] - -Specifically, the table contains the 'source_document' column, which is the id of the source document that \ -contains the core information about the entity. Make sure that you do not return more documents, i.e. if there \ -is a limit on source documents in the original SQL statement, the new SQL statement needs to have \ -the same limit. - -CRITICAL NOTES: - - Only return source documents and nothing else! - -Your task is then to create a new SQL statement that returns the source documents that are relevant to what the \ -original SQL statement is returning. So the source document of every row used in the original SQL statement should \ -be included in the result of the new SQL statement, and then you should apply a 'distinct'. - -Here is the *original* SQL statement: -{SEPARATOR_LINE} ----original_sql_statement--- -{SEPARATOR_LINE} - -Please structure your answer using , ,, start and end tags as in: - -[think very briefly through the problem step by step, not more than 2-3 sentences] \ -[the new SQL statement that returns the source documents involved in the original SQL statement] -""".strip() # noqa: S608 - LLM prompt template, not executed SQL - - -ENTITY_TABLE_DESCRIPTION = f"""\ - - Table name: entity_table - - Columns: - - entity (str): The name of the ENTITY, combining the nature of the entity and the id of the entity. \ -It is of the form :: [example: ACCOUNT::625482894]. - - entity_type (str): the type of the entity [example: ACCOUNT]. - - entity_attributes (json): the attributes of the entity [example: {{"priority": "high", "status": "active"}}] - - source_document (str): the id of the document that contains the entity. Note that the combination of \ -id_name and source_document IS UNIQUE! - - source_date (timestamp): the 'event' date of the source document [example: 2025-04-25 21:43:31.054741+00] - -{SEPARATOR_LINE} - -Importantly, here are the entity (node) types that you can use, with a short description of what they mean. You may need to \ -identify the proper entity type through its description. Also notice the allowed attributes for each entity type and \ -their values, if provided. Of particular importance is the 'subtype' attribute, if provided, as this is how \ -the entity type may also often be referred to. -{SEPARATOR_LINE} ----entity_types--- -{SEPARATOR_LINE} -""" - -RELATIONSHIP_TABLE_DESCRIPTION = f"""\ - - Table name: relationship_table - - Columns: - - relationship (str): The name of the RELATIONSHIP, combining the nature of the relationship and the names of the entities. \ -It is of the form \ -::____:: \ -[example: ACCOUNT::Nike__has__CONCERN::performance]. Note that this is NOT UNIQUE! - - source_entity (str): the id of the source ENTITY/NODE in the relationship [example: ACCOUNT::Nike] - - source_entity_attributes (json): the attributes of the source entity/node [example: {{"account_type": "customer"}}] - - target_entity (str): the id of the target ENTITY/NODE in the relationship [example: CONCERN::performance] - - target_entity_attributes (json): the attributes of the target entity/node [example: {{"degree": "severe"}}] - - source_entity_type (str): the type of the source entity/node [example: ACCOUNT]. Only the entity types provided \ - below are valid. - - target_entity_type (str): the type of the target entity/node [example: CONCERN]. Only the entity types provided \ - below are valid. - - relationship_type (str): the type of the relationship, formatted as \ -____. So the explicit entity_names have \ -been removed. [example: ACCOUNT__has__CONCERN] - - source_document (str): the id of the document that contains the relationship. Note that the combination of \ -id_name and source_document IS UNIQUE! - - source_date (timestamp): the 'event' date of the source document [example: 2025-04-25 21:43:31.054741+00] - -{SEPARATOR_LINE} - -Importantly, here are the entity (node) types that you can use, with a short description of what they mean. You may need to \ -identify the proper entity type through its description. Also notice the allowed attributes for each entity type and \ -their values, if provided. Of particular importance is the 'subtype' attribute, if provided, as this is how \ -the entity type may also often be referred to. -{SEPARATOR_LINE} ----entity_types--- -{SEPARATOR_LINE} - -Here are the relationship types that are in the table, denoted as ____. -In the table, the actual relationships are not quite of this form, but each is followed by '::' \ -in the relationship id as shown above. -{SEPARATOR_LINE} ----relationship_types--- -{SEPARATOR_LINE} -""" - - -SIMPLE_SQL_PROMPT = f""" -You are an expert in generating a SQL statement that only uses ONE TABLE that captures RELATIONSHIPS \ -between TWO ENTITIES. The table has the following structure: - -{SEPARATOR_LINE} -{RELATIONSHIP_TABLE_DESCRIPTION} - -Here is the question you are supposed to translate into a SQL statement: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -To help you, we already have identified the entities and relationships that the SQL statement likely *should* use (but note the \ -exception below!). The entities also contain the list of attributes and attribute values that should specify the entity. \ -The format is ::--[:, \ -:, ...]. -{SEPARATOR_LINE} -Identified entities with attributes in query: - ----query_entities_with_attributes--- - -These are the entities that should be used in the SQL statement. However, \ -note that these are the entities (with potential attributes) that were *matches* of Knowledge Graph identified with the \ -entities originally identified in the original question. A such, they may have id names that may not mean much by themselves, \ -eg ACCOUNT::a74f332. Here is the mapping of entities originally identified (whose role in the query should be obvious) with \ -the entities that were matched to them in the Knowledge Graph: - ----entity_explanation_string--- - --- - -Here are relationships that were identified as explicitly or implicitly referred to in the question: - ----query_relationships--- - -(Again, if applicable, the entities contained in the relationships are the same as the entities in the \ -query_entities_with_attributes, and those are the correct ones to use in the SQL statement.) - -{SEPARATOR_LINE} - -CRITICAL SPECIAL CASE: - - if an identified entity is of the form ::*, or an identified relationship contains an \ -entity of this form, this refers to *any* entity of that type. Correspondingly, the SQL query should use the *entity type*, \ -and possibly the relationship type, but not the entity with the * itself. \ -Example: if you see 'ACCOUNT::*', that means any account matches. So if you are supposed to count the 'ACCOUNT::*', \ -you should count the entities of entity_type 'ACCOUNT'. - - -IMPORTANT NOTES: -- The id_name of each relationship has the format \ -____. -- The relationship id_names are NOT UNIQUE, only the combinations of relationship id_name and source_document_id are unique. \ -That is because each relationship is extracted from a document. So make sure you use the proper 'distinct's! -- If the SQL contains a 'SELECT DISTINCT' clause and an ORDER BY clause, then you MUST include the columns from the ORDER BY \ -clause ALSO IN THE SELECT DISTINCT CLAUSE! This is very important! (This is a postgres db., so this is a MUST!). \ -You MUST NOT have a column in the ORDER BY clause that is not ALSO in the SELECT DISTINCT clause! -- If you join the relationship table on itself using the source_node or target_node, you need to make sure that you also \ -join on the source_document_id. -- The id_name of each node/entity has the format ::, where 'entity_type_id_name' \ -and 'name' are columns and \ - the values and can be used for filtering. -- The table can be joined on itself on source nodes and/or target nodes if needed. -- the SQL statement MUST ultimately only return NODES/ENTITIES (not relationships!), or aggregations of \ -entities/nodes(count, avg, max, min, etc.). \ -Again, DO NOT compose a SQL statement that returns id_name of relationships. -- You CAN ONLY return ENTITIES or COUNTS (or other aggregations) of ENTITIES, or you can return \ -source_date (but only if the question asks for event dates or times). DO NOT return \ -source documents or counts of source documents, or relationships or counts of relationships! \ -Those can only appear in where clauses, ordering etc., but they cannot be returned or ultimately \ -counted here! source_date and date operations can appear in select statements, particularly if \ -there is time ordering or grouping involved. -- ENTITIES can be target_entity or source_entity. Think about the allowed relationships and the \ -question to decide which one you want! -- It is ok to generate nested SQL as long as it is correct postgres syntax! -- Attributes are stored in the attributes json field. As this is postgres, querying for those must be done as \ -"attributes ->> '' = ''". -- The SELECT clause MUST only contain entities or aggregations/counts of entities, or, in cases the \ -question was about dates or times, then it can also include source_date. But source_document MUST NEVER appear \ -in the SELECT clause! -- Again, NEVER count or retrieve source documents in SELECT CLAUSE, whether it is in combination with \ -entities, with a distinct, etc. NO source_document in SELECT CLAUSE! So NEVER produce a \ -'SELECT COUNT(source_entity, source_document)...' -- Please think about whether you are interested in source entities or target entities! For that purpose, \ -consider the allowed relationship types to make sure you select or count the correct one! -- Again, ALWAYS make sure that EACH COLUMN in an ORDER-BY clause IS ALSO IN THE SELECT CLAUSE! Remind yourself \ -of that in the reasoning. -- Be careful with dates! Often a date will refer to the source data, which is the date when \ -an underlying piece of information was updated. However, if the attributes of an entity contain \ -time information as well (like 'started_at', 'completed_at', etc.), then you should really look at \ -the wording to see whether you should use a date in the attributes or the event date. -- Dates are ALWAYS in string format of the form YYYY-MM-DD, for source date as well as for date-like the attributes! \ -So please use that format, particularly if you use data comparisons (>, <, ...) -- Again, NO 'relationship' or 'source_document' in the SELECT CLAUSE, be it as direct columns are in aggregations! -- Careful with SORT! Really think in which order you want to sort if you have multiple columns you \ -want to sort by. If the sorting is time-based and there is a limit for example, then you do want to have a suitable date \ -variable as the first column to sort by. -- When doing a SORT on an attribute value of an entity, you MUST also apply a WHERE clause to filter \ -for entities that have the attribute value set. For example, if you want to sort the target entity \ -by the attribute 'created_date', you must also have a WHERE clause that checks whether the target \ -entity attribute contains 'created_date'. This is vital for proper ordering with null values. -- Usually, you will want to retrieve or count entities, maybe with attributes. But you almost always want to \ -have entities involved in the SELECT clause. -- Questions like 'What did Paul work on last week?' should generally be handled by finding all entities \ -that reasonably relate to 'work entities' that are i) related to Paul, and ii) that were created or \ -updated (by him) last week. So this would likely be a UNION of multiple queries. -- If you do joins consider the possibility that the second entity does not exist for all examples. \ -Therefore joins should generally be LEFT joins (or RIGHT joins) as appropriate. Think about which \ -entities you are interested in, and which ones provides attributes. -Another important note: - - For questions that really try to explore what a certain entity was involved in like 'what did Paul Smith do \ -in the last 3 months?', and Paul Smith has been extracted ie as an entity of type 'EMPLOYEE', you will \ -want to consider all entities that Paul Smith may be related to that satisfy any potential other conditions. -- Joins should always be made on entities, not source documents! -- Try to be as efficient as possible. - -APPROACH: -Please think through this step by step. Make sure that you include all columns in the ORDER BY clause \ -also in the SELECT DISTINCT clause, \ -if applicable! And again, joins should generally be LEFT JOINS! - -Also, in case it is important, today is ---today_date--- and the user/employee asking is ---user_name---. - -Please structure your answer using , , , start and end tags as in: - -[think through the logic but do so extremely briefly! Not more than 3-4 sentences.] -[the SQL statement that you generate to satisfy the task] -""".strip() # noqa: S608 - LLM prompt template, not executed SQL - -# TODO: remove following before merging after enough testing -SIMPLE_SQL_CORRECTION_PROMPT = f""" -You are an expert in reviewing and fixing SQL statements. - -Here is a draft SQL statement that you should consider as generally capturing the information intended. \ -However, it may or may not be syntactically 100% for our postgresql database. - -Guidance: - - Think about whether attributes should be numbers or strings. You may need to convert them. - - If we use SELECT DISTINCT we need to have the ORDER BY columns in the \ -SELECT statement as well! And it needs to be in the EXACT FORM! So if a \ -conversion took place, make sure to include the conversion in the SELECT and the ORDER BY clause! - - never should 'source_document' be in the SELECT clause! Remove if present! - - if there are joins, they must be on entities, never source documents - - if there are joins, consider the possibility that the second entity does not exist for all examples.\ - Therefore consider using LEFT joins (or RIGHT joins) as appropriate. - -Draft SQL: -{SEPARATOR_LINE} ----draft_sql--- -{SEPARATOR_LINE} - -Please structure your answer using , , , start and end tags as in: - -[think briefly through the problem step by step] -[the corrected (or original one, if correct) SQL statement] -""".strip() - -SIMPLE_ENTITY_SQL_PROMPT = f""" -You are an expert in generating a SQL statement that only uses ONE TABLE that captures ENTITIES \ -and their attributes and other data. The table has the following structure: - -{SEPARATOR_LINE} -{ENTITY_TABLE_DESCRIPTION} - -Here is the question you are supposed to translate into a SQL statement: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -To help you, we already have identified the entities that the SQL statement likely *should* use (but note the \ -exception below!). The entities as written below also contain the list of attributes and attribute values \ -that should specify the entity. \ -The format is ::--[:, \ -:, ...]. -{SEPARATOR_LINE} -Identified entities with attributes in query: - ----query_entities_with_attributes--- - -These are the entities that should be used in the SQL statement. However, \ -note that these are the entities (with potential attributes) that were *matches* of Knowledge Graph identified with the \ -entities originally identified in the original question. As such, they may have id names that may not mean much by themselves, \ -eg ACCOUNT::a74f332. Here is the mapping of entities originally identified (whose role in the query should be obvious) with \ -the entities that were matched to them in the Knowledge Graph: - ----entity_explanation_string--- - --- - - -{SEPARATOR_LINE} - -CRITICAL SPECIAL CASE: - - if an identified entity is of the form ::*, or an identified relationship contains an \ -entity of this form, this refers to *any* entity of that type. Correspondingly, the SQL query should use the *entity type*, \ -but not the entity with the * itself. \ -Example: if you see 'ACCOUNT::*', that means any account matches. So if you are supposed to count the 'ACCOUNT::*', \ -you should count the entities of entity_type 'ACCOUNT'. - - -IMPORTANT NOTES: -- The entities are unique in the table. -- If the SQL contains a 'SELECT DISTINCT' clause and an ORDER BY clause, then you MUST include the columns from the ORDER BY \ -clause ALSO IN THE SELECT DISTINCT CLAUSE! This is very important! (This is a postgres db., so this is a MUST!). \ -You MUST NOT have a column in the ORDER BY clause that is not ALSO in the SELECT DISTINCT clause! -- The table cannot be joined on itself. -- You CAN ONLY return ENTITIES or COUNTS (or other aggregations) of ENTITIES, or you can return \ -source_date (but only if the question asks for event dates or times, and then the \ -corresponding entity must also be returned). -- Generally, the query can only return ENTITIES or aggregations of ENTITIES: - - if individual entities are returned, then you MUST also return the source_document. \ -If the source date was requested, you can return that too. - - if aggregations of entities are returned, then you can only aggregate the entities. -- Attributes are stored in the attributes json field. As this is postgres, querying for those must be done as \ -"attributes ->> '' = ''". -- Again, ALWAYS make sure that EACH COLUMN in an ORDER-BY clause IS ALSO IN THE SELECT CLAUSE! Remind yourself \ -of that in the reasoning. -- Be careful with dates! Often a date will refer to the source data, which is the date when \ -an underlying piece of information was updated. However, if the attributes of an entity may contain \ -time information as well (like 'started_at', 'completed_at', etc.), then you should really look at \ -the wording to see whether you should use a date in the attributes or the event date. -- Dates are ALWAYS in string format of the form YYYY-MM-DD, for source date as well as for date-like the attributes! \ -So please use that format, particularly if you use data comparisons (>, <, ...) -- Careful with SORT! Really think in which order you want to sort if you have multiple columns you \ -want to sort by. If the sorting is time-based and there is a limit for example, then you do want to have a suitable date \ -variable as the first column to sort by. -- When doing a SORT on an attribute value of an entity, you MUST also apply a WHERE clause to filter \ -for entities that have the attribute value set. For example, if you want to sort the target entity \ -by the attribute 'created_date', you must also have a WHERE clause that checks whether the target \ -entity attribute contains 'created_date'. This is vital for proper ordering with null values. -- Usually, you will want to retrieve or count entities, maybe with attributes. But you almost always want to \ -have entities involved in the SELECT clause. -- You MUST ONLY rely on the entity attributes provided! This is essential! Do not assume \ -other attributes exist...they don't! Note that there will often be a search using the results \ -of this query. So if there is information in the question that does not fit the provided attributes, \ -you should not use it here but rely on the later search! -- Try to be as efficient as possible. - -APPROACH: -Please think through this step by step. Make sure that you include all columns in the ORDER BY clause \ -also in the SELECT DISTINCT clause, \ -if applicable! - -Also, in case it is important, today is ---today_date--- and the user/employee asking is ---user_name---. - -Please structure your answer using , , , start and end tags as in: - -[think through the logic but do so extremely briefly! Not more than 3-4 sentences.] -[the SQL statement that you generate to satisfy the task] -""".strip() # noqa: S608 - LLM prompt template, not executed SQL - -SIMPLE_SQL_ERROR_FIX_PROMPT = f""" -You are an expert at fixing SQL statements. You will be provided with a SQL statement that aims to address \ -a question, but it contains an error. Your task is to fix the SQL statement, based on the error message. - -Here is the description of the table that the SQL statement is supposed to use: ----table_description--- - -Here is the question you are supposed to translate into a SQL statement: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -Here is the SQL statement that you should fix: -{SEPARATOR_LINE} ----sql_statement--- -{SEPARATOR_LINE} - -Here is the error message that was returned: -{SEPARATOR_LINE} ----error_message--- -{SEPARATOR_LINE} - -Note that in the case the error states the sql statement did not return any results, it is possible that the \ -sql statement is correct, but the question is not addressable with the information in the knowledge graph. \ -If you are absolutely certain that is the case, you may return the original sql statement. - -Here are a couple common errors that you may encounter: -- source_document is in the SELECT clause -> remove it -- columns used in ORDER BY must also appear in the SELECT DISTINCT clause -- consider carefully the type of the columns you are using, especially for attributes. You may need to cast them -- dates are ALWAYS in string format of the form YYYY-MM-DD, for source date as well as for date-like the attributes! \ -So please use that format, particularly if you use data comparisons (>, <, ...) -- attributes are stored in the attributes json field. As this is postgres, querying for those must be done as \ -"attributes ->> '' = ''" (or "attributes ? ''" to check for existence). -- if you are using joins and the sql returned no joins, make sure you are using the appropriate join type (LEFT, RIGHT, etc.) \ -it is possible that the second entity does not exist for all examples. -- (ignore if using entity_table) if using the relationship_table and the sql returned no results, make sure you are \ -selecting the correct column! Use the available relationship types to determine whether to use the source or target entity. - -APPROACH: -Please think through this step by step. Please also bear in mind that the sql statement is written in postgres syntax. - -Also, in case it is important, today is ---today_date--- and the user/employee asking is ---user_name---. - -Please structure your answer using , , , start and end tags as in: - -[think through the logic but do so extremely briefly! Not more than 3-4 sentences.] -[the SQL statement that you generate to satisfy the task] -""" - - -SEARCH_FILTER_CONSTRUCTION_PROMPT = f""" -You need to prepare a search across text segments that contain the information necessary to \ -answer a question. The text segments have tags that can be used to filter for the relevant segments. \ -Key are suitable entities and relationships of a knowledge graph, as well as underlying source documents. - -Your overall task is to find the filters and structures that are needed to filtering a database to \ -properly address a user question. - -You will be given: - - the user question - - a description of all of the potential entity types involved - - a list of 'global' entities and relationships that should be filtered by, given the question - - the structure of a schema that was used to derive additional entity filters - - a SQL statement that was generated to derive those filters - - the results that were generated using the SQL statement. This can have multiple rows, \ -and those will be the 'local' filters (which will later mean that each retrieved result will \ -need to match at least one of the conditions that you will generate). - - the results of another query that asked for the underlying source documents that resulted \ -in the answers of the SQL statement - - -Here is the information: - -1) The overall user question -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -2) Here is a description of all of the entity types: -{SEPARATOR_LINE} ----entity_type_descriptions--- -{SEPARATOR_LINE} - -3) Here are the lists of entity and relationship filters that were derived from the question: -{SEPARATOR_LINE} -Entity filters: - ----entity_filters--- - --- - -Relationship filters: - ----relationship_filters--- - -{SEPARATOR_LINE} - -4) Here are the columns of a table in a database that has a lot of knowledge about the \ -data: -{SEPARATOR_LINE} - - relationship (str): The name of the RELATIONSHIP, combining the nature of the relationship and the names of the entities. \ -It is of the form \ -::____:: \ -[example: ACCOUNT::Nike__has__CONCERN::performance]. Note that this is NOT UNIQUE! - - source_entity (str): the id of the source ENTITY/NODE in the relationship [example: ACCOUNT::Nike] - - source_entity_attributes (json): the attributes of the source entity/node [example: {{"account_type": "customer"}}] - - target_entity (str): the id of the target ENTITY/NODE in the relationship [example: CONCERN::performance] - - target_entity_attributes (json): the attributes of the target entity/node [example: {{"degree": "severe"}}] - - source_entity_type (str): the type of the source entity/node [example: ACCOUNT]. Only the entity types provided \ - below are valid. - - target_entity_type (str): the type of the target entity/node [example: CONCERN]. Only the entity types provided \ - below are valid. - - relationship_type (str): the type of the relationship, formatted as \ -____. So the explicit entity_names have \ -been removed. [example: ACCOUNT__has__CONCERN] - - source_document (str): the id of the document that contains the relationship. Note that the combination of \ -id_name and source_document IS UNIQUE! - - source_date (str): the 'event' date of the source document [example: 2021-01-01] - -{SEPARATOR_LINE} - -5) Here is a query that was generated for that table to provide additional filters: -{SEPARATOR_LINE} ----sql_query--- -{SEPARATOR_LINE} - -6) Here are the results of that SQL query. (Consider the schema description and the \ -structure of the entities to interpret the results) -{SEPARATOR_LINE} ----sql_results--- -{SEPARATOR_LINE} - -7) Here are the results of the other query that provided the underlying source documents \ -using the schema: -{SEPARATOR_LINE} ----source_document_results--- -{SEPARATOR_LINE} - -Here is the detailed set of tasks that you should perform, including the proper output format for you: - -Please reply as a json dictionary in this form: - -{{ - "global_entity_filters": , - "global_relationship_filters": , - "local_entity_filters": , - "source_document_filters": , - "structure": -}} - -Again - DO NOT FORGET - here is the user question that motivates this whole task: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -Your json dictionary answer: -""".strip() - -OUTPUT_FORMAT_NO_EXAMPLES_PROMPT = f""" -You need to format an answer to a research question. \ -You will see what the desired output is, the original question, and the unformatted answer to the research question. \ -Your purpose is to generate the answer respecting the desired format. - -Notes: - - Note that you are a language model and that answers may or may not be perfect. To communicate \ -this to the user, consider phrases like 'I found [10 accounts]...', or 'Here are a number of [goals] that \ -I found...] -- Please DO NOT mention the explicit output format in your answer. Just use it to inform the formatting. - -Here is the unformatted answer to the research question: -{SEPARATOR_LINE} ----introductory_answer--- -{SEPARATOR_LINE} - -Here is the original question: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -And finally, here is the desired output format: -{SEPARATOR_LINE} ----output_format--- -{SEPARATOR_LINE} - -Please start generating the answer, without any explanation. There should be no real modifications to \ -the text, after all, all you need to do here is formatting. \ - -Your Answer: -""".strip() - - -OUTPUT_FORMAT_PROMPT = f""" -You need to format the answers to a research question that was generated using one or more objects. \ -An overall introductory answer may be provided to you, as well as the research results for each individual object. \ -You will also be provided with the original question as background, and the desired format. \ - -Your purpose is to generate a consolidated and FORMATTED answer that starts of with the introductory \ -answer, and then formats the research results for each individual object in the desired format. \ -Do not add any other text please! - -Notes: - - Note that you are a language model and that answers may or may not be perfect. To communicate \ -this to the user, consider phrases like 'I found [10 accounts]...', or 'Here are a number of [goals] that \ -I found...] -- Please DO NOT mention the explicit output format in your answer. Just use it to inform the formatting. -- DO NOT add any content to the introductory answer! - - -Here is the original question for your background: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -Here is the desired output format: -{SEPARATOR_LINE} ----output_format--- -{SEPARATOR_LINE} - -Here is the introductory answer: -{SEPARATOR_LINE} ----introductory_answer--- -{SEPARATOR_LINE} - -Here are the research results that you should - respecting the target format- return in a formatted way: -{SEPARATOR_LINE} ----research_results--- -{SEPARATOR_LINE} - -Please start generating the answer, without any explanation. After all, all you need to do here is formatting. \ - - -Your Answer: -""".strip() - -OUTPUT_FORMAT_NO_OVERALL_ANSWER_PROMPT = f""" -You need to format the return of research on multiple objects. The research results will be given \ -to you as a string. You will also see what the desired output is, as well as the original question. \ -Your purpose is to generate the answer respecting the desired format. - -Notes: - - Note that you are a language model and that answers may or may not be perfect. To communicate \ -this to the user, consider phrases like 'I found [10 accounts]...', or 'Here are a number of [goals] that \ -I found...] -- Please DO NOT mention the explicit output format in your answer. Just use it to inform the formatting. - - Often, you are also provided with a list of explicit examples. If - AND ONLY IF - the list is not \ -empty, then these should be listed at the end with the text: -'... -Here are some examples of what I found: - -...' - - Again if the list of examples is an empty string then skip this section! Do not use the \ -results data for this purpose instead! (They will already be handled in the answer.) -- Even if the desired output format is 'text', make sure that you keep the individual research results \ -separated by bullet points, and mention the object name first, followed by a new line. The object name \ -is at the beginning of the research result, and should be in the format ::. - - -Here is the original question: -{SEPARATOR_LINE} ----question--- -{SEPARATOR_LINE} - -And finally, here is the desired output format: -{SEPARATOR_LINE} ----output_format--- -{SEPARATOR_LINE} - -Here are the research results that you should properly format: -{SEPARATOR_LINE} ----research_results--- -{SEPARATOR_LINE} - -Please start generating the answer, without any explanation. After all, all you need to do here is formatting. \ - - -Your Answer: -""".strip() - -KG_OBJECT_SOURCE_RESEARCH_PROMPT = f""" -You are an expert in extracting relevant structured information from a list of documents that \ -should relate to one object. You are presented with a list of documents that have been determined to be \ -relevant to the task of interest. Your goal is to extract the information asked around these topics: -You should look at the documents - in no particular order! - and extract the information that relates \ -to a question: -{SEPARATOR_LINE} -{{question}} -{SEPARATOR_LINE} - -Here are the documents you are supposed to search through: --- -{{document_text}} -{SEPARATOR_LINE} -Note: in this case, please do NOT cite your sources. This is very important! - -Please now generate the answer to the question given the documents: -""".strip() - -KG_SEARCH_PROMPT = f""" -You are an expert in extracting relevant structured information from a list of documents that \ -should relate to one object. You are presented with a list of documents that have been determined to be \ -relevant to the task of interest. Your goal is to extract the information asked around these topics: -You should look at the documents and extract the information that relates \ -to a question: -{SEPARATOR_LINE} -{{question}} -{SEPARATOR_LINE} - -Here are the documents you are supposed to search through: --- -{{document_text}} -{SEPARATOR_LINE} -Note: in this case, please DO cite your sources. This is very important! \ -Use the format []. Ie, use [1], [2], and NOT [1,2] if \ -there are two documents to cite, etc. \ - - -Please now generate the answer to the question given the documents: -""".strip() - # KG Beta Assistant System Prompt KG_BETA_ASSISTANT_SYSTEM_PROMPT = """"You are a knowledge graph assistant that helps users explore and \ understand relationships between entities.""" KG_BETA_ASSISTANT_TASK_PROMPT = """"Help users explore and understand the knowledge graph by answering \ questions about entities and their relationships.""" - - -# Just in case, for best practice, send a system message with key rules. -# (The db user permissions executing the SQL will avoid issues anyway, -# but it does not hurt to to put multiple checks in place.) -SQL_INSTRUCTIONS_RELATIONSHIP_PROMPT = """ -You are an expert at generating SQL queries to answer questions about a knowledge graph. - -You will be given a lot of instructions later, but here rules that MUST BE FOLLOWED: - - the SQL generated MUST only use the table one table named 'relationship_table'. \ -This table is not a table that can be defined or overwritten by the user and the resulting SQL \ -statement, it MUST be seen as an existing table in the database. - - self-joins of the 'relationship_table' are allowed, as well as common table expressions \ - that reference only the 'relationship_table'. - - no other table or view can in any way or shape be \ -involved in the generated SQL. - - no other database operations can be generated except for those that query the 'relationship_table'. \ -(WHERE, GROUP BY, etc. are certainly allowed, but no other database table can be used in the generated SQL.) -""" - -SQL_INSTRUCTIONS_ENTITY_PROMPT = """ -You are an expert at generating SQL queries to answer questions about a knowledge graph. - -You will be given a lot of instructions later, but here rules that MUST BE FOLLOWED: - - the SQL generated MUST only use the table one table named 'entity_table'. \ -This table is not a table that can be defined or overwritten by the user and the resulting SQL \ -statement, it MUST be seen as an existing table in the database. - - common table expressions that reference only the 'entity_table' are allowed. - - no other table or view of a potential underlying schema can in any way or shape be \ -involved in the generated SQL. - - no other database operations can be generated except for those that query the 'entity_table'. \ -(WHERE, GROUP BY, etc. are certainly allowed, but no other database table can be used in the generated SQL.) -""" diff --git a/backend/onyx/server/documents/document.py b/backend/onyx/server/documents/document.py index ac96ca6ab6..06442fa564 100644 --- a/backend/onyx/server/documents/document.py +++ b/backend/onyx/server/documents/document.py @@ -11,7 +11,7 @@ from onyx.db.enums import Permission 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 -from onyx.document_index.interfaces_new import DocumentSectionRequest +from onyx.document_index.interfaces 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, DocumentInfo @@ -29,8 +29,7 @@ def get_document_info( db_session: Session = Depends(get_session), ) -> DocumentInfo: search_settings = get_current_search_settings(db_session) - # This flow is for search so we do not get all indices. - document_index = get_default_document_index(search_settings, None, db_session) + document_index = get_default_document_index(search_settings, None) user_acl_filters = build_access_filters_for_user(user, db_session) inference_chunks = document_index.id_based_retrieval( @@ -74,8 +73,7 @@ def get_chunk_info( db_session: Session = Depends(get_session), ) -> ChunkInfo: search_settings = get_current_search_settings(db_session) - # This flow is for search so we do not get all indices. - document_index = get_default_document_index(search_settings, None, db_session) + document_index = get_default_document_index(search_settings, None) user_acl_filters = build_access_filters_for_user(user, db_session) chunk_request = DocumentSectionRequest( 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 133bd9abae..7a9e09bd3f 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 @@ -344,7 +344,7 @@ _COMPOSE_INTERNAL_HOSTNAMES = { "indexing_model_server", "inference_model_server", "web_server", - "vespa", + "opensearch", } diff --git a/backend/onyx/server/features/hierarchy/api.py b/backend/onyx/server/features/hierarchy/api.py index 63d10819cf..4e1c31f58e 100644 --- a/backend/onyx/server/features/hierarchy/api.py +++ b/backend/onyx/server/features/hierarchy/api.py @@ -1,9 +1,8 @@ -from fastapi import APIRouter, Depends, HTTPException, Query +from fastapi import APIRouter, Depends, Query from sqlalchemy.orm import Session from onyx.access.hierarchy_access import get_user_external_group_ids from onyx.auth.permissions import require_permission -from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX 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 @@ -13,7 +12,6 @@ from onyx.db.hierarchy import ( 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, HIERARCHY_NODE_DOCUMENTS_PATH, @@ -35,32 +33,9 @@ from onyx.server.features.hierarchy.models import ( HierarchyNodeSummary, ) -OPENSEARCH_NOT_ENABLED_MESSAGE = "Per-source knowledge selection is coming soon in v3.0! OpenSearch indexing must be enabled to use this feature." - -MIGRATION_STATUS_MESSAGE = ( - "Our records indicate that the transition to OpenSearch is still in progress. " - "OpenSearch retrieval is necessary to use this feature. " - "You can still use Document Sets, though! " - "If you would like to manually switch to OpenSearch, " - 'Go to the "Document Index Migration" section in the Admin panel.' -) - router = APIRouter(prefix=HIERARCHY_NODES_PREFIX) -def _require_opensearch(db_session: Session) -> None: - if not ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: - raise HTTPException( - status_code=403, - detail=OPENSEARCH_NOT_ENABLED_MESSAGE, - ) - if not get_opensearch_retrieval_state(db_session): - raise HTTPException( - status_code=403, - detail=MIGRATION_STATUS_MESSAGE, - ) - - def _get_user_access_info(user: User, db_session: Session) -> tuple[str, list[str]]: return user.email, get_user_external_group_ids(db_session, user) @@ -71,7 +46,6 @@ def list_accessible_hierarchy_nodes( user: User = Depends(require_permission(Permission.BASIC_ACCESS)), db_session: Session = Depends(get_session), ) -> HierarchyNodesResponse: - _require_opensearch(db_session) user_email, external_group_ids = _get_user_access_info(user, db_session) nodes = get_accessible_hierarchy_nodes_for_source( db_session=db_session, @@ -99,7 +73,6 @@ def list_accessible_hierarchy_node_documents( user: User = Depends(require_permission(Permission.BASIC_ACCESS)), db_session: Session = Depends(get_session), ) -> HierarchyNodeDocumentsResponse: - _require_opensearch(db_session) user_email, external_group_ids = _get_user_access_info(user, db_session) cursor = documents_request.cursor sort_field = documents_request.sort_field @@ -156,7 +129,6 @@ def search_hierarchy_nodes( user: User = Depends(require_permission(Permission.BASIC_ACCESS)), db_session: Session = Depends(get_session), ) -> HierarchyNodeSearchResponse: - _require_opensearch(db_session) user_email, external_group_ids = _get_user_access_info(user, db_session) nodes = search_accessible_hierarchy_nodes( db_session=db_session, diff --git a/backend/onyx/server/features/search/api.py b/backend/onyx/server/features/search/api.py index c1fcf4b8ca..aebc7ada6b 100644 --- a/backend/onyx/server/features/search/api.py +++ b/backend/onyx/server/features/search/api.py @@ -156,7 +156,7 @@ def search( # 4. Get document index search_settings = get_current_search_settings(db_session) - document_index = get_default_document_index(search_settings, None, db_session) + document_index = get_default_document_index(search_settings, None) # 5. Get tool_id all_tools = get_tools(db_session) diff --git a/backend/onyx/server/manage/get_state.py b/backend/onyx/server/manage/get_state.py index 6f232d0486..9d6acf9418 100644 --- a/backend/onyx/server/manage/get_state.py +++ b/backend/onyx/server/manage/get_state.py @@ -242,6 +242,12 @@ def get_version() -> VersionResponse: return VersionResponse(backend_version=__version__) +# The document index image the docker compose deployment runs. Keep in sync +# with the `opensearch` service in +# deployment/docker_compose/docker-compose.template.yml. +DOCUMENT_INDEX_IMAGE = "opensearchproject/opensearch:3.6.0" + + @router.get("/versions", tags=PUBLIC_API_TAGS) def get_versions() -> AllVersions: """ @@ -325,19 +331,19 @@ def get_versions() -> AllVersions: stable=ContainerVersions( onyx=latest_stable_version, relational_db="postgres:15.2-alpine", - index="vespaengine/vespa:8.277.17", + index=DOCUMENT_INDEX_IMAGE, nginx="nginx:1.25.5-alpine", ), dev=ContainerVersions( onyx=latest_dev_version, relational_db="postgres:15.2-alpine", - index="vespaengine/vespa:8.277.17", + index=DOCUMENT_INDEX_IMAGE, nginx="nginx:1.25.5-alpine", ), migration=ContainerVersions( onyx="airgapped-intfloat-nomic-migration", relational_db="postgres:15.2-alpine", - index="vespaengine/vespa:8.277.17", + index=DOCUMENT_INDEX_IMAGE, nginx="nginx:1.25.5-alpine", ), ) diff --git a/backend/onyx/server/manage/opensearch_migration/api.py b/backend/onyx/server/manage/opensearch_migration/api.py deleted file mode 100644 index 78457dd3b4..0000000000 --- a/backend/onyx/server/manage/opensearch_migration/api.py +++ /dev/null @@ -1,66 +0,0 @@ -from fastapi import APIRouter, Depends -from sqlalchemy.orm import Session - -from onyx.auth.permissions import require_permission -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, - get_opensearch_retrieval_state, - set_enable_opensearch_retrieval_with_commit, -) -from onyx.server.manage.opensearch_migration.models import ( - OpenSearchMigrationStatusResponse, - OpenSearchRetrievalStatusRequest, - OpenSearchRetrievalStatusResponse, -) - -admin_router = APIRouter(prefix="/admin/opensearch-migration") - - -@admin_router.get("/status") -def get_opensearch_migration_status( - _: User = Depends(require_permission(Permission.FULL_ADMIN_PANEL_ACCESS)), - db_session: Session = Depends(get_session), -) -> OpenSearchMigrationStatusResponse: - ( - total_chunks_migrated, - created_at, - migration_completed_at, - approx_chunk_count_in_vespa, - ) = get_opensearch_migration_state(db_session) - return OpenSearchMigrationStatusResponse( - total_chunks_migrated=total_chunks_migrated, - created_at=created_at, - migration_completed_at=migration_completed_at, - approx_chunk_count_in_vespa=approx_chunk_count_in_vespa, - ) - - -@admin_router.get("/retrieval") -def get_opensearch_retrieval_status( - _: User = Depends(require_permission(Permission.FULL_ADMIN_PANEL_ACCESS)), - db_session: Session = Depends(get_session), -) -> OpenSearchRetrievalStatusResponse: - enable_opensearch_retrieval = get_opensearch_retrieval_state(db_session) - return OpenSearchRetrievalStatusResponse( - enable_opensearch_retrieval=enable_opensearch_retrieval, - toggling_retrieval_is_disabled=ONYX_DISABLE_VESPA, - ) - - -@admin_router.put("/retrieval") -def set_opensearch_retrieval_status( - request: OpenSearchRetrievalStatusRequest, - _: User = Depends(require_permission(Permission.FULL_ADMIN_PANEL_ACCESS)), - db_session: Session = Depends(get_session), -) -> OpenSearchRetrievalStatusResponse: - set_enable_opensearch_retrieval_with_commit( - db_session, request.enable_opensearch_retrieval - ) - return OpenSearchRetrievalStatusResponse( - enable_opensearch_retrieval=request.enable_opensearch_retrieval, - toggling_retrieval_is_disabled=ONYX_DISABLE_VESPA, - ) diff --git a/backend/onyx/server/manage/opensearch_migration/models.py b/backend/onyx/server/manage/opensearch_migration/models.py deleted file mode 100644 index bdaa7b612f..0000000000 --- a/backend/onyx/server/manage/opensearch_migration/models.py +++ /dev/null @@ -1,22 +0,0 @@ -from datetime import datetime - -from pydantic import BaseModel - - -class OpenSearchMigrationStatusResponse(BaseModel): - model_config = {"frozen": True} - total_chunks_migrated: int - created_at: datetime | None - migration_completed_at: datetime | None - approx_chunk_count_in_vespa: int | None - - -class OpenSearchRetrievalStatusRequest(BaseModel): - model_config = {"frozen": True} - enable_opensearch_retrieval: bool - - -class OpenSearchRetrievalStatusResponse(BaseModel): - model_config = {"frozen": True} - enable_opensearch_retrieval: bool - toggling_retrieval_is_disabled: bool = False diff --git a/backend/onyx/server/manage/search_settings.py b/backend/onyx/server/manage/search_settings.py index df6ff69536..80db90db51 100644 --- a/backend/onyx/server/manage/search_settings.py +++ b/backend/onyx/server/manage/search_settings.py @@ -13,7 +13,7 @@ from onyx.background.celery.tasks.port.tasks import ( from onyx.background.celery.versioned_apps.client import app as client_app from onyx.configs.app_configs import ( DISABLE_INDEX_UPDATE_ON_SWAP, - ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + DISABLE_VECTOR_DB, OLD_INDEX_RECLAIM_ENABLED, ) from onyx.context.search.models import ( @@ -64,11 +64,8 @@ from onyx.db.search_settings import ( update_current_search_settings, update_search_settings_status, ) -from onyx.document_index.factory import ( - get_all_document_indices, - get_default_document_index, -) -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.factory import get_default_document_index +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.client import OpenSearchClient from onyx.document_index.opensearch.constants import LUCENE_SCALAR_QUANTIZATION from onyx.document_index.opensearch.index_reclaim import ( @@ -221,14 +218,13 @@ def set_new_search_settings( commit=False, ) - # Ensure the document indices have the new index immediately. - document_indices = get_all_document_indices(search_settings, new_search_settings) - for document_index in document_indices: - # Pair instances already know about their secondary search settings via - # the factory; only the primary embedding info needs to be passed in. - document_index.verify_and_create_index_if_necessary( - embedding_dim=search_settings.final_embedding_dim, - ) + # Ensure the document index has the new index immediately. The pair already + # knows about its secondary search settings via the factory; only the primary + # embedding info needs to be passed in. + document_index = get_default_document_index(search_settings, new_search_settings) + document_index.verify_and_create_index_if_necessary( + embedding_dim=search_settings.final_embedding_dim, + ) # Pause index attempts for the currently in-use index to preserve resources. if DISABLE_INDEX_UPDATE_ON_SWAP: @@ -289,7 +285,7 @@ def _validate_vector_quantization_supported( its OpenSearch version is not checked. """ lucene_scalar_quantization = LUCENE_SCALAR_QUANTIZATION.get(vector_quantization) - if lucene_scalar_quantization is None or not ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: + if lucene_scalar_quantization is None or DISABLE_VECTOR_DB: return with OpenSearchClient() as opensearch_client: cluster_version = opensearch_client.get_opensearch_version() @@ -494,9 +490,7 @@ def cancel_new_embedding( clear_reclaim_intent__no_commit(db_session, primary_search_settings.id) db_session.commit() - document_index = get_default_document_index( - primary_search_settings, None, db_session - ) + document_index = get_default_document_index(primary_search_settings, None) document_index.verify_and_create_index_if_necessary( embedding_dim=primary_search_settings.final_embedding_dim, ) diff --git a/backend/onyx/server/metrics/indexing_pipeline.py b/backend/onyx/server/metrics/indexing_pipeline.py index da021a53a0..714a2bcda4 100644 --- a/backend/onyx/server/metrics/indexing_pipeline.py +++ b/backend/onyx/server/metrics/indexing_pipeline.py @@ -68,7 +68,6 @@ _QUEUE_LABEL_MAP: dict[str, str] = { OnyxCeleryQueues.USER_FILE_DELETE: "user_file_delete", OnyxCeleryQueues.MONITORING: "monitoring", OnyxCeleryQueues.SANDBOX: "sandbox", - OnyxCeleryQueues.OPENSEARCH_MIGRATION: "opensearch_migration", } # Queues where prefetched (unacked) task counts are meaningful diff --git a/backend/onyx/server/onyx_api/ingestion.py b/backend/onyx/server/onyx_api/ingestion.py index c1686d328d..6d193e1009 100644 --- a/backend/onyx/server/onyx_api/ingestion.py +++ b/backend/onyx/server/onyx_api/ingestion.py @@ -32,7 +32,7 @@ from onyx.db.search_settings import ( get_secondary_search_settings, ) from onyx.db.user_file import get_owned_file_ids, get_user_file_by_id -from onyx.document_index.factory import get_all_document_indices +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.indexing.adapters.document_indexing_adapter import ( @@ -186,12 +186,7 @@ def upsert_ingestion_doc( # Need to index for both the primary and secondary index if possible active_search_settings = get_active_search_settings(db_session) - # This flow is for indexing so we get all indices. - document_indices = get_all_document_indices( - active_search_settings.primary, - None, - None, - ) + document_index = get_default_document_index(active_search_settings.primary, None) search_settings = get_current_search_settings(db_session) @@ -212,7 +207,7 @@ def upsert_ingestion_doc( indexing_pipeline_result = run_indexing_pipeline( embedder=index_embedding_model, - document_indices=document_indices, + document_index=document_index, ignore_time_skip=True, db_session=db_session, tenant_id=tenant_id, @@ -235,14 +230,13 @@ def upsert_ingestion_doc( search_settings=sec_search_settings ) - # This flow is for indexing so we get all indices. - sec_document_indices = get_all_document_indices( - active_search_settings.secondary, None, None + sec_document_index = get_default_document_index( + active_search_settings.secondary, None ) run_indexing_pipeline( embedder=new_index_embedding_model, - document_indices=sec_document_indices, + document_index=sec_document_index, ignore_time_skip=True, # FUTURE write: skip content_hash dedup, else the primary run's hash # suppresses this into a no-op. @@ -301,18 +295,15 @@ def delete_ingestion_doc( if recorded_ids: db_session.commit() - # This flow is for deletion so we get all indices. - document_indices = get_all_document_indices( + document_index = get_default_document_index( active_search_settings.primary, active_search_settings.secondary, - None, ) try: - for document_index in document_indices: - document_index.delete( - document_id, - chunk_count=document.chunk_count, - ) + document_index.delete( + document_id, + chunk_count=document.chunk_count, + ) delete_documents_complete(db_session, [document_id]) except Exception: # A failed DB delete leaves the session aborted; roll back before the candidate diff --git a/backend/onyx/server/query_and_chat/query_backend.py b/backend/onyx/server/query_and_chat/query_backend.py index 87e2229395..acf04b4746 100644 --- a/backend/onyx/server/query_and_chat/query_backend.py +++ b/backend/onyx/server/query_and_chat/query_backend.py @@ -55,8 +55,7 @@ def admin_search( tenant_id=tenant_id, ) search_settings = get_current_search_settings(db_session) - # This flow is for search so we do not get all indices. - document_index = get_default_document_index(search_settings, None, db_session) + document_index = get_default_document_index(search_settings, None) if not query or query.strip() == "": matching_chunks = document_index.random_retrieval(filters=final_filters) diff --git a/backend/onyx/server/settings/models.py b/backend/onyx/server/settings/models.py index a11a69704e..3e2fbdf0af 100644 --- a/backend/onyx/server/settings/models.py +++ b/backend/onyx/server/settings/models.py @@ -107,8 +107,9 @@ class Settings(BaseModel): seat_count: int | None = None used_seats: int | None = None - # OpenSearch migration - opensearch_indexing_enabled: bool = False + # Deprecated: OpenSearch is the only document index, so this is always True. + # Kept for API clients that still read it (e.g. the Terraform provider). + opensearch_indexing_enabled: bool = True class UserSettings(Settings): @@ -124,7 +125,7 @@ class UserSettings(Settings): # user's sandbox pod's opencode-serve logs. Gated by the # ENABLE_OPENCODE_DEBUGGING env var; never set in prod. opencode_debugging_enabled: bool = False - # True when a vector database (Vespa/OpenSearch) is available. + # True when a vector database (OpenSearch) is available. # False when DISABLE_VECTOR_DB is set — connectors, RAG search, and # document sets are unavailable. vector_db_enabled: bool = True diff --git a/backend/onyx/server/settings/store.py b/backend/onyx/server/settings/store.py index 65388317be..65cb9ad577 100644 --- a/backend/onyx/server/settings/store.py +++ b/backend/onyx/server/settings/store.py @@ -7,7 +7,6 @@ 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, @@ -88,7 +87,7 @@ def load_settings(raise_on_error: bool = False) -> Settings: settings.user_knowledge_enabled = False settings.show_extra_connectors = SHOW_EXTRA_CONNECTORS - settings.opensearch_indexing_enabled = ENABLE_OPENSEARCH_INDEXING_FOR_ONYX + settings.opensearch_indexing_enabled = True settings.hide_query_history_from_admin_panel = HIDE_QUERY_HISTORY_FROM_ADMIN_PANEL # Resolve context-aware defaults for token threshold. diff --git a/backend/onyx/setup.py b/backend/onyx/setup.py index 95f9bf1d2b..7da4bef2a1 100644 --- a/backend/onyx/setup.py +++ b/backend/onyx/setup.py @@ -5,17 +5,10 @@ from sqlalchemy.orm import Session from onyx.configs.app_configs import ( DISABLE_INDEX_UPDATE_ON_SWAP, DISABLE_VECTOR_DB, - ENABLE_OPENSEARCH_INDEXING_FOR_ONYX, + DOCUMENT_INDEX_NUM_ATTEMPTS_ON_STARTUP, 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, - 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, create_initial_default_connector @@ -42,17 +35,14 @@ from onyx.db.search_settings import ( 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.factory import get_default_document_index +from onyx.document_index.interfaces import DocumentIndex from onyx.document_index.opensearch.client import ( OpenSearchClient, OpenSearchIndexWriteBlockedError, 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, -) from onyx.indexing.models import IndexingSetting from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError @@ -74,7 +64,6 @@ from onyx.server.settings.store import ( from onyx.utils.gpu_utils import gpu_status_request from onyx.utils.logger import setup_logger from shared_configs.configs import ( - ALT_INDEX_SUFFIX, MODEL_SERVER_HOST, MODEL_SERVER_PORT, MULTI_TENANT, @@ -100,15 +89,6 @@ def setup_onyx( search_settings = active_search_settings.primary secondary_search_settings = active_search_settings.secondary - # search_settings = get_current_search_settings(db_session) - # multipass_config_1 = get_multipass_config(search_settings) - - # secondary_large_chunks_enabled: bool | None = None - # secondary_search_settings = get_secondary_search_settings(db_session) - # if secondary_search_settings: - # multipass_config_2 = get_multipass_config(secondary_search_settings) - # secondary_large_chunks_enabled = multipass_config_2.enable_large_chunks - # Break bad state for thrashing indexes if secondary_search_settings and DISABLE_INDEX_UPDATE_ON_SWAP: expire_index_attempts( @@ -143,19 +123,17 @@ def setup_onyx( "DISABLE_VECTOR_DB is set — skipping document index setup and embedding model warm-up." ) else: - # Ensure the document indices are setup correctly. This step is - # relatively near the end because Vespa takes a bit of time to start up. - logger.notice("Verifying Document Index(s) is/are available.") - # This flow is for setting up the document index so we get all indices - # here. - document_indices = get_all_document_indices( + # Ensure the document index is set up correctly. This step is + # relatively near the end because the document index can take a bit of + # time to start up. + logger.notice("Verifying the document index is available.") + document_index = get_default_document_index( search_settings, secondary_search_settings, - None, ) - success = setup_document_indices( - document_indices, + success = setup_document_index( + document_index, IndexingSetting.from_db_model(search_settings), ) if not success: @@ -202,71 +180,64 @@ def mark_reindex_flag(db_session: Session) -> None: kv_store.store(KV_REINDEX_KEY, False) -def setup_document_indices( - document_indices: list[DocumentIndex], +def setup_document_index( + document_index: DocumentIndex, index_setting: IndexingSetting, - num_attempts: int = VESPA_NUM_ATTEMPTS_ON_STARTUP, + num_attempts: int = DOCUMENT_INDEX_NUM_ATTEMPTS_ON_STARTUP, ) -> bool: - """Sets up all input document indices. + """Sets up the input document index. - If any document index setup fails, the function will return False. Otherwise + Returns False if the setup does not succeed within num_attempts, otherwise returns True. """ - for document_index in document_indices: - # Document index startup is a bit slow, so give it a few seconds. - WAIT_SECONDS = 5 - document_index_setup_success = False - for x in range(num_attempts): - try: - logger.notice( - "Setting up document index %s (attempt %s/%s)...", - document_index.__class__.__name__, - x + 1, - num_attempts, - ) - document_index.verify_and_create_index_if_necessary( - embedding_dim=index_setting.final_embedding_dim, - ) - - logger.notice( - "Document index %s setup complete.", - document_index.__class__.__name__, - ) - document_index_setup_success = True - break - except OpenSearchIndexWriteBlockedError as e: - # The index exists but is write-blocked (typically the - # read_only_allow_delete block applied at the disk flood-stage - # watermark). It is still readable, so start up degraded rather - # than crash-loop until the block clears. A missing index or - # blocked creation raises a different error and still fails. - logger.error( - "Document index %s is write-blocked; continuing startup without " - "the mapping refresh. Search still works, but indexing will fail " - "until the block is cleared (usually by freeing disk space below " - "the flood-stage watermark). Error: %s", - document_index.__class__.__name__, - e, - ) - document_index_setup_success = True - break - except Exception: - logger.exception( - "Document index %s setup did not succeed. The relevant service may not be ready yet. Retrying in %s seconds.", - document_index.__class__.__name__, - WAIT_SECONDS, - ) - time.sleep(WAIT_SECONDS) - - if not document_index_setup_success: - logger.error( - "Document index %s setup did not succeed. Attempt limit reached. (%s)", + # Document index startup is a bit slow, so give it a few seconds. + WAIT_SECONDS = 5 + for x in range(num_attempts): + try: + logger.notice( + "Setting up document index %s (attempt %s/%s)...", document_index.__class__.__name__, + x + 1, num_attempts, ) - return False + document_index.verify_and_create_index_if_necessary( + embedding_dim=index_setting.final_embedding_dim, + ) - return True + logger.notice( + "Document index %s setup complete.", + document_index.__class__.__name__, + ) + return True + except OpenSearchIndexWriteBlockedError as e: + # The index exists but is write-blocked (typically the + # read_only_allow_delete block applied at the disk flood-stage + # watermark). It is still readable, so start up degraded rather + # than crash-loop until the block clears. A missing index or + # blocked creation raises a different error and still fails. + logger.error( + "Document index %s is write-blocked; continuing startup without " + "the mapping refresh. Search still works, but indexing will fail " + "until the block is cleared (usually by freeing disk space below " + "the flood-stage watermark). Error: %s", + document_index.__class__.__name__, + e, + ) + return True + except Exception: + logger.exception( + "Document index %s setup did not succeed. The relevant service may not be ready yet. Retrying in %s seconds.", + document_index.__class__.__name__, + WAIT_SECONDS, + ) + time.sleep(WAIT_SECONDS) + + logger.error( + "Document index %s setup did not succeed. Attempt limit reached. (%s)", + document_index.__class__.__name__, + num_attempts, + ) + return False def setup_postgres(db_session: Session) -> None: @@ -350,50 +321,12 @@ def update_default_multipass_indexing(db_session: Session) -> None: def setup_multitenant_onyx() -> None: if DISABLE_VECTOR_DB: - logger.notice("DISABLE_VECTOR_DB is set — skipping multitenant Vespa setup.") + logger.notice( + "DISABLE_VECTOR_DB is set — skipping multitenant document index setup." + ) return - if ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: - opensearch_client = OpenSearchClient() - if not wait_for_opensearch_with_timeout(client=opensearch_client): - raise RuntimeError("Failed to connect to OpenSearch.") - set_cluster_state(opensearch_client) - - # For Managed Vespa, the schema is sent over via the Vespa Console manually. - # NOTE: Pretty sure this code is never hit in any production environment. - if not MANAGED_VESPA and not ONYX_DISABLE_VESPA: - setup_vespa_multitenant(SUPPORTED_EMBEDDING_MODELS) - - -def setup_vespa_multitenant(supported_indices: list[SupportedEmbeddingModel]) -> bool: - # TODO(andrei): We don't yet support OpenSearch for multi-tenant instances - # so this function remains unchanged. - # This is for local testing - WAIT_SECONDS = 5 - VESPA_ATTEMPTS = 5 - for x in range(VESPA_ATTEMPTS): - try: - logger.notice("Setting up Vespa (attempt %s/%s)...", x + 1, VESPA_ATTEMPTS) - register_multitenant_vespa_indices( - indices=[index.index_name for index in supported_indices] - + [ - f"{index.index_name}{ALT_INDEX_SUFFIX}" - for index in supported_indices - ], - embedding_dims=[index.dim for index in supported_indices] - + [index.dim for index in supported_indices], - ) - - logger.notice("Vespa setup complete.") - return True - except Exception: - logger.notice( - "Vespa setup did not succeed. The Vespa service may not be ready yet. Retrying in %s seconds.", - WAIT_SECONDS, - ) - time.sleep(WAIT_SECONDS) - - logger.error( - "Vespa setup did not succeed. Attempt limit reached. (%s)", VESPA_ATTEMPTS - ) - return False + opensearch_client = OpenSearchClient() + if not wait_for_opensearch_with_timeout(client=opensearch_client): + raise RuntimeError("Failed to connect to OpenSearch.") + set_cluster_state(opensearch_client) diff --git a/backend/onyx/tools/tool_constructor.py b/backend/onyx/tools/tool_constructor.py index 9854035a95..95adaaf417 100644 --- a/backend/onyx/tools/tool_constructor.py +++ b/backend/onyx/tools/tool_constructor.py @@ -63,7 +63,7 @@ def _disambiguate_mcp_tool_names(tools: list[Tool]) -> None: class SearchToolConfig(BaseModel): user_selected_filters: BaseFilters | None = None - # Vespa metadata filters for overflowing user files. These are NOT the + # Document index metadata filters for overflowing user files. These are NOT the # IDs of the current project/persona — they are only set when the # project's/persona's user files didn't fit in the LLM context window and # must be found via vector DB search instead. @@ -214,8 +214,7 @@ def _construct_tools_impl( user_oauth_token: str | None = user.live_oauth_token search_settings = get_current_search_settings(db_session) - # This flow is for search so we do not get all indices. - document_index = get_default_document_index(search_settings, None, db_session) + document_index = get_default_document_index(search_settings, None) def _build_search_tool(tool_id: int, config: SearchToolConfig) -> SearchTool: persona_search_info = PersonaSearchInfo( 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 7b73d71f68..2f90266596 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 @@ -23,7 +23,7 @@ from onyx.context.search.utils import ( 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, DocumentSectionRequest +from onyx.document_index.interfaces import DocumentIndex, DocumentSectionRequest from onyx.server.query_and_chat.placement import Placement from onyx.server.query_and_chat.streaming_models import ( OpenUrlDocuments, diff --git a/backend/onyx/tools/tool_implementations/search/search_tool.py b/backend/onyx/tools/tool_implementations/search/search_tool.py index 5a2326bd19..4edd9dae29 100644 --- a/backend/onyx/tools/tool_implementations/search/search_tool.py +++ b/backend/onyx/tools/tool_implementations/search/search_tool.py @@ -77,7 +77,7 @@ from onyx.db.federated import ( 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.document_index.interfaces import DocumentIndex from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError from onyx.federated_connectors.federated_retrieval import ( @@ -284,7 +284,7 @@ class SearchTool(Tool[SearchToolOverrideKwargs]): document_index: DocumentIndex, # Respecting user selections user_selected_filters: BaseFilters | None, - # Vespa metadata filters for overflowing user files. NOT the raw IDs + # Document index metadata filters for overflowing user files. NOT the raw IDs # of the current project/persona — only set when user files couldn't # fit in the LLM context and need to be searched via vector DB. project_id_filter: int | None, @@ -1030,8 +1030,8 @@ class SearchTool(Tool[SearchToolOverrideKwargs]): ) search_weights.append(weight) - # Add Slack federated search (runs once in parallel with all Vespa queries) - # This avoids the query multiplication problem where each Vespa query + # Add Slack federated search (runs once in parallel with all index queries) + # This avoids the query multiplication problem where each index query # would trigger a separate Slack search. # Only run if pre-fetch found a valid Slack access token. if slack_access_token and override_kwargs.original_query: @@ -1050,7 +1050,7 @@ class SearchTool(Tool[SearchToolOverrideKwargs]): # Use same weight as original query for Slack results search_weights.append(ORIGINAL_QUERY_WEIGHT) - # Run all searches in parallel (Vespa queries + Slack) + # Run all searches in parallel (index queries + Slack) all_search_results = run_functions_tuples_in_parallel(search_functions) if not all_search_results: all_search_results = [] @@ -1083,7 +1083,7 @@ class SearchTool(Tool[SearchToolOverrideKwargs]): ) # Enrich chunks with `Document.file_id` (Postgres-only metadata not - # stored in Vespa). + # stored in the document index). with get_session_with_current_tenant() as enrichment_session: populate_file_ids_on_sections(top_sections, enrichment_session) diff --git a/backend/onyx/tools/tool_implementations/search/search_utils.py b/backend/onyx/tools/tool_implementations/search/search_utils.py index 8798554b8b..3dbe1c01f3 100644 --- a/backend/onyx/tools/tool_implementations/search/search_utils.py +++ b/backend/onyx/tools/tool_implementations/search/search_utils.py @@ -9,10 +9,7 @@ from onyx.context.search.models import ( InferenceSection, ) from onyx.context.search.utils import inference_section_from_chunks -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.document_index.interfaces import DocumentIndex, DocumentSectionRequest 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 @@ -169,7 +166,7 @@ def _retrieve_adjacent_chunks( above_max = min_chunk_id - 1 above_request = DocumentSectionRequest( - document_id=replace_invalid_doc_id_characters(document_id), + document_id=document_id, min_chunk_ind=above_min, max_chunk_ind=above_max, ) @@ -191,7 +188,7 @@ def _retrieve_adjacent_chunks( below_max = max_chunk_id + num_chunks_below below_request = DocumentSectionRequest( - document_id=replace_invalid_doc_id_characters(document_id), + document_id=document_id, min_chunk_ind=below_min, max_chunk_ind=below_max, ) diff --git a/backend/onyx/tracing/flows.py b/backend/onyx/tracing/flows.py index fdda5e16e4..c177b1c928 100644 --- a/backend/onyx/tracing/flows.py +++ b/backend/onyx/tracing/flows.py @@ -46,10 +46,6 @@ class LLMFlow(StrEnum): CONTEXTUAL_RAG_CHUNK_CONTEXT = "contextual_rag_chunk_context" IMAGE_SUMMARIZATION = "image_summarization" - # Knowledge graph - KG_DOCUMENT_CLASSIFICATION = "kg_document_classification" - KG_DEEP_EXTRACTION = "kg_deep_extraction" - # Image generation IMAGE_GENERATION = "image_generation" IMAGE_EDIT = "image_edit" @@ -79,7 +75,5 @@ SYSTEM_TEXT_GENERATION_FLOWS: frozenset[LLMFlow] = frozenset( LLMFlow.CONTEXTUAL_RAG_DOC_SUMMARY, LLMFlow.CONTEXTUAL_RAG_CHUNK_CONTEXT, LLMFlow.IMAGE_SUMMARIZATION, - LLMFlow.KG_DOCUMENT_CLASSIFICATION, - LLMFlow.KG_DEEP_EXTRACTION, } ) diff --git a/backend/scripts/debugging/onyx_redis.py b/backend/scripts/debugging/onyx_redis.py index da63dc237b..23e5fb7cdd 100644 --- a/backend/scripts/debugging/onyx_redis.py +++ b/backend/scripts/debugging/onyx_redis.py @@ -52,7 +52,6 @@ class OnyxRedisCommand(Enum): purge_documentset_taskset = "purge_documentset_taskset" purge_usergroup_taskset = "purge_usergroup_taskset" purge_locks_blocking_deletion = "purge_locks_blocking_deletion" - purge_vespa_syncing = "purge_vespa_syncing" purge_pidbox = "purge_pidbox" get_user_token = "get_user_token" delete_user_token = "delete_user_token" @@ -141,10 +140,6 @@ def onyx_redis( f"{tenant_id}:{redis_connector.external_group_sync.fence_key}", dry_run, r ) return 0 - elif command == OnyxRedisCommand.purge_vespa_syncing: - return purge_by_match_and_type( - "*connectorsync:vespa_syncing*", "string", batch, dry_run, r - ) elif command == OnyxRedisCommand.purge_pidbox: return purge_by_match_and_type( "*reply.celery.pidbox", "list", batch, dry_run, r diff --git a/backend/scripts/debugging/onyx_vespa_schemas.py b/backend/scripts/debugging/onyx_vespa_schemas.py deleted file mode 100644 index 013eb5e752..0000000000 --- a/backend/scripts/debugging/onyx_vespa_schemas.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Tool to generate all supported schema variations for Onyx Cloud's Vespa database. - -Usage: - -``` -PYTHONPATH=. python scripts/debugging/onyx_vespa_schemas.py -``` - -Then, paste them into the existing vespa schema downloaded from the Vespa console, -and then re-zip. -""" - -import argparse -import os -from pathlib import Path - -import jinja2 - -from onyx.configs.embedding_configs import SUPPORTED_EMBEDDING_MODELS -from onyx.utils.logger import setup_logger - -logger = setup_logger() - - -def write_schema( - index_name: str, - dim: int, - template: jinja2.Template, - output_path: Path, -) -> None: - # Create schemas directory if it doesn't exist - schemas_dir = output_path / "schemas" - schemas_dir.mkdir(parents=True, exist_ok=True) - - index_filename = schemas_dir / (index_name + ".sd") - - schema = template.render( - multi_tenant=True, - schema_name=index_name, - dim=dim, - ) - - with open(index_filename, "w", encoding="utf-8") as f: - f.write(schema) - - logger.info("Wrote %s", index_filename) - - -def generate_document_entries() -> str: - """Generate document entries for all supported embedding models.""" - document_entries = [] - - for model in SUPPORTED_EMBEDDING_MODELS: - # Add regular index - document_entries.append( - f' ' - ) - # Add alt index - document_entries.append( - f' ' - ) - - return "\n".join(document_entries) - - -def write_cloud_services(cloud_services_template_path: str, output_path: Path) -> None: - """Generate and write the cloud-services.xml file.""" - # Create output directory if it doesn't exist - output_path.mkdir(parents=True, exist_ok=True) - - jinja_env = jinja2.Environment() # noqa: S701 — renders Vespa schema files, not HTML - - with open(cloud_services_template_path, "r", encoding="utf-8") as f: - template_str = f.read() - - template = jinja_env.from_string(template_str) - document_entries = generate_document_entries() - - services_xml = template.render(document_elements=document_entries) - - services_file = output_path / "services.xml" - with open(services_file, "w", encoding="utf-8") as f: - f.write(services_xml) - - logger.info("Wrote %s", services_file) - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Generate multi tenant Vespa schemas and services configuration" - ) - parser.add_argument( - "--template", - help="The Jinja template to use for schemas", - default="onyx/document_index/vespa/app_config/schemas/danswer_chunk.sd.jinja", - ) - parser.add_argument( - "--cloud-services-template", - help="The cloud-services.xml.jinja template path", - default="ee/onyx/document_index/vespa/app_config/cloud-services.xml.jinja", - ) - parser.add_argument( - "--output-path", - help="Output directory path (defaults to current directory)", - default=".", - ) - args = parser.parse_args() - - # Convert output path to Path object - output_path = Path(args.output_path) - - jinja_env = jinja2.Environment() # noqa: S701 — renders Vespa schema files, not HTML - - # Generate schema files - with open(args.template, "r", encoding="utf-8") as f: - template_str = f.read() - - template = jinja_env.from_string(template_str) - - num_indexes = 0 - for model in SUPPORTED_EMBEDDING_MODELS: - write_schema( - model.index_name, - model.dim, - template, - output_path, - ) - write_schema( - model.index_name + "__danswer_alt_index", - model.dim, - template, - output_path, - ) - num_indexes += 2 - - logger.info("Wrote %s indexes.", num_indexes) - - # Generate cloud services configuration if template is provided - if args.cloud_services_template: - if os.path.exists(args.cloud_services_template): - write_cloud_services(args.cloud_services_template, output_path) - else: - logger.error( - "Cloud services template not found: %s", args.cloud_services_template - ) - - -if __name__ == "__main__": - main() diff --git a/backend/scripts/debugging/opensearch/benchmark_retrieval.py b/backend/scripts/debugging/opensearch/benchmark_retrieval.py index 3c886da77c..1bedd58be7 100644 --- a/backend/scripts/debugging/opensearch/benchmark_retrieval.py +++ b/backend/scripts/debugging/opensearch/benchmark_retrieval.py @@ -20,7 +20,7 @@ from onyx.context.search.enums import QueryType from onyx.context.search.models import IndexFilters from onyx.db.engine.sql_engine import SqlEngine, get_session_with_current_tenant from onyx.db.search_settings import get_current_search_settings -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, ) diff --git a/backend/scripts/dev_run_background_jobs.py b/backend/scripts/dev_run_background_jobs.py index 09745a87cb..729b07e46f 100644 --- a/backend/scripts/dev_run_background_jobs.py +++ b/backend/scripts/dev_run_background_jobs.py @@ -44,7 +44,7 @@ def run_jobs() -> None: "--loglevel=INFO", "--hostname=light@%n", "-Q", - "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration", + "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim", ] cmd_worker_docprocessing = [ diff --git a/backend/scripts/env_inventory_baseline.txt b/backend/scripts/env_inventory_baseline.txt index beed3d2f40..a440dee9ed 100644 --- a/backend/scripts/env_inventory_baseline.txt +++ b/backend/scripts/env_inventory_baseline.txt @@ -135,7 +135,6 @@ DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY DEFAULT_REDIS_PREFIX -DISABLE_OPENSEARCH_MIGRATION_TASK DISABLE_USER_KNOWLEDGE DISALLOWED_SLACK_BOT_TENANT_IDS DISPOSABLE_EMAIL_DOMAINS_URL @@ -203,36 +202,10 @@ JIRA_PERMISSION_GROUP_SYNC_FREQUENCY JIRA_SLIM_PAGE_SIZE JOB_TIMEOUT_SECONDS K8S_CONTEXT -KG_CLUSTERING_RETRIEVE_THRESHOLD -KG_CLUSTERING_THRESHOLD KG_DEFAULT_MAX_PARENT_RECURSION_DEPTH -KG_ENTITY_EXTRACTION_TIMEOUT -KG_FILTERED_SEARCH_TIMEOUT -KG_FILTER_CONSTRUCTION_TIMEOUT -KG_MAX_DECOMPOSITION_SEGMENTS -KG_MAX_DEEP_SEARCH_RESULTS -KG_MAX_SEARCH_DOCUMENTS -KG_MAX_TOKENS_ANSWER_GENERATION -KG_METADATA_TRACKING_THRESHOLD -KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT -KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT -KG_NORMALIZATION_RERANK_THRESHOLD -KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT -KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT -KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT -KG_OBJECT_SOURCE_RESEARCH_TIMEOUT -KG_RELATIONSHIP_EXTRACTION_TIMEOUT -KG_RESEARCH_NUM_RETRIEVED_DOCS -KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES -KG_SQL_GENERATION_MAX_TOKENS -KG_SQL_GENERATION_TIMEOUT -KG_SQL_GENERATION_TIMEOUT_OVERRIDE -KG_STRATEGY_GENERATION_TIMEOUT KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX KG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX KG_TEMP_KG_RELATIONSHIPS_VIEW_NAME_PREFIX -KG_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION -KG_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION LANGFUSE_HOST LANGFUSE_PUBLIC_KEY LANGFUSE_SECRET_KEY @@ -246,7 +219,6 @@ LLM_FIRST_CHUNK_MAX_RETRIES LOG_FILE_NAME LOG_FORMAT LOG_TO_FILE -MANAGED_VESPA MARKETING_POSTHOG_API_KEY MASK_CREDENTIAL_PREFIX MAX_CHUNKS_PER_DOC_BATCH @@ -308,7 +280,6 @@ ONYX_DB_SHARD_MAP_VERSION_POLL_SECONDS ONYX_DB_SHARD_OVERRIDES ONYX_DB_SHARD_POOL_OVERFLOW ONYX_DB_SHARD_POOL_SIZE -ONYX_DISABLE_VESPA ONYX_SANDBOX_PUSH_PRIVATE_KEY ONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH ONYX_VERSION @@ -334,7 +305,6 @@ OPENSEARCH_EXPLAIN_ENABLED OPENSEARCH_INDEX_NUM_REPLICAS OPENSEARCH_INDEX_NUM_SHARDS OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED -OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE OPENSEARCH_PROFILING_DISABLED OPENSEARCH_TEXT_ANALYZER OPENSEARCH_USE_SSL @@ -376,7 +346,6 @@ RECAPTCHA_ENTERPRISE_PROJECT_ID RECAPTCHA_HOSTNAME_ALLOWLIST RECAPTCHA_SCORE_THRESHOLD RECAPTCHA_SITE_KEY -RECENCY_BIAS_MULTIPLIER REDIS_AUTH_EXPIRE_TIME_SECONDS REDIS_DB_NUMBER REDIS_DB_NUMBER_CELERY @@ -393,7 +362,6 @@ REDIS_SSL_CERTFILE REDIS_SSL_KEYFILE REQUESTS_CA_BUNDLE REQUEST_TIMEOUT_SECONDS -RERANK_COUNT S3_FILE_STORE_PREFIX S3_GENERATE_LOCAL_CHECKSUM S3_VERIFY_SSL @@ -475,16 +443,6 @@ VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT VERTEXAI_DEFAULT_CREDENTIALS VERTEXAI_DEFAULT_LOCATION VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE -VESPA_CLOUD_CERT_PATH -VESPA_CLOUD_KEY_PATH -VESPA_CLOUD_URL -VESPA_CONFIG_SERVER_HOST -VESPA_HOST -VESPA_MIGRATION_REQUEST_TIMEOUT_S -VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT -VESPA_PORT -VESPA_REQUEST_TIMEOUT -VESPA_TENANT_PORT VOICE_DISABLE_STREAMING_FALLBACK VOICE_DISABLE_STREAMING_STT WEB_CONNECTOR_IGNORED_CLASSES diff --git a/backend/scripts/force_delete_connector_by_id.py b/backend/scripts/force_delete_connector_by_id.py index 5c3ae89f1f..e31fbd925c 100755 --- a/backend/scripts/force_delete_connector_by_id.py +++ b/backend/scripts/force_delete_connector_by_id.py @@ -40,8 +40,8 @@ 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 -from onyx.document_index.interfaces_new import DocumentIndex +from onyx.document_index.factory import get_default_document_index +from onyx.document_index.interfaces import DocumentIndex from onyx.file_store.file_store import get_default_file_store from onyx.utils.logger import setup_logger @@ -56,7 +56,7 @@ _DELETION_BATCH_SIZE = 1000 def _unsafe_deletion( db_session: Session, - document_indices: list[DocumentIndex], + document_index: DocumentIndex, cc_pair: ConnectorCredentialPair, pair_id: int, ) -> int: @@ -77,11 +77,10 @@ def _unsafe_deletion( break for document in documents: - for document_index in document_indices: - document_index.delete( - document.id, - chunk_count=document.chunk_count, - ) + document_index.delete( + document.id, + chunk_count=document.chunk_count, + ) delete_documents_complete__no_commit( db_session=db_session, @@ -212,18 +211,16 @@ def _delete_connector(cc_pair_id: int, db_session: Session) -> None: else [] ) try: - logger.notice("Deleting information from Vespa and Postgres") + logger.notice("Deleting information from the document index and Postgres") active_search_settings = get_active_search_settings(db_session) - # This flow is for deletion so we get all indices. - document_indices = get_all_document_indices( + document_index = get_default_document_index( active_search_settings.primary, active_search_settings.secondary, - None, ) files_deleted_count = _unsafe_deletion( db_session=db_session, - document_indices=document_indices, + document_index=document_index, cc_pair=cc_pair, pair_id=cc_pair_id, ) diff --git a/backend/scripts/orphan_doc_cleanup_script.py b/backend/scripts/orphan_doc_cleanup_script.py index c3f7656abe..50072517b7 100644 --- a/backend/scripts/orphan_doc_cleanup_script.py +++ b/backend/scripts/orphan_doc_cleanup_script.py @@ -5,29 +5,19 @@ import sys from sqlalchemy import text from sqlalchemy.orm import Session -from onyx.document_index.document_index_utils import get_multipass_config - # makes it so `PYTHONPATH=.` is not required when running this script 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 ( # 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.search_settings import get_active_search_settings # noqa: E402 from onyx.db.tag import delete_orphan_tags_batched # 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, -) -from shared_configs.configs import MULTI_TENANT # noqa: E402 -from shared_configs.contextvars import get_current_tenant_id # noqa: E402 +from onyx.document_index.factory import get_default_document_index # noqa: E402 +from onyx.document_index.interfaces import DocumentIndex # noqa: E402 BATCH_SIZE = 100 @@ -61,54 +51,31 @@ def main() -> None: ) return - # Setup Vespa index - search_settings = get_current_search_settings(db_session) - multipass_config = get_multipass_config(search_settings) - index_name = search_settings.index_name - tenant_state = TenantState( - tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT - ) - vespa_index = VespaDocumentIndex( - index_name=index_name, - tenant_state=tenant_state, - large_chunks_enabled=multipass_config.enable_large_chunks, + # Include the secondary index so an orphan's chunks are also + # removed from the future index during an index swap. + active_search_settings = get_active_search_settings(db_session) + document_index = get_default_document_index( + active_search_settings.primary, active_search_settings.secondary ) - # Delete chunks from Vespa first - print("Deleting orphaned document chunks from Vespa") - successfully_vespa_deleted_doc_ids: list[str] = [] + # Delete chunks from the document index first + print("Deleting orphaned document chunks from the document index") + successfully_index_deleted_doc_ids: list[str] = [] # Process documents in parallel using ThreadPoolExecutor with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor: def process_doc( - doc_id: str, vespa_index: VespaDocumentIndex = vespa_index + doc_id: str, document_index: DocumentIndex = document_index ) -> str | None: document = get_document(doc_id, db_session) if not document: return None - # Check if document exists in Vespa first + # Delete without a lookup first: lookups read only the + # primary index, and delete is a no-op for a missing + # document. try: - chunks = vespa_index.id_based_retrieval( - chunk_requests=[ - DocumentSectionRequest( - document_id=doc_id, max_chunk_ind=2 - ) - ], - filters=IndexFilters(access_control_list=None), - batch_retrieval=True, - ) - if not chunks: - print(f"Document {doc_id} not found in Vespa") - return doc_id - except Exception as e: - print( - f"Error checking if document {doc_id} exists in Vespa: {e}" - ) - return None - - try: - print(f"Deleting document {doc_id} in Vespa") - chunks_deleted = vespa_index.delete( + print(f"Deleting document {doc_id} in the document index") + chunks_deleted = document_index.delete( doc_id, chunk_count=document.chunk_count, ) @@ -119,7 +86,7 @@ def main() -> None: return doc_id except Exception as e: print( - f"Error deleting document {doc_id} in Vespa and will not delete from Postgres: {e}" + f"Error deleting document {doc_id} in the document index and will not delete from Postgres: {e}" ) return None @@ -130,13 +97,21 @@ def main() -> None: for future in concurrent.futures.as_completed(futures): doc_id = future.result() if doc_id: - successfully_vespa_deleted_doc_ids.append(doc_id) + successfully_index_deleted_doc_ids.append(doc_id) + + if not successfully_index_deleted_doc_ids: + # The next query would return the same documents, so stop + # instead of retrying them forever. + print( + "Could not delete any orphaned document in this batch from the document index. Stopping." + ) + break # Delete documents from Postgres print("Deleting orphaned documents from Postgres") try: delete_documents_complete__no_commit( - db_session, successfully_vespa_deleted_doc_ids + db_session, successfully_index_deleted_doc_ids ) db_session.commit() delete_orphan_tags_batched(db_session) @@ -144,9 +119,9 @@ def main() -> None: print(f"Error deleting documents from Postgres: {e}") break - total_processed += len(successfully_vespa_deleted_doc_ids) + total_processed += len(successfully_index_deleted_doc_ids) print( - f"Successfully cleaned up {len(successfully_vespa_deleted_doc_ids)} orphaned documents in this batch" + f"Successfully cleaned up {len(successfully_index_deleted_doc_ids)} orphaned documents in this batch" ) print(f"Total documents processed so far: {total_processed}") diff --git a/backend/scripts/query_time_check/seed_dummy_docs.py b/backend/scripts/query_time_check/seed_dummy_docs.py deleted file mode 100644 index 59f2ce3013..0000000000 --- a/backend/scripts/query_time_check/seed_dummy_docs.py +++ /dev/null @@ -1,185 +0,0 @@ -""" -launch: -- api server -- postgres -- vespa -- model server (this is only needed so the api server can startup, no embedding is done) - -Run this script to seed the database with dummy documents. -Then run test_query_times.py to test query times. -""" - -import random -from datetime import datetime - -from onyx.access.models import DocumentAccess -from onyx.configs.constants import DocumentSource -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, TenantState -from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex -from onyx.indexing.models import ChunkEmbedding, DocMetadataAwareIndexChunk, IndexChunk -from onyx.utils.timing import log_function_time -from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA -from shared_configs.model_server_models import Embedding - -TOTAL_DOC_SETS = 8 -TOTAL_ACL_ENTRIES_PER_CATEGORY = 80 - - -def generate_random_embedding(dim: int) -> Embedding: - return [random.uniform(-1, 1) for _ in range(dim)] - - -def generate_random_identifier() -> str: - return f"dummy_doc_{random.randint(1, 1000)}" - - -def generate_dummy_chunk( - doc_id: str, - chunk_id: int, - embedding_dim: int, - number_of_acl_entries: int, - number_of_document_sets: int, -) -> DocMetadataAwareIndexChunk: - document = Document( - id=doc_id, - source=DocumentSource.GOOGLE_DRIVE, - sections=[], - metadata={}, - semantic_identifier=generate_random_identifier(), - ) - - chunk = IndexChunk( - chunk_id=chunk_id, - blurb=f"Blurb for chunk {chunk_id} of document {doc_id}.", - content=f"Content for chunk {chunk_id} of document {doc_id}. This is dummy text for testing purposes.", - source_links={}, - section_continuation=False, - source_document=document, - title_prefix=f"Title prefix for doc {doc_id}", - metadata_suffix_semantic="", - metadata_suffix_keyword="", - doc_summary="", - chunk_context="", - mini_chunk_texts=None, - contextual_rag_reserved_tokens=0, - embeddings=ChunkEmbedding( - full_embedding=generate_random_embedding(embedding_dim), - mini_chunk_embeddings=[], - ), - title_embedding=generate_random_embedding(embedding_dim), - large_chunk_id=None, - large_chunk_reference_ids=[], - image_file_id=None, - ) - - document_set_names = [f"Document Set {i}" for i in range(number_of_document_sets)] - - user_emails: list[str | None] = [] - user_groups: list[str] = [] - external_user_emails: list[str] = [] - external_user_group_ids: list[str] = [] - for i in range(number_of_acl_entries): - user_emails.append(f"user_{i}@example.com") - user_groups.append(f"group_{i}") - external_user_emails.append(f"external_user_{i}@example.com") - external_user_group_ids.append(f"external_group_{i}") - - return DocMetadataAwareIndexChunk.from_index_chunk( - index_chunk=chunk, - user_project=[], - personas=[], - access=DocumentAccess.build( - user_emails=user_emails, - user_groups=user_groups, - external_user_emails=external_user_emails, - external_user_group_ids=external_user_group_ids, - is_public=random.choice([True, False]), - ), - document_sets=set(document_set_names), - boost=random.randint(-1, 1), - aggregated_chunk_boost_factor=random.random(), - tenant_id=POSTGRES_DEFAULT_SCHEMA, - ) - - -@log_function_time() -def do_insertion( - vespa_index: VespaDocumentIndex, all_chunks: list[DocMetadataAwareIndexChunk] -) -> None: - insertion_records = vespa_index.index( - chunks=all_chunks, - indexing_metadata=IndexingMetadata(doc_id_to_chunk_cnt_diff={}), - ) - print(f"Indexed {len(insertion_records)} documents.") - print( - f"New documents: {sum(1 for record in insertion_records if not record.already_existed)}" - ) - print( - f"Existing documents updated: {sum(1 for record in insertion_records if record.already_existed)}" - ) - - -@log_function_time() -def seed_dummy_docs( - number_of_document_sets: int, - number_of_acl_entries: int, - num_docs: int = 1000, - chunks_per_doc: int = 5, - batch_size: int = 100, -) -> None: - with get_session_with_current_tenant() as db_session: - search_settings = get_current_search_settings(db_session) - multipass_config = get_multipass_config(search_settings) - index_name = search_settings.index_name - embedding_dim = search_settings.final_embedding_dim - - vespa_index = VespaDocumentIndex( - index_name=index_name, - tenant_state=TenantState( - tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=MULTI_TENANT - ), - large_chunks_enabled=multipass_config.enable_large_chunks, - ) - print(index_name) - - all_chunks = [] - chunk_count = 0 - for doc_num in range(num_docs): - doc_id = f"dummy_doc_{doc_num}_{datetime.now().isoformat()}" - for chunk_num in range(chunks_per_doc): - chunk = generate_dummy_chunk( - doc_id=doc_id, - chunk_id=chunk_num, - embedding_dim=embedding_dim, - number_of_acl_entries=number_of_acl_entries, - number_of_document_sets=number_of_document_sets, - ) - all_chunks.append(chunk) - chunk_count += 1 - - if len(all_chunks) >= chunks_per_doc * batch_size: - do_insertion(vespa_index, all_chunks) - print( - f"Indexed {chunk_count} chunks out of {num_docs * chunks_per_doc}." - ) - print( - f"percentage: {chunk_count / (num_docs * chunks_per_doc) * 100:.2f}% \n" - ) - all_chunks = [] - - if all_chunks: - do_insertion(vespa_index, all_chunks) - - -if __name__ == "__main__": - seed_dummy_docs( - number_of_document_sets=TOTAL_DOC_SETS, - number_of_acl_entries=TOTAL_ACL_ENTRIES_PER_CATEGORY, - num_docs=100000, - chunks_per_doc=5, - batch_size=1000, - ) diff --git a/backend/scripts/query_time_check/test_query_times.py b/backend/scripts/query_time_check/test_query_times.py deleted file mode 100644 index ea4251dd70..0000000000 --- a/backend/scripts/query_time_check/test_query_times.py +++ /dev/null @@ -1,132 +0,0 @@ -# """ -# RUN THIS AFTER SEED_DUMMY_DOCS.PY -# """ - -# import random -# import time - -# from onyx.agents.agent_search.shared_graph_utils.models import QueryExpansionType -# from onyx.configs.constants import DocumentSource -# from onyx.configs.model_configs import DOC_EMBEDDING_DIM -# from onyx.context.search.models import IndexFilters -# 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.vespa.index import VespaIndex -# from scripts.query_time_check.seed_dummy_docs import TOTAL_ACL_ENTRIES_PER_CATEGORY -# from scripts.query_time_check.seed_dummy_docs import TOTAL_DOC_SETS -# from shared_configs.model_server_models import Embedding - -# # make sure these are smaller than TOTAL_ACL_ENTRIES_PER_CATEGORY and TOTAL_DOC_SETS, respectively -# NUMBER_OF_ACL_ENTRIES_PER_QUERY = 6 -# NUMBER_OF_DOC_SETS_PER_QUERY = 2 - - -# def get_slowest_99th_percentile(results: list[float]) -> float: -# return sorted(results)[int(0.99 * len(results))] - - -# # Generate random filters -# def _random_filters() -> IndexFilters: -# """ -# Generate random filters for the query containing: -# - NUMBER_OF_ACL_ENTRIES_PER_QUERY user emails -# - NUMBER_OF_ACL_ENTRIES_PER_QUERY groups -# - NUMBER_OF_ACL_ENTRIES_PER_QUERY external groups -# - NUMBER_OF_DOC_SETS_PER_QUERY document sets -# """ -# access_control_list = [ -# f"user_email:user_{random.randint(0, TOTAL_ACL_ENTRIES_PER_CATEGORY - 1)}@example.com", -# ] -# acl_indices = random.sample( -# range(TOTAL_ACL_ENTRIES_PER_CATEGORY), NUMBER_OF_ACL_ENTRIES_PER_QUERY -# ) -# for i in acl_indices: -# access_control_list.append(f"group:group_{acl_indices[i]}") -# access_control_list.append(f"external_group:external_group_{acl_indices[i]}") - -# doc_sets = [] -# doc_set_indices = random.sample( -# range(TOTAL_DOC_SETS), NUMBER_OF_ACL_ENTRIES_PER_QUERY -# ) -# for i in doc_set_indices: -# doc_sets.append(f"document_set:Document Set {doc_set_indices[i]}") - -# return IndexFilters( -# source_type=[DocumentSource.GOOGLE_DRIVE], -# document_set=doc_sets, -# tags=[], -# access_control_list=access_control_list, -# ) - - -# def test_hybrid_retrieval_times( -# number_of_queries: int, -# ) -> None: -# with get_session_with_current_tenant() as db_session: -# search_settings = get_current_search_settings(db_session) -# multipass_config = get_multipass_config(search_settings) -# index_name = search_settings.index_name - -# vespa_index = VespaIndex( -# index_name=index_name, -# secondary_index_name=None, -# large_chunks_enabled=multipass_config.enable_large_chunks, -# secondary_large_chunks_enabled=None, -# ) - -# # Generate random queries -# queries = [f"Random Query {i}" for i in range(number_of_queries)] - -# # Generate random embeddings -# embeddings = [ -# Embedding([random.random() for _ in range(DOC_EMBEDDING_DIM)]) -# for _ in range(number_of_queries) -# ] - -# total_time = 0.0 -# results = [] -# for i in range(number_of_queries): -# start_time = time.time() - -# vespa_index.hybrid_retrieval( -# query=queries[i], -# query_embedding=embeddings[i], -# final_keywords=None, -# filters=_random_filters(), -# hybrid_alpha=0.5, -# time_decay_multiplier=1.0, -# num_to_retrieve=50, -# ranking_profile_type=QueryExpansionType.SEMANTIC, -# offset=0, -# title_content_ratio=0.5, -# ) - -# end_time = time.time() -# query_time = end_time - start_time -# total_time += query_time -# results.append(query_time) - -# print(f"Query {i+1}: {query_time:.4f} seconds") - -# avg_time = total_time / number_of_queries -# fast_time = min(results) -# slow_time = max(results) -# ninety_ninth_percentile = get_slowest_99th_percentile(results) -# # Write results to a file -# _OUTPUT_PATH = "query_times_results_large_more.txt" -# with open(_OUTPUT_PATH, "w") as f: -# f.write(f"Average query time: {avg_time:.4f} seconds\n") -# f.write(f"Fastest query: {fast_time:.4f} seconds\n") -# f.write(f"Slowest query: {slow_time:.4f} seconds\n") -# f.write(f"99th percentile: {ninety_ninth_percentile:.4f} seconds\n") -# print(f"Results written to {_OUTPUT_PATH}") - -# print(f"\nAverage query time: {avg_time:.4f} seconds") -# print(f"Fastest query: {fast_time:.4f} seconds") -# print(f"Slowest query: {max(results):.4f} seconds") -# print(f"99th percentile: {get_slowest_99th_percentile(results):.4f} seconds") - - -# if __name__ == "__main__": -# test_hybrid_retrieval_times(number_of_queries=1000) diff --git a/backend/scripts/reset_indexes.py b/backend/scripts/reset_indexes.py deleted file mode 100644 index 59f489c715..0000000000 --- a/backend/scripts/reset_indexes.py +++ /dev/null @@ -1,76 +0,0 @@ -# This file is purely for development use, not included in any builds -import os -import sys -from time import sleep - -import requests -from requests.exceptions import RequestException - -# makes it so `PYTHONPATH=.` is not required when running this script -parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(parent_dir) - -from onyx.configs.app_configs import DOCUMENT_INDEX_NAME # noqa: E402 -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT # noqa: E402 -from onyx.utils.logger import setup_logger # noqa: E402 - -logger = setup_logger() - - -def wipe_vespa_index() -> bool: - """ - Wipes the Vespa index by deleting all documents. - """ - continuation = None - should_continue = True - RETRIES = 3 - - endpoint = DOCUMENT_ID_ENDPOINT.format(index_name=DOCUMENT_INDEX_NAME) - - while should_continue: - params = {"selection": "true", "cluster": DOCUMENT_INDEX_NAME} - if continuation: - params["continuation"] = continuation - - for attempt in range(RETRIES): - try: - response = requests.delete(endpoint, params=params) - response.raise_for_status() - - response_json = response.json() - logger.info("Response: %s", response_json) - - continuation = response_json.get("continuation") - should_continue = bool(continuation) - break # Exit the retry loop if the request is successful - - except RequestException: - logger.exception("Request failed") - sleep(2**attempt) # Exponential backoff - else: - logger.error("Max retries (%s) exceeded. Exiting.", RETRIES) - return False - - return True - - -def main() -> int: - """ - Main function to execute the script. - """ - try: - succeeded = wipe_vespa_index() - except Exception: - logger.exception("wipe_vespa_index exceptioned.") - return 1 - - if not succeeded: - logger.info("Vespa index wipe failed.") - return 0 - - logger.info("Vespa index wiped successfully.") - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/backend/scripts/restart_containers.sh b/backend/scripts/restart_containers.sh index f388b16eda..eec2076fb7 100755 --- a/backend/scripts/restart_containers.sh +++ b/backend/scripts/restart_containers.sh @@ -9,6 +9,8 @@ COMPOSE_FILE="$SCRIPT_DIR/../../deployment/docker_compose/docker-compose.yml" COMPOSE_DEV_FILE="$SCRIPT_DIR/../../deployment/docker_compose/docker-compose.dev.yml" stop_and_remove_containers() { + # onyx_vespa was started by older versions of this script; remove it if it + # is still around. docker stop onyx_postgres onyx_vespa onyx_redis onyx_minio onyx_code_interpreter 2>/dev/null || true docker rm onyx_postgres onyx_vespa onyx_redis onyx_minio onyx_code_interpreter 2>/dev/null || true docker compose -f "$COMPOSE_FILE" -f "$COMPOSE_DEV_FILE" --profile opensearch-enabled stop opensearch 2>/dev/null || true @@ -24,8 +26,8 @@ cleanup() { trap 'echo "Error occurred on line $LINENO. Exiting script." >&2; cleanup' ERR # Usage of the script with optional volume arguments -# ./restart_containers.sh [vespa_volume] [postgres_volume] [redis_volume] -# [minio_volume] [--keep-opensearch-data] +# ./restart_containers.sh [postgres_volume] [redis_volume] [minio_volume] +# [--keep-opensearch-data] KEEP_OPENSEARCH_DATA=false POSITIONAL_ARGS=() @@ -37,10 +39,17 @@ for arg in "$@"; do fi done -VESPA_VOLUME=${POSITIONAL_ARGS[0]:-""} -POSTGRES_VOLUME=${POSITIONAL_ARGS[1]:-""} -REDIS_VOLUME=${POSITIONAL_ARGS[2]:-""} -MINIO_VOLUME=${POSITIONAL_ARGS[3]:-""} +# The old first argument (vespa_volume) is gone. Stop on the old four-argument +# form so a Vespa volume is never mounted as the Postgres data volume. +if [[ ${#POSITIONAL_ARGS[@]} -gt 3 ]]; then + echo "Usage: $0 [postgres_volume] [redis_volume] [minio_volume] [--keep-opensearch-data]" >&2 + echo "The vespa_volume argument was removed. Pass the Postgres volume first." >&2 + exit 1 +fi + +POSTGRES_VOLUME=${POSITIONAL_ARGS[0]:-""} +REDIS_VOLUME=${POSITIONAL_ARGS[1]:-""} +MINIO_VOLUME=${POSITIONAL_ARGS[2]:-""} # Stop and remove the existing containers echo "Stopping and removing existing containers..." @@ -54,14 +63,6 @@ else docker run -p 5432:5432 --name onyx_postgres -e POSTGRES_PASSWORD=password -d postgres -c max_connections=250 fi -# Start the Vespa container with optional volume -echo "Starting Vespa container..." -if [[ -n "$VESPA_VOLUME" ]]; then - docker run --detach --name onyx_vespa --hostname vespa-container --publish 8081:8081 --publish 19071:19071 -v "$VESPA_VOLUME":/opt/vespa/var vespaengine/vespa:8 -else - docker run --detach --name onyx_vespa --hostname vespa-container --publish 8081:8081 --publish 19071:19071 vespaengine/vespa:8 -fi - # If OPENSEARCH_ADMIN_PASSWORD is not already set, try loading it from # .vscode/.env so existing dev setups that stored it there aren't silently # broken. diff --git a/backend/scripts/save_load_state.py b/backend/scripts/save_load_state.py deleted file mode 100644 index 1ddfe41da4..0000000000 --- a/backend/scripts/save_load_state.py +++ /dev/null @@ -1,141 +0,0 @@ -# This file is purely for development use, not included in any builds -# Remember to first to send over the schema information (run API Server) -import argparse -import json -import os -import subprocess - -import requests -from alembic import command -from alembic.config import Config - -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 - -logger = setup_logger() - - -def save_postgres(filename: str, container_name: str) -> None: - logger.notice("Attempting to take Postgres snapshot") - cmd = f"docker exec {container_name} pg_dump -U {POSTGRES_USER} -h {POSTGRES_HOST} -p {POSTGRES_PORT} -W -F t {POSTGRES_DB}" - with open(filename, "w") as file: - subprocess.run( - cmd, - shell=True, - check=True, - stdout=file, - text=True, - input=f"{POSTGRES_PASSWORD}\n", - ) - - -def load_postgres(filename: str, container_name: str) -> None: - logger.notice("Attempting to load Postgres snapshot") - try: - alembic_cfg = Config("alembic.ini") - command.upgrade(alembic_cfg, "head") - except Exception as e: - logger.error("Alembic upgrade failed: %s", e) - - host_file_path = os.path.abspath(filename) - - copy_cmd = f"docker cp {host_file_path} {container_name}:/tmp/" - subprocess.run(copy_cmd, shell=True, check=True) - - container_file_path = f"/tmp/{os.path.basename(filename)}" - - restore_cmd = ( - f"docker exec {container_name} pg_restore --clean -U {POSTGRES_USER} " - f"-h localhost -p {POSTGRES_PORT} -d {POSTGRES_DB} -1 -F t {container_file_path}" - ) - subprocess.run(restore_cmd, shell=True, check=True) - - -def save_vespa(filename: str) -> None: - logger.notice("Attempting to take Vespa snapshot") - endpoint = DOCUMENT_ID_ENDPOINT.format(index_name=DOCUMENT_INDEX_NAME) - continuation = "" - params = {} - doc_jsons: list[dict] = [] - while continuation is not None: - if continuation: - params = {"continuation": continuation} - response = requests.get(endpoint, params=params) - response.raise_for_status() - found = response.json() - continuation = found.get("continuation") - docs = found["documents"] - for doc in docs: - doc_json = {"update": doc["id"], "create": True, "fields": doc["fields"]} - doc_jsons.append(doc_json) - - with open(filename, "w") as jsonl_file: - for doc in doc_jsons: - json_str = json.dumps(doc) - jsonl_file.write(json_str + "\n") - - -def load_vespa(filename: str) -> None: - headers = {"Content-Type": "application/json"} - endpoint = DOCUMENT_ID_ENDPOINT.format(index_name=DOCUMENT_INDEX_NAME) - with open(filename, "r") as f: - for line in f: - new_doc = json.loads(line.strip()) - doc_id = new_doc["update"].split("::")[-1] - response = requests.post( - endpoint + "/" + doc_id, - headers=headers, - json=new_doc, - ) - response.raise_for_status() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Onyx checkpoint saving and loading.") - parser.add_argument( - "--save", action="store_true", help="Save Onyx state to directory." - ) - parser.add_argument( - "--load", action="store_true", help="Load Onyx state from save directory." - ) - parser.add_argument( - "--postgres_container_name", - type=str, - default="onyx-relational_db-1", - help="Name of the postgres container to dump", - ) - parser.add_argument( - "--checkpoint_dir", - type=str, - default=os.path.join("..", "onyx_checkpoint"), - help="A directory to store temporary files to.", - ) - - args = parser.parse_args() - checkpoint_dir = args.checkpoint_dir - postgres_container = args.postgres_container_name - - if not os.path.exists(checkpoint_dir): - os.makedirs(checkpoint_dir) - - if not args.save and not args.load: - raise ValueError("Must specify --save or --load") - - if args.load: - load_postgres( - os.path.join(checkpoint_dir, "postgres_snapshot.tar"), postgres_container - ) - load_vespa(os.path.join(checkpoint_dir, "vespa_snapshot.jsonl")) - else: - save_postgres( - os.path.join(checkpoint_dir, "postgres_snapshot.tar"), postgres_container - ) - save_vespa(os.path.join(checkpoint_dir, "vespa_snapshot.jsonl")) diff --git a/backend/supervisord.conf b/backend/supervisord.conf index 30a30131fa..8b99fe2f69 100644 --- a/backend/supervisord.conf +++ b/backend/supervisord.conf @@ -26,7 +26,7 @@ supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface # This means workers will not be able take advantage of multiple CPU cores # on a system, but this should be okay for now since all our celery tasks are # relatively compute-light (e.g. they tend to just make a bunch of requests to -# Vespa / Postgres) +# OpenSearch / Postgres) [program:celery_worker_primary] command=celery -A onyx.background.celery.versioned_apps.primary worker --hostname=primary@%%n @@ -44,7 +44,7 @@ stopasgroup=true [program:celery_worker_light] command=celery -A onyx.background.celery.versioned_apps.light worker --hostname=light@%%n - -Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration,chat_ttl_deletion + -Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,chat_ttl_deletion stdout_logfile=/var/log/celery_worker_light.log stdout_logfile_maxbytes=16MB redirect_stderr=true diff --git a/backend/tests/README.md b/backend/tests/README.md index 3b223de394..1f287dc37a 100644 --- a/backend/tests/README.md +++ b/backend/tests/README.md @@ -15,7 +15,7 @@ pytest -xv backend/tests/unit ### External Dependency Unit Tests (`tests/external_dependency_unit/`) -Real Postgres, Redis, MinIO, and Vespa available. Real OpenAI key when set. Real +Real Postgres, Redis, MinIO, and OpenSearch available. Real OpenAI key when set. Real Docker daemon when a Docker-backend test opts into it. Onyx application processes (API server, Celery workers) are **not** running. Tests import and call functions directly and can mock selectively. diff --git a/backend/tests/api/test_api.py b/backend/tests/api/test_api.py index 29d620d026..844967411b 100644 --- a/backend/tests/api/test_api.py +++ b/backend/tests/api/test_api.py @@ -152,7 +152,7 @@ def test_versions_endpoint(client: TestClient) -> None: # Verify migration has expected values assert migration["onyx"] == "airgapped-intfloat-nomic-migration" assert migration["relational_db"] == "postgres:15.2-alpine" - assert migration["index"] == "vespaengine/vespa:8.277.17" + assert migration["index"] == "opensearchproject/opensearch:3.6.0" assert migration["nginx"] == "nginx:1.25.5-alpine" # Verify versions are different between stable and dev diff --git a/backend/tests/daily/conftest.py b/backend/tests/daily/conftest.py index b441bbec85..e46e198945 100644 --- a/backend/tests/daily/conftest.py +++ b/backend/tests/daily/conftest.py @@ -1,7 +1,7 @@ import os # Daily tests run without a live backend; EE code paths that depend on -# Redis/Vespa/etc are not available, so disable enforcement before any +# Redis/OpenSearch/etc are not available, so disable enforcement before any # module-level imports below pull in EE versioned implementations. os.environ["LICENSE_ENFORCEMENT_ENABLED"] = "false" diff --git a/backend/tests/evals/connector_filter_eval/conftest.py b/backend/tests/evals/connector_filter_eval/conftest.py index 52d4f82835..33c9d91cb3 100644 --- a/backend/tests/evals/connector_filter_eval/conftest.py +++ b/backend/tests/evals/connector_filter_eval/conftest.py @@ -27,7 +27,6 @@ from tests.external_dependency_unit.answer.conftest import ( # noqa: F401 # no mock_file_store, # noqa: F401 mock_gpu_status, # noqa: F401 mock_nlp_embeddings_post, - mock_vespa_query, ) from tests.external_dependency_unit.conftest import ( db_session, # noqa: F401 diff --git a/backend/tests/external_dependency_unit/answer/conftest.py b/backend/tests/external_dependency_unit/answer/conftest.py index 3653477464..b6f0f0e45f 100644 --- a/backend/tests/external_dependency_unit/answer/conftest.py +++ b/backend/tests/external_dependency_unit/answer/conftest.py @@ -83,15 +83,6 @@ def mock_gpu_status() -> Iterator[None]: yield -@pytest.fixture -def mock_vespa_query() -> Iterator[None]: - """Stub Vespa query to a safe empty response to avoid CI flakiness.""" - with patch( - "onyx.document_index.vespa.vespa_document_index.query_vespa", return_value=[] - ): - yield - - @pytest.fixture def mock_file_store() -> Iterator[None]: """Mock the file store to avoid S3/storage dependencies in tests.""" @@ -118,7 +109,6 @@ def mock_file_store() -> Iterator[None]: def mock_external_deps( mock_nlp_embeddings_post: None, # noqa: ARG001 mock_gpu_status: None, # noqa: ARG001 - mock_vespa_query: None, # noqa: ARG001 mock_file_store: None, # noqa: ARG001 ) -> Iterator[None]: """Convenience fixture to enable all common external dependency mocks.""" diff --git a/backend/tests/external_dependency_unit/celery/test_document_sync_no_pg_txn_during_index_io.py b/backend/tests/external_dependency_unit/celery/test_document_sync_no_pg_txn_during_index_io.py index 86fd7a3635..50cc15b0ac 100644 --- a/backend/tests/external_dependency_unit/celery/test_document_sync_no_pg_txn_during_index_io.py +++ b/backend/tests/external_dependency_unit/celery/test_document_sync_no_pg_txn_during_index_io.py @@ -1,17 +1,16 @@ """ Regression coverage for document_index_metadata_sync_task holding a Postgres -transaction across the document-index (Vespa/OpenSearch) HTTP round-trip. +transaction across the document-index HTTP round-trip. Under bulk-deletion fan-out, every docprocessing-sync worker slot used to pin a -DB connection in state idle-in-transaction for the duration of the index call -(up to RetryDocumentIndex.STOP_AFTER seconds of tenacity retries). Those -transactions blocked `UPDATE document SET last_synced` writers cluster-wide. +DB connection in state idle-in-transaction for the duration of the index call. +Those transactions blocked `UPDATE document SET last_synced` writers cluster-wide. The task is now split into three phases: read DB state and close the session, do index I/O with no connection held, then reopen a fresh session to mark the document synced. Uses real PostgreSQL for Document rows and search settings. Only -RetryDocumentIndex.update is mocked (no index network I/O), so index-client +OpenSearchIndexPair.update is mocked (no index network I/O), so index-client construction runs for real against the detached-from-session SearchSettings — guarding the phase boundary itself. """ @@ -26,11 +25,13 @@ import pytest from sqlalchemy.orm import Session from sqlalchemy.pool import QueuePool -from onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex from onyx.background.celery.tasks.vespa import tasks as vespa_tasks from onyx.db.engine.sql_engine import SqlEngine, get_session_with_current_tenant from onyx.db.models import Document as DbDocument -from onyx.document_index.interfaces_new import SecondaryIndexDocumentMissingError +from onyx.document_index.interfaces import SecondaryIndexDocumentMissingError +from onyx.document_index.opensearch.opensearch_document_index import ( + OpenSearchIndexPair, +) from onyx.kg.models import KGStage from shared_configs.configs import ( POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE as TEST_TENANT_ID, @@ -65,13 +66,13 @@ def stale_document( def _run_task(doc_id: str, index_update_side_effect: object = None) -> bool: - """Run the sync task with RetryDocumentIndex.update patched. + """Run the sync task with OpenSearchIndexPair.update patched. Index clients are still built by the real factory from search settings that are detached from their (closed) phase-1 session. """ with patch.object( - RetryDocumentIndex, "update", side_effect=index_update_side_effect + OpenSearchIndexPair, "update", side_effect=index_update_side_effect ): result: bool = vespa_tasks.document_index_metadata_sync_task.apply( args=(doc_id,), kwargs={"tenant_id": TEST_TENANT_ID} @@ -92,7 +93,7 @@ def test_no_db_connection_held_during_index_io( result = _run_task(stale_document, index_update_side_effect=_record_pool_state) assert result is True - # the index update ran (once per configured index)... + # the index update ran... assert len(checked_out_during_update) >= 1 # ...with no DB connection checked out by the task at any of those moments assert all(count == 0 for count in checked_out_during_update) @@ -151,3 +152,31 @@ def test_port_missing_doc_still_marked_synced_in_fresh_session( assert row is not None assert row.last_synced == _LAST_MODIFIED assert row.secondary_only_sync_pending is False + + +def test_index_failure_retries_then_fails_and_leaves_doc_stale( + stale_document: str, db_session: Session +) -> None: + """Any index error is retryable: the task retries up to max_retries, then + fails without marking the doc synced, so the next beat pass re-syncs it.""" + index_update_calls: list[object] = [] + + def _fail(update_requests: object) -> None: + index_update_calls.append(update_requests) + raise RuntimeError("document index unavailable") + + with patch.object(OpenSearchIndexPair, "update", side_effect=_fail): + eager_result = vespa_tasks.document_index_metadata_sync_task.apply( + args=(stale_document,), kwargs={"tenant_id": TEST_TENANT_ID} + ) + + assert eager_result.failed() + assert isinstance(eager_result.result, RuntimeError) + max_retries = vespa_tasks.document_index_metadata_sync_task.max_retries + assert max_retries is not None + assert len(index_update_calls) == max_retries + 1 + + db_session.expire_all() + row = db_session.get(DbDocument, stale_document) + assert row is not None + assert row.last_synced is None 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 70e79db96e..bc45084f5b 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 @@ -8,13 +8,13 @@ Validates that: 2. The process_single_user_file_project_sync worker task reads persona associations from the DB, passes persona_ids to the document index via - VespaDocumentUserFields, and clears needs_persona_sync afterwards. + a MetadataUpdateRequest, and clears needs_persona_sync afterwards. 3. upsert_persona correctly marks affected UserFiles with needs_persona_sync=True when file associations change. -Uses real Redis and PostgreSQL. Document index (Vespa) calls are mocked -since we only need to verify the arguments passed to update_single. +Uses real Redis and PostgreSQL. Document index calls are mocked since we +only need to verify the arguments passed to update. """ from collections.abc import Generator @@ -36,7 +36,7 @@ from onyx.background.celery.tasks.user_file_processing.tasks import ( from onyx.db.enums import UserFileStatus from onyx.db.models import Persona, Persona__UserFile, User, UserFile from onyx.db.persona import upsert_persona -from onyx.document_index.interfaces_new import ( +from onyx.document_index.interfaces import ( MetadataUpdateRequest, SecondaryIndexDocumentMissingError, ) @@ -250,11 +250,8 @@ class TestCheckSweepIncludesPersonaSync: _PATCH_GET_SETTINGS = ( "onyx.background.celery.tasks.user_file_processing.tasks.get_active_search_settings" ) -_PATCH_GET_INDICES = ( - "onyx.background.celery.tasks.user_file_processing.tasks.get_all_document_indices" -) -_PATCH_HTTPX_INIT = ( - "onyx.background.celery.tasks.user_file_processing.tasks.httpx_init_vespa_pool" +_PATCH_GET_INDEX = ( + "onyx.background.celery.tasks.user_file_processing.tasks.get_default_document_index" ) _PATCH_DISABLE_VDB = ( "onyx.background.celery.tasks.user_file_processing.tasks.DISABLE_VECTOR_DB" @@ -312,9 +309,8 @@ class TestSyncTaskWritesPersonaIds: with ( patch(_PATCH_DISABLE_VDB, False), - patch(_PATCH_HTTPX_INIT), patch(_PATCH_GET_SETTINGS, return_value=mock_search_settings), - patch(_PATCH_GET_INDICES, return_value=[mock_doc_index]), + patch(_PATCH_GET_INDEX, return_value=mock_doc_index), ): process_single_user_file_project_sync.run( user_file_id=str(uf.id), # ty: ignore[invalid-argument-type] @@ -392,9 +388,8 @@ class TestSyncTaskWritesPersonaIds: with ( patch(_PATCH_DISABLE_VDB, False), - patch(_PATCH_HTTPX_INIT), patch(_PATCH_GET_SETTINGS, return_value=mock_search_settings), - patch(_PATCH_GET_INDICES, return_value=[mock_doc_index]), + patch(_PATCH_GET_INDEX, return_value=mock_doc_index), ): process_single_user_file_project_sync.run( user_file_id=str(uf.id), # ty: ignore[invalid-argument-type] @@ -421,7 +416,7 @@ class TestSyncTaskWritesPersonaIds: db_session: Session, tenant_context: None, # noqa: ARG002 ) -> None: - """A soft-deleted persona should NOT appear in the persona_ids sent to Vespa.""" + """A soft-deleted persona should NOT appear in the persona_ids sent to the index.""" user = create_test_user(db_session, "sync_deleted") uf = _create_completed_user_file(db_session, user, needs_persona_sync=True) persona = _create_test_persona(db_session, user) @@ -444,9 +439,8 @@ class TestSyncTaskWritesPersonaIds: with ( patch(_PATCH_DISABLE_VDB, False), - patch(_PATCH_HTTPX_INIT), patch(_PATCH_GET_SETTINGS, return_value=mock_search_settings), - patch(_PATCH_GET_INDICES, return_value=[mock_doc_index]), + patch(_PATCH_GET_INDEX, return_value=mock_doc_index), ): process_single_user_file_project_sync.run( user_file_id=str(uf.id), # ty: ignore[invalid-argument-type] @@ -610,9 +604,8 @@ def _run_sync( patches: list[Any] = [ patch(_PATCH_DISABLE_VDB, False), - patch(_PATCH_HTTPX_INIT), patch(_PATCH_GET_SETTINGS, return_value=mock_search_settings), - patch(_PATCH_GET_INDICES, return_value=[mock_doc_index]), + patch(_PATCH_GET_INDEX, return_value=mock_doc_index), patch(_PATCH_ACTIVE_SECONDARY, return_value=port_target), ] if supply is not None: @@ -826,10 +819,9 @@ def _run_index_pass(uf_id: str, port_target: Any, index_secondary: MagicMock) -> current.status.is_current.return_value = True with ( - patch(_PATCH_HTTPX_INIT), patch(_PATCH_GET_SETTINGS_LIST, return_value=[current]), patch(_PATCH_EMBEDDER), - patch(_PATCH_GET_INDICES, return_value=[MagicMock()]), + patch(_PATCH_GET_INDEX, return_value=MagicMock()), patch(_PATCH_RUN_PIPELINE, return_value=present_result), patch(_PATCH_ACTIVE_SECONDARY, return_value=port_target), patch(_PATCH_INDEX_SECONDARY, index_secondary), diff --git a/backend/tests/external_dependency_unit/conftest.py b/backend/tests/external_dependency_unit/conftest.py index 00ecdfab41..e5b7716b55 100644 --- a/backend/tests/external_dependency_unit/conftest.py +++ b/backend/tests/external_dependency_unit/conftest.py @@ -54,8 +54,8 @@ def db_session() -> Generator[Session, None, None]: def full_deployment_setup() -> Generator[None, None, None]: """Optional fixture to perform full deployment-like setup on demand. - Import and call tests.external_dependency_unit.startup.full_setup.ensure_full_deployment_setup - to initialize Postgres defaults, Vespa indices, and seed initial docs. + Import and call tests.external_dependency_unit.full_setup.ensure_full_deployment_setup + to initialize Postgres defaults and the document index. """ ensure_full_deployment_setup() yield 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 c10047e4e3..875d7df3a1 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 @@ -42,7 +42,7 @@ from onyx.db.models import ( ) from onyx.db.port_attempt import get_port_attempt 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.interfaces 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 ( @@ -305,7 +305,7 @@ def test_port_flow_end_to_end( # check_and_perform_index_swap fetches the FUTURE row via # get_secondary_search_settings (a stale FUTURE row exists in this dev - # DB), and get_all_document_indices would touch real Vespa/OpenSearch + # DB), and get_default_document_index would touch real OpenSearch # provisioning -> patch both in the swap_index namespace. The deferred # metadata-sync backlog is 0 in this dev DB; if it weren't, the swap # criterion would block, so we don't need to patch the count. @@ -323,7 +323,7 @@ def test_port_flow_end_to_end( "fetch_indexable_standard_connector_credential_pair_ids", lambda *_, **__: [pair.id], ), - patch.object(swap_index, "get_all_document_indices", return_value=[]), + patch.object(swap_index, "get_default_document_index"), ): old_settings = swap_index.check_and_perform_index_swap(db_session) 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 f3fc7def2a..7e38edd118 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 @@ -47,7 +47,7 @@ from onyx.db.port_attempt import ( create_port_attempt, mark_port_in_progress, ) -from onyx.document_index.interfaces_new import ( +from onyx.document_index.interfaces import ( MetadataUpdateRequest, SecondaryIndexDocumentMissingError, TenantState, 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 ea0b863ce0..f1e12c2abe 100644 --- a/backend/tests/external_dependency_unit/db/test_reindex_port.py +++ b/backend/tests/external_dependency_unit/db/test_reindex_port.py @@ -100,7 +100,7 @@ from onyx.db.port_orphan_candidate import ( ) 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.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState 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 @@ -1656,7 +1656,7 @@ def test_port_target_settings_id() -> None: def test_delete_port_written_chunks_query() -> None: """Filters written_by_port + doc-ids; adds the tenant term only in multitenant mode (single-tenant has no tenant_id field, so it would match zero docs).""" - from onyx.document_index.interfaces_new import TenantState + from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.schema import ( DOCUMENT_ID_FIELD_NAME, TENANT_ID_FIELD_NAME, diff --git a/backend/tests/external_dependency_unit/document_index/conftest.py b/backend/tests/external_dependency_unit/document_index/conftest.py index 0da454ae21..847150d8e6 100644 --- a/backend/tests/external_dependency_unit/document_index/conftest.py +++ b/backend/tests/external_dependency_unit/document_index/conftest.py @@ -12,7 +12,7 @@ from onyx.access.models import DocumentAccess from onyx.configs.constants import DocumentSource from onyx.connectors.models import Document from onyx.db.enums import VectorQuantization -from onyx.document_index.interfaces_new import IndexingMetadata, TenantState +from onyx.document_index.interfaces 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, 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 749c4bdc79..f1b60bd985 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 @@ -1,4 +1,4 @@ -"""External dependency tests for the new DocumentIndex interface. +"""External dependency tests for the DocumentIndex interface. These tests assume OpenSearch is running. """ @@ -15,8 +15,8 @@ from onyx.access.utils import prefix_user_email from onyx.configs.constants import PUBLIC_DOC_PAT from onyx.context.search.models import IndexFilters, InferenceChunk from onyx.db.enums import VectorQuantization -from onyx.document_index.interfaces_new import DocumentIndex as DocumentIndexNew -from onyx.document_index.interfaces_new import ( +from onyx.document_index.interfaces import ( + DocumentIndex, DocumentSectionRequest, MetadataUpdateRequest, TenantState, @@ -34,7 +34,7 @@ from tests.external_dependency_unit.document_index.conftest import ( def _retrieve_chunks_with_expected_boost( - document_index: DocumentIndexNew, + document_index: DocumentIndex, document_id: str, expected_chunk_count: int, expected_boost: int, @@ -89,10 +89,10 @@ def opensearch_document_index( @pytest.fixture(scope="module") -def document_indices( +def document_index( opensearch_document_index: OpenSearchDocumentIndex, -) -> Generator[list[DocumentIndexNew], None, None]: - yield [opensearch_document_index] +) -> Generator[DocumentIndex, None, None]: + yield opensearch_document_index # ------------------------------------------------------------------------------ @@ -100,14 +100,14 @@ def document_indices( # ------------------------------------------------------------------------------ -class TestDocumentIndexNew: +class TestDocumentIndex: """ - Tests the new DocumentIndex interface against a real OpenSearch. + Tests the DocumentIndex interface against a real OpenSearch. """ def test_index_single_new_doc( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -115,22 +115,21 @@ class TestDocumentIndexNew: already_existed=False. """ # Precondition. - for document_index in document_indices: - doc_id = f"test_single_new_{uuid.uuid4().hex[:8]}" - chunk = make_chunk(doc_id) - metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[1]) + doc_id = f"test_single_new_{uuid.uuid4().hex[:8]}" + chunk = make_chunk(doc_id) + metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[1]) - # Under test. - results = document_index.index(chunks=[chunk], indexing_metadata=metadata) + # Under test. + results = document_index.index(chunks=[chunk], indexing_metadata=metadata) - # Postcondition. - assert len(results) == 1 - assert results[0].document_id == doc_id - assert results[0].already_existed is False + # Postcondition. + assert len(results) == 1 + assert results[0].document_id == doc_id + assert results[0].already_existed is False def test_index_existing_doc_already_existed_true( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -138,36 +137,35 @@ class TestDocumentIndexNew: already_existed=True. """ # Precondition. - for document_index in document_indices: - doc_id = f"test_existing_{uuid.uuid4().hex[:8]}" - chunk = make_chunk(doc_id) + doc_id = f"test_existing_{uuid.uuid4().hex[:8]}" + chunk = make_chunk(doc_id) - # First index — brand new document. - metadata_first = make_indexing_metadata( - [doc_id], old_counts=[0], new_counts=[1] - ) - document_index.index(chunks=[chunk], indexing_metadata=metadata_first) + # First index — brand new document. + metadata_first = make_indexing_metadata( + [doc_id], old_counts=[0], new_counts=[1] + ) + document_index.index(chunks=[chunk], indexing_metadata=metadata_first) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Re-index — old_chunk_cnt=1 signals the document already existed. - metadata_second = make_indexing_metadata( - [doc_id], old_counts=[1], new_counts=[1] - ) + # Re-index — old_chunk_cnt=1 signals the document already existed. + metadata_second = make_indexing_metadata( + [doc_id], old_counts=[1], new_counts=[1] + ) - # Under test. - results = document_index.index( - chunks=[chunk], indexing_metadata=metadata_second - ) + # Under test. + results = document_index.index( + chunks=[chunk], indexing_metadata=metadata_second + ) - # Postcondition. - assert len(results) == 1 - assert results[0].already_existed is True + # Postcondition. + assert len(results) == 1 + assert results[0].already_existed is True def test_index_multiple_docs( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -175,30 +173,29 @@ class TestDocumentIndexNew: document. """ # Precondition. - for document_index in document_indices: - doc1 = f"test_multi_1_{uuid.uuid4().hex[:8]}" - doc2 = f"test_multi_2_{uuid.uuid4().hex[:8]}" - chunks = [ - make_chunk(doc1, chunk_id=0), - make_chunk(doc1, chunk_id=1), - make_chunk(doc2, chunk_id=0), - ] - metadata = make_indexing_metadata( - [doc1, doc2], old_counts=[0, 0], new_counts=[2, 1] - ) + doc1 = f"test_multi_1_{uuid.uuid4().hex[:8]}" + doc2 = f"test_multi_2_{uuid.uuid4().hex[:8]}" + chunks = [ + make_chunk(doc1, chunk_id=0), + make_chunk(doc1, chunk_id=1), + make_chunk(doc2, chunk_id=0), + ] + metadata = make_indexing_metadata( + [doc1, doc2], old_counts=[0, 0], new_counts=[2, 1] + ) - # Under test. - results = document_index.index(chunks=chunks, indexing_metadata=metadata) + # Under test. + results = document_index.index(chunks=chunks, indexing_metadata=metadata) - # Postcondition. - result_map = {r.document_id: r.already_existed for r in results} - assert len(result_map) == 2 - assert result_map[doc1] is False - assert result_map[doc2] is False + # Postcondition. + result_map = {r.document_id: r.already_existed for r in results} + assert len(result_map) == 2 + assert result_map[doc1] is False + assert result_map[doc2] is False def test_index_deduplicates_doc_ids_in_results( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -206,21 +203,20 @@ class TestDocumentIndexNew: DocumentInsertionRecord. """ # Precondition. - for document_index in document_indices: - doc_id = f"test_dedup_{uuid.uuid4().hex[:8]}" - chunks = [make_chunk(doc_id, chunk_id=i) for i in range(5)] - metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[5]) + doc_id = f"test_dedup_{uuid.uuid4().hex[:8]}" + chunks = [make_chunk(doc_id, chunk_id=i) for i in range(5)] + metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[5]) - # Under test. - results = document_index.index(chunks=chunks, indexing_metadata=metadata) + # Under test. + results = document_index.index(chunks=chunks, indexing_metadata=metadata) - # Postcondition. - assert len(results) == 1 - assert results[0].document_id == doc_id + # Postcondition. + assert len(results) == 1 + assert results[0].document_id == doc_id def test_index_mixed_new_and_existing_docs( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -228,64 +224,60 @@ class TestDocumentIndexNew: correct already_existed flag for each. """ # Precondition. - for document_index in document_indices: - existing_doc = f"test_mixed_exist_{uuid.uuid4().hex[:8]}" - new_doc = f"test_mixed_new_{uuid.uuid4().hex[:8]}" + existing_doc = f"test_mixed_exist_{uuid.uuid4().hex[:8]}" + new_doc = f"test_mixed_new_{uuid.uuid4().hex[:8]}" - # Pre-index the existing document. - pre_chunk = make_chunk(existing_doc) - pre_metadata = make_indexing_metadata( - [existing_doc], old_counts=[0], new_counts=[1] - ) - document_index.index(chunks=[pre_chunk], indexing_metadata=pre_metadata) + # Pre-index the existing document. + pre_chunk = make_chunk(existing_doc) + pre_metadata = make_indexing_metadata( + [existing_doc], old_counts=[0], new_counts=[1] + ) + document_index.index(chunks=[pre_chunk], indexing_metadata=pre_metadata) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Now index a batch with the existing doc and a new doc. - chunks = [ - make_chunk(existing_doc, chunk_id=0), - make_chunk(new_doc, chunk_id=0), - ] - metadata = make_indexing_metadata( - [existing_doc, new_doc], old_counts=[1, 0], new_counts=[1, 1] - ) + # Now index a batch with the existing doc and a new doc. + chunks = [ + make_chunk(existing_doc, chunk_id=0), + make_chunk(new_doc, chunk_id=0), + ] + metadata = make_indexing_metadata( + [existing_doc, new_doc], old_counts=[1, 0], new_counts=[1, 1] + ) - # Under test. - results = document_index.index(chunks=chunks, indexing_metadata=metadata) + # Under test. + results = document_index.index(chunks=chunks, indexing_metadata=metadata) - # Postcondition. - result_map = {r.document_id: r.already_existed for r in results} - assert len(result_map) == 2 - assert result_map[existing_doc] is True - assert result_map[new_doc] is False + # Postcondition. + result_map = {r.document_id: r.already_existed for r in results} + assert len(result_map) == 2 + assert result_map[existing_doc] is True + assert result_map[new_doc] is False def test_index_accepts_generator( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ Tests that index() accepts a generator (any iterable), not just a list. """ # Precondition. - for document_index in document_indices: - doc_id = f"test_gen_{uuid.uuid4().hex[:8]}" - metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[3]) + doc_id = f"test_gen_{uuid.uuid4().hex[:8]}" + metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[3]) - def chunk_gen(doc_id: str = doc_id) -> Iterator[DocMetadataAwareIndexChunk]: - for i in range(3): - yield make_chunk(doc_id, chunk_id=i) + def chunk_gen(doc_id: str = doc_id) -> Iterator[DocMetadataAwareIndexChunk]: + for i in range(3): + yield make_chunk(doc_id, chunk_id=i) - # Under test. - results = document_index.index( - chunks=chunk_gen(), indexing_metadata=metadata - ) + # Under test. + results = document_index.index(chunks=chunk_gen(), indexing_metadata=metadata) - # Postcondition. - assert len(results) == 1 - assert results[0].document_id == doc_id - assert results[0].already_existed is False + # Postcondition. + assert len(results) == 1 + assert results[0].document_id == doc_id + assert results[0].already_existed is False def test_mt_cloud_opensearch_index_verification_only_happens_once( self, @@ -327,7 +319,7 @@ class TestDocumentIndexNew: def test_update_changes_boost_across_multiple_docs_in_single_request( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -336,61 +328,60 @@ class TestDocumentIndexNew: doc. """ # Precondition. - for document_index in document_indices: - doc1 = f"test_update_boost_multi_1_{uuid.uuid4().hex[:8]}" - doc2 = f"test_update_boost_multi_2_{uuid.uuid4().hex[:8]}" - chunks = [ - make_chunk(doc1, chunk_id=0), - make_chunk(doc1, chunk_id=1), - make_chunk(doc1, chunk_id=2), - make_chunk(doc2, chunk_id=0), - make_chunk(doc2, chunk_id=1), - ] - metadata = make_indexing_metadata( - [doc1, doc2], old_counts=[0, 0], new_counts=[3, 2] - ) - document_index.index(chunks=chunks, indexing_metadata=metadata) + doc1 = f"test_update_boost_multi_1_{uuid.uuid4().hex[:8]}" + doc2 = f"test_update_boost_multi_2_{uuid.uuid4().hex[:8]}" + chunks = [ + make_chunk(doc1, chunk_id=0), + make_chunk(doc1, chunk_id=1), + make_chunk(doc1, chunk_id=2), + make_chunk(doc2, chunk_id=0), + make_chunk(doc2, chunk_id=1), + ] + metadata = make_indexing_metadata( + [doc1, doc2], old_counts=[0, 0], new_counts=[3, 2] + ) + document_index.index(chunks=chunks, indexing_metadata=metadata) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Under test. - update_request = MetadataUpdateRequest( - document_ids=[doc1, doc2], - doc_id_to_chunk_cnt={doc1: 3, doc2: 2}, - boost=7, - ) - document_index.update([update_request]) + # Under test. + update_request = MetadataUpdateRequest( + document_ids=[doc1, doc2], + doc_id_to_chunk_cnt={doc1: 3, doc2: 2}, + boost=7, + ) + document_index.update([update_request]) - # Postcondition. Poll until the eventually-consistent indexes - # reflect the updates rather than racing a fixed sleep against - # OpenSearch's ~1s refresh window. - filters = IndexFilters( - access_control_list=[PUBLIC_DOC_PAT], - tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, - ) - retrieved_doc1 = _retrieve_chunks_with_expected_boost( - document_index=document_index, - document_id=doc1, - expected_chunk_count=3, - expected_boost=7, - filters=filters, - ) - retrieved_doc2 = _retrieve_chunks_with_expected_boost( - document_index=document_index, - document_id=doc2, - expected_chunk_count=2, - expected_boost=7, - filters=filters, - ) - assert len(retrieved_doc1) == 3 - assert len(retrieved_doc2) == 2 - for chunk in retrieved_doc1 + retrieved_doc2: - assert chunk.boost == 7 + # Postcondition. Poll until the eventually-consistent indexes + # reflect the updates rather than racing a fixed sleep against + # OpenSearch's ~1s refresh window. + filters = IndexFilters( + access_control_list=[PUBLIC_DOC_PAT], + tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ) + retrieved_doc1 = _retrieve_chunks_with_expected_boost( + document_index=document_index, + document_id=doc1, + expected_chunk_count=3, + expected_boost=7, + filters=filters, + ) + retrieved_doc2 = _retrieve_chunks_with_expected_boost( + document_index=document_index, + document_id=doc2, + expected_chunk_count=2, + expected_boost=7, + filters=filters, + ) + assert len(retrieved_doc1) == 3 + assert len(retrieved_doc2) == 2 + for chunk in retrieved_doc1 + retrieved_doc2: + assert chunk.boost == 7 def test_update_applies_each_request_independently( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -398,65 +389,64 @@ class TestDocumentIndexNew: each apply their own fields to their own documents. """ # Precondition. - for document_index in document_indices: - doc1 = f"test_update_indep_1_{uuid.uuid4().hex[:8]}" - doc2 = f"test_update_indep_2_{uuid.uuid4().hex[:8]}" - chunks = [ - make_chunk(doc1, chunk_id=0), - make_chunk(doc1, chunk_id=1), - make_chunk(doc2, chunk_id=0), - ] - metadata = make_indexing_metadata( - [doc1, doc2], old_counts=[0, 0], new_counts=[2, 1] - ) - document_index.index(chunks=chunks, indexing_metadata=metadata) + doc1 = f"test_update_indep_1_{uuid.uuid4().hex[:8]}" + doc2 = f"test_update_indep_2_{uuid.uuid4().hex[:8]}" + chunks = [ + make_chunk(doc1, chunk_id=0), + make_chunk(doc1, chunk_id=1), + make_chunk(doc2, chunk_id=0), + ] + metadata = make_indexing_metadata( + [doc1, doc2], old_counts=[0, 0], new_counts=[2, 1] + ) + document_index.index(chunks=chunks, indexing_metadata=metadata) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Under test - two separate requests, each updating a different doc. - req1 = MetadataUpdateRequest( - document_ids=[doc1], - doc_id_to_chunk_cnt={doc1: 2}, - boost=3, - ) - req2 = MetadataUpdateRequest( - document_ids=[doc2], - doc_id_to_chunk_cnt={doc2: 1}, - boost=9, - ) - document_index.update([req1, req2]) + # Under test - two separate requests, each updating a different doc. + req1 = MetadataUpdateRequest( + document_ids=[doc1], + doc_id_to_chunk_cnt={doc1: 2}, + boost=3, + ) + req2 = MetadataUpdateRequest( + document_ids=[doc2], + doc_id_to_chunk_cnt={doc2: 1}, + boost=9, + ) + document_index.update([req1, req2]) - # Postcondition. Poll until the eventually-consistent indexes - # reflect the updates rather than racing a fixed sleep against - # OpenSearch's ~1s refresh window. - filters = IndexFilters( - access_control_list=[PUBLIC_DOC_PAT], - tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, - ) - retrieved_doc1 = _retrieve_chunks_with_expected_boost( - document_index=document_index, - document_id=doc1, - expected_chunk_count=2, - expected_boost=3, - filters=filters, - ) - retrieved_doc2 = _retrieve_chunks_with_expected_boost( - document_index=document_index, - document_id=doc2, - expected_chunk_count=1, - expected_boost=9, - filters=filters, - ) - assert len(retrieved_doc1) == 2 - assert len(retrieved_doc2) == 1 - for chunk in retrieved_doc1: - assert chunk.boost == 3 - assert retrieved_doc2[0].boost == 9 + # Postcondition. Poll until the eventually-consistent indexes + # reflect the updates rather than racing a fixed sleep against + # OpenSearch's ~1s refresh window. + filters = IndexFilters( + access_control_list=[PUBLIC_DOC_PAT], + tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ) + retrieved_doc1 = _retrieve_chunks_with_expected_boost( + document_index=document_index, + document_id=doc1, + expected_chunk_count=2, + expected_boost=3, + filters=filters, + ) + retrieved_doc2 = _retrieve_chunks_with_expected_boost( + document_index=document_index, + document_id=doc2, + expected_chunk_count=1, + expected_boost=9, + filters=filters, + ) + assert len(retrieved_doc1) == 2 + assert len(retrieved_doc2) == 1 + for chunk in retrieved_doc1: + assert chunk.boost == 3 + assert retrieved_doc2[0].boost == 9 def test_update_access_revokes_public( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -465,68 +455,67 @@ class TestDocumentIndexNew: users in its new ACL. """ # Precondition. - for document_index in document_indices: - doc_id = f"test_update_revoke_public_{uuid.uuid4().hex[:8]}" - user_email = "revoke_public_user@example.com" - chunks = [make_chunk(doc_id, chunk_id=0), make_chunk(doc_id, chunk_id=1)] - metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[2]) - document_index.index(chunks=chunks, indexing_metadata=metadata) + doc_id = f"test_update_revoke_public_{uuid.uuid4().hex[:8]}" + user_email = "revoke_public_user@example.com" + chunks = [make_chunk(doc_id, chunk_id=0), make_chunk(doc_id, chunk_id=1)] + metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[2]) + document_index.index(chunks=chunks, indexing_metadata=metadata) - public_filters = IndexFilters( - access_control_list=[PUBLIC_DOC_PAT], - tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, - ) - _retrieve_chunks_with_expected_boost( - document_index=document_index, - document_id=doc_id, - expected_chunk_count=2, - expected_boost=0, + public_filters = IndexFilters( + access_control_list=[PUBLIC_DOC_PAT], + tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ) + _retrieve_chunks_with_expected_boost( + document_index=document_index, + document_id=doc_id, + expected_chunk_count=2, + expected_boost=0, + filters=public_filters, + ) + + # Under test. + private_access = DocumentAccess.build( + user_emails=[user_email], + user_groups=[], + external_user_emails=[], + external_user_group_ids=[], + is_public=False, + ) + document_index.update( + [ + MetadataUpdateRequest( + document_ids=[doc_id], + doc_id_to_chunk_cnt={doc_id: 2}, + access=private_access, + ) + ] + ) + + # Postcondition. + deadline = time.time() + 10.0 + public_retrieved: list[InferenceChunk] = [] + while time.time() < deadline: + public_retrieved = document_index.id_based_retrieval( + chunk_requests=[DocumentSectionRequest(document_id=doc_id)], filters=public_filters, ) + if not public_retrieved: + break + time.sleep(0.25) + assert public_retrieved == [] - # Under test. - private_access = DocumentAccess.build( - user_emails=[user_email], - user_groups=[], - external_user_emails=[], - external_user_group_ids=[], - is_public=False, - ) - document_index.update( - [ - MetadataUpdateRequest( - document_ids=[doc_id], - doc_id_to_chunk_cnt={doc_id: 2}, - access=private_access, - ) - ] - ) - - # Postcondition. - deadline = time.time() + 10.0 - public_retrieved: list[InferenceChunk] = [] - while time.time() < deadline: - public_retrieved = document_index.id_based_retrieval( - chunk_requests=[DocumentSectionRequest(document_id=doc_id)], - filters=public_filters, - ) - if not public_retrieved: - break - time.sleep(0.25) - assert public_retrieved == [] - - user_retrieved = document_index.id_based_retrieval( - chunk_requests=[DocumentSectionRequest(document_id=doc_id)], - filters=IndexFilters( - access_control_list=[PUBLIC_DOC_PAT, prefix_user_email(user_email)], - tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, - ), - ) - assert len(user_retrieved) == 2 + user_retrieved = document_index.id_based_retrieval( + chunk_requests=[DocumentSectionRequest(document_id=doc_id)], + filters=IndexFilters( + access_control_list=[PUBLIC_DOC_PAT, prefix_user_email(user_email)], + tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ), + ) + assert len(user_retrieved) == 2 def test_update_with_no_fields_does_not_modify_chunks( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -534,41 +523,40 @@ class TestDocumentIndexNew: no-op and the chunks remain retrievable with their original values. """ # Precondition. - for document_index in document_indices: - doc_id = f"test_update_noop_{uuid.uuid4().hex[:8]}" - chunks = [make_chunk(doc_id, chunk_id=0), make_chunk(doc_id, chunk_id=1)] - metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[2]) - document_index.index(chunks=chunks, indexing_metadata=metadata) + doc_id = f"test_update_noop_{uuid.uuid4().hex[:8]}" + chunks = [make_chunk(doc_id, chunk_id=0), make_chunk(doc_id, chunk_id=1)] + metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[2]) + document_index.index(chunks=chunks, indexing_metadata=metadata) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Under test - no fields set. - update_request = MetadataUpdateRequest( - document_ids=[doc_id], - doc_id_to_chunk_cnt={doc_id: 2}, - ) - document_index.update([update_request]) + # Under test - no fields set. + update_request = MetadataUpdateRequest( + document_ids=[doc_id], + doc_id_to_chunk_cnt={doc_id: 2}, + ) + document_index.update([update_request]) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Postcondition - chunks still retrievable with their default boost. - filters = IndexFilters( - access_control_list=[PUBLIC_DOC_PAT], - tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, - ) - retrieved = document_index.id_based_retrieval( - chunk_requests=[DocumentSectionRequest(document_id=doc_id)], - filters=filters, - ) - assert len(retrieved) == 2 - for chunk in retrieved: - assert chunk.boost == 0 + # Postcondition - chunks still retrievable with their default boost. + filters = IndexFilters( + access_control_list=[PUBLIC_DOC_PAT], + tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ) + retrieved = document_index.id_based_retrieval( + chunk_requests=[DocumentSectionRequest(document_id=doc_id)], + filters=filters, + ) + assert len(retrieved) == 2 + for chunk in retrieved: + assert chunk.boost == 0 def test_update_skips_doc_with_unknown_chunk_count( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -578,50 +566,47 @@ class TestDocumentIndexNew: """ # Precondition - index one real doc; pair it with a phantom doc whose # chunk count is unknown (absent from doc_id_to_chunk_cnt -> -1). - for document_index in document_indices: - real_doc = f"test_update_unknown_real_{uuid.uuid4().hex[:8]}" - phantom_doc = f"test_update_unknown_phantom_{uuid.uuid4().hex[:8]}" - chunks = [ - make_chunk(real_doc, chunk_id=0), - make_chunk(real_doc, chunk_id=1), - ] - metadata = make_indexing_metadata( - [real_doc], old_counts=[0], new_counts=[2] - ) - document_index.index(chunks=chunks, indexing_metadata=metadata) + real_doc = f"test_update_unknown_real_{uuid.uuid4().hex[:8]}" + phantom_doc = f"test_update_unknown_phantom_{uuid.uuid4().hex[:8]}" + chunks = [ + make_chunk(real_doc, chunk_id=0), + make_chunk(real_doc, chunk_id=1), + ] + metadata = make_indexing_metadata([real_doc], old_counts=[0], new_counts=[2]) + document_index.index(chunks=chunks, indexing_metadata=metadata) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Under test - phantom_doc has no entry in doc_id_to_chunk_cnt, so - # its chunk count resolves to -1 (unknown). This must not raise. - update_request = MetadataUpdateRequest( - document_ids=[real_doc, phantom_doc], - doc_id_to_chunk_cnt={real_doc: 2}, - boost=5, - ) - document_index.update([update_request]) + # Under test - phantom_doc has no entry in doc_id_to_chunk_cnt, so + # its chunk count resolves to -1 (unknown). This must not raise. + update_request = MetadataUpdateRequest( + document_ids=[real_doc, phantom_doc], + doc_id_to_chunk_cnt={real_doc: 2}, + boost=5, + ) + document_index.update([update_request]) - # Postcondition - the real doc is still updated despite the phantom - # doc being skipped. - filters = IndexFilters( - access_control_list=[PUBLIC_DOC_PAT], - tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, - ) - retrieved = _retrieve_chunks_with_expected_boost( - document_index=document_index, - document_id=real_doc, - expected_chunk_count=2, - expected_boost=5, - filters=filters, - ) - assert len(retrieved) == 2 - for chunk in retrieved: - assert chunk.boost == 5 + # Postcondition - the real doc is still updated despite the phantom + # doc being skipped. + filters = IndexFilters( + access_control_list=[PUBLIC_DOC_PAT], + tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ) + retrieved = _retrieve_chunks_with_expected_boost( + document_index=document_index, + document_id=real_doc, + expected_chunk_count=2, + expected_boost=5, + filters=filters, + ) + assert len(retrieved) == 2 + for chunk in retrieved: + assert chunk.boost == 5 def test_update_skips_doc_with_zero_chunk_count( self, - document_indices: list[DocumentIndexNew], + document_index: DocumentIndex, tenant_context: None, # noqa: ARG002 ) -> None: """ @@ -631,43 +616,40 @@ class TestDocumentIndexNew: """ # Precondition - index one real doc; pair it with a phantom doc whose # chunk count is explicitly 0. - for document_index in document_indices: - real_doc = f"test_update_zero_real_{uuid.uuid4().hex[:8]}" - phantom_doc = f"test_update_zero_phantom_{uuid.uuid4().hex[:8]}" - chunks = [ - make_chunk(real_doc, chunk_id=0), - make_chunk(real_doc, chunk_id=1), - ] - metadata = make_indexing_metadata( - [real_doc], old_counts=[0], new_counts=[2] - ) - document_index.index(chunks=chunks, indexing_metadata=metadata) + real_doc = f"test_update_zero_real_{uuid.uuid4().hex[:8]}" + phantom_doc = f"test_update_zero_phantom_{uuid.uuid4().hex[:8]}" + chunks = [ + make_chunk(real_doc, chunk_id=0), + make_chunk(real_doc, chunk_id=1), + ] + metadata = make_indexing_metadata([real_doc], old_counts=[0], new_counts=[2]) + document_index.index(chunks=chunks, indexing_metadata=metadata) - # Allow OpenSearch refresh interval to settle. - time.sleep(1) + # Allow OpenSearch refresh interval to settle. + time.sleep(1) - # Under test - phantom_doc has a chunk count of 0. This must not - # raise. - update_request = MetadataUpdateRequest( - document_ids=[real_doc, phantom_doc], - doc_id_to_chunk_cnt={real_doc: 2, phantom_doc: 0}, - boost=6, - ) - document_index.update([update_request]) + # Under test - phantom_doc has a chunk count of 0. This must not + # raise. + update_request = MetadataUpdateRequest( + document_ids=[real_doc, phantom_doc], + doc_id_to_chunk_cnt={real_doc: 2, phantom_doc: 0}, + boost=6, + ) + document_index.update([update_request]) - # Postcondition - the real doc is still updated despite the phantom - # doc being skipped. - filters = IndexFilters( - access_control_list=[PUBLIC_DOC_PAT], - tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, - ) - retrieved = _retrieve_chunks_with_expected_boost( - document_index=document_index, - document_id=real_doc, - expected_chunk_count=2, - expected_boost=6, - filters=filters, - ) - assert len(retrieved) == 2 - for chunk in retrieved: - assert chunk.boost == 6 + # Postcondition - the real doc is still updated despite the phantom + # doc being skipped. + filters = IndexFilters( + access_control_list=[PUBLIC_DOC_PAT], + tenant_id=POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ) + retrieved = _retrieve_chunks_with_expected_boost( + document_index=document_index, + document_id=real_doc, + expected_chunk_count=2, + expected_boost=6, + filters=filters, + ) + assert len(retrieved) == 2 + for chunk in retrieved: + assert chunk.boost == 6 diff --git a/backend/tests/external_dependency_unit/document_index/test_index_pair_search_and_permissions.py b/backend/tests/external_dependency_unit/document_index/test_index_pair_search_and_permissions.py new file mode 100644 index 0000000000..175a1d7e60 --- /dev/null +++ b/backend/tests/external_dependency_unit/document_index/test_index_pair_search_and_permissions.py @@ -0,0 +1,471 @@ +"""External dependency tests for the search and permission paths of the +document index returned by `get_default_document_index`. + +Every write and read goes through an `OpenSearchIndexPair`: one primary index +and, during a re-index, one secondary index. These tests run against a real +OpenSearch and check that: +- metadata updates (access, document sets, hidden) reach both indices and that + retrieval enforces them, +- deletes reach both indices, and writes and reads use the primary only, +- search-result expansion finds the neighbors of documents whose IDs contain + characters the former Vespa backend rewrote (for example `'`). + +These tests assume OpenSearch is running. +""" + +import time +import uuid +from collections.abc import Callable, Generator +from typing import TypeVar +from unittest.mock import MagicMock + +import pytest + +from onyx.access.models import DocumentAccess +from onyx.access.utils import prefix_user_email +from onyx.configs.constants import PUBLIC_DOC_PAT +from onyx.context.search.models import IndexFilters, InferenceChunk, InferenceSection +from onyx.db.enums import VectorQuantization +from onyx.document_index.interfaces import ( + DocumentSectionRequest, + MetadataUpdateRequest, + TenantState, +) +from onyx.document_index.opensearch.client import ( + OpenSearchIndexClient, + wait_for_opensearch_with_timeout, +) +from onyx.document_index.opensearch.opensearch_document_index import ( + OpenSearchDocumentIndex, + OpenSearchIndexPair, +) +from onyx.indexing.models import DocMetadataAwareIndexChunk +from onyx.tools.tool_implementations.search.search_utils import ( + _retrieve_adjacent_chunks, + expand_section_with_context, +) +from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE +from tests.external_dependency_unit.document_index.conftest import ( + EMBEDDING_DIM, + make_chunk, + make_indexing_metadata, +) + +T = TypeVar("T") + +TENANT_ID = POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE +OWNER = "owner@example.com" +OTHER_USER = "other@example.com" + + +# ------------------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------------------ + + +def _wait_for( + read: Callable[[], T], done: Callable[[T], bool], timeout_s: float = 10.0 +) -> T: + """Polls `read` until `done` holds. OpenSearch makes writes visible to + search only after a refresh.""" + deadline = time.monotonic() + timeout_s + value = read() + while not done(value) and time.monotonic() < deadline: + time.sleep(0.25) + value = read() + return value + + +def _filters(acl: list[str], document_sets: list[str] | None = None) -> IndexFilters: + return IndexFilters( + access_control_list=acl, + document_set=document_sets, + tenant_id=TENANT_ID, + ) + + +def _owner_filters(document_sets: list[str] | None = None) -> IndexFilters: + return _filters([PUBLIC_DOC_PAT, prefix_user_email(OWNER)], document_sets) + + +def _other_user_filters() -> IndexFilters: + return _filters([PUBLIC_DOC_PAT, prefix_user_email(OTHER_USER)]) + + +def _private_access(user_email: str) -> DocumentAccess: + return DocumentAccess.build( + user_emails=[user_email], + user_groups=[], + external_user_emails=[], + external_user_group_ids=[], + is_public=False, + ) + + +def _chunks( + doc_id: str, count: int, access: DocumentAccess | None = None +) -> list[DocMetadataAwareIndexChunk]: + chunks = [] + for chunk_id in range(count): + chunk = make_chunk( + doc_id, + chunk_id=chunk_id, + content=f"quarterly zebra roadmap part {chunk_id}", + ) + if access is not None: + chunk.access = access + chunks.append(chunk) + return chunks + + +def _index( + document_index: OpenSearchDocumentIndex | OpenSearchIndexPair, + chunks: list[DocMetadataAwareIndexChunk], +) -> None: + doc_id = chunks[0].source_document.id + document_index.index( + chunks=chunks, + indexing_metadata=make_indexing_metadata( + [doc_id], old_counts=[0], new_counts=[len(chunks)] + ), + ) + + +def _chunk_ids( + document_index: OpenSearchDocumentIndex | OpenSearchIndexPair, + doc_id: str, + filters: IndexFilters, +) -> list[int]: + chunks = document_index.id_based_retrieval( + chunk_requests=[DocumentSectionRequest(document_id=doc_id)], + filters=filters, + batch_retrieval=True, + ) + return sorted(chunk.chunk_id for chunk in chunks) + + +def _wait_for_chunk_ids( + document_index: OpenSearchDocumentIndex | OpenSearchIndexPair, + doc_id: str, + filters: IndexFilters, + expected: list[int], +) -> list[int]: + return _wait_for( + lambda: _chunk_ids(document_index, doc_id, filters), + lambda ids: ids == expected, + ) + + +def _new_index(name: str) -> OpenSearchDocumentIndex: + document_index = OpenSearchDocumentIndex( + tenant_state=TenantState(tenant_id=TENANT_ID, multitenant=False), + index_name=name, + embedding_dim=EMBEDDING_DIM, + vector_quantization=VectorQuantization.NONE, + ) + document_index.verify_and_create_index_if_necessary(embedding_dim=EMBEDDING_DIM) + return document_index + + +# ------------------------------------------------------------------------------ +# Fixtures +# ------------------------------------------------------------------------------ + + +@pytest.fixture(scope="module") +def primary_and_secondary( + tenant_context: None, # noqa: ARG001 +) -> Generator[tuple[OpenSearchDocumentIndex, OpenSearchDocumentIndex], None, None]: + if not wait_for_opensearch_with_timeout(): + pytest.fail("OpenSearch is not available.") + suffix = uuid.uuid4().hex[:8] + names = [f"test_pair_primary_{suffix}", f"test_pair_secondary_{suffix}"] + primary, secondary = (_new_index(name) for name in names) + try: + yield primary, secondary + finally: + for name in names: + OpenSearchIndexClient(index_name=name).delete_index() + + +@pytest.fixture(scope="module") +def pair( + primary_and_secondary: tuple[OpenSearchDocumentIndex, OpenSearchDocumentIndex], +) -> OpenSearchIndexPair: + primary, secondary = primary_and_secondary + return OpenSearchIndexPair( + primary=primary, + secondary=secondary, + secondary_embedding_dim=EMBEDDING_DIM, + ) + + +@pytest.fixture(scope="module") +def primary( + primary_and_secondary: tuple[OpenSearchDocumentIndex, OpenSearchDocumentIndex], +) -> OpenSearchDocumentIndex: + return primary_and_secondary[0] + + +@pytest.fixture(scope="module") +def secondary( + primary_and_secondary: tuple[OpenSearchDocumentIndex, OpenSearchDocumentIndex], +) -> OpenSearchDocumentIndex: + return primary_and_secondary[1] + + +def _doc_id(label: str) -> str: + return f"{label}_{uuid.uuid4().hex[:8]}" + + +def _index_in_both( + pair: OpenSearchIndexPair, + secondary: OpenSearchDocumentIndex, + chunks: list[DocMetadataAwareIndexChunk], +) -> None: + """Writes go to the primary only; a re-index fills the secondary on its own, + which is what the direct secondary write stands in for.""" + _index(pair, chunks) + _index(secondary, chunks) + + +# ------------------------------------------------------------------------------ +# Writes and reads +# ------------------------------------------------------------------------------ + + +class TestPairRouting: + def test_index_writes_the_primary_only( + self, + pair: OpenSearchIndexPair, + primary: OpenSearchDocumentIndex, + secondary: OpenSearchDocumentIndex, + ) -> None: + doc_id = _doc_id("routing_index") + _index(pair, _chunks(doc_id, 2)) + + assert _wait_for_chunk_ids(primary, doc_id, _owner_filters(), [0, 1]) == [ + 0, + 1, + ] + assert _chunk_ids(secondary, doc_id, _owner_filters()) == [] + + def test_retrieval_reads_the_primary_only( + self, + pair: OpenSearchIndexPair, + secondary: OpenSearchDocumentIndex, + ) -> None: + doc_id = _doc_id("routing_read") + _index(secondary, _chunks(doc_id, 2)) + _wait_for_chunk_ids(secondary, doc_id, _owner_filters(), [0, 1]) + + assert _chunk_ids(pair, doc_id, _owner_filters()) == [] + + def test_delete_removes_the_document_from_both_indices( + self, + pair: OpenSearchIndexPair, + primary: OpenSearchDocumentIndex, + secondary: OpenSearchDocumentIndex, + ) -> None: + doc_id = _doc_id("routing_delete") + _index_in_both(pair, secondary, _chunks(doc_id, 3)) + _wait_for_chunk_ids(primary, doc_id, _owner_filters(), [0, 1, 2]) + _wait_for_chunk_ids(secondary, doc_id, _owner_filters(), [0, 1, 2]) + + deleted = pair.delete(doc_id, chunk_count=3) + + assert deleted == 6 + for document_index in (primary, secondary): + assert ( + _wait_for_chunk_ids(document_index, doc_id, _owner_filters(), []) == [] + ) + + +# ------------------------------------------------------------------------------ +# Permission and metadata updates +# ------------------------------------------------------------------------------ + + +class TestPairMetadataUpdates: + def test_access_update_reaches_both_indices_and_is_enforced( + self, + pair: OpenSearchIndexPair, + primary: OpenSearchDocumentIndex, + secondary: OpenSearchDocumentIndex, + ) -> None: + """Moving a document from one user to another must revoke the first + user's access and grant the second user access in both indices.""" + doc_id = _doc_id("acl_move") + _index_in_both(pair, secondary, _chunks(doc_id, 2, _private_access(OWNER))) + for document_index in (primary, secondary): + _wait_for_chunk_ids(document_index, doc_id, _owner_filters(), [0, 1]) + assert _chunk_ids(document_index, doc_id, _other_user_filters()) == [] + + pair.update( + [ + MetadataUpdateRequest( + document_ids=[doc_id], + doc_id_to_chunk_cnt={doc_id: 2}, + access=_private_access(OTHER_USER), + ) + ] + ) + + for document_index in (primary, secondary): + assert ( + _wait_for_chunk_ids(document_index, doc_id, _owner_filters(), []) == [] + ) + assert _wait_for_chunk_ids( + document_index, doc_id, _other_user_filters(), [0, 1] + ) == [0, 1] + + def test_access_update_is_enforced_by_keyword_search( + self, + pair: OpenSearchIndexPair, + secondary: OpenSearchDocumentIndex, + ) -> None: + doc_id = _doc_id("acl_keyword") + _index_in_both(pair, secondary, _chunks(doc_id, 1, _private_access(OWNER))) + + def _hit_ids(filters: IndexFilters) -> set[str]: + chunks: list[InferenceChunk] = pair.keyword_retrieval( + query="quarterly zebra roadmap", + filters=filters, + num_to_retrieve=50, + ) + return {chunk.document_id for chunk in chunks} + + assert doc_id in _wait_for( + lambda: _hit_ids(_owner_filters()), lambda ids: doc_id in ids + ) + assert doc_id not in _hit_ids(_other_user_filters()) + + pair.update( + [ + MetadataUpdateRequest( + document_ids=[doc_id], + doc_id_to_chunk_cnt={doc_id: 1}, + access=_private_access(OTHER_USER), + ) + ] + ) + + assert doc_id in _wait_for( + lambda: _hit_ids(_other_user_filters()), lambda ids: doc_id in ids + ) + assert doc_id not in _wait_for( + lambda: _hit_ids(_owner_filters()), lambda ids: doc_id not in ids + ) + + def test_document_set_update_reaches_both_indices( + self, + pair: OpenSearchIndexPair, + primary: OpenSearchDocumentIndex, + secondary: OpenSearchDocumentIndex, + ) -> None: + doc_id = _doc_id("doc_set") + _index_in_both(pair, secondary, _chunks(doc_id, 2)) + engineering = _owner_filters(document_sets=["Engineering"]) + for document_index in (primary, secondary): + _wait_for_chunk_ids(document_index, doc_id, _owner_filters(), [0, 1]) + assert _chunk_ids(document_index, doc_id, engineering) == [] + + pair.update( + [ + MetadataUpdateRequest( + document_ids=[doc_id], + doc_id_to_chunk_cnt={doc_id: 2}, + document_sets={"Engineering"}, + ) + ] + ) + + legal = _owner_filters(document_sets=["Legal"]) + for document_index in (primary, secondary): + assert _wait_for_chunk_ids(document_index, doc_id, engineering, [0, 1]) == [ + 0, + 1, + ] + assert _chunk_ids(document_index, doc_id, legal) == [] + + def test_hidden_update_reaches_both_indices( + self, + pair: OpenSearchIndexPair, + primary: OpenSearchDocumentIndex, + secondary: OpenSearchDocumentIndex, + ) -> None: + doc_id = _doc_id("hidden") + _index_in_both(pair, secondary, _chunks(doc_id, 2)) + for document_index in (primary, secondary): + _wait_for_chunk_ids(document_index, doc_id, _owner_filters(), [0, 1]) + + pair.update( + [ + MetadataUpdateRequest( + document_ids=[doc_id], + doc_id_to_chunk_cnt={doc_id: 2}, + hidden=True, + ) + ] + ) + + for document_index in (primary, secondary): + assert ( + _wait_for_chunk_ids(document_index, doc_id, _owner_filters(), []) == [] + ) + + +# ------------------------------------------------------------------------------ +# Search-result expansion +# ------------------------------------------------------------------------------ + + +class TestSectionExpansion: + @pytest.mark.parametrize( + "doc_id", + [ + "plain_document", + "https://example.com/o'brien's notes?page=1", + ], + ) + def test_expansion_finds_neighbors( + self, + pair: OpenSearchIndexPair, + doc_id: str, + ) -> None: + doc_id = f"{doc_id}#{uuid.uuid4().hex[:8]}" + _index(pair, _chunks(doc_id, 6)) + _wait_for_chunk_ids(pair, doc_id, _owner_filters(), list(range(6))) + (center,) = pair.id_based_retrieval( + chunk_requests=[ + DocumentSectionRequest( + document_id=doc_id, min_chunk_ind=2, max_chunk_ind=2 + ) + ], + filters=_owner_filters(), + batch_retrieval=True, + ) + section = InferenceSection( + center_chunk=center, chunks=[center], combined_content=center.content + ) + + above, below = _retrieve_adjacent_chunks( + section=section, + document_index=pair, + num_chunks_above=2, + num_chunks_below=2, + ) + assert [chunk.chunk_id for chunk in above] == [0, 1] + assert [chunk.chunk_id for chunk in below] == [3, 4] + + llm = MagicMock() + expanded = expand_section_with_context( + section=section, + user_query="roadmap", + llm=llm, + document_index=pair, + expand_override=True, + ) + assert expanded is not None + assert [chunk.chunk_id for chunk in expanded.chunks] == list(range(6)) + llm.invoke.assert_not_called() diff --git a/backend/tests/external_dependency_unit/document_index/test_multi_source_opensearch_e2e.py b/backend/tests/external_dependency_unit/document_index/test_multi_source_opensearch_e2e.py index 66737101de..358074e89b 100644 --- a/backend/tests/external_dependency_unit/document_index/test_multi_source_opensearch_e2e.py +++ b/backend/tests/external_dependency_unit/document_index/test_multi_source_opensearch_e2e.py @@ -3,7 +3,7 @@ import uuid from onyx.configs.constants import PUBLIC_DOC_PAT, DocumentSource from onyx.context.search.models import IndexFilters, InferenceChunk -from onyx.document_index.interfaces_new import ( +from onyx.document_index.interfaces import ( DocumentSectionRequest, MetadataUpdateRequest, ) diff --git a/backend/tests/external_dependency_unit/document_index/test_vector_quantization.py b/backend/tests/external_dependency_unit/document_index/test_vector_quantization.py index ab667861f3..f418d77cb4 100644 --- a/backend/tests/external_dependency_unit/document_index/test_vector_quantization.py +++ b/backend/tests/external_dependency_unit/document_index/test_vector_quantization.py @@ -15,7 +15,7 @@ from onyx.configs.constants import PUBLIC_DOC_PAT from onyx.context.search.enums import QueryType from onyx.context.search.models import IndexFilters from onyx.db.enums import VectorQuantization -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.client import ( OpenSearchIndexClient, wait_for_opensearch_with_timeout, diff --git a/backend/tests/external_dependency_unit/document_index/test_verify_write_blocked.py b/backend/tests/external_dependency_unit/document_index/test_verify_write_blocked.py index f9ff79769a..a51def50da 100644 --- a/backend/tests/external_dependency_unit/document_index/test_verify_write_blocked.py +++ b/backend/tests/external_dependency_unit/document_index/test_verify_write_blocked.py @@ -5,7 +5,7 @@ Regression tests for api-server pods crash-looping on startup while the index was read_only_allow_delete-blocked: the mapping refresh is a metadata write, so it is rejected while the block is active. verify_and_create_index_if_necessary still raises (callers like embedding-model swaps must not silently continue); -the tolerant call sites — startup's setup_document_indices and the multitenant +the tolerant call sites — startup's setup_document_index and the multitenant DocumentIndex init — catch the block error and proceed degraded. """ @@ -15,7 +15,7 @@ from unittest.mock import patch import pytest from onyx.db.enums import VectorQuantization -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch import ( opensearch_document_index as opensearch_document_index_module, ) @@ -28,7 +28,7 @@ from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, ) from onyx.indexing.models import IndexingSetting -from onyx.setup import setup_document_indices +from onyx.setup import setup_document_index from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from tests.external_dependency_unit.document_index.conftest import EMBEDDING_DIM @@ -68,15 +68,15 @@ def test_verify_raises_typed_error_under_write_block( assert is_cluster_block_error(cause) -def test_setup_document_indices_succeeds_under_write_block( +def test_setup_document_index_succeeds_under_write_block( write_blocked_index: OpenSearchDocumentIndex, ) -> None: """Startup must survive an existing, readable index that is merely write-blocked instead of crash-looping.""" index_setting = IndexingSetting.model_construct(model_dim=EMBEDDING_DIM) - assert setup_document_indices( - document_indices=[write_blocked_index], + assert setup_document_index( + document_index=write_blocked_index, index_setting=index_setting, num_attempts=1, ) diff --git a/backend/tests/external_dependency_unit/full_setup.py b/backend/tests/external_dependency_unit/full_setup.py index df281c4fa9..b1784e1340 100644 --- a/backend/tests/external_dependency_unit/full_setup.py +++ b/backend/tests/external_dependency_unit/full_setup.py @@ -6,13 +6,10 @@ from typing import Optional from onyx.db.engine.sql_engine import SqlEngine, get_session_with_current_tenant from onyx.db.search_settings import get_active_search_settings -from onyx.document_index.factory import ( - get_all_document_indices, - get_default_document_index, -) +from onyx.document_index.factory import 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, setup_postgres +from onyx.setup import setup_document_index, 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 @@ -22,15 +19,14 @@ _SETUP_COMPLETE: bool = False def ensure_full_deployment_setup( tenant_id: Optional[str] = None, - opensearch_available: bool = False, ) -> None: """Initialize test environment to mirror a real deployment, on demand. - Initializes DB engine and sets tenant context - Skips model warm-ups during setup - - Runs setup_onyx (Postgres defaults, Vespa indices) + - Runs setup_postgres (Postgres defaults) - Initializes file store (best-effort) - - Ensures Vespa indices exist + - Ensures the document index exists """ global _SETUP_COMPLETE if _SETUP_COMPLETE: @@ -62,29 +58,19 @@ def ensure_full_deployment_setup( except Exception: pass - # Also ensure indices exist explicitly (no-op if already created) + # Also ensure the index exists explicitly (no-op if already created) with get_session_with_current_tenant() as db_session: active = get_active_search_settings(db_session) - if opensearch_available: - # We use this special bool here instead of just relying on - # ENABLE_OPENSEARCH_INDEXING_FOR_ONYX because not all testing - # infra is configured for OpenSearch. - document_indices = get_all_document_indices( + ok = setup_document_index( + document_index=get_default_document_index( active.primary, active.secondary - ) - else: - document_indices = [ - get_default_document_index( - active.primary, active.secondary, db_session - ) - ] - ok = setup_document_indices( - document_indices=document_indices, + ), index_setting=IndexingSetting.from_db_model(active.primary), ) if not ok: raise RuntimeError( - "Vespa did not initialize within the specified timeout." + "The document index did not initialize within the specified " + "timeout." ) _SETUP_COMPLETE = True 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 f3a43d5746..5a3d3923aa 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 @@ -182,12 +182,9 @@ class TestDeleteIngestionDoc: _index_doc(db_session, doc, attempt_metadata) assert get_filerecord(db_session, file_id) is not None - # Patch out Vespa — we're testing the file cleanup, not the document - # index integration. - with patch( - "onyx.server.onyx_api.ingestion.get_all_document_indices", - return_value=[], - ): + # Patch out the document index — we're testing the file cleanup, not + # the document index integration. + with patch("onyx.server.onyx_api.ingestion.get_default_document_index"): delete_ingestion_doc( document_id=doc.id, # a real admin: the body now feeds this to the GATE 2 cc_pair check @@ -217,11 +214,10 @@ class TestDocumentByCcPairCleanupTask: assert get_filerecord(db_session, file_id) is not None - # Patch out Vespa interaction — no chunks were ever written, and we're + # Patch out the document index — no chunks were ever written, and we're # not testing the document index here. with patch( - "onyx.background.celery.tasks.shared.tasks.get_all_document_indices", - return_value=[], + "onyx.background.celery.tasks.shared.tasks.get_default_document_index" ): result = document_by_cc_pair_cleanup_task.apply( args=( @@ -261,8 +257,7 @@ class TestDocumentByCcPairCleanupTask: db_session.commit() with patch( - "onyx.background.celery.tasks.shared.tasks.get_all_document_indices", - return_value=[], + "onyx.background.celery.tasks.shared.tasks.get_default_document_index" ): results = [ document_by_cc_pair_cleanup_task.apply( @@ -313,16 +308,12 @@ class TestDocumentByCcPairCleanupTask: assert get_document_connector_count(db_session, doc.id) == 1 assert stored_doc.last_modified > indexed_at - with ( - patch( - "onyx.background.celery.tasks.shared.tasks.get_all_document_indices", - return_value=[object()], - ), - patch( - "onyx.background.celery.tasks.shared.tasks.RetryDocumentIndex.update", - side_effect=assert_relationship_removed, - ), - ): + with patch( + "onyx.background.celery.tasks.shared.tasks.get_default_document_index" + ) as mock_get_document_index: + mock_get_document_index.return_value.update.side_effect = ( + assert_relationship_removed + ) result = document_by_cc_pair_cleanup_task.apply( args=( doc.id, @@ -333,3 +324,52 @@ class TestDocumentByCcPairCleanupTask: ) assert result.successful(), result.traceback + + def test_index_failure_hands_the_doc_to_reconciliation_on_last_retry( + self, + db_session: Session, + cc_pair: ConnectorCredentialPair, + attempt_metadata: IndexAttemptMetadata, + full_deployment_setup: None, # noqa: ARG002 + ) -> None: + """When the index delete keeps failing, the task retries, and on the + last attempt it drops the cc_pair link and marks the doc modified so + stale-document reconciliation removes it later. The file stays, because + the document row still exists.""" + file_id = stage_file() + doc = make_doc(f"doc-{uuid4().hex[:8]}", file_id=file_id) + _index_doc(db_session, doc, attempt_metadata) + indexed_doc = get_doc_row(db_session, doc.id) + assert indexed_doc is not None + assert indexed_doc.last_modified is not None + indexed_at = indexed_doc.last_modified + + with patch( + "onyx.background.celery.tasks.shared.tasks.get_default_document_index" + ) as mock_get_document_index: + mock_get_document_index.return_value.delete.side_effect = RuntimeError( + "document index unavailable" + ) + result = document_by_cc_pair_cleanup_task.apply( + args=( + doc.id, + cc_pair.connector_id, + cc_pair.credential_id, + POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE, + ), + ) + delete_calls = mock_get_document_index.return_value.delete.call_count + + assert result.successful(), result.traceback + assert result.result is False + max_retries = document_by_cc_pair_cleanup_task.max_retries + assert max_retries is not None + assert delete_calls == max_retries + 1 + + db_session.expire_all() + assert get_document_connector_count(db_session, doc.id) == 0 + stored_doc = get_doc_row(db_session, doc.id) + assert stored_doc is not None + assert stored_doc.last_modified is not None + assert stored_doc.last_modified > indexed_at + assert get_filerecord(db_session, file_id) is not None 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 0c10ad2d2a..41829c36a8 100644 --- a/backend/tests/external_dependency_unit/indexing/test_port_reembed.py +++ b/backend/tests/external_dependency_unit/indexing/test_port_reembed.py @@ -32,7 +32,7 @@ from onyx.db.models import SearchSettings from onyx.document_index.chunk_content_enrichment import ( generate_enriched_content_for_chunk_embedding, ) -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces 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 diff --git a/backend/tests/external_dependency_unit/mock_search_pipeline.py b/backend/tests/external_dependency_unit/mock_search_pipeline.py index 55c38aadee..3705829a32 100644 --- a/backend/tests/external_dependency_unit/mock_search_pipeline.py +++ b/backend/tests/external_dependency_unit/mock_search_pipeline.py @@ -14,7 +14,7 @@ from onyx.context.search.models import ( SearchDoc, ) from onyx.db.models import SearchSettings, User -from onyx.document_index.interfaces_new import DocumentIndex +from onyx.document_index.interfaces import DocumentIndex from onyx.federated_connectors.federated_retrieval import FederatedRetrievalInfo from onyx.llm.interfaces import LLM from onyx.natural_language_processing.search_nlp_models import EmbeddingModel 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 5b7166da0f..98bcb03aba 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 @@ -9,7 +9,7 @@ hierarchy nodes, document sets, and persona/project user files. from typing import Any from onyx.configs.constants import DocumentSource -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.schema import ( ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME, DOCUMENT_ID_FIELD_NAME, diff --git a/backend/tests/external_dependency_unit/opensearch/test_forced_document_set_filter.py b/backend/tests/external_dependency_unit/opensearch/test_forced_document_set_filter.py index 4f1d0568b3..e63a4c66c6 100644 --- a/backend/tests/external_dependency_unit/opensearch/test_forced_document_set_filter.py +++ b/backend/tests/external_dependency_unit/opensearch/test_forced_document_set_filter.py @@ -10,7 +10,7 @@ enforcement intact. from typing import Any from onyx.configs.constants import DocumentSource -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.schema import ( ACCESS_CONTROL_LIST_FIELD_NAME, DOCUMENT_SETS_FIELD_NAME, 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 89db753147..6447b7ecd3 100644 --- a/backend/tests/external_dependency_unit/opensearch/test_opensearch_client.py +++ b/backend/tests/external_dependency_unit/opensearch/test_opensearch_client.py @@ -21,7 +21,7 @@ from onyx.access.utils import prefix_user_email from onyx.configs.constants import DocumentSource from onyx.context.search.models import IndexFilters, TimeRange from onyx.db.enums import VectorQuantization -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.client import ( OpenSearchDocumentMissingError, OpenSearchIndexClient, diff --git a/backend/tests/external_dependency_unit/scripts/test_orphan_doc_cleanup_script.py b/backend/tests/external_dependency_unit/scripts/test_orphan_doc_cleanup_script.py new file mode 100644 index 0000000000..b3d471bbb3 --- /dev/null +++ b/backend/tests/external_dependency_unit/scripts/test_orphan_doc_cleanup_script.py @@ -0,0 +1,264 @@ +"""External dependency tests for `scripts/orphan_doc_cleanup_script.py`. + +An orphan is a document with no connector/credential pair. The script deletes +its chunks from the document index, then deletes the document from Postgres. +During an index swap, the chunks can be in the primary index, the secondary +(future) index, or both, so the script must clean both indices. + +These tests assume Postgres and OpenSearch are running. +""" + +import time +import uuid +from collections.abc import Callable, Generator +from unittest.mock import MagicMock, patch + +import pytest +from scripts import orphan_doc_cleanup_script +from sqlalchemy import select +from sqlalchemy.orm import Session + +from onyx.context.search.models import IndexFilters +from onyx.db.enums import IndexModelStatus, SwitchoverType, VectorQuantization +from onyx.db.models import Document as DBDocument +from onyx.db.models import SearchSettings +from onyx.db.search_settings import ActiveSearchSettings +from onyx.document_index.interfaces import DocumentSectionRequest, TenantState +from onyx.document_index.opensearch.client import ( + OpenSearchIndexClient, + wait_for_opensearch_with_timeout, +) +from onyx.document_index.opensearch.opensearch_document_index import ( + OpenSearchDocumentIndex, +) +from onyx.kg.models import KGStage +from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE +from tests.external_dependency_unit.document_index.conftest import ( + EMBEDDING_DIM, + make_chunk, + make_indexing_metadata, +) + +TENANT_ID = POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE +CHUNKS_PER_DOC = 3 + + +def _search_settings(index_name: str, status: IndexModelStatus) -> SearchSettings: + """Search settings for an index that the test creates. Not saved, so other + tests never see a FUTURE search settings row.""" + return SearchSettings( + model_name="test-orphan-cleanup-model", + model_dim=EMBEDDING_DIM, + normalize=True, + query_prefix="", + passage_prefix="", + provider_type=None, + index_name=index_name, + multipass_indexing=False, + reduced_dimension=None, + vector_quantization=VectorQuantization.NONE, + switchover_type=SwitchoverType.REINDEX, + enable_contextual_rag=False, + status=status, + ) + + +def _new_index(name: str) -> OpenSearchDocumentIndex: + document_index = OpenSearchDocumentIndex( + tenant_state=TenantState(tenant_id=TENANT_ID, multitenant=False), + index_name=name, + embedding_dim=EMBEDDING_DIM, + vector_quantization=VectorQuantization.NONE, + ) + document_index.verify_and_create_index_if_necessary(embedding_dim=EMBEDDING_DIM) + return document_index + + +def _index_doc(document_index: OpenSearchDocumentIndex, doc_id: str) -> None: + document_index.index( + chunks=[ + make_chunk(doc_id, chunk_id=chunk_id) for chunk_id in range(CHUNKS_PER_DOC) + ], + indexing_metadata=make_indexing_metadata( + [doc_id], old_counts=[0], new_counts=[CHUNKS_PER_DOC] + ), + ) + + +def _chunk_count(document_index: OpenSearchDocumentIndex, doc_id: str) -> int: + return len( + document_index.id_based_retrieval( + chunk_requests=[DocumentSectionRequest(document_id=doc_id)], + filters=IndexFilters(access_control_list=None, tenant_id=TENANT_ID), + batch_retrieval=True, + ) + ) + + +def _wait_for_chunk_count( + document_index: OpenSearchDocumentIndex, doc_id: str, expected: int +) -> int: + """OpenSearch makes writes and deletes visible only after a refresh.""" + deadline = time.monotonic() + 10 + count = _chunk_count(document_index, doc_id) + while count != expected and time.monotonic() < deadline: + time.sleep(0.25) + count = _chunk_count(document_index, doc_id) + return count + + +@pytest.fixture +def index_names( + tenant_context: None, # noqa: ARG001 +) -> Generator[tuple[str, str], None, None]: + """Creates a primary and a secondary index and yields their names.""" + if not wait_for_opensearch_with_timeout(): + pytest.fail("OpenSearch is not available.") + suffix = uuid.uuid4().hex[:8] + names = [f"test_orphan_primary_{suffix}", f"test_orphan_secondary_{suffix}"] + for name in names: + _new_index(name) + try: + yield names[0], names[1] + finally: + for name in names: + OpenSearchIndexClient(index_name=name).delete_index() + + +def _add_orphans(db_session: Session, doc_ids: list[str]) -> None: + """Adds documents with no connector/credential pair.""" + for doc_id in doc_ids: + db_session.add( + DBDocument( + id=doc_id, + semantic_id=doc_id, + kg_stage=KGStage.NOT_STARTED, + chunk_count=CHUNKS_PER_DOC, + ) + ) + db_session.commit() + + +def _remaining_docs(db_session: Session, doc_ids: list[str]) -> list[str]: + db_session.expire_all() + return sorted( + db_session.scalars( + select(DBDocument.id).where(DBDocument.id.in_(doc_ids)) + ).all() + ) + + +def _delete_docs(db_session: Session, doc_ids: list[str]) -> None: + db_session.rollback() + for document in db_session.scalars( + select(DBDocument).where(DBDocument.id.in_(doc_ids)) + ).all(): + db_session.delete(document) + db_session.commit() + + +def _patched_orphan_query( + doc_ids: list[str], +) -> tuple[Callable[[Session, int], list[str]], list[int]]: + """Runs the real orphan query, but keeps only this test's documents so the + script leaves orphans from other tests alone. Fails instead of hanging if + the script keeps asking for the same batch.""" + find_orphans = orphan_doc_cleanup_script._get_orphaned_document_ids + calls: list[int] = [] + + def _find_test_orphans(session: Session, limit: int) -> list[str]: + calls.append(1) + assert len(calls) <= 5, "the script keeps retrying the same batch" + return [ + doc_id + for doc_id in find_orphans(session, max(limit, 10_000)) + if doc_id in doc_ids + ] + + return _find_test_orphans, calls + + +def test_cleanup_deletes_orphans_from_both_indices_and_postgres( + db_session: Session, + index_names: tuple[str, str], +) -> None: + primary_name, secondary_name = index_names + primary, secondary = _new_index(primary_name), _new_index(secondary_name) + suffix = uuid.uuid4().hex[:8] + in_both = f"orphan_in_both_{suffix}" + # During an index swap, a document can exist in the future index only. + in_secondary_only = f"orphan_in_secondary_only_{suffix}" + not_indexed = f"orphan_not_indexed_{suffix}" + doc_ids = [in_both, in_secondary_only, not_indexed] + + _add_orphans(db_session, doc_ids) + _index_doc(primary, in_both) + _index_doc(secondary, in_both) + _index_doc(secondary, in_secondary_only) + assert _wait_for_chunk_count(primary, in_both, CHUNKS_PER_DOC) == CHUNKS_PER_DOC + for doc_id in (in_both, in_secondary_only): + assert ( + _wait_for_chunk_count(secondary, doc_id, CHUNKS_PER_DOC) == CHUNKS_PER_DOC + ) + + active_search_settings = ActiveSearchSettings( + primary=_search_settings(primary_name, IndexModelStatus.PRESENT), + secondary=_search_settings(secondary_name, IndexModelStatus.FUTURE), + ) + find_test_orphans, _ = _patched_orphan_query(doc_ids) + try: + with ( + patch.object( + orphan_doc_cleanup_script, + "get_active_search_settings", + return_value=active_search_settings, + ), + patch.object( + orphan_doc_cleanup_script, + "_get_orphaned_document_ids", + side_effect=find_test_orphans, + ), + ): + orphan_doc_cleanup_script.main() + + for doc_id in (in_both, in_secondary_only): + assert _wait_for_chunk_count(primary, doc_id, 0) == 0 + assert _wait_for_chunk_count(secondary, doc_id, 0) == 0 + assert _remaining_docs(db_session, doc_ids) == [] + finally: + _delete_docs(db_session, doc_ids) + + +def test_cleanup_stops_when_no_orphan_in_a_batch_can_be_deleted( + db_session: Session, + tenant_context: None, # noqa: ARG001 +) -> None: + """A batch where every index delete fails would come back unchanged from + the next query. The script must stop and keep the documents in Postgres.""" + suffix = uuid.uuid4().hex[:8] + doc_ids = [f"orphan_delete_fails_{i}_{suffix}" for i in range(2)] + _add_orphans(db_session, doc_ids) + failing_index = MagicMock() + failing_index.delete.side_effect = RuntimeError("document index unavailable") + find_test_orphans, calls = _patched_orphan_query(doc_ids) + try: + with ( + patch.object(orphan_doc_cleanup_script, "get_active_search_settings"), + patch.object( + orphan_doc_cleanup_script, + "get_default_document_index", + return_value=failing_index, + ), + patch.object( + orphan_doc_cleanup_script, + "_get_orphaned_document_ids", + side_effect=find_test_orphans, + ), + ): + orphan_doc_cleanup_script.main() + + assert len(calls) == 1 + assert failing_index.delete.call_count == len(doc_ids) + assert _remaining_docs(db_session, doc_ids) == sorted(doc_ids) + finally: + _delete_docs(db_session, doc_ids) 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 86e544690d..5d7ec4ddb4 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 @@ -5,8 +5,9 @@ calls `delete_all_documents_for_connector_credential_pair` for each cc_pair. This test exercises that full workflow end-to-end and asserts that the attached `Document.file_id`s are also reaped — not just the document rows. -Mocks Vespa (`get_all_document_indices`) since this is testing the postgres + -file_store side effects of the swap, not the document index integration. +Mocks the document index (`get_default_document_index`) since this is testing +the postgres + file_store side effects of the swap, not the document index +integration. """ from collections.abc import Generator @@ -132,12 +133,9 @@ class TestInstantIndexSwap: status=IndexModelStatus.FUTURE, ) - # Vespa is patched out — we're testing the postgres + file_store - # side effects, not the document-index integration. - with patch( - "onyx.db.swap_index.get_all_document_indices", - return_value=[], - ): + # The document index is patched out — we're testing the postgres + + # file_store side effects, not the document-index integration. + with patch("onyx.db.swap_index.get_default_document_index"): old_settings = check_and_perform_index_swap(db_session) assert old_settings is not None, "INSTANT swap should have executed" @@ -178,10 +176,7 @@ class TestInstantIndexSwap: status=IndexModelStatus.FUTURE, ) - with patch( - "onyx.db.swap_index.get_all_document_indices", - return_value=[], - ): + with patch("onyx.db.swap_index.get_default_document_index"): old_settings = check_and_perform_index_swap(db_session) assert old_settings is not None 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 6b5d4bed51..5b2b9bd737 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 @@ -142,7 +142,7 @@ def _run_indexing_pipeline_with_mocks( document_batch=[], request_id=None, embedder=MagicMock(), - document_indices=[], + document_index=MagicMock(), db_session=db_session, tenant_id="public", adapter=MagicMock(), @@ -302,11 +302,9 @@ def test_contextual_model_update_rejects_unknown_model( assert str(unknown_model_configuration_id) in exc.value.detail -@patch("onyx.server.manage.search_settings.get_all_document_indices") @patch("onyx.server.manage.search_settings.get_default_document_index") def test_port_seed_excludes_invalid_cc_pair( mock_get_default_doc_index: MagicMock, # noqa: ARG001 - mock_get_all_doc_indices: MagicMock, baseline_search_settings: None, # noqa: ARG001 db_session: Session, ) -> None: @@ -314,7 +312,6 @@ def test_port_seed_excludes_invalid_cc_pair( INVALID cc_pair (excluded by the port's indexable_statuses scope) must NOT get a synthetic seed — otherwise its backlog is never ported while the seed cursor claims "already done", so its docs vanish from the live index once it's fixed.""" - mock_get_all_doc_indices.return_value = [] active_pair = make_cc_pair(db_session) invalid_pair = make_cc_pair(db_session) @@ -364,8 +361,7 @@ def test_port_seed_excludes_invalid_cc_pair( "onyx.db.swap_index.fetch_indexable_standard_connector_credential_pair_ids", new=lambda *_a, **_k: [], ) -@patch("onyx.db.swap_index.get_all_document_indices") -@patch("onyx.server.manage.search_settings.get_all_document_indices") +@patch("onyx.db.swap_index.get_default_document_index") @patch("onyx.server.manage.search_settings.get_default_document_index") @patch("onyx.indexing.indexing_pipeline.get_contextual_rag_llm_for_search_settings") @patch("onyx.indexing.indexing_pipeline.index_doc_batch_with_handler") @@ -373,8 +369,7 @@ def test_indexing_pipeline_uses_contextual_rag_settings_from_create( mock_index_handler: MagicMock, mock_get_llm: MagicMock, mock_get_doc_index: MagicMock, # noqa: ARG001 - mock_get_all_doc_indices_search_settings: MagicMock, # noqa: ARG001 - mock_get_all_doc_indices: MagicMock, + mock_get_swap_doc_index: MagicMock, # noqa: ARG001 baseline_search_settings: None, # noqa: ARG001 db_session: Session, ) -> None: @@ -400,7 +395,6 @@ def test_indexing_pipeline_uses_contextual_rag_settings_from_create( # Swap FUTURE → PRESENT. New settings use the port flow, whose swap gate waits # for each portable cc_pair's port; none require porting here (patched empty), # so the swap proceeds immediately. - mock_get_all_doc_indices.return_value = [] old_settings = check_and_perform_index_swap(db_session) assert old_settings is not None, "Swap should have occurred" @@ -422,8 +416,7 @@ def test_indexing_pipeline_uses_contextual_rag_settings_from_create( "onyx.db.swap_index.fetch_indexable_standard_connector_credential_pair_ids", new=lambda *_a, **_k: [], ) -@patch("onyx.db.swap_index.get_all_document_indices") -@patch("onyx.server.manage.search_settings.get_all_document_indices") +@patch("onyx.db.swap_index.get_default_document_index") @patch("onyx.server.manage.search_settings.get_default_document_index") @patch("onyx.indexing.indexing_pipeline.get_contextual_rag_llm_for_search_settings") @patch("onyx.indexing.indexing_pipeline.index_doc_batch_with_handler") @@ -431,8 +424,7 @@ def test_indexing_pipeline_uses_updated_contextual_rag_settings( mock_index_handler: MagicMock, mock_get_llm: MagicMock, mock_get_doc_index: MagicMock, # noqa: ARG001 - mock_get_all_doc_indices_search_settings: MagicMock, # noqa: ARG001 - mock_get_all_doc_indices: MagicMock, + mock_get_swap_doc_index: MagicMock, # noqa: ARG001 baseline_search_settings: None, # noqa: ARG001 db_session: Session, ) -> None: @@ -464,7 +456,6 @@ def test_indexing_pipeline_uses_updated_contextual_rag_settings( # Swap FUTURE → PRESENT. New settings use the port flow, whose swap gate waits # for each portable cc_pair's port; none require porting here (patched empty), # so the swap proceeds immediately. - mock_get_all_doc_indices.return_value = [] old_settings = check_and_perform_index_swap(db_session) assert old_settings is not None, "Swap should have occurred" @@ -516,7 +507,6 @@ def test_indexing_pipeline_uses_updated_contextual_rag_settings( assert called_settings.contextual_rag_model_configuration_id == updated_mc_id -@patch("onyx.server.manage.search_settings.get_all_document_indices") @patch("onyx.server.manage.search_settings.get_default_document_index") @patch("onyx.indexing.indexing_pipeline.get_contextual_rag_llm_for_search_settings") @patch("onyx.indexing.indexing_pipeline.index_doc_batch_with_handler") @@ -524,7 +514,6 @@ def test_indexing_pipeline_skips_llm_when_contextual_rag_disabled( mock_index_handler: MagicMock, mock_get_llm: MagicMock, mock_get_doc_index: MagicMock, # noqa: ARG001 - mock_get_all_doc_indices_search_settings: MagicMock, # noqa: ARG001 baseline_search_settings: None, # noqa: ARG001 db_session: Session, ) -> None: 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 e7fbc242be..3e06486324 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 @@ -472,12 +472,12 @@ class TestSlackBotFederatedSearch: @patch("onyx.utils.gpu_utils.fast_gpu_status_request", return_value=False) @patch( - "onyx.document_index.vespa.vespa_document_index.VespaDocumentIndex.hybrid_retrieval", + "onyx.document_index.opensearch.opensearch_document_index.OpenSearchIndexPair.hybrid_retrieval", return_value=[], ) def test_slack_bot_public_channel_filtering( self, - mock_vespa: Mock, # noqa: ARG002 + mock_hybrid_retrieval: Mock, # noqa: ARG002 mock_gpu_status: Mock, # noqa: ARG002 db_session: Session, ) -> None: @@ -532,12 +532,12 @@ class TestSlackBotFederatedSearch: @patch("onyx.utils.gpu_utils.fast_gpu_status_request", return_value=False) @patch( - "onyx.document_index.vespa.vespa_document_index.VespaDocumentIndex.hybrid_retrieval", + "onyx.document_index.opensearch.opensearch_document_index.OpenSearchIndexPair.hybrid_retrieval", return_value=[], ) def test_slack_bot_private_channel_filtering( self, - mock_vespa: Mock, # noqa: ARG002 + mock_hybrid_retrieval: Mock, # noqa: ARG002 mock_gpu_status: Mock, # noqa: ARG002 db_session: Session, ) -> None: @@ -592,12 +592,12 @@ class TestSlackBotFederatedSearch: @patch("onyx.utils.gpu_utils.fast_gpu_status_request", return_value=False) @patch( - "onyx.document_index.vespa.vespa_document_index.VespaDocumentIndex.hybrid_retrieval", + "onyx.document_index.opensearch.opensearch_document_index.OpenSearchIndexPair.hybrid_retrieval", return_value=[], ) def test_slack_bot_dm_filtering( self, - mock_vespa: Mock, # noqa: ARG002 + mock_hybrid_retrieval: Mock, # noqa: ARG002 mock_gpu_status: Mock, # noqa: ARG002 db_session: Session, ) -> None: diff --git a/backend/tests/integration/README.md b/backend/tests/integration/README.md index 8e59ad425d..28fd7e5d87 100644 --- a/backend/tests/integration/README.md +++ b/backend/tests/integration/README.md @@ -162,7 +162,7 @@ The current testing coverage should be checked by reading the comments at the to This could make test writing a bit cleaner by preventing test writers from having to pass around objects into functions that the objects have a 1:1 relationship with. -### Rework VespaClient +### Rework DocumentIndexClient Right now, its used a fixture and has to be passed around between manager classes. Could just be built where its used diff --git a/backend/tests/integration/common_utils/document_index.py b/backend/tests/integration/common_utils/document_index.py new file mode 100644 index 0000000000..4b9cd6b374 --- /dev/null +++ b/backend/tests/integration/common_utils/document_index.py @@ -0,0 +1,63 @@ +from typing import Any + +from opensearchpy import OpenSearch +from opensearchpy.exceptions import NotFoundError + +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 + + +class DocumentIndexClient: + """Test client for inspecting the chunks stored in the document index. + + The current index name is resolved lazily on each call rather than at + construction time. The docprocessing worker performs an in-flight swap from + ``danswer_chunk`` to ``danswer_chunk_`` the first time it indexes + after a Postgres reset; resolving the name eagerly would cache the pre-swap + value and query a non-existent index. + """ + + def __init__(self) -> None: + self._client = OpenSearch( + hosts=[{"host": OPENSEARCH_HOST, "port": OPENSEARCH_REST_API_PORT}], + http_auth=(OPENSEARCH_ADMIN_USERNAME, OPENSEARCH_ADMIN_PASSWORD), + use_ssl=OPENSEARCH_USE_SSL, + verify_certs=False, + ssl_show_warn=False, + ) + + @property + def index_name(self) -> str: + with get_session_with_current_tenant() as db_session: + return get_current_search_settings(db_session).index_name + + def get_chunks_by_document_id( + self, document_ids: list[str], wanted_chunk_count: int = 1_000 + ) -> list[dict[str, Any]]: + """Returns the stored fields (``_source``) of every chunk that belongs + to one of the given documents.""" + index_name = self.index_name + # Refresh first so chunks indexed just before the call are visible. + try: + self._client.indices.refresh(index=index_name) + except NotFoundError: + return [] + + body: dict[str, Any] = { + "size": wanted_chunk_count, + "query": {"terms": {"document_id": document_ids}}, + } + try: + result = self._client.search(index=index_name, body=body) + except NotFoundError: + return [] + + hits = result.get("hits", {}).get("hits", []) + return [dict(hit.get("_source", {})) for hit in hits] diff --git a/backend/tests/integration/common_utils/managers/cc_pair.py b/backend/tests/integration/common_utils/managers/cc_pair.py index f0a6527462..84bb4210fe 100644 --- a/backend/tests/integration/common_utils/managers/cc_pair.py +++ b/backend/tests/integration/common_utils/managers/cc_pair.py @@ -546,8 +546,8 @@ class CCPairManager: number_of_updated_docs: int = 0, # Sometimes waiting for a group sync is not necessary should_wait_for_group_sync: bool = True, - # Sometimes waiting for a vespa sync is not necessary - should_wait_for_vespa_sync: bool = True, + # Sometimes waiting for a document index sync is not necessary + should_wait_for_index_sync: bool = True, ) -> None: """after: The task register time must be after this time.""" doc_synced = False @@ -592,11 +592,12 @@ class CCPairManager: # this shouldnt be necessary but something is off with the timing for the sync jobs time.sleep(5) - if not should_wait_for_vespa_sync: + if not should_wait_for_index_sync: return - print("waiting for vespa sync") - # wait for the vespa sync to complete once the permission sync is complete + print("waiting for document index sync") + # wait for the document index sync to complete once the permission sync is + # complete start = time.monotonic() while True: doc_sync_statuses = CCPairManager.get_doc_sync_statuses( @@ -621,11 +622,11 @@ class CCPairManager: elapsed = time.monotonic() - start if elapsed > timeout: raise TimeoutError( - f"Vespa sync was not completed within {timeout} seconds" + f"Document index sync was not completed within {timeout} seconds" ) print( - f"Waiting for vespa sync to complete. elapsed={elapsed:.2f} timeout={timeout}" + f"Waiting for document index sync to complete. elapsed={elapsed:.2f} timeout={timeout}" ) time.sleep(5) diff --git a/backend/tests/integration/common_utils/managers/document.py b/backend/tests/integration/common_utils/managers/document.py index 265e55ffdf..789401cd18 100644 --- a/backend/tests/integration/common_utils/managers/document.py +++ b/backend/tests/integration/common_utils/managers/document.py @@ -8,6 +8,7 @@ from onyx.configs.constants import DocumentSource from onyx.db.enums import AccessType from onyx.db.models import ConnectorCredentialPair, DocumentByConnectorCredentialPair from tests.integration.common_utils.constants import API_SERVER_URL, NUM_DOCS +from tests.integration.common_utils.document_index import DocumentIndexClient 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 ( @@ -15,7 +16,6 @@ from tests.integration.common_utils.test_models import ( DATestUser, SimpleTestDocument, ) -from tests.integration.common_utils.vespa import vespa_fixture def _verify_document_permissions( @@ -25,13 +25,13 @@ def _verify_document_permissions( doc_set_names: list[str] | None = None, group_names: list[str] | None = None, ) -> None: - acl_keys = set(retrieved_doc.get("access_control_list", {}).keys()) + acl_keys = set(retrieved_doc.get("access_control_list") or []) print(f"ACL keys: {acl_keys}") if cc_pair.access_type == AccessType.PUBLIC: - if "PUBLIC" not in acl_keys: + if not retrieved_doc.get("public"): raise ValueError( - f"Document {retrieved_doc['document_id']} is public but does not have the PUBLIC ACL key" + f"Document {retrieved_doc['document_id']} is public but is not marked public in the index" ) if f"user_email:{doc_creating_user.email}" not in acl_keys: @@ -51,7 +51,7 @@ def _verify_document_permissions( ) if doc_set_names is not None: - found_doc_set_names = set(retrieved_doc.get("document_sets", {}).keys()) + found_doc_set_names = set(retrieved_doc.get("document_sets") or []) if found_doc_set_names != set(doc_set_names): raise ValueError( f"Document set names mismatch. \nFound: {found_doc_set_names}, \nExpected: {set(doc_set_names)}" @@ -190,7 +190,7 @@ class DocumentManager: @staticmethod def verify( - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, cc_pair: DATestCCPair, doc_creating_user: DATestUser, # If None, will not check doc sets or groups @@ -200,11 +200,9 @@ class DocumentManager: verify_deleted: bool = False, ) -> None: doc_ids = [document.id for document in cc_pair.documents] - retrieved_docs_dict = vespa_client.get_documents_by_id(doc_ids)["documents"] + retrieved_chunks = document_index_client.get_chunks_by_document_id(doc_ids) - retrieved_docs = { - doc["fields"]["document_id"]: doc["fields"] for doc in retrieved_docs_dict - } + retrieved_docs = {chunk["document_id"]: chunk for chunk in retrieved_chunks} # NOTE(rkuo): too much log spam # Left this here for debugging purposes. @@ -244,7 +242,7 @@ class DocumentManager: def fetch_documents_for_cc_pair( cc_pair_id: int, db_session: Session, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> list[SimpleTestDocument]: stmt = ( select(DocumentByConnectorCredentialPair) @@ -264,16 +262,15 @@ class DocumentManager: return [] doc_ids = [document.id for document in documents] - retrieved_docs_dict = vespa_client.get_documents_by_id(doc_ids)["documents"] + retrieved_chunks = document_index_client.get_chunks_by_document_id(doc_ids) final_docs: list[SimpleTestDocument] = [] - # NOTE: they are really chunks, but we're assuming that for these tests - # we only have one chunk per document for now - for doc_dict in retrieved_docs_dict: - doc_id = doc_dict["fields"]["document_id"] - doc_content = doc_dict["fields"]["content"] - # still called `image_file_name` in Vespa for backwards compatibility - image_file_id = doc_dict["fields"].get("image_file_name", None) + # NOTE: we're assuming that for these tests we only have one chunk per + # document for now + for chunk in retrieved_chunks: + doc_id = chunk["document_id"] + doc_content = chunk["content"] + image_file_id = chunk.get("image_file_id") final_docs.append( SimpleTestDocument( id=doc_id, content=doc_content, image_file_id=image_file_id diff --git a/backend/tests/integration/common_utils/reset.py b/backend/tests/integration/common_utils/reset.py index 676dc227ff..01d145f37a 100644 --- a/backend/tests/integration/common_utils/reset.py +++ b/backend/tests/integration/common_utils/reset.py @@ -296,9 +296,7 @@ def reset_postgres( _seed_dev_license_if_set(db_session) # Promote the FUTURE search-settings row (danswer_chunk_) to # PRESENT so secondary_search_settings is None and the api_server - # doesn't have to perform the swap mid-request. Previously this - # lived in reset_vespa(); when Vespa was deprecated the swap call - # needs to stay. + # doesn't have to perform the swap mid-request. check_and_perform_index_swap(db_session) diff --git a/backend/tests/integration/common_utils/vespa.py b/backend/tests/integration/common_utils/vespa.py deleted file mode 100644 index 611d80c332..0000000000 --- a/backend/tests/integration/common_utils/vespa.py +++ /dev/null @@ -1,89 +0,0 @@ -from typing import Any - -from opensearchpy import OpenSearch -from opensearchpy.exceptions import NotFoundError - -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 - - -class vespa_fixture: - """Test fixture for inspecting the document index. - - Kept named ``vespa_fixture`` for backwards compatibility with the many - existing integration tests that take it as a parameter. Internally it is now - backed by OpenSearch, and it reshapes hits into the dict-of-keys layout that - the legacy Vespa assertions expect (``access_control_list`` and - ``document_sets`` as dicts; ``image_file_name`` mirrored from OpenSearch's - ``image_file_id``; the ``public`` boolean folded back into the ACL as the - ``"PUBLIC"`` entry). - - The current index name is resolved lazily on each call rather than at - construction time. The docprocessing worker performs an in-flight swap from - ``danswer_chunk`` to ``danswer_chunk_`` the first time it indexes - after a Postgres reset; resolving the name eagerly in the fixture would - cache the pre-swap value and query a non-existent index. - """ - - def __init__(self, index_name: str | None = None) -> None: - # index_name is accepted for backwards compat with the prior Vespa - # fixture signature but ignored — see class docstring. - del index_name - self._client = OpenSearch( - hosts=[{"host": OPENSEARCH_HOST, "port": OPENSEARCH_REST_API_PORT}], - http_auth=(OPENSEARCH_ADMIN_USERNAME, OPENSEARCH_ADMIN_PASSWORD), - use_ssl=OPENSEARCH_USE_SSL, - verify_certs=False, - ssl_show_warn=False, - ) - - @property - def index_name(self) -> str: - with get_session_with_current_tenant() as db_session: - return get_current_search_settings(db_session).index_name - - def get_documents_by_id( - self, document_ids: list[str], wanted_doc_count: int = 1_000 - ) -> dict[str, Any]: - index_name = self.index_name - # Refresh first so chunks indexed just before the call are visible. - try: - self._client.indices.refresh(index=index_name) - except NotFoundError: - return {"documents": []} - - body: dict[str, Any] = { - "size": wanted_doc_count, - "query": {"terms": {"document_id": document_ids}}, - } - try: - result = self._client.search(index=index_name, body=body) - except NotFoundError: - return {"documents": []} - - hits = result.get("hits", {}).get("hits", []) - documents: list[dict[str, Any]] = [] - for hit in hits: - source: dict[str, Any] = dict(hit.get("_source", {})) - - acl_entries: set[str] = set(source.get("access_control_list") or []) - if source.get("public"): - acl_entries.add("PUBLIC") - source["access_control_list"] = dict.fromkeys(acl_entries, 1) - - source["document_sets"] = dict.fromkeys( - source.get("document_sets") or [], 1 - ) - - if "image_file_id" in source: - source["image_file_name"] = source["image_file_id"] - - documents.append({"fields": source}) - return {"documents": documents} diff --git a/backend/tests/integration/conftest.py b/backend/tests/integration/conftest.py index 2b92c3a538..10dc51d5c3 100644 --- a/backend/tests/integration/conftest.py +++ b/backend/tests/integration/conftest.py @@ -57,7 +57,6 @@ from onyx.db.engine.sql_engine import ( # noqa: E402 SqlEngine, get_session_with_current_tenant, ) -from onyx.db.search_settings import get_current_search_settings # noqa: E402 from onyx.utils.variable_functionality import ( # noqa: E402 fetch_versioned_implementation, ) @@ -67,6 +66,9 @@ from tests.integration.common_utils.constants import ( # noqa: E402 ADMIN_USER_NAME, GENERAL_HEADERS, ) +from tests.integration.common_utils.document_index import ( # noqa: E402 + DocumentIndexClient, +) from tests.integration.common_utils.managers.api_key import APIKeyManager # noqa: E402 from tests.integration.common_utils.managers.document import ( # noqa: E402 DocumentManager, @@ -101,7 +103,6 @@ from tests.integration.common_utils.test_models import ( # noqa: E402 DATestUser, SimpleTestDocument, ) -from tests.integration.common_utils.vespa import vespa_fixture # noqa: E402 from tests.integration.mock_services.mock_llm_server.server import ( # noqa: E402 run_in_thread, ) @@ -149,7 +150,7 @@ _CELERY_WORKER_PROGRAMS: list[tuple[str, str]] = [ ( "light", "vespa_metadata_sync,connector_deletion,doc_permissions_upsert," - "checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration", + "checkpoint_cleanup,index_attempt_cleanup,index_reclaim", ), ( "heavy", @@ -352,10 +353,8 @@ instantiate the session directly within the test. @pytest.fixture -def vespa_client() -> vespa_fixture: - with get_session_with_current_tenant() as db_session: - search_settings = get_current_search_settings(db_session) - return vespa_fixture(index_name=search_settings.index_name) +def document_index_client() -> DocumentIndexClient: + return DocumentIndexClient() @pytest.fixture 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 8a7ae42869..0315c51956 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 @@ -76,7 +76,8 @@ def test_github_private_repo_permission_sync( # Use a longer timeout for GitHub permission sync operations # GitHub API operations can be slow, especially with rate limiting - # This accounts for document sync, group sync, and vespa sync operations + # This accounts for document sync, group sync, and document index sync + # operations CCPairManager.wait_for_sync( cc_pair=github_cc_pair, user_performing_action=admin_user, 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 b097445af6..8861809287 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 @@ -28,7 +28,6 @@ from tests.integration.common_utils.test_models import ( DATestCredential, DATestUser, ) -from tests.integration.common_utils.vespa import vespa_fixture from tests.integration.connector_job_tests.google.google_drive_api_utils import ( GoogleDriveManager, ) @@ -108,7 +107,6 @@ def google_drive_test_env_setup() -> Generator[ @pytest.mark.xfail(reason="Needs to be tested for flakiness") def test_google_permission_sync( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, # noqa: ARG001 google_drive_test_env_setup: tuple[ GoogleDriveService, str, DATestCCPair, DATestUser, DATestUser, DATestUser ], @@ -278,8 +276,9 @@ def test_google_permission_sync( number_of_updated_docs=2, user_performing_action=admin_user, # if we are only updating the group definition for this test we use this varaiable, - # since it doesn't result in a vespa sync so we don't want to wait for it - should_wait_for_vespa_sync=False, + # since it doesn't result in a document index sync so we don't want to + # wait for it + should_wait_for_index_sync=False, ) # Verify user 1 can access both documents 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 f22fe751ee..cb1326f738 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 @@ -23,7 +23,6 @@ from tests.integration.common_utils.test_models import ( DATestCredential, DATestUser, ) -from tests.integration.common_utils.vespa import vespa_fixture from tests.integration.connector_job_tests.slack.conftest import ( SLACK_ADMIN_EMAIL, SLACK_TEST_USER_1_EMAIL, @@ -40,7 +39,6 @@ from tests.integration.connector_job_tests.slack.slack_api_utils import SlackMan ) def test_slack_permission_sync( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, # noqa: ARG001 slack_perm_sync_test_setup: tuple[ChannelType, ChannelType], ) -> None: public_channel, private_channel = slack_perm_sync_test_setup @@ -146,7 +144,7 @@ def test_slack_permission_sync( number_of_updated_docs=0, user_performing_action=admin_user, should_wait_for_group_sync=False, - should_wait_for_vespa_sync=False, + should_wait_for_index_sync=False, ) # Verify admin can see messages from both channels @@ -213,7 +211,6 @@ def test_slack_permission_sync( ) def test_slack_group_permission_sync( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, # noqa: ARG001 slack_perm_sync_test_setup: tuple[ChannelType, ChannelType], ) -> None: """ @@ -320,7 +317,7 @@ def test_slack_group_permission_sync( number_of_updated_docs=0, user_performing_action=admin_user, should_wait_for_group_sync=False, - should_wait_for_vespa_sync=False, + should_wait_for_index_sync=False, ) # Verify admin can see the message 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 6b381780e4..15f29019ee 100644 --- a/backend/tests/integration/connector_job_tests/slack/test_prune.py +++ b/backend/tests/integration/connector_job_tests/slack/test_prune.py @@ -22,14 +22,12 @@ from tests.integration.common_utils.test_models import ( DATestCredential, DATestUser, ) -from tests.integration.common_utils.vespa import vespa_fixture from tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager @pytest.mark.xfail(reason="flaky - see DAN-986 for details", strict=False) def test_slack_prune( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, # noqa: ARG001 slack_test_setup: tuple[ChannelType, ChannelType], ) -> None: public_channel, private_channel = slack_test_setup diff --git a/backend/tests/integration/tests/connector/test_connector_deletion.py b/backend/tests/integration/tests/connector/test_connector_deletion.py index 3d22c6920e..94229133c2 100644 --- a/backend/tests/integration/tests/connector/test_connector_deletion.py +++ b/backend/tests/integration/tests/connector/test_connector_deletion.py @@ -1,7 +1,7 @@ """ This file contains tests for the following: - Ensuring deletion of a connector also: - - deletes the documents in vespa for that connector + - deletes the documents in the document index for that connector - updates the document sets and user groups to remove the connector - Ensure that deleting a connector that is part of an overlapping document set and/or user group works as expected """ @@ -19,6 +19,7 @@ from onyx.db.models import IndexAttempt from onyx.db.search_settings import get_current_search_settings from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.constants import NUM_DOCS +from tests.integration.common_utils.document_index import DocumentIndexClient 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 @@ -30,12 +31,11 @@ from tests.integration.common_utils.test_models import ( DATestUser, DATestUserGroup, ) -from tests.integration.common_utils.vespa import vespa_fixture def test_connector_deletion( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> None: user_group_1: DATestUserGroup user_group_2: DATestUserGroup @@ -178,9 +178,9 @@ def test_connector_deletion( cc_pair_id=cc_pair_1.id, user_performing_action=admin_user ) - # validate vespa documents + # validate documents in the document index DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, doc_set_names=[], group_names=[], @@ -195,7 +195,7 @@ def test_connector_deletion( ] DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, doc_set_names=[doc_set_2.name], group_names=cc_pair_2_group_name_expected, @@ -238,7 +238,7 @@ def test_connector_deletion( def test_connector_deletion_for_overlapping_connectors( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> None: """Checks to make sure that connectors with overlapping documents work properly. Specifically, that the overlapping document (1) still exists and (2) has the right document set / group post-deletion of one of the connectors. @@ -283,16 +283,17 @@ def test_connector_deletion_for_overlapping_connectors( api_key=api_key, ) - # verify vespa document exists and that it is not in any document sets or groups + # verify the indexed document exists and that it is not in any document sets + # or groups DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, doc_set_names=[], group_names=[], doc_creating_user=admin_user, ) DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, doc_set_names=[], group_names=[], @@ -312,15 +313,15 @@ def test_connector_deletion_for_overlapping_connectors( print("Document set 1 created and synced") - # verify vespa document is in the document set + # verify the indexed document is in the document set DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, doc_set_names=[doc_set_1.name], doc_creating_user=admin_user, ) DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, doc_creating_user=admin_user, ) @@ -354,15 +355,15 @@ def test_connector_deletion_for_overlapping_connectors( print("User group 2 created and synced") - # verify vespa document is in the user group + # verify the indexed document is in the user group DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, group_names=[user_group_1.name, user_group_2.name], doc_creating_user=admin_user, ) DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, group_names=[user_group_1.name, user_group_2.name], doc_creating_user=admin_user, @@ -406,7 +407,7 @@ def test_connector_deletion_for_overlapping_connectors( ] DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, doc_set_names=[], group_names=group_names_expected, diff --git a/backend/tests/integration/tests/document_set/test_syncing.py b/backend/tests/integration/tests/document_set/test_syncing.py index 66f09dfa64..1228cc2dbc 100644 --- a/backend/tests/integration/tests/document_set/test_syncing.py +++ b/backend/tests/integration/tests/document_set/test_syncing.py @@ -2,17 +2,17 @@ from uuid import uuid4 from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.constants import NUM_DOCS +from tests.integration.common_utils.document_index import DocumentIndexClient 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.managers.user_group import UserGroupManager from tests.integration.common_utils.test_models import DATestAPIKey, DATestUser -from tests.integration.common_utils.vespa import vespa_fixture def test_multiple_document_sets_syncing_same_connnector( - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: # create api key with admin-group scope so it can hit the ingestion API @@ -62,7 +62,7 @@ def test_multiple_document_sets_syncing_same_connnector( # make sure documents are as expected DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, doc_set_names=[doc_set_1.name, doc_set_2.name], doc_creating_user=admin_user, @@ -70,7 +70,7 @@ def test_multiple_document_sets_syncing_same_connnector( def test_removing_connector( - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: # create api key with admin-group scope so it can hit the ingestion API @@ -122,7 +122,7 @@ def test_removing_connector( # make sure cc_pair_1 docs are doc_set_1 only DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, doc_set_names=[doc_set_1.name], doc_creating_user=admin_user, @@ -130,7 +130,7 @@ def test_removing_connector( # make sure cc_pair_2 docs are doc_set_1 only DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, doc_set_names=[doc_set_1.name], doc_creating_user=admin_user, @@ -149,7 +149,7 @@ def test_removing_connector( # make sure cc_pair_1 docs are doc_set_1 only DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, doc_set_names=[doc_set_1.name], doc_creating_user=admin_user, @@ -157,7 +157,7 @@ def test_removing_connector( # make sure cc_pair_2 docs have no doc set DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, doc_set_names=[], doc_creating_user=admin_user, @@ -165,7 +165,7 @@ def test_removing_connector( def test_renaming_document_set( - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: api_key: DATestAPIKey = APIKeyManager.create( @@ -210,7 +210,7 @@ def test_renaming_document_set( ) DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair, doc_set_names=[new_name], doc_creating_user=admin_user, 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 8648b34250..817edf9e7d 100644 --- a/backend/tests/integration/tests/image_indexing/test_indexing_images.py +++ b/backend/tests/integration/tests/image_indexing/test_indexing_images.py @@ -7,6 +7,7 @@ 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.server.documents.models import DocumentSource +from tests.integration.common_utils.document_index import DocumentIndexClient 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 @@ -15,7 +16,6 @@ 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, DATestUser -from tests.integration.common_utils.vespa import vespa_fixture FILE_NAME = "Sample.pdf" FILE_PATH = "tests/integration/common_utils/test_files" @@ -25,7 +25,7 @@ DOCX_FILE_NAME = "three_images.docx" def test_image_indexing( reset: None, # noqa: ARG001 admin_user: DATestUser, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> None: os.makedirs(FILE_PATH, exist_ok=True) test_file_path = os.path.join(FILE_PATH, FILE_NAME) @@ -103,12 +103,13 @@ def test_image_indexing( ) with get_session_with_current_tenant() as db_session: - # really gets the chunks from Vespa, which is why there are two; + # really gets the chunks from the document index, which is why there + # are two; # one for the raw text and one for the summarized image. documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(documents) == 2 @@ -126,7 +127,7 @@ def test_image_indexing( def test_docx_image_indexing( reset: None, # noqa: ARG001 admin_user: DATestUser, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> None: """Test that images from docx files are correctly extracted and indexed.""" os.makedirs(FILE_PATH, exist_ok=True) @@ -207,11 +208,12 @@ def test_docx_image_indexing( ) with get_session_with_current_tenant() as db_session: - # Fetch documents from Vespa - expect text content plus 3 images + # Fetch documents from the document index - expect text content plus + # 3 images documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) # Should have documents for text content plus 3 images 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 8764a22267..b93fd4cf6f 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 @@ -15,7 +15,6 @@ from tests.integration.common_utils.managers.credential import CredentialManager from tests.integration.common_utils.managers.file import FileManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.vespa import vespa_fixture # This is a placeholder - you'll need to create this zip file with actual test files TEST_FILES_BASE = "tests/integration/tests/indexing/file_connector/test_files" @@ -33,7 +32,6 @@ TEST_METADATA_FILE = f"{TEST_FILES_BASE}/.onyx_metadata.json" ) def test_zip_metadata_handling( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, # noqa: ARG001 zip_path: str, has_metadata: bool, ) -> None: diff --git a/backend/tests/integration/tests/indexing/test_checkpointing.py b/backend/tests/integration/tests/indexing/test_checkpointing.py index 8fd00f5c06..a484446512 100644 --- a/backend/tests/integration/tests/indexing/test_checkpointing.py +++ b/backend/tests/integration/tests/indexing/test_checkpointing.py @@ -12,6 +12,7 @@ from tests.integration.common_utils.constants import ( MOCK_CONNECTOR_SERVER_HOST, MOCK_CONNECTOR_SERVER_PORT, ) +from tests.integration.common_utils.document_index import DocumentIndexClient 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 @@ -20,12 +21,11 @@ from tests.integration.common_utils.test_document_utils import ( create_test_document_failure, ) from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.vespa import vespa_fixture def test_mock_connector_basic_flow( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: """Test that the mock connector can successfully process documents and failures""" @@ -85,7 +85,7 @@ def test_mock_connector_basic_flow( chunks = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(chunks) == 1 assert chunks[0].id == test_doc.id @@ -99,7 +99,7 @@ def test_mock_connector_basic_flow( def test_mock_connector_with_failures( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: """Test that the mock connector processes both successes and failures properly.""" @@ -158,7 +158,7 @@ def test_mock_connector_with_failures( documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(documents) == 1 assert documents[0].id == doc1.id @@ -175,7 +175,7 @@ def test_mock_connector_with_failures( def test_mock_connector_failure_recovery( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: """Test that a failed document can be successfully indexed in a subsequent attempt @@ -249,7 +249,7 @@ def test_mock_connector_failure_recovery( documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(documents) == 1 assert documents[0].id == doc1.id @@ -313,7 +313,7 @@ def test_mock_connector_failure_recovery( documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(documents) == 2 document_ids = {doc.id for doc in documents} @@ -335,7 +335,7 @@ def test_mock_connector_failure_recovery( def test_mock_connector_checkpoint_recovery( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: """Test that checkpointing works correctly when an unhandled exception occurs @@ -435,7 +435,7 @@ def test_mock_connector_checkpoint_recovery( documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) # This is no longer guaranteed because docfetching and docprocessing are decoupled! # Some batches may not be processed when docfetching fails, but they should still stick around @@ -510,7 +510,7 @@ def test_mock_connector_checkpoint_recovery( documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(documents) == 102 # 100 docs from first batch + doc2 + doc3 document_ids = {doc.id for doc in documents} 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 cef0debe94..72556781a7 100644 --- a/backend/tests/integration/tests/indexing/test_initial_permission_sync.py +++ b/backend/tests/integration/tests/indexing/test_initial_permission_sync.py @@ -21,12 +21,12 @@ from tests.integration.common_utils.constants import ( MOCK_CONNECTOR_SERVER_HOST, MOCK_CONNECTOR_SERVER_PORT, ) +from tests.integration.common_utils.document_index import DocumentIndexClient 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, DATestUser -from tests.integration.common_utils.vespa import vespa_fixture def _setup_mock_connector( @@ -89,7 +89,7 @@ def _setup_mock_connector( ) def test_mock_connector_initial_permission_sync( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: """Test that the MockConnector fetches and sets permissions during initial indexing @@ -101,7 +101,7 @@ def test_mock_connector_initial_permission_sync( documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(documents) == 1 assert documents[0].id == test_doc.id @@ -140,7 +140,7 @@ def test_mock_connector_initial_permission_sync( number_of_updated_docs=1, user_performing_action=admin_user, should_wait_for_group_sync=False, - should_wait_for_vespa_sync=False, + should_wait_for_index_sync=False, ) updated_cc_pair_info = CCPairManager.get_single( @@ -156,7 +156,6 @@ def test_mock_connector_initial_permission_sync( ) def test_permission_sync_attempt_tracking_integration( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, # noqa: ARG001 admin_user: DATestUser, ) -> None: """Test that permission sync attempts are properly tracked during real sync workflows.""" @@ -175,7 +174,7 @@ def test_permission_sync_attempt_tracking_integration( number_of_updated_docs=1, user_performing_action=admin_user, should_wait_for_group_sync=False, - should_wait_for_vespa_sync=False, + should_wait_for_index_sync=False, ) with get_session_with_current_tenant() as db_session: @@ -203,7 +202,6 @@ def test_permission_sync_attempt_tracking_integration( ) def test_permission_sync_attempt_status_success( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, # noqa: ARG001 admin_user: DATestUser, ) -> None: """Test that permission sync attempts are marked as SUCCESS when sync completes without errors.""" @@ -222,7 +220,7 @@ def test_permission_sync_attempt_status_success( number_of_updated_docs=1, user_performing_action=admin_user, should_wait_for_group_sync=False, - should_wait_for_vespa_sync=False, + should_wait_for_index_sync=False, ) with get_session_with_current_tenant() as db_session: 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 a8ff8f6999..64d42ef775 100644 --- a/backend/tests/integration/tests/indexing/test_repeated_error_state.py +++ b/backend/tests/integration/tests/indexing/test_repeated_error_state.py @@ -16,17 +16,17 @@ from tests.integration.common_utils.constants import ( MOCK_CONNECTOR_SERVER_HOST, MOCK_CONNECTOR_SERVER_PORT, ) +from tests.integration.common_utils.document_index import DocumentIndexClient 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 DATestUser -from tests.integration.common_utils.vespa import vespa_fixture def test_repeated_error_state_detection_and_recovery( mock_server_client: httpx.Client, - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, admin_user: DATestUser, ) -> None: """Test that a connector is marked as in a repeated error state after @@ -189,7 +189,7 @@ def test_repeated_error_state_detection_and_recovery( documents = DocumentManager.fetch_documents_for_cc_pair( cc_pair_id=cc_pair.id, db_session=db_session, - vespa_client=vespa_client, + document_index_client=document_index_client, ) assert len(documents) == 1 assert documents[0].id == test_doc.id diff --git a/backend/tests/integration/tests/ingestion/test_ingestion_api.py b/backend/tests/integration/tests/ingestion/test_ingestion_api.py index abc30937e8..6104e8fbde 100644 --- a/backend/tests/integration/tests/ingestion/test_ingestion_api.py +++ b/backend/tests/integration/tests/ingestion/test_ingestion_api.py @@ -2,17 +2,17 @@ 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.models import Document +from tests.integration.common_utils.document_index import DocumentIndexClient 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 IngestionManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.vespa import vespa_fixture def test_ingestion_api_crud( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> None: """Test create, list, and delete via the ingestion API.""" admin_user: DATestUser = UserManager.create(email="admin@onyx.app") @@ -43,8 +43,8 @@ def test_ingestion_api_crud( assert doc_db is not None assert doc_db.from_ingestion_api is True - vespa_docs = vespa_client.get_documents_by_id([doc.id])["documents"] - assert len(vespa_docs) == 1 + chunks = document_index_client.get_chunks_by_document_id([doc.id]) + assert len(chunks) == 1 # LIST docs_list = IngestionManager.list_all_ingestion_docs(api_key=api_key) @@ -57,5 +57,5 @@ def test_ingestion_api_crud( doc_db = db_session.query(Document).filter(Document.id == doc.id).first() assert doc_db is None - vespa_docs = vespa_client.get_documents_by_id([doc.id])["documents"] - assert len(vespa_docs) == 0 + chunks = document_index_client.get_chunks_by_document_id([doc.id]) + assert len(chunks) == 0 diff --git a/backend/tests/integration/tests/migrations/conftest.py b/backend/tests/integration/tests/migrations/conftest.py index 6aaec70403..bee80f4231 100644 --- a/backend/tests/integration/tests/migrations/conftest.py +++ b/backend/tests/integration/tests/migrations/conftest.py @@ -31,8 +31,8 @@ from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA # Migration tests only need Postgres (provided by the workflow) and the # pytest-alembic fixtures below — they must NOT pre-migrate the schema # (would break pytest-alembic's test_upgrade) and must NOT start the -# FastAPI app (its lifespan calls setup_onyx() which requires Vespa, and -# the database-tests workflow doesn't start Vespa). +# FastAPI app (its lifespan calls setup_onyx() which requires OpenSearch, +# and the database-tests workflow doesn't start OpenSearch). @pytest.fixture(scope="session", autouse=True) diff --git a/backend/tests/integration/tests/migrations/test_drop_opensearch_migration_tables.py b/backend/tests/integration/tests/migrations/test_drop_opensearch_migration_tables.py new file mode 100644 index 0000000000..ace72dd373 --- /dev/null +++ b/backend/tests/integration/tests/migrations/test_drop_opensearch_migration_tables.py @@ -0,0 +1,218 @@ +"""Tests for migration 3067343245d1, which drops the tables that tracked the +Vespa to OpenSearch migration.""" + +import threading +import time +from collections.abc import Generator + +import pytest +from sqlalchemy import Engine, create_engine, text + +from onyx.configs.app_configs import ( + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, +) +from onyx.db.engine.sql_engine import SYNC_DB_API, build_connection_string +from tests.integration.common_utils.reset import downgrade_postgres, upgrade_postgres + +PREVIOUS_REVISION = "b7c2e4f1a9d3" +DROP_REVISION = "3067343245d1" +DOCUMENT_RECORD_TABLE = "opensearch_document_migration_record" +TENANT_RECORD_TABLE = "opensearch_tenant_migration_record" + + +def _engine() -> Engine: + return create_engine( + build_connection_string( + db="postgres", + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + host=POSTGRES_HOST, + port=POSTGRES_PORT, + db_api=SYNC_DB_API, + ) + ) + + +def _migration_tables(engine: Engine) -> set[str]: + with engine.connect() as conn: + rows = conn.execute( + text( + "SELECT table_name FROM information_schema.tables " + "WHERE table_schema = 'public' AND table_name IN (:a, :b)" + ), + {"a": DOCUMENT_RECORD_TABLE, "b": TENANT_RECORD_TABLE}, + ) + return {row[0] for row in rows} + + +def _current_revision(engine: Engine) -> str: + with engine.connect() as conn: + return conn.execute( + text("SELECT version_num FROM alembic_version") + ).scalar_one() + + +def _insert_document(engine: Engine, doc_id: str) -> None: + with engine.begin() as conn: + conn.execute( + text( + "INSERT INTO document (id, semantic_id, boost, hidden, " + "from_ingestion_api, kg_stage) " + "VALUES (:id, :id, 0, false, false, 'NOT_STARTED')" + ), + {"id": doc_id}, + ) + + +@pytest.fixture(scope="module") +def engine() -> Generator[Engine, None, None]: + downgrade_postgres( + database="postgres", config_name="alembic", revision="base", clear_data=True + ) + upgrade_postgres( + database="postgres", config_name="alembic", revision=PREVIOUS_REVISION + ) + engine = _engine() + try: + yield engine + finally: + engine.dispose() + upgrade_postgres(database="postgres", config_name="alembic", revision="head") + + +@pytest.fixture +def at_previous_revision(engine: Engine) -> Engine: + """Puts the database at the revision just before the drop, with empty + migration tables.""" + if _current_revision(engine) == DROP_REVISION: + downgrade_postgres( + database="postgres", config_name="alembic", revision=PREVIOUS_REVISION + ) + assert _current_revision(engine) == PREVIOUS_REVISION + with engine.begin() as conn: + conn.execute(text(f"DELETE FROM {DOCUMENT_RECORD_TABLE}")) + conn.execute(text(f"DELETE FROM {TENANT_RECORD_TABLE}")) + return engine + + +def test_upgrade_drops_populated_tables_and_keeps_documents( + at_previous_revision: Engine, +) -> None: + engine = at_previous_revision + _insert_document(engine, "drop-migration-doc") + with engine.begin() as conn: + conn.execute( + text( + f"INSERT INTO {DOCUMENT_RECORD_TABLE} (document_id) " + "VALUES ('drop-migration-doc')" + ) + ) + conn.execute( + text( + f"INSERT INTO {TENANT_RECORD_TABLE} " + "(id, enable_opensearch_retrieval, migration_completed_at) " + "VALUES (1, true, now())" + ) + ) + + upgrade_postgres(database="postgres", config_name="alembic", revision=DROP_REVISION) + + assert _migration_tables(engine) == set() + with engine.connect() as conn: + remaining = conn.execute( + text("SELECT count(*) FROM document WHERE id = 'drop-migration-doc'") + ).scalar_one() + assert remaining == 1 + + +def test_downgrade_recreates_the_tables(at_previous_revision: Engine) -> None: + engine = at_previous_revision + upgrade_postgres(database="postgres", config_name="alembic", revision=DROP_REVISION) + + downgrade_postgres( + database="postgres", config_name="alembic", revision=PREVIOUS_REVISION + ) + + assert _migration_tables(engine) == {DOCUMENT_RECORD_TABLE, TENANT_RECORD_TABLE} + _insert_document(engine, "recreated-table-doc") + with engine.begin() as conn: + conn.execute( + text( + f"INSERT INTO {DOCUMENT_RECORD_TABLE} (document_id) " + "VALUES ('recreated-table-doc')" + ) + ) + conn.execute(text(f"INSERT INTO {TENANT_RECORD_TABLE} (id) VALUES (1)")) + status = conn.execute( + text(f"SELECT status FROM {DOCUMENT_RECORD_TABLE}") + ).scalar_one() + assert status == "pending" + # The singleton index allows only one tenant record. + with pytest.raises(Exception, match="idx_opensearch_tenant_migration_singleton"): + with engine.begin() as conn: + conn.execute(text(f"INSERT INTO {TENANT_RECORD_TABLE} (id) VALUES (2)")) + + +def test_upgrade_waits_out_a_long_lock_on_document( + at_previous_revision: Engine, +) -> None: + """A transaction that holds a lock on `document` (for example an indexing + batch) must delay the drop, not fail the migration.""" + engine = at_previous_revision + lock_held = threading.Event() + hold_lock_s = 7.0 + + def _hold_lock() -> None: + with engine.begin() as conn: + conn.execute(text("LOCK TABLE document IN ROW SHARE MODE")) + lock_held.set() + time.sleep(hold_lock_s) + + holder = threading.Thread(target=_hold_lock) + holder.start() + assert lock_held.wait(timeout=10) + start = time.monotonic() + + upgrade_postgres(database="postgres", config_name="alembic", revision=DROP_REVISION) + + elapsed = time.monotonic() - start + holder.join() + assert _current_revision(engine) == DROP_REVISION + assert _migration_tables(engine) == set() + # The first attempt timed out behind the lock, so the drop was retried. + assert elapsed >= 5.0 + + +def test_upgrade_fails_fast_on_an_error_that_is_not_a_lock_timeout( + at_previous_revision: Engine, +) -> None: + """Only a lock timeout is retried. A permanent error, here a view that + depends on the table, must fail the migration on the first attempt.""" + engine = at_previous_revision + view_name = "drop_migration_test_dependent_view" + with engine.begin() as conn: + conn.execute( + text( + f"CREATE VIEW {view_name} AS SELECT document_id " + f"FROM {DOCUMENT_RECORD_TABLE}" + ) + ) + try: + start = time.monotonic() + with pytest.raises(Exception, match="depend"): + upgrade_postgres( + database="postgres", config_name="alembic", revision=DROP_REVISION + ) + elapsed = time.monotonic() - start + finally: + with engine.begin() as conn: + conn.execute(text(f"DROP VIEW IF EXISTS {view_name}")) + + # A retry would add at least one 5 s delay. + assert elapsed < 5.0 + # The migration ran in one transaction, so nothing was dropped. + assert _current_revision(engine) == PREVIOUS_REVISION + assert _migration_tables(engine) == {DOCUMENT_RECORD_TABLE, TENANT_RECORD_TABLE} diff --git a/backend/tests/integration/tests/permissions_membership/test_group_manager_membership.py b/backend/tests/integration/tests/permissions_membership/test_group_manager_membership.py index e2ffadab8f..a0fb863ad0 100644 --- a/backend/tests/integration/tests/permissions_membership/test_group_manager_membership.py +++ b/backend/tests/integration/tests/permissions_membership/test_group_manager_membership.py @@ -200,7 +200,7 @@ def _set_cc_pair_creator(cc_pair_id: int, user_id: str) -> None: def _insert_stale_cc_pair_junction(group_id: int, cc_pair_id: int) -> None: """Simulate a removed-but-not-yet-swept cc_pair: an is_current=False junction row - (removals mark the row stale rather than delete it, until the Vespa sync runs).""" + (removals mark the row stale rather than delete it, until the document index sync runs).""" with get_session_with_current_tenant() as db_session: db_session.add( UserGroup__ConnectorCredentialPair( diff --git a/backend/tests/integration/tests/pruning/test_pruning.py b/backend/tests/integration/tests/pruning/test_pruning.py index 2f29117a75..421b796031 100644 --- a/backend/tests/integration/tests/pruning/test_pruning.py +++ b/backend/tests/integration/tests/pruning/test_pruning.py @@ -15,11 +15,11 @@ from fastapi.staticfiles import StaticFiles from onyx.server.documents.models import DocumentSource from onyx.utils.logger import setup_logger +from tests.integration.common_utils.document_index import DocumentIndexClient 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.user import UserManager from tests.integration.common_utils.test_models import DATestUser -from tests.integration.common_utils.vespa import vespa_fixture logger = setup_logger() @@ -99,7 +99,7 @@ def http_server_context( def test_web_pruning( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> None: # Creating an admin user (first user created is automatically an admin) admin_user: DATestUser = UserManager.create(name="admin_user") @@ -168,24 +168,22 @@ def test_web_pruning( assert selected_cc_pair is not None, "cc_pair not found after pruning!" assert selected_cc_pair.docs_indexed == 12 - # check vespa + # check the document index root_id = f"http://{hostname}:{port}/" index_id = f"http://{hostname}:{port}/index.html" about_id = f"http://{hostname}:{port}/about.html" courses_id = f"http://{hostname}:{port}/courses.html" doc_ids = [root_id, index_id, about_id, courses_id] - retrieved_docs_dict = vespa_client.get_documents_by_id(doc_ids)["documents"] - retrieved_docs = { - doc["fields"]["document_id"]: doc["fields"] - for doc in retrieved_docs_dict - } + retrieved_chunks = document_index_client.get_chunks_by_document_id(doc_ids) + retrieved_docs = {chunk["document_id"]: chunk for chunk in retrieved_chunks} - # verify root exists in Vespa + # verify root exists in the document index retrieved_doc = retrieved_docs.get(root_id) assert retrieved_doc - # verify index.html does not exist in Vespa since it is a duplicate of root + # verify index.html does not exist in the document index since it is a + # duplicate of root retrieved_doc = retrieved_docs.get(index_id) assert not retrieved_doc 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 dcb0f0639f..94ac84d020 100644 --- a/backend/tests/integration/tests/usergroup/test_user_group_deletion.py +++ b/backend/tests/integration/tests/usergroup/test_user_group_deletion.py @@ -29,7 +29,6 @@ from tests.integration.common_utils.test_models import ( DATestUser, DATestUserGroup, ) -from tests.integration.common_utils.vespa import vespa_fixture @pytest.mark.skipif( @@ -38,7 +37,6 @@ from tests.integration.common_utils.vespa import vespa_fixture ) def test_user_group_deletion( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, # noqa: ARG001 ) -> None: # Creating an admin user (first user created is automatically an admin) admin_user: DATestUser = UserManager.create(name="admin_user") diff --git a/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py b/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py index 8f60f49310..2432d609fe 100644 --- a/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py +++ b/backend/tests/integration/tests/usergroup/test_usergroup_syncing.py @@ -4,6 +4,7 @@ import pytest from onyx.server.documents.models import DocumentSource from tests.integration.common_utils.constants import NUM_DOCS +from tests.integration.common_utils.document_index import DocumentIndexClient 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 @@ -14,7 +15,6 @@ from tests.integration.common_utils.test_models import ( DATestUser, DATestUserGroup, ) -from tests.integration.common_utils.vespa import vespa_fixture @pytest.mark.skipif( @@ -23,7 +23,7 @@ from tests.integration.common_utils.vespa import vespa_fixture ) def test_removing_connector( reset: None, # noqa: ARG001 - vespa_client: vespa_fixture, + document_index_client: DocumentIndexClient, ) -> None: # Creating an admin user (first user created is automatically an admin) admin_user: DATestUser = UserManager.create(name="admin_user") @@ -77,7 +77,7 @@ def test_removing_connector( # make sure cc_pair_1 docs are user_group_1 only DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, group_names=[user_group_1.name], doc_creating_user=admin_user, @@ -85,7 +85,7 @@ def test_removing_connector( # make sure cc_pair_2 docs are user_group_1 only DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, group_names=[user_group_1.name], doc_creating_user=admin_user, @@ -104,7 +104,7 @@ def test_removing_connector( # make sure cc_pair_1 docs are user_group_1 only DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_1, group_names=[user_group_1.name], doc_creating_user=admin_user, @@ -112,7 +112,7 @@ def test_removing_connector( # make sure cc_pair_2 docs have no user group DocumentManager.verify( - vespa_client=vespa_client, + document_index_client=document_index_client, cc_pair=cc_pair_2, group_names=[], doc_creating_user=admin_user, diff --git a/backend/tests/regression/answer_quality/README.md b/backend/tests/regression/answer_quality/README.md index 0937528a6b..e8b1cd6f0d 100644 --- a/backend/tests/regression/answer_quality/README.md +++ b/backend/tests/regression/answer_quality/README.md @@ -68,7 +68,7 @@ Edit `search_test_config.yaml` to set: - output_folder - This is the folder where the folders for each test will go - - These folders will contain the postgres/vespa data as well as the results for each test + - These folders will contain the postgres data as well as the results for each test - zipped_documents_file - The path to the zip file containing the files you'd like to test against - questions_file @@ -82,7 +82,7 @@ Edit `search_test_config.yaml` to set: - launch_web_ui - Set this to true if you want to use the UI during/after the testing process - only_state - - Whether to only run Vespa and Postgres + - Whether to only run OpenSearch and Postgres - only_retrieve_docs - Set true to only retrieve documents, not LLM response - This is to save on API costs diff --git a/backend/tests/regression/answer_quality/cli_utils.py b/backend/tests/regression/answer_quality/cli_utils.py index 35bb8b9713..9cb3491d36 100644 --- a/backend/tests/regression/answer_quality/cli_utils.py +++ b/backend/tests/regression/answer_quality/cli_utils.py @@ -3,7 +3,6 @@ import os import socket import subprocess import sys -import time from datetime import datetime from threading import Thread from typing import IO @@ -101,7 +100,6 @@ def manage_data_directories(env_name: str, base_path: str, use_cloud_gpu: bool) target_path = os.path.join(os.path.expanduser(base_path), env_name) directories = { "DANSWER_POSTGRES_DATA_DIR": os.path.join(target_path, "postgres/"), - "DANSWER_VESPA_DATA_DIR": os.path.join(target_path, "vespa/"), } if not use_cloud_gpu: directories["DANSWER_INDEX_MODEL_CACHE_DIR"] = os.path.join( @@ -157,7 +155,7 @@ def start_docker_compose( command += " --force-recreate" if only_state: - command += " index relational_db" + command += " opensearch relational_db" else: if use_cloud_gpu: command += " --scale indexing_model_server=0" @@ -287,33 +285,11 @@ def get_api_server_host_port(env_name: str) -> str: return matching_ports[0] -# Added function to restart Vespa container -def restart_vespa_container(env_name: str) -> None: - print(f"Restarting Vespa container for env_name: {env_name}") - - # Find the Vespa container - stdout, _ = _run_command( - f"docker ps -a --format '{{{{.Names}}}}' | awk '/index-1/ && /{env_name}/'" - ) - container_name = stdout.strip() - - if not container_name: - raise RuntimeError(f"No Vespa container found with env_name: {env_name}") - - # Restart the container - _run_command(f"docker restart {container_name}") - - print(f"Vespa container '{container_name}' has begun restarting") - - time.sleep(30) - print(f"Vespa container '{container_name}' has been restarted") - - if __name__ == "__main__": """ Running this just cleans up the docker environment for the container indicated by environment_name If no environment_name is indicated, will just clean up all onyx docker containers/volumes/networks - Note: vespa/postgres mounts are not deleted + Note: postgres mounts are not deleted """ current_dir = os.path.dirname(os.path.abspath(__file__)) config_path = os.path.join(current_dir, "search_test_config.yaml") diff --git a/backend/tests/regression/search_quality/utils.py b/backend/tests/regression/search_quality/utils.py index 631a67d4d1..4e2f51d4b1 100644 --- a/backend/tests/regression/search_quality/utils.py +++ b/backend/tests/regression/search_quality/utils.py @@ -23,7 +23,7 @@ from onyx.db.engine.sql_engine import get_session_with_tenant 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.document_index.interfaces 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 ( @@ -89,7 +89,7 @@ def get_doc_contents( ) -> dict[tuple[str, int], str]: with get_session_with_tenant(tenant_id=tenant_id) as db_session: search_settings = get_current_search_settings(db_session) - document_index = get_default_document_index(search_settings, None, db_session) + document_index = get_default_document_index(search_settings, None) filters = IndexFilters(access_control_list=None, tenant_id=tenant_id) 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 cb39e46a4e..16a274e37b 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 @@ -367,9 +367,8 @@ class TestDeleteImplNoVectorDb: mock_get_file_store.return_value = MagicMock() with ( - patch(f"{TASKS_MODULE}.get_all_document_indices") as mock_get_indices, + patch(f"{TASKS_MODULE}.get_default_document_index") as mock_get_index, patch(f"{TASKS_MODULE}.get_active_search_settings") as mock_get_ss, - patch(f"{TASKS_MODULE}.httpx_init_vespa_pool") as mock_vespa_pool, ): delete_user_file_impl( user_file_id=str(uf.id), @@ -377,9 +376,8 @@ class TestDeleteImplNoVectorDb: redis_locking=False, ) - mock_get_indices.assert_not_called() + mock_get_index.assert_not_called() mock_get_ss.assert_not_called() - mock_vespa_pool.assert_not_called() session.delete.assert_called_once_with(uf) session.commit.assert_called_once() @@ -433,9 +431,8 @@ class TestProjectSyncImplNoVectorDb: f"{TASKS_MODULE}.fetch_user_files_with_access_relationships", return_value=[uf], ), - patch(f"{TASKS_MODULE}.get_all_document_indices") as mock_get_indices, + patch(f"{TASKS_MODULE}.get_default_document_index") as mock_get_index, patch(f"{TASKS_MODULE}.get_active_search_settings") as mock_get_ss, - patch(f"{TASKS_MODULE}.httpx_init_vespa_pool") as mock_vespa_pool, ): project_sync_user_file_impl( user_file_id=str(uf.id), @@ -443,9 +440,8 @@ class TestProjectSyncImplNoVectorDb: redis_locking=False, ) - mock_get_indices.assert_not_called() + mock_get_index.assert_not_called() mock_get_ss.assert_not_called() - mock_vespa_pool.assert_not_called() @patch(f"{TASKS_MODULE}.DISABLE_VECTOR_DB", True) @patch(f"{TASKS_MODULE}.get_session_with_current_tenant") 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 b6d9f02244..98fa87f774 100644 --- a/backend/tests/unit/onyx/db/engine/test_engine_disposal.py +++ b/backend/tests/unit/onyx/db/engine/test_engine_disposal.py @@ -8,7 +8,7 @@ and the api server hung. These tests exercise each disposal function directly. A higher-level integration check would re-run the FastAPI lifespan, but the lifespan touches -auth, telemetry, vespa, and the file store — far more scaffolding than the +auth, telemetry, OpenSearch, and the file store — far more scaffolding than the behavior we care about here, which is that ``dispose()`` is in fact called and the cached engine references are released. """ diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_chunk_presence_lookup.py b/backend/tests/unit/onyx/document_index/opensearch/test_chunk_presence_lookup.py index 77e1f3df54..2277ffb845 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_chunk_presence_lookup.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_chunk_presence_lookup.py @@ -8,7 +8,7 @@ from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, ) diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_get_doc_chunk_id.py b/backend/tests/unit/onyx/document_index/opensearch/test_get_doc_chunk_id.py index 31098d4fea..ee8c816452 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_get_doc_chunk_id.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_get_doc_chunk_id.py @@ -1,6 +1,6 @@ import pytest -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE from onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id from onyx.document_index.opensearch.string_filtering import ( 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 665c032ea9..e863b4df88 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,7 +18,7 @@ from unittest.mock import MagicMock import pytest -from onyx.document_index.interfaces_new import ( +from onyx.document_index.interfaces import ( MetadataUpdateRequest, SecondaryIndexDocumentMissingError, TenantState, diff --git a/backend/tests/unit/onyx/document_index/opensearch/test_multi_source_opensearch.py b/backend/tests/unit/onyx/document_index/opensearch/test_multi_source_opensearch.py index 69e8ca8e56..83464cc595 100644 --- a/backend/tests/unit/onyx/document_index/opensearch/test_multi_source_opensearch.py +++ b/backend/tests/unit/onyx/document_index/opensearch/test_multi_source_opensearch.py @@ -5,7 +5,7 @@ import pytest from onyx.configs.constants import DocumentSource from onyx.context.search.models import IndexFilters, SearchDoc -from onyx.document_index.interfaces_new import MetadataUpdateRequest, TenantState +from onyx.document_index.interfaces import MetadataUpdateRequest, TenantState from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, convert_retrieved_opensearch_chunk_to_inference_chunk_uncleaned, 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 10d092c631..70b0af3e84 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 @@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch from onyx.access.models import DocumentAccess from onyx.configs.constants import DocumentSource from onyx.connectors.models import Document, TextSection -from onyx.document_index.interfaces_new import IndexingMetadata, TenantState +from onyx.document_index.interfaces import IndexingMetadata, TenantState from onyx.document_index.opensearch.opensearch_document_index import ( OpenSearchDocumentIndex, ) 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 6e474d75e8..06bd5c72d0 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 @@ -9,7 +9,7 @@ from unittest.mock import MagicMock, patch from pydantic import BaseModel -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState 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 cbc34852c5..6384035277 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 @@ -8,7 +8,7 @@ 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.interfaces import TenantState from onyx.document_index.opensearch.constants import ASSUMED_DOCUMENT_AGE_DAYS from onyx.document_index.opensearch.schema import ( CREATED_AT_FIELD_NAME, 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 adf9cc217c..32a65c6f4f 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,7 +12,7 @@ import pytest from onyx.context.search.enums import QueryType from onyx.context.search.models import IndexFilters from onyx.document_index.disabled import VECTOR_DB_DISABLED_ERROR, DisabledDocumentIndex -from onyx.document_index.interfaces_new import IndexingMetadata, MetadataUpdateRequest +from onyx.document_index.interfaces import IndexingMetadata, MetadataUpdateRequest ESCAPED_ERROR = re.escape(VECTOR_DB_DISABLED_ERROR) diff --git a/backend/tests/unit/onyx/document_index/test_document_index_factory.py b/backend/tests/unit/onyx/document_index/test_document_index_factory.py new file mode 100644 index 0000000000..da9ef0f1b5 --- /dev/null +++ b/backend/tests/unit/onyx/document_index/test_document_index_factory.py @@ -0,0 +1,140 @@ +"""Tests for the document index factory. + +`get_default_document_index` is the only way callers get a document index, for +both retrieval and indexing. These tests pin how it maps search settings to the +OpenSearch primary/secondary pair. +""" + +from collections.abc import Iterator +from typing import Any, cast +from unittest.mock import MagicMock, patch + +import pytest + +from onyx.db.enums import VectorQuantization +from onyx.db.models import SearchSettings +from onyx.document_index import factory +from onyx.document_index.disabled import DisabledDocumentIndex +from onyx.document_index.opensearch.opensearch_document_index import ( + OpenSearchIndexPair, +) + +_FACTORY = "onyx.document_index.factory" + + +def _search_settings(index_name: str) -> MagicMock: + search_settings = MagicMock(spec=SearchSettings) + search_settings.index_name = index_name + return search_settings + + +def _indexing_setting(embedding_dim: int) -> MagicMock: + indexing_setting = MagicMock() + indexing_setting.final_embedding_dim = embedding_dim + indexing_setting.vector_quantization = VectorQuantization.NONE + return indexing_setting + + +@pytest.fixture +def mock_opensearch_index() -> Iterator[MagicMock]: + """Replaces the OpenSearch index class so no client is built.""" + with ( + patch(f"{_FACTORY}.OpenSearchDocumentIndex") as index_cls, + patch(f"{_FACTORY}.get_current_tenant_id", return_value="tenant_a"), + patch(f"{_FACTORY}.MULTI_TENANT", False), + patch(f"{_FACTORY}.DISABLE_VECTOR_DB", False), + ): + # One distinct instance per construction, so primary != secondary. + index_cls.side_effect = lambda **_kwargs: MagicMock() + yield index_cls + + +def test_disable_vector_db_returns_disabled_index() -> None: + with patch(f"{_FACTORY}.DISABLE_VECTOR_DB", True): + document_index = factory.get_default_document_index( + _search_settings("primary_index"), None + ) + + assert isinstance(document_index, DisabledDocumentIndex) + + +def test_primary_only_builds_pair_without_secondary( + mock_opensearch_index: MagicMock, +) -> None: + primary_settings = _search_settings("primary_index") + with patch( + f"{_FACTORY}.IndexingSetting.from_db_model", + return_value=_indexing_setting(768), + ): + document_index = factory.get_default_document_index(primary_settings, None) + + assert isinstance(document_index, OpenSearchIndexPair) + assert document_index.secondary is None + mock_opensearch_index.assert_called_once() + kwargs = mock_opensearch_index.call_args.kwargs + assert kwargs["index_name"] == "primary_index" + assert kwargs["embedding_dim"] == 768 + assert kwargs["vector_quantization"] == VectorQuantization.NONE + assert kwargs["tenant_state"].tenant_id == "tenant_a" + assert kwargs["tenant_state"].multitenant is False + + +def test_secondary_settings_build_secondary_index( + mock_opensearch_index: MagicMock, +) -> None: + primary_settings = _search_settings("primary_index") + secondary_settings = _search_settings("secondary_index") + + def _from_db_model(search_settings: MagicMock) -> MagicMock: + return _indexing_setting(768 if search_settings is primary_settings else 1024) + + with patch(f"{_FACTORY}.IndexingSetting.from_db_model", side_effect=_from_db_model): + document_index = factory.get_default_document_index( + primary_settings, secondary_settings + ) + + assert isinstance(document_index, OpenSearchIndexPair) + assert document_index.secondary is not None + assert document_index.primary is not document_index.secondary + index_names = [ + call.kwargs["index_name"] for call in mock_opensearch_index.call_args_list + ] + assert index_names == ["primary_index", "secondary_index"] + embedding_dims = [ + call.kwargs["embedding_dim"] for call in mock_opensearch_index.call_args_list + ] + assert embedding_dims == [768, 1024] + + +@pytest.mark.parametrize("with_secondary", [False, True]) +@pytest.mark.parametrize("backfill", [False, True]) +def test_primary_backfill_flag_reaches_the_pair( + mock_opensearch_index: MagicMock, # noqa: ARG001 + with_secondary: bool, + backfill: bool, +) -> None: + """The INSTANT reindex-port flag must reach the pair on both the + primary-only and the primary+secondary paths.""" + secondary_settings = _search_settings("secondary_index") if with_secondary else None + with patch( + f"{_FACTORY}.IndexingSetting.from_db_model", + return_value=_indexing_setting(768), + ): + document_index = factory.get_default_document_index( + _search_settings("primary_index"), + secondary_settings, + primary_backfill_in_progress=backfill, + ) + + assert isinstance(document_index, OpenSearchIndexPair) + assert document_index._primary_backfill_in_progress is backfill + + +def test_primary_backfill_flag_is_keyword_only() -> None: + """A positional third argument used to be the DB session. Keeping the flag + keyword-only makes a stale positional call fail loudly instead of passing a + truthy session as the flag.""" + # Typed as Any so the type checker lets the stale call through to runtime. + get_default_document_index = cast(Any, factory.get_default_document_index) + with pytest.raises(TypeError): + get_default_document_index(_search_settings("primary_index"), None, True) diff --git a/backend/tests/unit/onyx/document_index/test_tenant_scoping.py b/backend/tests/unit/onyx/document_index/test_tenant_scoping.py index 24432c394c..9fe1b1973f 100644 --- a/backend/tests/unit/onyx/document_index/test_tenant_scoping.py +++ b/backend/tests/unit/onyx/document_index/test_tenant_scoping.py @@ -1,13 +1,11 @@ -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from opensearchpy import NotFoundError -from onyx.context.search.models import IndexFilters -from onyx.document_index.interfaces_new import TenantState +from onyx.document_index.interfaces import TenantState from onyx.document_index.opensearch.client import OpenSearchIndexClient from onyx.document_index.opensearch.schema import TENANT_ID_FIELD_NAME -from onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex @pytest.mark.parametrize("expired", [False, True]) @@ -42,36 +40,3 @@ def test_pit_scan_keeps_tenant_filter_on_every_request(expired: bool) -> None: assert {"term": {TENANT_ID_FIELD_NAME: {"value": "tenant-a"}}} in call.kwargs[ "body" ]["query"]["bool"]["filter"] - - -@pytest.mark.parametrize("batch_retrieval", [False, True]) -@pytest.mark.parametrize("supplied_tenant", [None, "tenant-b"]) -def test_vespa_id_retrieval_pins_tenant_without_mutating_filters( - batch_retrieval: bool, supplied_tenant: str | None -) -> None: - index = VespaDocumentIndex( - index_name="test-index", - tenant_state=TenantState(tenant_id="tenant-a", multitenant=True), - large_chunks_enabled=False, - ) - filters = IndexFilters( - access_control_list=["user:reader"], tenant_id=supplied_tenant - ) - with ( - patch( - "onyx.document_index.vespa.vespa_document_index.batch_search_api_retrieval", - return_value=[], - ) as batch, - patch( - "onyx.document_index.vespa.vespa_document_index.parallel_visit_api_retrieval", - return_value=[], - ) as visit, - ): - assert ( - index.id_based_retrieval([], filters, batch_retrieval=batch_retrieval) == [] - ) - retrieval: MagicMock = batch if batch_retrieval else visit - effective_filters = retrieval.call_args.kwargs["filters"] - assert effective_filters.tenant_id == "tenant-a" - assert effective_filters.access_control_list == ["user:reader"] - assert filters.tenant_id == supplied_tenant 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 deleted file mode 100644 index 8f08c0228e..0000000000 --- a/backend/tests/unit/onyx/document_index/vespa/test_vespa_batch_flush.py +++ /dev/null @@ -1,147 +0,0 @@ -"""Unit tests for VespaDocumentIndex.index(). - -These tests mock all external I/O (HTTP calls, thread pools) and verify -the streaming logic, ID cleaning/mapping, and DocumentInsertionRecord -construction. -""" - -from unittest.mock import MagicMock, patch - -from onyx.access.models import DocumentAccess -from onyx.configs.constants import DocumentSource -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, DocMetadataAwareIndexChunk, IndexChunk - - -def _make_chunk( - doc_id: str, - chunk_id: int = 0, - content: str = "test content", -) -> DocMetadataAwareIndexChunk: - doc = Document( - id=doc_id, - semantic_identifier="test_doc", - sections=[TextSection(text=content, link=None)], - source=DocumentSource.NOT_APPLICABLE, - metadata={}, - ) - index_chunk = IndexChunk( - chunk_id=chunk_id, - blurb=content[:50], - content=content, - source_links=None, - image_file_id=None, - section_continuation=False, - source_document=doc, - title_prefix="", - metadata_suffix_semantic="", - metadata_suffix_keyword="", - contextual_rag_reserved_tokens=0, - doc_summary="", - chunk_context="", - mini_chunk_texts=None, - large_chunk_id=None, - embeddings=ChunkEmbedding( - full_embedding=[0.1] * 10, - mini_chunk_embeddings=[], - ), - title_embedding=None, - ) - access = DocumentAccess.build( - user_emails=[], - user_groups=[], - external_user_emails=[], - external_user_group_ids=[], - is_public=True, - ) - return DocMetadataAwareIndexChunk.from_index_chunk( - index_chunk=index_chunk, - access=access, - document_sets=set(), - user_project=[], - personas=[], - boost=0, - aggregated_chunk_boost_factor=1.0, - tenant_id="test_tenant", - ) - - -def _make_indexing_metadata( - doc_ids: list[str], - old_counts: list[int], - new_counts: list[int], -) -> IndexingMetadata: - return IndexingMetadata( - doc_id_to_chunk_cnt_diff={ - doc_id: IndexingMetadata.ChunkCounts( - old_chunk_cnt=old, - new_chunk_cnt=new, - ) - for doc_id, old, new in zip(doc_ids, old_counts, new_counts, strict=True) - } - ) - - -def _stub_enrich( - doc_id: str, - old_chunk_cnt: int, -) -> EnrichedDocumentIndexingInfo: - """Build an EnrichedDocumentIndexingInfo that says 'no chunks to delete' - when old_chunk_cnt == 0, or 'has existing chunks' otherwise.""" - return EnrichedDocumentIndexingInfo( - doc_id=doc_id, - chunk_start_index=0, - old_version=False, - chunk_end_index=old_chunk_cnt, - ) - - -@patch("onyx.document_index.vespa.vespa_document_index.batch_index_vespa_chunks") -@patch("onyx.document_index.vespa.vespa_document_index.delete_vespa_chunks") -@patch( - "onyx.document_index.vespa.vespa_document_index.get_document_chunk_ids", - return_value=[], -) -@patch("onyx.document_index.vespa.vespa_document_index._enrich_basic_chunk_info") -@patch( - "onyx.document_index.vespa.vespa_document_index.BATCH_SIZE", - 3, -) -def test_index_respects_batch_size( - mock_enrich: MagicMock, - mock_get_chunk_ids: MagicMock, # noqa: ARG001 - mock_delete: MagicMock, # noqa: ARG001 - mock_batch_index: MagicMock, -) -> None: - """When chunks exceed BATCH_SIZE, batch_index_vespa_chunks is called - multiple times with correctly sized batches.""" - mock_enrich.return_value = _stub_enrich("doc1", old_chunk_cnt=0) - - index = VespaDocumentIndex( - index_name="test_index", - tenant_state=TenantState(tenant_id="test_tenant", multitenant=False), - large_chunks_enabled=False, - httpx_client=MagicMock(), - ) - - chunks = [_make_chunk("doc1", chunk_id=i) for i in range(7)] - metadata = _make_indexing_metadata(["doc1"], old_counts=[0], new_counts=[7]) - - results = index.index(chunks=chunks, indexing_metadata=metadata) - - assert len(results) == 1 - - # With BATCH_SIZE=3 and 7 chunks: batches of 3, 3, 1 - assert mock_batch_index.call_count == 3 - batch_sizes = [len(c.kwargs["chunks"]) for c in mock_batch_index.call_args_list] - assert batch_sizes == [3, 3, 1] - - # Verify all chunks are accounted for and in order - all_indexed = [ - chunk for c in mock_batch_index.call_args_list for chunk in c.kwargs["chunks"] - ] - assert len(all_indexed) == 7 - assert [c.chunk_id for c in all_indexed] == list(range(7)) diff --git a/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py b/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py index 7c65519c24..ce702b7447 100644 --- a/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py +++ b/backend/tests/unit/onyx/indexing/test_indexing_pipeline.py @@ -315,7 +315,7 @@ def _make_cc_pair(is_public: bool) -> MagicMock: def _make_insertion_records(doc_ids: list[str]) -> list[Any]: - from onyx.document_index.interfaces_new import DocumentInsertionRecord + from onyx.document_index.interfaces import DocumentInsertionRecord return [ DocumentInsertionRecord(document_id=d, already_existed=False) for d in doc_ids @@ -698,7 +698,7 @@ def test_run_pipeline_owns_llm_enrichment_trace() -> None: document_batch=[document], request_id=None, embedder=MagicMock(), - document_indices=[], + document_index=MagicMock(), db_session=MagicMock(), tenant_id="tenant", adapter=MagicMock(), @@ -789,7 +789,7 @@ def test_unavailable_vision_llm_does_not_enable_spend_gate() -> None: document_batch=[document], request_id=None, embedder=MagicMock(), - document_indices=[], + document_index=MagicMock(), db_session=MagicMock(), tenant_id="tenant", adapter=MagicMock(), @@ -887,7 +887,7 @@ def test_index_batch_returns_spend_limit_failures_before_contextual_rag() -> Non document_batch=[document], chunker=chunker, embedder=MagicMock(), - document_indices=[], + document_index=MagicMock(), request_id=None, tenant_id="tenant", adapter=adapter, diff --git a/backend/tests/unit/onyx/indexing/test_vespa.py b/backend/tests/unit/onyx/indexing/test_vespa.py deleted file mode 100644 index 6f1c70ea5b..0000000000 --- a/backend/tests/unit/onyx/indexing/test_vespa.py +++ /dev/null @@ -1,95 +0,0 @@ -from http import HTTPStatus -from typing import Any - -import httpx -import pytest -from sqlalchemy.orm import Session - -from onyx.db.engine.sql_engine import get_sqlalchemy_engine -from onyx.document_index.document_index_utils import get_both_index_properties -from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT - - -@pytest.mark.skip() -def test_vespa_update() -> None: - """This Test exercises some ambiguous Vespa behavior and - shows exactly what happens. - """ - - doc_id = "test-vespa-update" - - with Session(get_sqlalchemy_engine()) as db_session: - primary_index_name, _, _, _ = get_both_index_properties(db_session) - endpoint = ( - f"{DOCUMENT_ID_ENDPOINT.format(index_name=primary_index_name)}/{doc_id}" - ) - with httpx.Client(http2=True) as http_client: - payload: dict[str, Any] = {} - - # always delete to set up the test, should always be OK - res = http_client.delete(endpoint) - assert HTTPStatus.OK == res.status_code - - # Verify the document is not found - res = http_client.get(endpoint) - assert HTTPStatus.NOT_FOUND == res.status_code - - # Attempt to update a nonexistent test document. Should return OK - payload["fields"] = {} - payload["fields"]["title"] = {"assign": "Best of Bob Dylan"} - - res = http_client.put( - endpoint, - headers={"Content-Type": "application/json"}, - json=payload, - ) - assert HTTPStatus.OK == res.status_code - - # when we look for it, should be NOT_FOUND - res = http_client.get(endpoint) - assert HTTPStatus.NOT_FOUND == res.status_code - - # POST/Put new document - payload = {} - payload["fields"] = {} - payload["fields"]["document_id"] = doc_id - payload["fields"]["title"] = "A Head Full of Dreams" - - res = http_client.post( - endpoint, - headers={"Content-Type": "application/json"}, - json=payload, - ) - assert HTTPStatus.OK == res.status_code - - # when we look for it, now we should find it - res = http_client.get(endpoint) - assert HTTPStatus.OK == res.status_code - d = res.json() - - assert payload["fields"]["title"] == d["fields"]["title"] - - # Attempt to update the document that we know exists. Should return OK - payload["fields"] = {} - payload["fields"]["title"] = {"assign": "Remember The Name"} - - res = http_client.put( - endpoint, - headers={"Content-Type": "application/json"}, - json=payload, - ) - assert HTTPStatus.OK == res.status_code - - # verify the change - res = http_client.get(endpoint) - assert HTTPStatus.OK == res.status_code - d = res.json() - assert payload["fields"]["title"]["assign"] == d["fields"]["title"] - - # always delete to clean up the test, should always be OK - res = http_client.delete(endpoint) - assert HTTPStatus.OK == res.status_code - - # Verify the document is not found - res = http_client.get(endpoint) - assert HTTPStatus.NOT_FOUND == res.status_code diff --git a/backend/tests/unit/onyx/kg/test_kg_extraction_tracing.py b/backend/tests/unit/onyx/kg/test_kg_extraction_tracing.py deleted file mode 100644 index f8552c0b25..0000000000 --- a/backend/tests/unit/onyx/kg/test_kg_extraction_tracing.py +++ /dev/null @@ -1,37 +0,0 @@ -from contextlib import nullcontext -from unittest.mock import MagicMock, patch - -from onyx.kg.utils.extraction_utils import ( - KG_DOCUMENT_PROCESSING_TRACE_NAME, - kg_deep_extraction, -) -from onyx.tracing.framework.traces import TraceContentMode - -_EXTRACTION_UTILS = "onyx.kg.utils.extraction_utils" - - -def test_kg_deep_extraction_owns_document_trace() -> None: - expected_result = MagicMock() - with ( - patch( - f"{_EXTRACTION_UTILS}._kg_deep_extraction", - return_value=expected_result, - ), - patch( - f"{_EXTRACTION_UTILS}.ensure_trace", return_value=nullcontext() - ) as ensure, - ): - result = kg_deep_extraction( - document_id="document", - metadata=MagicMock(), - implied_extraction=MagicMock(), - tenant_id="tenant", - index_name="index", - kg_config_settings=MagicMock(), - ) - - assert result is expected_result - ensure.assert_called_once_with( - KG_DOCUMENT_PROCESSING_TRACE_NAME, - content_mode=TraceContentMode.METADATA_ONLY, - ) diff --git a/backend/tests/unit/onyx/server/manage/test_get_versions.py b/backend/tests/unit/onyx/server/manage/test_get_versions.py new file mode 100644 index 0000000000..f041e60bf4 --- /dev/null +++ b/backend/tests/unit/onyx/server/manage/test_get_versions.py @@ -0,0 +1,49 @@ +"""Tests for the public /versions endpoint, with DockerHub mocked out.""" + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +from onyx.server.manage.get_state import DOCUMENT_INDEX_IMAGE, get_versions + +_COMPOSE_TEMPLATE = ( + Path(__file__).resolve().parents[6] + / "deployment" + / "docker_compose" + / "docker-compose.template.yml" +) + + +def _dockerhub_page(_url: str, timeout: int) -> MagicMock: # noqa: ARG001 + response = MagicMock() + page: dict[str, Any] = { + "results": [ + {"name": "v3.1.0"}, + {"name": "v3.2.0"}, + {"name": "v3.3.0-beta.1"}, + {"name": "v3.3.0-beta.2"}, + {"name": "latest"}, + ], + "next": None, + } + response.json.return_value = page + return response + + +def test_versions_report_the_opensearch_index_image() -> None: + with patch( + "onyx.server.manage.get_state.requests.get", side_effect=_dockerhub_page + ): + versions = get_versions() + + assert versions.stable.onyx == "v3.2.0" + assert versions.dev.onyx == "v3.3.0-beta.2" + for container_versions in (versions.stable, versions.dev, versions.migration): + assert container_versions.index == DOCUMENT_INDEX_IMAGE + assert "vespa" not in container_versions.index + + +def test_index_image_matches_the_docker_compose_deployment() -> None: + """The /versions index image must be the one docker compose deploys.""" + template = _COMPOSE_TEMPLATE.read_text() + assert f"/{DOCUMENT_INDEX_IMAGE}\n" in template diff --git a/backend/tests/unit/onyx/server/test_settings_store.py b/backend/tests/unit/onyx/server/test_settings_store.py index 5717186bd2..498be7ef60 100644 --- a/backend/tests/unit/onyx/server/test_settings_store.py +++ b/backend/tests/unit/onyx/server/test_settings_store.py @@ -168,3 +168,28 @@ def test_load_settings_default_clamped_to_max( settings = settings_store.load_settings() assert settings.user_file_max_upload_size_mb == 50 + + +@pytest.mark.parametrize("stored_value", [None, False, True]) +@pytest.mark.parametrize("vector_db_disabled", [False, True]) +def test_load_settings_reports_opensearch_indexing_enabled( + monkeypatch: pytest.MonkeyPatch, + stored_value: bool | None, + vector_db_disabled: bool, +) -> None: + """`opensearch_indexing_enabled` is deprecated and always True. API clients + such as the Terraform provider still read it, so a stale stored value must + not leak through.""" + stored = ( + None + if stored_value is None + else Settings(opensearch_indexing_enabled=stored_value).model_dump() + ) + monkeypatch.setattr(settings_store, "get_kv_store", lambda: _FakeKvStore(stored)) + monkeypatch.setattr(settings_store, "get_cache_backend", lambda: _FakeCache()) + monkeypatch.setattr(settings_store, "DISABLE_VECTOR_DB", vector_db_disabled) + + settings = settings_store.load_settings() + + assert settings.opensearch_indexing_enabled is True + assert settings.model_dump()["opensearch_indexing_enabled"] is True diff --git a/backend/tests/unit/onyx/test_setup_document_index.py b/backend/tests/unit/onyx/test_setup_document_index.py new file mode 100644 index 0000000000..f7d0716608 --- /dev/null +++ b/backend/tests/unit/onyx/test_setup_document_index.py @@ -0,0 +1,69 @@ +"""Tests for `setup_document_index`, the startup check that the document index +exists and matches the expected schema.""" + +from collections.abc import Iterator +from unittest.mock import MagicMock, patch + +import pytest + +from onyx.document_index.opensearch.client import OpenSearchIndexWriteBlockedError +from onyx.setup import setup_document_index + + +@pytest.fixture(autouse=True) +def _no_sleep() -> Iterator[MagicMock]: + with patch("onyx.setup.time.sleep") as sleep: + yield sleep + + +def _index_setting(embedding_dim: int = 768) -> MagicMock: + index_setting = MagicMock() + index_setting.final_embedding_dim = embedding_dim + return index_setting + + +def test_success_on_first_attempt() -> None: + document_index = MagicMock() + + assert setup_document_index(document_index, _index_setting(768), num_attempts=3) + + document_index.verify_and_create_index_if_necessary.assert_called_once_with( + embedding_dim=768 + ) + + +def test_retries_until_the_index_is_ready(_no_sleep: MagicMock) -> None: + document_index = MagicMock() + document_index.verify_and_create_index_if_necessary.side_effect = [ + ConnectionError("not ready"), + ConnectionError("not ready"), + None, + ] + + assert setup_document_index(document_index, _index_setting(), num_attempts=5) + + assert document_index.verify_and_create_index_if_necessary.call_count == 3 + assert _no_sleep.call_count == 2 + + +def test_gives_up_after_the_attempt_limit() -> None: + document_index = MagicMock() + document_index.verify_and_create_index_if_necessary.side_effect = ConnectionError( + "down" + ) + + assert not setup_document_index(document_index, _index_setting(), num_attempts=3) + + assert document_index.verify_and_create_index_if_necessary.call_count == 3 + + +def test_write_blocked_index_starts_degraded() -> None: + """A readable but write-blocked index must not stop startup.""" + document_index = MagicMock() + document_index.verify_and_create_index_if_necessary.side_effect = ( + OpenSearchIndexWriteBlockedError("read_only_allow_delete") + ) + + assert setup_document_index(document_index, _index_setting(), num_attempts=3) + + document_index.verify_and_create_index_if_necessary.assert_called_once() diff --git a/backend/tests/unit/onyx/tools/test_search_utils.py b/backend/tests/unit/onyx/tools/test_search_utils.py index a54f6237ef..6538b52d3d 100644 --- a/backend/tests/unit/onyx/tools/test_search_utils.py +++ b/backend/tests/unit/onyx/tools/test_search_utils.py @@ -1,11 +1,15 @@ """Unit tests for search utility functions.""" from typing import NamedTuple +from unittest.mock import MagicMock import pytest +from onyx.configs.constants import DocumentSource +from onyx.context.search.models import InferenceChunk, InferenceSection from onyx.tools.tool_implementations.search.search_tool import deduplicate_queries from onyx.tools.tool_implementations.search.search_utils import ( + _retrieve_adjacent_chunks, weighted_reciprocal_rank_fusion, ) @@ -570,3 +574,123 @@ class TestDeduplicateQueries: assert len(result) == 1 assert result[0][0] == "Café" assert result[0][1] == 4.5 + + +# ============================================================================= +# Tests for _retrieve_adjacent_chunks +# ============================================================================= + + +def _inference_chunk(document_id: str, chunk_id: int) -> InferenceChunk: + return InferenceChunk( + document_id=document_id, + chunk_id=chunk_id, + content=f"content-{chunk_id}", + source_type=DocumentSource.MOCK_CONNECTOR, + semantic_identifier=document_id, + title=document_id, + boost=1, + score=0.5, + hidden=False, + metadata={}, + match_highlights=[], + doc_summary="", + chunk_context="", + updated_at=None, + image_file_id=None, + source_links={}, + section_continuation=False, + blurb=f"blurb-{chunk_id}", + ) + + +def _section(document_id: str, chunk_ids: list[int]) -> InferenceSection: + chunks = [_inference_chunk(document_id, chunk_id) for chunk_id in chunk_ids] + return InferenceSection( + center_chunk=chunks[0], + chunks=chunks, + combined_content=" ".join(chunk.content for chunk in chunks), + ) + + +class TestRetrieveAdjacentChunks: + """`_retrieve_adjacent_chunks` expands a search hit with its neighbors.""" + + def test_passes_the_raw_document_id(self) -> None: + """OpenSearch stores document IDs unchanged, so the lookup must not + rewrite them. The former Vespa backend replaced `'` with `_`, which made + expansion miss every document whose ID has an apostrophe.""" + document_id = "https://example.com/o'brien's-notes" + document_index = MagicMock() + document_index.id_based_retrieval.return_value = [] + + _retrieve_adjacent_chunks( + section=_section(document_id, [3, 4]), + document_index=document_index, + num_chunks_above=2, + num_chunks_below=2, + ) + + requested_ids = { + call.kwargs["chunk_requests"][0].document_id + for call in document_index.id_based_retrieval.call_args_list + } + assert requested_ids == {document_id} + + def test_requests_the_ranges_around_the_section(self) -> None: + document_index = MagicMock() + document_index.id_based_retrieval.return_value = [] + + _retrieve_adjacent_chunks( + section=_section("doc", [3, 4]), + document_index=document_index, + num_chunks_above=2, + num_chunks_below=3, + ) + + ranges = [ + ( + call.kwargs["chunk_requests"][0].min_chunk_ind, + call.kwargs["chunk_requests"][0].max_chunk_ind, + ) + for call in document_index.id_based_retrieval.call_args_list + ] + assert ranges == [(1, 2), (5, 7)] + for call in document_index.id_based_retrieval.call_args_list: + # Permissions were already enforced when the section was found. + assert call.kwargs["filters"].access_control_list is None + assert call.kwargs["batch_retrieval"] is True + + def test_skips_above_request_at_document_start(self) -> None: + document_index = MagicMock() + document_index.id_based_retrieval.return_value = [] + + _retrieve_adjacent_chunks( + section=_section("doc", [0, 1]), + document_index=document_index, + num_chunks_above=2, + num_chunks_below=1, + ) + + assert document_index.id_based_retrieval.call_count == 1 + request = document_index.id_based_retrieval.call_args.kwargs["chunk_requests"][ + 0 + ] + assert (request.min_chunk_ind, request.max_chunk_ind) == (2, 2) + + def test_sorts_results_and_tolerates_index_errors(self) -> None: + document_index = MagicMock() + document_index.id_based_retrieval.side_effect = [ + [_inference_chunk("doc", 2), _inference_chunk("doc", 1)], + RuntimeError("index unavailable"), + ] + + above, below = _retrieve_adjacent_chunks( + section=_section("doc", [3]), + document_index=document_index, + num_chunks_above=2, + num_chunks_below=2, + ) + + assert [chunk.chunk_id for chunk in above] == [1, 2] + assert below == [] diff --git a/backend/tests/unit/onyx/document_index/vespa/shared_utils/test_utils.py b/backend/tests/unit/onyx/utils/test_text_processing.py similarity index 100% rename from backend/tests/unit/onyx/document_index/vespa/shared_utils/test_utils.py rename to backend/tests/unit/onyx/utils/test_text_processing.py diff --git a/backend/tests/unit/onyx/utils/test_vespa_query.py b/backend/tests/unit/onyx/utils/test_vespa_query.py deleted file mode 100644 index 0a2acc4f16..0000000000 --- a/backend/tests/unit/onyx/utils/test_vespa_query.py +++ /dev/null @@ -1,363 +0,0 @@ -from datetime import datetime, timedelta, timezone - -from onyx.configs.constants import INDEX_SEPARATOR, DocumentSource -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, - DOCUMENT_SETS, - HIDDEN, - METADATA_LIST, - PERSONAS, - SOURCE_TYPE, - TENANT_ID, - USER_PROJECT, -) -from shared_configs.configs import MULTI_TENANT - - -class TestBuildVespaFilters: - def test_empty_filters(self) -> None: - """Test with empty filters object.""" - filters = IndexFilters(access_control_list=[]) - result = build_vespa_filters(filters) - assert result == f"!({HIDDEN}=true) and " - - # With trailing AND removed - result = build_vespa_filters(filters, remove_trailing_and=True) - assert result == f"!({HIDDEN}=true)" - - def test_include_hidden(self) -> None: - """Test with include_hidden flag.""" - filters = IndexFilters(access_control_list=[]) - result = build_vespa_filters(filters, include_hidden=True) - assert result == "" # No filters applied when including hidden - - # With some other filter to ensure proper AND chaining - filters = IndexFilters(access_control_list=[], source_type=[DocumentSource.WEB]) - result = build_vespa_filters(filters, include_hidden=True) - assert result == f'({SOURCE_TYPE} contains "web") and ' - - def test_acl(self) -> None: - """Test with acls — uses weightedSet operator for efficient matching.""" - # Single ACL - filters = IndexFilters(access_control_list=["user1"]) - result = build_vespa_filters(filters) - assert ( - result - == f'!({HIDDEN}=true) and weightedSet(access_control_list, {{"user1":1}}) and ' - ) - - # Multiple ACL's - filters = IndexFilters(access_control_list=["user2", "group2"]) - result = build_vespa_filters(filters) - assert ( - result - == f'!({HIDDEN}=true) and weightedSet(access_control_list, {{"user2":1, "group2":1}}) and ' - ) - - def test_tenant_filter(self) -> None: - """Test tenant ID filtering.""" - # With tenant ID - if MULTI_TENANT: - filters = IndexFilters(access_control_list=[], tenant_id="tenant1") - result = build_vespa_filters(filters) - assert ( - f'!({HIDDEN}=true) and ({TENANT_ID} contains "tenant1") and ' == result - ) - - # No tenant ID - filters = IndexFilters(access_control_list=[], tenant_id=None) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result - - def test_source_type_filter(self) -> None: - """Test source type filtering.""" - # Single source type - filters = IndexFilters(access_control_list=[], source_type=[DocumentSource.WEB]) - result = build_vespa_filters(filters) - assert f'!({HIDDEN}=true) and ({SOURCE_TYPE} contains "web") and ' == result - - # Multiple source types - filters = IndexFilters( - access_control_list=[], - source_type=[DocumentSource.WEB, DocumentSource.JIRA], - ) - result = build_vespa_filters(filters) - assert ( - f'!({HIDDEN}=true) and ({SOURCE_TYPE} contains "web" or {SOURCE_TYPE} contains "jira") and ' - == result - ) - - # Empty source type list - filters = IndexFilters(access_control_list=[], source_type=[]) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result - - def test_tag_filters(self) -> None: - """Test tag filtering.""" - # Single tag - filters = IndexFilters( - access_control_list=[], tags=[Tag(tag_key="color", tag_value="red")] - ) - result = build_vespa_filters(filters) - assert ( - f'!({HIDDEN}=true) and ({METADATA_LIST} contains "color{INDEX_SEPARATOR}red") and ' - == result - ) - - # Multiple tags - filters = IndexFilters( - access_control_list=[], - tags=[ - Tag(tag_key="color", tag_value="red"), - Tag(tag_key="size", tag_value="large"), - ], - ) - result = build_vespa_filters(filters) - expected = ( - f'!({HIDDEN}=true) and ({METADATA_LIST} contains "color{INDEX_SEPARATOR}red" ' - f'or {METADATA_LIST} contains "size{INDEX_SEPARATOR}large") and ' - ) - assert expected == result - - # Empty tags list - filters = IndexFilters(access_control_list=[], tags=[]) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result - - def test_document_sets_filter(self) -> None: - """Test document sets filtering.""" - # Single document set - filters = IndexFilters(access_control_list=[], document_set=["set1"]) - result = build_vespa_filters(filters) - assert f'!({HIDDEN}=true) and ({DOCUMENT_SETS} contains "set1") and ' == result - - # Multiple document sets - filters = IndexFilters(access_control_list=[], document_set=["set1", "set2"]) - result = build_vespa_filters(filters) - assert ( - f'!({HIDDEN}=true) and ({DOCUMENT_SETS} contains "set1" or {DOCUMENT_SETS} contains "set2") and ' - == result - ) - - # Empty document sets - filters = IndexFilters(access_control_list=[], document_set=[]) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result - - def test_user_project_filter(self) -> None: - """Test user project filtering. - - project_id_filter alone does NOT trigger a knowledge scope restriction - (an agent with no explicit knowledge should search everything). - It only participates when explicit knowledge filters are present. - """ - # project_id_filter alone → no restriction - filters = IndexFilters(access_control_list=[], project_id_filter=789) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result - - # project_id_filter with document_set → both OR'd - filters = IndexFilters( - access_control_list=[], project_id_filter=789, document_set=["set1"] - ) - result = build_vespa_filters(filters) - assert ( - f'!({HIDDEN}=true) and (({DOCUMENT_SETS} contains "set1") or ({USER_PROJECT} contains "789")) and ' - == result - ) - - # No project id filter - filters = IndexFilters(access_control_list=[], project_id_filter=None) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result - - def test_time_cutoff_filter(self) -> None: - """Test time cutoff filtering.""" - # With cutoff time - cutoff_time = datetime(2023, 1, 1, tzinfo=timezone.utc) - filters = IndexFilters( - access_control_list=[], updated_at_range=TimeRange(start=cutoff_time) - ) - result = build_vespa_filters(filters) - cutoff_secs = int(cutoff_time.timestamp()) - assert ( - f"!({HIDDEN}=true) and !({DOC_UPDATED_AT} < {cutoff_secs}) and " == result - ) - - # No cutoff time - filters = IndexFilters(access_control_list=[], updated_at_range=None) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result - - # Test untimed logic (when cutoff is old enough) - old_cutoff = datetime.now(timezone.utc) - timedelta(days=100) - filters = IndexFilters( - access_control_list=[], updated_at_range=TimeRange(start=old_cutoff) - ) - result = build_vespa_filters(filters) - old_cutoff_secs = int(old_cutoff.timestamp()) - assert ( - f"!({HIDDEN}=true) and !({DOC_UPDATED_AT} < {old_cutoff_secs}) and " - == result - ) - - def test_time_range_filter(self) -> None: - """A bounded updated_at_range emits strict >= and <= clauses and - excludes untimed docs (no `!(... < ...)` form), even when the lower - bound is old.""" - start = datetime(2020, 1, 1, tzinfo=timezone.utc) - end = datetime(2020, 1, 31, 23, 59, 59, tzinfo=timezone.utc) - filters = IndexFilters( - access_control_list=[], - updated_at_range=TimeRange(start=start, end=end), - ) - result = build_vespa_filters(filters) - start_secs = int(start.timestamp()) - end_secs = int(end.timestamp()) - assert ( - f"!({HIDDEN}=true) and ({DOC_UPDATED_AT} >= {start_secs}) and " - f"({DOC_UPDATED_AT} <= {end_secs}) and " == result - ) - - def test_time_upper_bound_only(self) -> None: - """An upper bound alone (via updated_at_range) emits just a <= clause.""" - end = datetime(2023, 6, 1, tzinfo=timezone.utc) - filters = IndexFilters( - access_control_list=[], - updated_at_range=TimeRange(end=end), - ) - result = build_vespa_filters(filters) - end_secs = int(end.timestamp()) - assert f"!({HIDDEN}=true) and ({DOC_UPDATED_AT} <= {end_secs}) and " == result - - def test_combined_filters(self) -> None: - """Test combining multiple filter types. - - Knowledge-scope filters (document_set, project_id_filter, persona_id_filter) - are OR'd together, while all other filters are AND'd. - """ - filters = IndexFilters( - access_control_list=["user1", "group1"], - source_type=[DocumentSource.WEB], - tags=[Tag(tag_key="color", tag_value="red")], - document_set=["set1"], - project_id_filter=789, - persona_id_filter=42, - updated_at_range=TimeRange(start=datetime(2023, 1, 1, tzinfo=timezone.utc)), - ) - - result = build_vespa_filters(filters) - - expected = f"!({HIDDEN}=true) and " - expected += 'weightedSet(access_control_list, {"user1":1, "group1":1}) and ' - expected += f'({SOURCE_TYPE} contains "web") and ' - expected += f'({METADATA_LIST} contains "color{INDEX_SEPARATOR}red") and ' - # Knowledge scope filters are OR'd together - # (persona_id_filter is primary, project_id_filter is additive — order reflects this) - expected += ( - f'(({DOCUMENT_SETS} contains "set1")' - f' or ({PERSONAS} contains "42")' - f' or ({USER_PROJECT} contains "789")' - f") and " - ) - cutoff_secs = int(datetime(2023, 1, 1, tzinfo=timezone.utc).timestamp()) - expected += f"!({DOC_UPDATED_AT} < {cutoff_secs}) and " - - assert expected == result - - # With trailing AND removed - result_no_trailing = build_vespa_filters(filters, remove_trailing_and=True) - assert expected[:-5] == result_no_trailing # Remove trailing " and " - - def test_knowledge_scope_single_filter_not_wrapped(self) -> None: - """When only one knowledge-scope filter is present it should not - be wrapped in an extra OR group.""" - filters = IndexFilters(access_control_list=[], document_set=["set1"]) - result = build_vespa_filters(filters) - assert f'!({HIDDEN}=true) and ({DOCUMENT_SETS} contains "set1") and ' == result - - def test_persona_id_filter_is_primary_knowledge_scope(self) -> None: - """persona_id_filter alone should trigger a knowledge scope restriction - (a persona with user files IS explicit knowledge).""" - filters = IndexFilters(access_control_list=[], persona_id_filter=42) - result = build_vespa_filters(filters) - assert f'!({HIDDEN}=true) and ({PERSONAS} contains "42") and ' == result - - def test_persona_id_filter_with_project_id_filter(self) -> None: - """When persona_id_filter triggers the scope, project_id_filter should be - OR'd in additively.""" - filters = IndexFilters( - access_control_list=[], persona_id_filter=42, project_id_filter=789 - ) - result = build_vespa_filters(filters) - expected = ( - f"!({HIDDEN}=true) and " - f'(({PERSONAS} contains "42") or ({USER_PROJECT} contains "789")) and ' - ) - assert expected == result - - def test_knowledge_scope_document_set_and_persona_filter_ored(self) -> None: - """Document set filter and persona_id_filter must be OR'd so that - connector documents (in the set) and persona user files can - both be found.""" - filters = IndexFilters( - access_control_list=[], - document_set=["engineering"], - persona_id_filter=42, - ) - result = build_vespa_filters(filters) - expected = f'!({HIDDEN}=true) and (({DOCUMENT_SETS} contains "engineering") or ({PERSONAS} contains "42")) and ' - assert expected == result - - def test_acl_large_list_uses_weighted_set(self) -> None: - """Verify that large ACL lists produce a weightedSet clause - instead of OR-chained contains — this is what prevents Vespa - HTTP 400 errors for users with thousands of permission groups.""" - acl = [f"external_group:google_drive_{i}" for i in range(10_000)] - acl += ["user_email:user@example.com", "__PUBLIC__"] - filters = IndexFilters(access_control_list=acl) - result = build_vespa_filters(filters) - - assert "weightedSet(access_control_list, {" in result - # Must NOT contain OR-chained contains clauses - assert "access_control_list contains" not in result - # All entries should be present - assert '"external_group:google_drive_0":1' in result - assert '"external_group:google_drive_9999":1' in result - assert '"user_email:user@example.com":1' in result - assert '"__PUBLIC__":1' in result - - def test_acl_empty_strings_filtered(self) -> None: - """Empty strings in the ACL list should be filtered out.""" - filters = IndexFilters(access_control_list=["user1", "", "group1"]) - result = build_vespa_filters(filters) - assert ( - result - == f'!({HIDDEN}=true) and weightedSet(access_control_list, {{"user1":1, "group1":1}}) and ' - ) - - # All empty - filters = IndexFilters(access_control_list=["", ""]) - result = build_vespa_filters(filters) - assert result == f"!({HIDDEN}=true) and " - - def test_empty_or_none_values(self) -> None: - """Test with empty or None values in filter lists.""" - # Empty strings in document set - filters = IndexFilters( - access_control_list=[], document_set=["set1", "", "set2"] - ) - result = build_vespa_filters(filters) - assert ( - f'!({HIDDEN}=true) and ({DOCUMENT_SETS} contains "set1" or {DOCUMENT_SETS} contains "set2") and ' - == result - ) - - # All empty strings in document set - filters = IndexFilters(access_control_list=[], document_set=["", ""]) - result = build_vespa_filters(filters) - assert f"!({HIDDEN}=true) and " == result diff --git a/cli/internal/deploy/deployfiles/embedded/docker_compose/docker-compose.prod.yml b/cli/internal/deploy/deployfiles/embedded/docker_compose/docker-compose.prod.yml index 6c13d45c1c..16ae409ccd 100644 --- a/cli/internal/deploy/deployfiles/embedded/docker_compose/docker-compose.prod.yml +++ b/cli/internal/deploy/deployfiles/embedded/docker_compose/docker-compose.prod.yml @@ -318,7 +318,6 @@ services: - INDEXING_ONLY=True # Set to debug to get more fine-grained logs - LOG_LEVEL=${LOG_LEVEL:-info} - - VESPA_SEARCHER_THREADS=${VESPA_SEARCHER_THREADS:-1} volumes: # Not necessary, this is just to reduce download time during startup - indexing_huggingface_model_cache:/app/.cache/huggingface/ diff --git a/cli/internal/deploy/deployfiles/embedded/docker_compose/env.prod.template b/cli/internal/deploy/deployfiles/embedded/docker_compose/env.prod.template index fdafee9696..d3d93e66e7 100644 --- a/cli/internal/deploy/deployfiles/embedded/docker_compose/env.prod.template +++ b/cli/internal/deploy/deployfiles/embedded/docker_compose/env.prod.template @@ -73,10 +73,6 @@ POSTGRES_PASSWORD=password DB_READONLY_USER=db_readonly_user DB_READONLY_PASSWORD=password -# If setting the vespa language is required, set this ('en', 'de', etc.). -# See: https://docs.vespa.ai/en/linguistics.html -#VESPA_LANGUAGE_OVERRIDE= - # Show extra/uncommon connectors # See https://docs.onyx.app/admins/connectors/overview for a full list of connectors SHOW_EXTRA_CONNECTORS=False diff --git a/cli/internal/deploy/deployfiles/embedded/docker_compose/env.template b/cli/internal/deploy/deployfiles/embedded/docker_compose/env.template index 33bc800891..76ad183574 100644 --- a/cli/internal/deploy/deployfiles/embedded/docker_compose/env.template +++ b/cli/internal/deploy/deployfiles/embedded/docker_compose/env.template @@ -287,7 +287,6 @@ USE_IAM_AUTH=false # LOG_LEVEL inside the worker process. LOG_LEVEL=info LOG_ONYX_MODEL_INTERACTIONS=False -# LOG_VESPA_TIMING_INFORMATION= # LOG_ENDPOINT_LATENCY= # LOG_POSTGRES_LATENCY= # LOG_POSTGRES_CONN_COUNTS= @@ -390,7 +389,6 @@ LOG_ONYX_MODEL_INTERACTIONS=False # MIN_THREADS_ML_MODELS= ## Indexing Configuration -# VESPA_SEARCHER_THREADS= # ENABLED_CONNECTOR_TYPES= # DISABLE_INDEX_UPDATE_ON_SWAP= # OpenSearch Point-In-Time keep-alive for the reindex/port scan (default 5m) @@ -452,7 +450,6 @@ LOG_ONYX_MODEL_INTERACTIONS=False # APIs + recording keep working, unlike ONYX_QUERY_HISTORY_TYPE=disabled). # HIDE_QUERY_HISTORY_FROM_ADMIN_PANEL=true # CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS= -# VESPA_LANGUAGE_OVERRIDE= ## Frontend Configs # NEXT_PUBLIC_DISABLE_LOGOUT= diff --git a/ct.yaml b/ct.yaml index e530457707..f9176f1c74 100644 --- a/ct.yaml +++ b/ct.yaml @@ -6,7 +6,6 @@ chart-dirs: # must be kept in sync with Chart.yaml chart-repos: - - vespa=https://onyx-dot-app.github.io/vespa-helm-charts - opensearch=https://opensearch-project.github.io/helm-charts - ingress-nginx=https://kubernetes.github.io/ingress-nginx - postgresql=https://cloudnative-pg.github.io/charts diff --git a/deployment/aws_ecs_fargate/cloudformation/README.md b/deployment/aws_ecs_fargate/cloudformation/README.md index eb377e30fa..ce19d13dd8 100644 --- a/deployment/aws_ecs_fargate/cloudformation/README.md +++ b/deployment/aws_ecs_fargate/cloudformation/README.md @@ -47,7 +47,6 @@ The deployment follows this order: 2. Service stacks: - Postgres - Redis - - Vespa Engine - Model Server (Indexing) - Model Server (Inference) - Backend API Server diff --git a/deployment/aws_ecs_fargate/cloudformation/deploy.sh b/deployment/aws_ecs_fargate/cloudformation/deploy.sh index 51ddaa4c4a..835d3394c9 100755 --- a/deployment/aws_ecs_fargate/cloudformation/deploy.sh +++ b/deployment/aws_ecs_fargate/cloudformation/deploy.sh @@ -45,7 +45,6 @@ INFRA_ORDER=( SERVICE_ORDER=( "onyx_postgres_service_template.yaml" "onyx_redis_service_template.yaml" - "onyx_vespaengine_service_template.yaml" "onyx_model_server_indexing_service_template.yaml" "onyx_model_server_inference_service_template.yaml" "onyx_backend_api_server_service_template.yaml" diff --git a/deployment/aws_ecs_fargate/cloudformation/onyx_efs_template.yaml b/deployment/aws_ecs_fargate/cloudformation/onyx_efs_template.yaml index 3ff6d9b521..c5dc97b335 100644 --- a/deployment/aws_ecs_fargate/cloudformation/onyx_efs_template.yaml +++ b/deployment/aws_ecs_fargate/cloudformation/onyx_efs_template.yaml @@ -29,6 +29,9 @@ Resources: ReplicationOverwriteProtection: ENABLED ThroughputMode: elastic + # Legacy Vespa access points. Onyx no longer uses Vespa. They are kept so + # that stacks which still import them (an old vespaengine service stack) can + # update. Delete that service stack before you remove these. VespaEngineTmpEfsAccessPoint: Type: AWS::EFS::AccessPoint Properties: diff --git a/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_api_server_service_template.yaml b/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_api_server_service_template.yaml index 9fe71756ae..489dc04fe3 100644 --- a/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_api_server_service_template.yaml +++ b/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_api_server_service_template.yaml @@ -151,11 +151,6 @@ Resources: - "${Environment}-onyx-model-server-inference-service.${ImportedNamespace}" - ImportedNamespace: !ImportValue Fn::Sub: "${Environment}-onyx-cluster-OnyxNamespaceName" - - Name: VESPA_HOST - Value: !Sub - - "${Environment}-onyx-vespaengine-service.${ImportedNamespace}" - - ImportedNamespace: !ImportValue - Fn::Sub: "${Environment}-onyx-cluster-OnyxNamespaceName" - Name: POSTGRES_HOST Value: !Sub - "${Environment}-onyx-postgres-service.${ImportedNamespace}" diff --git a/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_background_server_service_template.yaml b/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_background_server_service_template.yaml index 8c5163e1ec..6b160f2c9b 100644 --- a/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_background_server_service_template.yaml +++ b/deployment/aws_ecs_fargate/cloudformation/services/onyx_backend_background_server_service_template.yaml @@ -150,11 +150,6 @@ Resources: - "${Environment}-onyx-model-server-inference-service.${ImportedNamespace}" - ImportedNamespace: !ImportValue Fn::Sub: "${Environment}-onyx-cluster-OnyxNamespaceName" - - Name: VESPA_HOST - Value: !Sub - - "${Environment}-onyx-vespaengine-service.${ImportedNamespace}" - - ImportedNamespace: !ImportValue - Fn::Sub: "${Environment}-onyx-cluster-OnyxNamespaceName" - Name: POSTGRES_HOST Value: !Sub - "${Environment}-onyx-postgres-service.${ImportedNamespace}" diff --git a/deployment/aws_ecs_fargate/cloudformation/services/onyx_model_server_indexing_service_template.yaml b/deployment/aws_ecs_fargate/cloudformation/services/onyx_model_server_indexing_service_template.yaml index 580c17bb14..5d237f1c68 100644 --- a/deployment/aws_ecs_fargate/cloudformation/services/onyx_model_server_indexing_service_template.yaml +++ b/deployment/aws_ecs_fargate/cloudformation/services/onyx_model_server_indexing_service_template.yaml @@ -131,8 +131,6 @@ Resources: Value: info - Name: INDEXING_ONLY Value: True - - Name: VESPA_SEARCHER_THREADS - Value: "1" MountPoints: - SourceVolume: efs-volume ContainerPath: /app/.cache/huggingface/ diff --git a/deployment/aws_ecs_fargate/cloudformation/services/onyx_vespaengine_service_template.yaml b/deployment/aws_ecs_fargate/cloudformation/services/onyx_vespaengine_service_template.yaml deleted file mode 100644 index 0ebe3d8e8d..0000000000 --- a/deployment/aws_ecs_fargate/cloudformation/services/onyx_vespaengine_service_template.yaml +++ /dev/null @@ -1,192 +0,0 @@ -AWSTemplateFormatVersion: "2010-09-09" -Description: CloudFormation template for Onyx Vespa Engine TaskDefinition -Parameters: - Environment: - Type: String - SubnetIDs: - Type: CommaDelimitedList - Description: "Comma-delimited list of at least two subnet IDs in different Availability Zones" - VpcID: - Type: String - Description: "The ID of the VPC to deploy into" - ServiceName: - Type: String - Default: onyx-vespaengine - TaskCpu: - Type: String - Default: "4096" - TaskMemory: - Type: String - Default: "16384" - TaskDesiredCount: - Type: Number - Default: 1 - -Resources: - - ECSService: - Type: AWS::ECS::Service - Properties: - Cluster: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-cluster-ECSClusterName" - CapacityProviderStrategy: - - CapacityProvider: FARGATE - Base: 0 - Weight: 1 - TaskDefinition: !Ref TaskDefinition - ServiceName: !Sub ${Environment}-${ServiceName}-service - SchedulingStrategy: REPLICA - DesiredCount: !Ref TaskDesiredCount - AvailabilityZoneRebalancing: ENABLED - NetworkConfiguration: - AwsvpcConfiguration: - AssignPublicIp: ENABLED - SecurityGroups: - - Ref: SecurityGroup - Subnets: !Ref SubnetIDs - PlatformVersion: LATEST - DeploymentConfiguration: - MaximumPercent: 200 - MinimumHealthyPercent: 100 - DeploymentCircuitBreaker: - Enable: true - Rollback: true - DeploymentController: - Type: ECS - ServiceConnectConfiguration: - Enabled: false - ServiceRegistries: - - RegistryArn: !GetAtt ServiceDiscoveryService.Arn - Tags: - - Key: app - Value: onyx - - Key: service - Value: !Ref ServiceName - - Key: env - Value: !Ref Environment - EnableECSManagedTags: true - - SecurityGroup: - Type: AWS::EC2::SecurityGroup - Properties: - GroupDescription: !Sub Onyx SecurityGroup access to EFS mount and ${ServiceName}. - GroupName: !Sub ${Environment}-ecs-${ServiceName} - VpcId: !Ref VpcID - SecurityGroupIngress: - - FromPort: 19071 - ToPort: 19071 - IpProtocol: tcp - CidrIp: 0.0.0.0/0 - - FromPort: 19071 - ToPort: 19071 - IpProtocol: tcp - CidrIpv6: "::/0" - - FromPort: 8081 - ToPort: 8081 - IpProtocol: tcp - CidrIp: 0.0.0.0/0 - - FromPort: 8081 - ToPort: 8081 - IpProtocol: tcp - CidrIpv6: "::/0" - - FromPort: 2049 - ToPort: 2049 - IpProtocol: tcp - SourceSecurityGroupId: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-efs-EFSSecurityGroupMountTargets" - - ServiceDiscoveryService: - Type: "AWS::ServiceDiscovery::Service" - Properties: - Name: !Sub ${Environment}-${ServiceName}-service - DnsConfig: - DnsRecords: - - Type: "A" - TTL: 15 - NamespaceId: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-cluster-OnyxNamespace" - HealthCheckCustomConfig: - FailureThreshold: 1 - - TaskDefinition: - Type: AWS::ECS::TaskDefinition - Properties: - Family: !Sub ${Environment}-${ServiceName}-TaskDefinition - TaskRoleArn: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-cluster-ECSTaskRole" - ExecutionRoleArn: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-cluster-ECSTaskExecutionRole" - NetworkMode: awsvpc - RequiresCompatibilities: - - FARGATE - Cpu: !Ref TaskCpu - Memory: !Ref TaskMemory - RuntimePlatform: - CpuArchitecture: ARM64 - OperatingSystemFamily: LINUX - ContainerDefinitions: - - Name: vespaengine - Image: vespaengine/vespa:8.609.39 - Cpu: 0 - Essential: true - PortMappings: - - Name: vespaengine_port - ContainerPort: 19071 - HostPort: 19071 - Protocol: tcp - AppProtocol: http - - Name: vespaengine_port2 - ContainerPort: 8081 - HostPort: 8081 - Protocol: tcp - AppProtocol: http - MountPoints: - - SourceVolume: efs-volume-data - ContainerPath: /opt/vespa/var - ReadOnly: false - - SourceVolume: efs-volume-tmp - ContainerPath: /var/tmp - ReadOnly: false - LogConfiguration: - LogDriver: awslogs - Options: - awslogs-group: /ecs/OnyxVespaEngineTaskDefinition - mode: non-blocking - awslogs-create-group: "true" - max-buffer-size: "25m" - awslogs-region: !Ref AWS::Region - awslogs-stream-prefix: ecs - User: "1000" - Environment: - - Name: VESPA_SKIP_UPGRADE_CHECK - Value: "true" - VolumesFrom: [] - SystemControls: [] - Volumes: - - Name: efs-volume-tmp - EFSVolumeConfiguration: - FilesystemId: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-efs-OnyxEfsId" - RootDirectory: "/" - TransitEncryption: ENABLED - AuthorizationConfig: - AccessPointId: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-efs-VespaEngineTmpEfsAccessPoint" - - Name: efs-volume-data - EFSVolumeConfiguration: - FilesystemId: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-efs-OnyxEfsId" - RootDirectory: "/" - TransitEncryption: ENABLED - AuthorizationConfig: - AccessPointId: - Fn::ImportValue: - Fn::Sub: "${Environment}-onyx-efs-VespaEngineDataEfsAccessPoint" diff --git a/deployment/aws_ecs_fargate/cloudformation/uninstall.sh b/deployment/aws_ecs_fargate/cloudformation/uninstall.sh index 154a7f67bc..f473344acf 100755 --- a/deployment/aws_ecs_fargate/cloudformation/uninstall.sh +++ b/deployment/aws_ecs_fargate/cloudformation/uninstall.sh @@ -27,6 +27,8 @@ STACK_NAMES=( "${ENVIRONMENT}-onyx-backend-api-server-service" "${ENVIRONMENT}-onyx-model-server-inference-service" "${ENVIRONMENT}-onyx-model-server-indexing-service" + # Legacy Vespa service stack, removed from deploy.sh. Kept so old installs + # can still be torn down. "${ENVIRONMENT}-onyx-vespaengine-service" "${ENVIRONMENT}-onyx-redis-service" "${ENVIRONMENT}-onyx-postgres-service" diff --git a/deployment/docker_compose/docker-compose.prod-no-letsencrypt.yml b/deployment/docker_compose/docker-compose.prod-no-letsencrypt.yml index fd73df2ab0..c58019461a 100644 --- a/deployment/docker_compose/docker-compose.prod-no-letsencrypt.yml +++ b/deployment/docker_compose/docker-compose.prod-no-letsencrypt.yml @@ -318,7 +318,6 @@ services: - INDEXING_ONLY=True # Set to debug to get more fine-grained logs - LOG_LEVEL=${LOG_LEVEL:-info} - - VESPA_SEARCHER_THREADS=${VESPA_SEARCHER_THREADS:-1} volumes: # Not necessary, this is just to reduce download time during startup - indexing_huggingface_model_cache:/app/.cache/huggingface/ diff --git a/deployment/docker_compose/docker-compose.prod.yml b/deployment/docker_compose/docker-compose.prod.yml index 6c13d45c1c..16ae409ccd 100644 --- a/deployment/docker_compose/docker-compose.prod.yml +++ b/deployment/docker_compose/docker-compose.prod.yml @@ -318,7 +318,6 @@ services: - INDEXING_ONLY=True # Set to debug to get more fine-grained logs - LOG_LEVEL=${LOG_LEVEL:-info} - - VESPA_SEARCHER_THREADS=${VESPA_SEARCHER_THREADS:-1} volumes: # Not necessary, this is just to reduce download time during startup - indexing_huggingface_model_cache:/app/.cache/huggingface/ diff --git a/deployment/docker_compose/docker-compose.resources.yml b/deployment/docker_compose/docker-compose.resources.yml index c9bb1415a0..835c72da57 100644 --- a/deployment/docker_compose/docker-compose.resources.yml +++ b/deployment/docker_compose/docker-compose.resources.yml @@ -35,16 +35,6 @@ services: # cpus: ${API_SERVER_CPU_RESERVATION} # memory: ${API_SERVER_MEM_RESERVATION} # - # index: - # deploy: - # resources: - # limits: - # cpus: ${VESPA_CPU_LIMIT} - # memory: ${VESPA_MEM_LIMIT} - # reservations: - # cpus: ${VESPA_CPU_RESERVATION} - # memory: ${VESPA_MEM_RESERVATION} - # inference_model_server: deploy: resources: diff --git a/deployment/docker_compose/docker-compose.search-testing.yml b/deployment/docker_compose/docker-compose.search-testing.yml index 8fedc460e5..880a4afa3a 100644 --- a/deployment/docker_compose/docker-compose.search-testing.yml +++ b/deployment/docker_compose/docker-compose.search-testing.yml @@ -139,7 +139,6 @@ services: - MIN_THREADS_ML_MODELS=${MIN_THREADS_ML_MODELS:-} - INDEXING_ONLY=True - LOG_LEVEL=${LOG_LEVEL:-debug} - - VESPA_SEARCHER_THREADS=${VESPA_SEARCHER_THREADS:-1} volumes: - inference_model_cache_huggingface:/app/.cache/huggingface/ logging: diff --git a/deployment/docker_compose/docker-compose.template.yml b/deployment/docker_compose/docker-compose.template.yml index dab82ad847..9009d4fdad 100644 --- a/deployment/docker_compose/docker-compose.template.yml +++ b/deployment/docker_compose/docker-compose.template.yml @@ -414,7 +414,6 @@ services: #!for prod,no-letsencrypt # Set to debug to get more fine-grained logs - LOG_LEVEL=${LOG_LEVEL:-info} - - VESPA_SEARCHER_THREADS=${VESPA_SEARCHER_THREADS:-1} #!endfor volumes: # Not necessary, this is just to reduce download time during startup diff --git a/deployment/docker_compose/env.prod.template b/deployment/docker_compose/env.prod.template index fdafee9696..d3d93e66e7 100644 --- a/deployment/docker_compose/env.prod.template +++ b/deployment/docker_compose/env.prod.template @@ -73,10 +73,6 @@ POSTGRES_PASSWORD=password DB_READONLY_USER=db_readonly_user DB_READONLY_PASSWORD=password -# If setting the vespa language is required, set this ('en', 'de', etc.). -# See: https://docs.vespa.ai/en/linguistics.html -#VESPA_LANGUAGE_OVERRIDE= - # Show extra/uncommon connectors # See https://docs.onyx.app/admins/connectors/overview for a full list of connectors SHOW_EXTRA_CONNECTORS=False diff --git a/deployment/docker_compose/env.template b/deployment/docker_compose/env.template index 33bc800891..76ad183574 100644 --- a/deployment/docker_compose/env.template +++ b/deployment/docker_compose/env.template @@ -287,7 +287,6 @@ USE_IAM_AUTH=false # LOG_LEVEL inside the worker process. LOG_LEVEL=info LOG_ONYX_MODEL_INTERACTIONS=False -# LOG_VESPA_TIMING_INFORMATION= # LOG_ENDPOINT_LATENCY= # LOG_POSTGRES_LATENCY= # LOG_POSTGRES_CONN_COUNTS= @@ -390,7 +389,6 @@ LOG_ONYX_MODEL_INTERACTIONS=False # MIN_THREADS_ML_MODELS= ## Indexing Configuration -# VESPA_SEARCHER_THREADS= # ENABLED_CONNECTOR_TYPES= # DISABLE_INDEX_UPDATE_ON_SWAP= # OpenSearch Point-In-Time keep-alive for the reindex/port scan (default 5m) @@ -452,7 +450,6 @@ LOG_ONYX_MODEL_INTERACTIONS=False # APIs + recording keep working, unlike ONYX_QUERY_HISTORY_TYPE=disabled). # HIDE_QUERY_HISTORY_FROM_ADMIN_PANEL=true # CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS= -# VESPA_LANGUAGE_OVERRIDE= ## Frontend Configs # NEXT_PUBLIC_DISABLE_LOGOUT= diff --git a/deployment/docker_compose/install.ps1 b/deployment/docker_compose/install.ps1 index cf35182919..cee441de1a 100644 --- a/deployment/docker_compose/install.ps1 +++ b/deployment/docker_compose/install.ps1 @@ -314,7 +314,7 @@ function Show-OnyxHelp { $help += "`nUsage: .\install.ps1 [OPTIONS]`n" $help += "`nOptions:" $help += "`n -IncludeCraft Enable Onyx Craft (AI-powered web app building)" - $help += "`n -Lite Deploy Onyx Lite (no Vespa, Redis, or model servers)" + $help += "`n -Lite Deploy Onyx Lite (no OpenSearch, Redis, or model servers)" $help += "`n -Local Use existing config files instead of downloading from GitHub" $help += "`n -Shutdown Stop (pause) Onyx containers" $help += "`n -DeleteData Remove all Onyx data (containers, volumes, and files)" @@ -1110,7 +1110,7 @@ function Main { Print-Info "The first user created will automatically have admin privileges" if ($script:LiteMode) { - Print-Info "Running in Lite mode - Vespa, Redis, model servers, and background workers are NOT started." + Print-Info "Running in Lite mode - OpenSearch, Redis, model servers, and background workers are NOT started." Print-Info "Connectors and RAG search are disabled. LLM chat, tools, Projects still work." } diff --git a/deployment/helm/MIGRATION.md b/deployment/helm/MIGRATION.md index 51aa915ed2..44cb2baa4c 100644 --- a/deployment/helm/MIGRATION.md +++ b/deployment/helm/MIGRATION.md @@ -111,62 +111,16 @@ New installs set `minio.enabled: false` and `objectStore.enabled: true`. ## chart 0.4.x → 0.5.x -### What changed +### Vespa removal Chart 0.5.0 removed the bundled `charts/vespa/` subchart. Earlier chart -versions installed Vespa as a `da-vespa` StatefulSet alongside the rest of -Onyx; the api-server connected to it on `localhost:19071` (Vespa application -deploy port) and the chart-managed PV held the indexed corpus. +versions installed Vespa as a `da-vespa` StatefulSet. Onyx no longer uses +Vespa. OpenSearch is the only document index. -The 0.5.x line assumes you are running Vespa **outside** the chart — either -managed Vespa, Vespa Cloud, or a separately-managed deployment in another -namespace. - -### Why this matters - -A naive `helm upgrade` from 0.4.x to 0.5.x will delete the `da-vespa` -StatefulSet (no template renders it anymore). The PV underneath the -StatefulSet's PVC will be orphaned but the data inside is unreachable -until you reattach it manually. Meanwhile the api-server will crash-loop -trying to deploy its Vespa application package: - -``` -ConnectionRefusedError: [Errno 111] Connection refused -HTTPConnection(host='localhost', port=19071): Failed to establish a new -connection -``` - -The chart now ships a guard that detects this situation at install/upgrade -time and fails fast with a clear message instead of silently breaking. See -`templates/legacy-vespa-check.yaml`. - -### How to upgrade safely - -1. **Stand up an external Vespa cluster.** Vespa Cloud or a self-hosted - deployment outside this chart, whichever fits your operational model. -2. **Re-index is automatic.** Vespa data does not roundtrip directly - between releases (chunk schemas have changed over time anyway). No - manual action here; Onyx connectors will reindex on their own once - the api-server can reach the new endpoint (after the upgrade in - step 5). -3. **Update your values** to point Onyx at the external Vespa endpoint - (the api-server respects `VESPA_HOST` / `VESPA_PORT` env vars; set - them through your `configMap:` block). -4. **Delete the old StatefulSet** once you no longer need it. The PV - reclaim policy determines whether the underlying disk goes with it — - verify before you delete anything. -5. **Run `helm upgrade`** with chart 0.5.x. If the old StatefulSet is - already gone the legacy check passes automatically. - -If you need to bypass the check (e.g. you've already migrated and only -have a stale PV lingering), set in your values: - -```yaml -legacyVespaCheck: - acknowledged: true -``` - -or disable the check entirely with `legacyVespaCheck.enabled: false`. +A `helm upgrade` from 0.4.x deletes the `da-vespa` StatefulSet. Its PVC +(`vespa-storage-da-vespa-0`) stays. Onyx does not read the Vespa data, so +re-index your connectors into OpenSearch after the upgrade. Delete the PVC +when you no longer need it. ### `celery-worker-scheduled-tasks` deployment diff --git a/deployment/helm/README.md b/deployment/helm/README.md index 918f5eeefa..468ed7b7f6 100644 --- a/deployment/helm/README.md +++ b/deployment/helm/README.md @@ -43,7 +43,7 @@ unless you specifically need OCI. # Recent chart changes (0.5.0) -If you are upgrading from an earlier 0.4.x release, **read [MIGRATION.md](./MIGRATION.md) first.** The 0.5.0 release dropped the bundled Vespa subchart; chart 0.5.6 ships a guard that fails `helm upgrade` if the legacy `da-vespa` StatefulSet is still in the namespace so you don't lose the indexed data silently. +If you are upgrading from an earlier 0.4.x release, **read [MIGRATION.md](./MIGRATION.md) first.** The 0.5.0 release dropped the bundled Vespa subchart. Onyx no longer uses Vespa. Other 0.5.0 changes: @@ -277,14 +277,10 @@ Other docker-compose-style values you should set deliberately: * helm install onyx . -n onyx --set postgresql.primary.persistence.enabled=false --set auth.opensearch.values.opensearch_admin_password='StrongPassword123!' * the postgres flag is to keep the storage ephemeral for testing. You probably don't want to set that in prod. * the OpenSearch admin password must be set on first install unless you are supplying `auth.opensearch.existingSecret`. - * no flag for ephemeral vespa storage yet, might be good for testing * kubectl -n onyx port-forward service/onyx-nginx 8080:80 * this will forward the local port 8080 to the installed chart for you to run tests, etc. * When you are finished * helm uninstall onyx -n onyx - * Vespa leaves behind a PVC. Delete it if you are completely done. - * k -n onyx get pvc - * k -n onyx delete pvc vespa-storage-da-vespa-0 * If you didn't disable Postgres persistence earlier, you may want to delete that PVC too. ## Run as non-root user @@ -295,14 +291,6 @@ By default, some onyx containers run as root. If you'd like to explicitly run th runAsNonRoot: true runAsUser: 1001 ``` - * `vespa` - ```yaml - podSecurityContext: - fsGroup: 1000 - securityContext: - privileged: false - runAsUser: 1000 - ``` ## Resourcing In the helm charts, we have resource suggestions for all Onyx-owned components. diff --git a/deployment/helm/charts/onyx/Chart.yaml b/deployment/helm/charts/onyx/Chart.yaml index e44ccfd318..2189782d09 100644 --- a/deployment/helm/charts/onyx/Chart.yaml +++ b/deployment/helm/charts/onyx/Chart.yaml @@ -5,7 +5,7 @@ home: https://www.onyx.app/ sources: - "https://github.com/onyx-dot-app/onyx" type: application -version: 0.9.0 +version: 0.9.1 appVersion: latest annotations: category: Productivity @@ -20,6 +20,10 @@ annotations: - name: background image: docker.io/onyxdotapp/onyx-backend:latest artifacthub.io/changes: | + - kind: removed + description: The legacyVespaCheck upgrade guard and the unused + ENABLE_OPENSEARCH_INDEXING_FOR_ONYX / ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX + env vars. Onyx no longer uses Vespa. - kind: added description: objectStore, an in-cluster SeaweedFS S3 store that replaces bundled MinIO, with a Job that copies MinIO's files while both stores serve. diff --git a/deployment/helm/charts/onyx/SIZING.md b/deployment/helm/charts/onyx/SIZING.md index 5818581ace..17882f6842 100644 --- a/deployment/helm/charts/onyx/SIZING.md +++ b/deployment/helm/charts/onyx/SIZING.md @@ -109,7 +109,8 @@ opensearch: persistence: size: 128Gi # existing PVCs must be expanded manually # Pin to the dedicated index node group if you provisioned one — see the - # placement note in the Large section. + # placement note in the Large section. The terraform module keeps the legacy + # `vespa` names for this group so existing state is not replaced. nodeSelector: eks.amazonaws.com/nodegroup: vespa-node-group tolerations: @@ -192,5 +193,5 @@ opensearch: If you use a managed OpenSearch domain (`opensearch.enabled: false` + `OPENSEARCH_HOST` in the configMap), skip every `opensearch:` block above and size the domain in terraform instead. In that case the terraform module's -dedicated index node group has nothing to run — shrink -`vespa_node_instance_types` to a small instance. +dedicated index node group has nothing to run — set +`vespa_node_enabled = false`. diff --git a/deployment/helm/charts/onyx/templates/celery-worker-light.yaml b/deployment/helm/charts/onyx/templates/celery-worker-light.yaml index 04d82cc32a..bb4742bb05 100644 --- a/deployment/helm/charts/onyx/templates/celery-worker-light.yaml +++ b/deployment/helm/charts/onyx/templates/celery-worker-light.yaml @@ -73,7 +73,7 @@ spec: {{- end }} "--hostname=light@%n", "-Q", - "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,opensearch_migration,chat_ttl_deletion", + "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,index_reclaim,chat_ttl_deletion", ] ports: - name: metrics diff --git a/deployment/helm/charts/onyx/templates/configmap.yaml b/deployment/helm/charts/onyx/templates/configmap.yaml index 84350b2f7a..8a14b8ea68 100755 --- a/deployment/helm/charts/onyx/templates/configmap.yaml +++ b/deployment/helm/charts/onyx/templates/configmap.yaml @@ -13,8 +13,6 @@ data: {{- if .Values.opensearch.enabled }} OPENSEARCH_HOST: {{ .Values.opensearch.clusterName }}-{{ .Values.opensearch.nodeGroup }}.{{ .Release.Namespace }}.svc.cluster.local OPENSEARCH_REST_API_PORT: "9200" - ENABLE_OPENSEARCH_INDEXING_FOR_ONYX: "true" - ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX: "false" {{- end }} {{- if .Values.redis.enabled }} # Standalone Redis CR produces a Service named after the CR (no `-master` diff --git a/deployment/helm/charts/onyx/templates/legacy-vespa-check.yaml b/deployment/helm/charts/onyx/templates/legacy-vespa-check.yaml deleted file mode 100644 index 3f83dee5b4..0000000000 --- a/deployment/helm/charts/onyx/templates/legacy-vespa-check.yaml +++ /dev/null @@ -1,36 +0,0 @@ -{{- /* -Detect upgrades from chart < 0.5.0, which bundled Vespa as a subchart. -Chart 0.5.0 dropped that subchart, so a naive `helm upgrade` from 0.4.x -would delete the `da-vespa` StatefulSet and orphan the Vespa data -volume. The api-server would then crash-loop trying to deploy its Vespa -application package against localhost:19071. - -Fail fast at install/upgrade time with a clear migration message. The -check is a `lookup` against the live cluster, so it only fires during -`helm install`/`helm upgrade` (not during local `helm template`). - -If the legacy StatefulSet has already been migrated away and only the -PV is hanging around, set `legacyVespaCheck.acknowledged=true` to -override the check. - -The hasKey guards below protect `helm upgrade --reuse-values` against a -release whose previously-stored values predate the `legacyVespaCheck` -key — direct nested access would nil-deref and silently bypass the -entire fail-fast path. We treat absent values as the "safe default" -(enabled, not acknowledged) so the protection is in force by default. -*/}} -{{- $cfg := default (dict) .Values.legacyVespaCheck }} -{{- $checkEnabled := true }} -{{- if hasKey $cfg "enabled" }} -{{- $checkEnabled = $cfg.enabled }} -{{- end }} -{{- $checkAck := false }} -{{- if hasKey $cfg "acknowledged" }} -{{- $checkAck = $cfg.acknowledged }} -{{- end }} -{{- if $checkEnabled }} -{{- $legacyVespa := lookup "apps/v1" "StatefulSet" .Release.Namespace "da-vespa" }} -{{- if and $legacyVespa (not $checkAck) }} -{{- fail (printf "\n\nERROR: legacy `da-vespa` StatefulSet detected in namespace %q.\n\nChart 0.5.0 removed the bundled Vespa subchart. Upgrading without first\nmigrating to external/managed Vespa will delete the StatefulSet and you\nwill lose access to the indexed data managed by it. The api-server will\nfail to start (connection refused on localhost:19071) until it can reach\nan external Vespa endpoint.\n\nWhat to do:\n 1. Provision an external Vespa cluster (Vespa Cloud, self-hosted, etc.).\n 2. Update your values to point Onyx at the new endpoint — set\n VESPA_HOST / VESPA_PORT under the `configMap:` block.\n 3. Delete the legacy `da-vespa` StatefulSet (check the PV reclaim\n policy first; the data is gone once the PV is reclaimed).\n 4. Re-run `helm upgrade`. Onyx connectors will re-index into the new\n cluster automatically once the api-server can reach it.\n\nOr set `legacyVespaCheck.acknowledged=true` to bypass this check\nexplicitly (only if the StatefulSet is already gone or is being kept\nintentionally).\n\nSee deployment/helm/MIGRATION.md (chart 0.4.x → 0.5.x section) for details.\n" .Release.Namespace) }} -{{- end }} -{{- end }} diff --git a/deployment/helm/charts/onyx/values.yaml b/deployment/helm/charts/onyx/values.yaml index 27c6491495..023361089b 100644 --- a/deployment/helm/charts/onyx/values.yaml +++ b/deployment/helm/charts/onyx/values.yaml @@ -1462,17 +1462,6 @@ externalSecret: # Typically set to the same value as externalSecret.secretName. extraEnvFromSecret: "" -# -- Safeguards against destructive upgrades from chart 0.4.x. Chart 0.5.0 -# removed the bundled Vespa subchart; the check below fails `helm upgrade` -# if the legacy `da-vespa` StatefulSet is still in the namespace, so -# the StatefulSet and its Vespa data volume are not silently deleted. -# Set `acknowledged: true` once you've migrated to external Vespa, or -# `enabled: false` to disable the check entirely. -# See deployment/helm/MIGRATION.md. -legacyVespaCheck: - enabled: true - acknowledged: false - # -- Governs all Secrets created or used by this chart. Values set by this chart will be base64 encoded in the k8s cluster. auth: postgresql: diff --git a/deployment/helm/dev/k8s-down.sh b/deployment/helm/dev/k8s-down.sh index 3dba535080..d16a217d1a 100755 --- a/deployment/helm/dev/k8s-down.sh +++ b/deployment/helm/dev/k8s-down.sh @@ -55,7 +55,7 @@ fi if [[ "$KEEP_CLUSTER" -eq 1 ]]; then # PVCs are intentionally left intact so the next install reuses postgres / - # opensearch / vespa / minio data. For a clean slate, omit --keep-cluster + # opensearch / minio data. For a clean slate, omit --keep-cluster # or run: kubectl -n $NAMESPACE delete pvc --all echo "uninstalling helm release 'onyx' in namespace '$NAMESPACE' ..." helm uninstall onyx -n "$NAMESPACE" 2>/dev/null || true diff --git a/deployment/helm/dev/k8s-up.sh b/deployment/helm/dev/k8s-up.sh index 1596d4921e..982ae6eab6 100755 --- a/deployment/helm/dev/k8s-up.sh +++ b/deployment/helm/dev/k8s-up.sh @@ -117,7 +117,6 @@ trap 'rm -rf "$HELM_DEV_HOME"' EXIT # Repo names must match the dep names in Chart.yaml. helm repo add cloudnative-pg https://cloudnative-pg.github.io/charts >/dev/null -helm repo add vespa https://onyx-dot-app.github.io/vespa-helm-charts >/dev/null helm repo add opensearch https://opensearch-project.github.io/helm-charts >/dev/null helm repo add ingress-nginx https://kubernetes.github.io/ingress-nginx >/dev/null helm repo add redis-ot https://ot-container-kit.github.io/helm-charts >/dev/null diff --git a/deployment/terraform/modules/aws/README.md b/deployment/terraform/modules/aws/README.md index 3307762f9d..3b77a77aeb 100644 --- a/deployment/terraform/modules/aws/README.md +++ b/deployment/terraform/modules/aws/README.md @@ -251,7 +251,7 @@ Key inputs include: - `vpc_id`, `subnet_ids` - `public_cluster_enabled` (default true), `private_cluster_enabled` (default true) - `cluster_endpoint_public_access_cidrs` (default `[]`). Empty denies all public API access. Set it when `public_cluster_enabled` is true and you need to reach the API server -- `eks_managed_node_groups` (defaults include a main and a vespa-dedicated group with GP3 volumes) +- `eks_managed_node_groups` (defaults include a main group and a dedicated document-index group, legacy key `vespa`, with GP3 volumes) - `s3_bucket_names` (optional list). If set, creates an IRSA role and Kubernetes service account for S3 access ### `postgres` diff --git a/deployment/terraform/modules/aws/eks/main.tf b/deployment/terraform/modules/aws/eks/main.tf index 4c32a438cf..897379ee95 100644 --- a/deployment/terraform/modules/aws/eks/main.tf +++ b/deployment/terraform/modules/aws/eks/main.tf @@ -160,7 +160,7 @@ module "eks" { v.instance_types ) }, - # Only add subnet_ids override for vespa node group if specified + # Only add subnet_ids override for the document-index node group if specified k == "vespa" && length(var.vespa_node_subnet_ids) > 0 ? { subnet_ids = var.vespa_node_subnet_ids } : {}, @@ -178,7 +178,7 @@ module "eks" { max_size = coalesce(var.main_node_max_size, v.max_size) desired_size = try(v.desired_size, coalesce(var.main_node_min_size, v.min_size)) } : {}, - # Disk override for the Vespa/document-index node; null keeps the map + # Disk override for the document-index node; null keeps the map # default. Merge preserves any other device mappings on the group. k == "vespa" && var.vespa_node_disk_size_gb != null ? { block_device_mappings = merge(try(v.block_device_mappings, {}), { diff --git a/deployment/terraform/modules/aws/eks/variables.tf b/deployment/terraform/modules/aws/eks/variables.tf index d27d743d72..a9ea44721b 100644 --- a/deployment/terraform/modules/aws/eks/variables.tf +++ b/deployment/terraform/modules/aws/eks/variables.tf @@ -45,19 +45,19 @@ variable "main_node_instance_types" { variable "vespa_node_enabled" { type = bool - description = "Whether to create the dedicated Vespa/document-index node group. Disable when the index runs off-cluster (managed OpenSearch) or fits on the main node group." + description = "Whether to create the dedicated document-index node group (legacy `vespa` name). Disable when the index runs off-cluster (managed OpenSearch) or fits on the main node group." default = true } variable "vespa_node_instance_types" { type = list(string) - description = "Instance types for the Vespa node group" + description = "Instance types for the dedicated document-index node group" default = ["m6i.2xlarge"] } variable "vespa_node_subnet_ids" { type = list(string) - description = "Subnet IDs for the Vespa node group (must be in same AZ as Vespa PV). If not specified, uses all cluster subnets." + description = "Subnet IDs for the dedicated document-index node group (must be in same AZ as the index PV). If not specified, uses all cluster subnets." default = [] } @@ -75,7 +75,7 @@ variable "main_node_min_size" { variable "vespa_node_disk_size_gb" { type = number - description = "Root EBS volume (GiB) for the Vespa/document-index node. Null keeps the node-group default." + description = "Root EBS volume (GiB) for the dedicated document-index node. Null keeps the node-group default." default = null } @@ -152,7 +152,7 @@ variable "eks_managed_node_groups" { type = map(any) description = "EKS managed node groups with EBS volume configuration" default = { - # Main node group for all pods except Vespa + # Main node group for all pods except the in-cluster document index main = { name = "main-node-group" instance_types = null # Will be set from var.main_node_instance_types @@ -175,13 +175,14 @@ variable "eks_managed_node_groups" { # No taints for main node group taints = [] } - # Vespa dedicated node group + # Dedicated document-index (OpenSearch) node group. The key, name, and taint + # keep their legacy `vespa` names so existing node groups are not replaced. vespa = { name = "vespa-node-group" instance_types = null # Will be set from var.vespa_node_instance_types min_size = 1 max_size = 1 - # Larger EBS volume for Vespa storage + # Larger EBS volume for document index storage block_device_mappings = { xvda = { device_name = "/dev/xvda" @@ -195,7 +196,7 @@ variable "eks_managed_node_groups" { } } } - # Taint to ensure only Vespa pods can schedule here + # Taint to ensure only document index pods can schedule here taints = [ { key = "vespa-dedicated" diff --git a/deployment/terraform/modules/aws/onyx/variables.tf b/deployment/terraform/modules/aws/onyx/variables.tf index adcdb807e9..aa8e81d1bd 100644 --- a/deployment/terraform/modules/aws/onyx/variables.tf +++ b/deployment/terraform/modules/aws/onyx/variables.tf @@ -279,19 +279,19 @@ variable "vespa_node_enabled" { variable "vespa_node_instance_types" { type = list(string) - description = "Instance types for the Vespa EKS node group Null uses the t-shirt size default." + description = "Instance types for the dedicated document-index EKS node group. Null uses the t-shirt size default." default = null } variable "vespa_node_disk_size_gb" { type = number - description = "Root EBS volume (GiB) for the Vespa/document-index node. Null keeps the node-group default (100 GiB). Null uses the t-shirt size default." + description = "Root EBS volume (GiB) for the dedicated document-index node. Null keeps the node-group default (100 GiB). Null uses the t-shirt size default." default = null } variable "vespa_node_subnet_ids" { type = list(string) - description = "Subnet IDs for the Vespa node group (must be in same AZ as Vespa PV)" + description = "Subnet IDs for the dedicated document-index node group (must be in same AZ as the index PV)" default = [] } diff --git a/docs/craft/dev/local-kubernetes.md b/docs/craft/dev/local-kubernetes.md index e6ed162c32..88969981f9 100644 --- a/docs/craft/dev/local-kubernetes.md +++ b/docs/craft/dev/local-kubernetes.md @@ -120,7 +120,7 @@ per cluster. New clusters use the `kindest/node:v1.33.1` node image so Craft's native init sidecar pod shape is supported. Existing clusters are not recreated; set `KIND_NODE_IMAGE` to override the default for a newly created cluster. -Watch pods (vespa and CNPG-postgres take a minute or two on first boot): +Watch pods (opensearch and CNPG-postgres take a minute or two on first boot): ```bash kubectl -n onyx get pods -w @@ -332,7 +332,7 @@ external-dependency-unit tests against a temp dir. See Run **`k8s: pause cluster`** (or `docker stop onyx-dev-control-plane`) to stop the kind node container. PVC data lives inside that container, so postgres, -redis, opensearch, vespa, and minio state all survive. Resume with +redis, opensearch, and minio state all survive. Resume with **`k8s: resume cluster`** — the kubelet reconciles pods automatically. Reach for **`k8s: cluster down (full teardown)`** only when you want a clean diff --git a/docs/craft/kubernetes/craft-eks-runbook.md b/docs/craft/kubernetes/craft-eks-runbook.md index 430bf8d965..c84623e8e6 100644 --- a/docs/craft/kubernetes/craft-eks-runbook.md +++ b/docs/craft/kubernetes/craft-eks-runbook.md @@ -14,7 +14,7 @@ Craft provisions a sandbox + snapshot/restore, with no manual kubectl/RBAC/node **Clean seam: Terraform = AWS infrastructure; Helm = everything inside Kubernetes. Terraform's outputs are Helm's inputs.** **Terraform** (`deployment/terraform/…`) creates ONLY AWS resources: -- VPC / subnets / NAT; EKS **cluster** + **node groups** (main, vespa, **sandbox** — labeled/tainted/IMDSv2) + the **OIDC provider**; EKS add-ons + gp3 storage class. +- VPC / subnets / NAT; EKS **cluster** + **node groups** (main, vespa — the legacy name of the document-index group, **sandbox** — labeled/tainted/IMDSv2) + the **OIDC provider**; EKS add-ons + gp3 storage class. - Managed data stores: **RDS** (Postgres), **ElastiCache** (Redis), **OpenSearch** domain; the main Onyx **S3 file-store bucket**; **WAF**. - IAM/IRSA **roles**: the workload role for Onyx application pods. - **Outputs** (the only thing Helm consumes): cluster name, RDS/Redis/OpenSearch endpoints, the file-store bucket name, the workload IRSA **role ARN**, the OIDC provider ARN/URL. @@ -208,7 +208,7 @@ sandboxProxy: ``` - ElastiCache: `REDIS_HOST`, `REDIS_SSL=true`, `REDIS_SSL_CERT_REQS=none`, `auth.redis.enabled=false` (no auth token). -- OpenSearch (v4.0 search backend): `ONYX_DISABLE_VESPA=true`, `ENABLE_OPENSEARCH_INDEXING/RETRIEVAL_FOR_ONYX=true`, `USING_AWS_MANAGED_OPENSEARCH=true`, `OPENSEARCH_REST_API_PORT=443`, `OPENSEARCH_USE_SSL=true`, `OPENSEARCH_ADMIN_USERNAME=admin`. +- OpenSearch (v4.0 search backend): `USING_AWS_MANAGED_OPENSEARCH=true`, `OPENSEARCH_REST_API_PORT=443`, `OPENSEARCH_USE_SSL=true`, `OPENSEARCH_ADMIN_USERNAME=admin`. - S3: `S3_FILE_STORE_BUCKET_NAME`, `S3_ENDPOINT_URL=""`, `AWS_REGION_NAME`. - Craft: `ENABLE_CRAFT=true`, `ONYX_SERVER_URL=http://onyx-api-service.onyx.svc.cluster.local:8080`, `auth.sandboxPushSecret.enabled=true`. (`SANDBOX_SERVICE_ACCOUNT_NAME`/`SANDBOX_CONTAINER_IMAGE` default correctly.) @@ -424,7 +424,7 @@ The original lead infra TODO list is accounted for as: ## 7. Notes / gotchas (condensed) - us-west-1 has only 2 AZs (→ the slice fix). Shared account near the EIP quota → use `single_nat_gateway=true`. -- The `vespa` node group is vestigial in v4.0 (OpenSearch replaced Vespa) — size it small or make it optional. +- The `vespa` node group (legacy name) only matters for an in-cluster OpenSearch. With a managed OpenSearch domain, set `vespa_node_enabled = false`. - `cluster_endpoint_public_access_cidrs=[]` causes a perpetual no-op diff AWS rejects — set explicitly (e.g. `["0.0.0.0/0"]`). - Codified chart/terraform changes live in the **local** chart, not the published `onyx/onyx` — install from `.` until released. - LLM provider: configure via admin UI (encrypted in DB) — never `GEN_AI_API_KEY` in a ConfigMap. diff --git a/docs/group-manager-scoped-permissions/00-index.md b/docs/group-manager-scoped-permissions/00-index.md index 647d37c548..7cecaa62e8 100644 --- a/docs/group-manager-scoped-permissions/00-index.md +++ b/docs/group-manager-scoped-permissions/00-index.md @@ -26,7 +26,7 @@ live) · D2 admins-only create groups · D3 admin-or-manager-of-that-group assig everything **except delete** · **D7 attaching an agent to a group is controlled by `manage:agents`** (standard GATE 2 keyed on `MANAGE_AGENTS` — admins/global holders self-share to their groups, scoped managers to managed groups; `add:agents`-only users can't group-share). The reviews confirmed PAT, -chat-runtime, and document/Vespa ACL are untouched, and refuted the backfill data-loss concern. Full +chat-runtime, and document-index ACL are untouched, and refuted the backfill data-loss concern. Full case-by-case coverage + the boot-bug prerequisite are in **[03 §11](03-detailed-design.md)** — the authoritative implementation checklist. diff --git a/docs/group-manager-scoped-permissions/02-high-level-design.md b/docs/group-manager-scoped-permissions/02-high-level-design.md index b8cb50dbe7..26071d9d6b 100644 --- a/docs/group-manager-scoped-permissions/02-high-level-design.md +++ b/docs/group-manager-scoped-permissions/02-high-level-design.md @@ -206,6 +206,6 @@ The gate therefore loads the resource's **current** groups in the same transacti - Dropping `role` / `UserRole` / `is_curator` and migrating the 3 residual `user.role==ADMIN` readers (`persona_sharing.py:53`, `build_session.py:638`, `search/api.py:104`) — deferred cleanup release. -- Any change to document-level Vespa ACL (`get_acl_for_user`) — unchanged; managers affect entity-level access +- Any change to document-level document-index ACL (`get_acl_for_user`) — unchanged; managers affect entity-level access only. - CE behavior — Group Manager is an EE capability (custom groups are EE); CE has only Basic + Admins. diff --git a/docs/group-manager-scoped-permissions/03-detailed-design.md b/docs/group-manager-scoped-permissions/03-detailed-design.md index 695330dd84..9b298d3b44 100644 --- a/docs/group-manager-scoped-permissions/03-detailed-design.md +++ b/docs/group-manager-scoped-permissions/03-detailed-design.md @@ -654,7 +654,7 @@ managed scope + PRIVATE); reject otherwise. Admins unaffected (global bypass). filter edits inside the `if get_editable:` block as correlated subqueries — do not perturb the shared joins above the split. - **Junction-only:** scope is purely `User__UserGroup.is_manager` + resource↔group junctions; - `permission_grant` stays global-only; Vespa ACL untouched. Backfill is safe — no upgrade migration nulls + `permission_grant` stays global-only; document-index ACL untouched. Backfill is safe — no upgrade migration nulls `role`, and `native_enum=False` stores `'CURATOR'`/`'GLOBAL_CURATOR'` literally. Ship the zero-managed-group migration report for the GLOBAL_CURATOR snapshot caveat. diff --git a/docs/group-manager-scoped-permissions/04-implementation-plan.md b/docs/group-manager-scoped-permissions/04-implementation-plan.md index 75bc8361c9..28236a0087 100644 --- a/docs/group-manager-scoped-permissions/04-implementation-plan.md +++ b/docs/group-manager-scoped-permissions/04-implementation-plan.md @@ -73,7 +73,7 @@ PRIVATE and strictly within their managed groups — enforced authoritatively at - **Corrections (§11.7):** feedback `db/feedback.py` = **no change** (admin-only; not in bundle); `recompute_user_permissions__no_commit` takes `(user_ids, db_session)` and must be extended to set `is_group_manager`. - - **Confirmed SAFE (§11.8):** PAT cap, chat runtime, and document/Vespa ACL are untouched — keep them so. + - **Confirmed SAFE (§11.8):** PAT cap, chat runtime, and document-index ACL are untouched — keep them so. ## Implementation Strategy diff --git a/mobile/src/components/chat/SourceRow.tsx b/mobile/src/components/chat/SourceRow.tsx index d44887c89a..6902edd77b 100644 --- a/mobile/src/components/chat/SourceRow.tsx +++ b/mobile/src/components/chat/SourceRow.tsx @@ -21,7 +21,7 @@ function metaLine(doc: SearchDoc): string { function snippet(doc: SearchDoc): string { const raw = doc.match_highlights?.[0] ?? doc.blurb ?? ""; - // Strip Vespa … highlight markup so it doesn't render literally. + // Strip … highlight markup so it doesn't render literally. const clean = raw.replace(/<\/?hi>/g, "").trim(); return clean.length > 200 ? `${clean.slice(0, 200)}…` : clean; } diff --git a/pyproject.toml b/pyproject.toml index 86877f775f..1c91a4f42c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -397,9 +397,7 @@ known-first-party = ["onyx", "ee", "tests", "shared_configs", "model_server"] "backend/onyx/db/persona.py" = ["C901"] # 49 "backend/onyx/deep_research/dr_loop.py" = ["C901"] # 26 "backend/onyx/document_index/opensearch/search.py" = ["C901"] # 48 -"backend/onyx/document_index/vespa/shared_utils/vespa_request_builders.py" = ["C901"] # 41 "backend/onyx/federated_connectors/federated_retrieval.py" = ["C901"] # 21 -"backend/onyx/kg/extractions/extraction_processing.py" = ["C901"] # 30 "backend/onyx/llm/multi_llm.py" = ["C901"] # 45 "backend/onyx/llm/utils.py" = ["C901"] # 30 "backend/onyx/onyxbot/slack/handlers/handle_message.py" = ["C901"] # 34 diff --git a/terraform-provider-onyx/README.md b/terraform-provider-onyx/README.md index e637e09586..76f24ad9c3 100644 --- a/terraform-provider-onyx/README.md +++ b/terraform-provider-onyx/README.md @@ -354,7 +354,7 @@ celery -A onyx.background.celery.versioned_apps.primary worker \ --pool=threads --concurrency=4 --loglevel=INFO --hostname=tfacc-primary@%n -Q celery & celery -A onyx.background.celery.versioned_apps.light worker \ --pool=threads --concurrency=8 --loglevel=INFO --hostname=tfacc-light@%n \ - -Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,opensearch_migration & + -Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup & ``` The primary worker picks up the deletion checks the API server dispatches; the light diff --git a/terraform-provider-onyx/docs/resources/settings.md b/terraform-provider-onyx/docs/resources/settings.md index 4f384c3f52..f28b205059 100644 --- a/terraform-provider-onyx/docs/resources/settings.md +++ b/terraform-provider-onyx/docs/resources/settings.md @@ -55,7 +55,7 @@ resource "onyx_settings" "workspace" { - `gpu_enabled` (Boolean) Whether the deployment has GPU support (read-only). - `hide_query_history_from_admin_panel` (Boolean) Whether the query history page is hidden in the admin panel. Read-only: controlled by the HIDE_QUERY_HISTORY_FROM_ADMIN_PANEL backend env var. - `id` (String) Always `"settings"`. -- `opensearch_indexing_enabled` (Boolean) OpenSearch migration flag. Read-only: controlled by the ENABLE_OPENSEARCH_INDEXING_FOR_ONYX backend env var. +- `opensearch_indexing_enabled` (Boolean) Deprecated: always true, because OpenSearch is the only document index. Read-only. - `seat_count` (Number) Licensed seat count (read-only). - `show_extra_connectors` (Boolean) Whether the extended connector catalog is shown. Read-only: controlled by the SHOW_EXTRA_CONNECTORS backend env var. - `tier` (String) Resolved license tier: `community`, `business`, or `enterprise` (read-only). diff --git a/terraform-provider-onyx/internal/provider/settings_resource.go b/terraform-provider-onyx/internal/provider/settings_resource.go index 426cfc0e1c..e62a724d79 100644 --- a/terraform-provider-onyx/internal/provider/settings_resource.go +++ b/terraform-provider-onyx/internal/provider/settings_resource.go @@ -36,8 +36,8 @@ type settingsResource struct { // // These fields are read-only. The license gives ApplicationStatus, Tier, // EEFeaturesEnabled, GPUEnabled, SeatCount and UsedSeats. Backend env vars -// overwrite HideQueryHistoryFromAdminPanel, ShowExtraConnectors and -// OpenSearchIndexingEnabled on every read. +// overwrite HideQueryHistoryFromAdminPanel and ShowExtraConnectors on every +// read. OpenSearchIndexingEnabled is deprecated and always true. type settingsResourceModel struct { ID types.String `tfsdk:"id"` CompanyName types.String `tfsdk:"company_name"` @@ -213,8 +213,8 @@ func (r *settingsResource) Schema(_ context.Context, _ resource.SchemaRequest, r }, "opensearch_indexing_enabled": schema.BoolAttribute{ Computed: true, - MarkdownDescription: "OpenSearch migration flag. Read-only: controlled by the " + - "ENABLE_OPENSEARCH_INDEXING_FOR_ONYX backend env var.", + MarkdownDescription: "Deprecated: always true, because OpenSearch is " + + "the only document index. Read-only.", }, }, } diff --git a/tools/loadtest/README.md b/tools/loadtest/README.md index ba54f56be0..21cdfa137b 100644 --- a/tools/loadtest/README.md +++ b/tools/loadtest/README.md @@ -84,7 +84,7 @@ approximating production traffic shape: | Scenario | Metrics | Weight | What it exercises | |---|---|---|---| | **BasicChatUser** | `chat:*` | 70 | single-turn chat, plain answer | -| **ChatWithSearchUser** | `search:*` | 20 | one `internal_search` tool call → query expansion, embedding model server, Vespa/OpenSearch retrieval (needs indexed docs) | +| **ChatWithSearchUser** | `search:*` | 20 | one `internal_search` tool call → query expansion, embedding model server, OpenSearch retrieval (needs indexed docs) | | **MultiToolUser** | `multitool:*` | 8 | up to 3 retrieval tools in parallel in one turn (`mock-tools3`) | | **DeepResearchUser** | `dr:*` | 2 | full DR turn — plan, parallel agents, reports; ~8+ LLM calls on one held stream | diff --git a/tools/loadtest/scenarios/chat_with_search.py b/tools/loadtest/scenarios/chat_with_search.py index d4a7e45329..d8e39e8b78 100644 --- a/tools/loadtest/scenarios/chat_with_search.py +++ b/tools/loadtest/scenarios/chat_with_search.py @@ -2,7 +2,7 @@ The `-tools1` model knob makes the mock LLM answer the first AUTO-tool-choice cycle with an `internal_search` tool call, so Onyx genuinely executes the -search tool: query expansion, embedding model server, Vespa/OpenSearch +search tool: query expansion, embedding model server, OpenSearch retrieval, document streaming. The follow-up LLM call (tool result in history) streams the final answer. diff --git a/tools/profiling/README.md b/tools/profiling/README.md index 5aa0c0d05d..217e14dd64 100644 --- a/tools/profiling/README.md +++ b/tools/profiling/README.md @@ -29,7 +29,7 @@ docker compose up -d | `onyx-docfetching-worker` | 9092 | Celery docfetching worker | | `onyx-docprocessing-worker`| 9093 | Celery docprocessing worker | | `onyx-heavy-worker` | 9094 | Celery heavy worker (pruning, perm sync, group sync) | -| `onyx-light-worker` | 9095 | Celery light worker (vespa sync, deletion, permissions upsert) | +| `onyx-light-worker` | 9095 | Celery light worker (document index sync, deletion, permissions upsert) | ## Environment variables diff --git a/web/playwright.config.ts b/web/playwright.config.ts index 48f9325936..6e4363591e 100644 --- a/web/playwright.config.ts +++ b/web/playwright.config.ts @@ -56,7 +56,7 @@ export default defineConfig({ workers: 1, }, { - // runs against the Onyx Lite stack (DISABLE_VECTOR_DB=true, no Vespa/Redis) + // runs against the Onyx Lite stack (DISABLE_VECTOR_DB=true, no OpenSearch/Redis) name: "lite", use: { ...devices["Desktop Chrome"], diff --git a/web/src/i18n/messages/ar.json b/web/src/i18n/messages/ar.json index 45d22ac18e..6b8424e600 100644 --- a/web/src/i18n/messages/ar.json +++ b/web/src/i18n/messages/ar.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "تلخيص الصور" }, - "kgDeepExtraction": { - "label": "استخراج الرسم البياني المعرفي" - }, - "kgDocumentClassification": { - "label": "تصنيف الرسم البياني المعرفي" - }, "other": { "label": "أخرى" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · استخدام النظام هو نشاط نماذج اللغة الكبيرة الذي يشغّله {appName} في الخلفية بدلًا من تشغيله لمستخدم محدد. تشمل الأمثلة تلخيص الصور ومعالجة RAG السياقية واستخراج الرسم البياني المعرفي.", - "emptyDescription": "استخدام النظام هو نشاط نماذج اللغة الكبيرة الذي يشغّله {appName} في الخلفية بدلًا من تشغيله لمستخدم محدد. تشمل الأمثلة تلخيص الصور ومعالجة RAG السياقية واستخراج الرسم البياني المعرفي.", + "description": "{start} – {end} · استخدام النظام هو نشاط نماذج اللغة الكبيرة الذي يشغّله {appName} في الخلفية بدلًا من تشغيله لمستخدم محدد. تشمل الأمثلة تلخيص الصور ومعالجة RAG السياقية.", + "emptyDescription": "استخدام النظام هو نشاط نماذج اللغة الكبيرة الذي يشغّله {appName} في الخلفية بدلًا من تشغيله لمستخدم محدد. تشمل الأمثلة تلخيص الصور ومعالجة RAG السياقية.", "error": { "title": "تعذر تحميل استخدام النظام لهذه الفترة." }, diff --git a/web/src/i18n/messages/de.json b/web/src/i18n/messages/de.json index 64573c7248..a0594cca63 100644 --- a/web/src/i18n/messages/de.json +++ b/web/src/i18n/messages/de.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "Bildzusammenfassung" }, - "kgDeepExtraction": { - "label": "Knowledge-Graph-Extraktion" - }, - "kgDocumentClassification": { - "label": "Knowledge-Graph-Klassifizierung" - }, "other": { "label": "Sonstige" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · Systemnutzung ist LLM-Aktivität, die {appName} im Hintergrund statt für einen bestimmten Benutzer ausführt. Beispiele sind Bildzusammenfassungen, kontextbezogene RAG-Verarbeitung und Knowledge-Graph-Extraktion.", - "emptyDescription": "Systemnutzung ist LLM-Aktivität, die {appName} im Hintergrund statt für einen bestimmten Benutzer ausführt. Beispiele sind Bildzusammenfassungen, kontextbezogene RAG-Verarbeitung und Knowledge-Graph-Extraktion.", + "description": "{start} – {end} · Systemnutzung ist LLM-Aktivität, die {appName} im Hintergrund statt für einen bestimmten Benutzer ausführt. Beispiele sind Bildzusammenfassungen und kontextbezogene RAG-Verarbeitung.", + "emptyDescription": "Systemnutzung ist LLM-Aktivität, die {appName} im Hintergrund statt für einen bestimmten Benutzer ausführt. Beispiele sind Bildzusammenfassungen und kontextbezogene RAG-Verarbeitung.", "error": { "title": "Die Systemnutzung für diesen Zeitraum konnte nicht geladen werden." }, diff --git a/web/src/i18n/messages/en.json b/web/src/i18n/messages/en.json index 2336b353e6..3104d7df04 100644 --- a/web/src/i18n/messages/en.json +++ b/web/src/i18n/messages/en.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "Image summarization" }, - "kgDeepExtraction": { - "label": "Knowledge graph extraction" - }, - "kgDocumentClassification": { - "label": "Knowledge graph classification" - }, "other": { "label": "Other" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · System usage is LLM activity that {appName} runs in the background rather than for a specific user. Examples include image summaries, contextual RAG processing, and knowledge graph extraction.", - "emptyDescription": "System usage is LLM activity that {appName} runs in the background rather than for a specific user. Examples include image summaries, contextual RAG processing, and knowledge graph extraction.", + "description": "{start} – {end} · System usage is LLM activity that {appName} runs in the background rather than for a specific user. Examples include image summaries and contextual RAG processing.", + "emptyDescription": "System usage is LLM activity that {appName} runs in the background rather than for a specific user. Examples include image summaries and contextual RAG processing.", "error": { "title": "Failed to load system usage for this period." }, diff --git a/web/src/i18n/messages/es.json b/web/src/i18n/messages/es.json index 260642abec..4d7ec0d2b9 100644 --- a/web/src/i18n/messages/es.json +++ b/web/src/i18n/messages/es.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "Resumen de imágenes" }, - "kgDeepExtraction": { - "label": "Extracción del grafo de conocimiento" - }, - "kgDocumentClassification": { - "label": "Clasificación del grafo de conocimiento" - }, "other": { "label": "Otros" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · El uso del sistema es actividad del LLM que {appName} ejecuta en segundo plano, no para un usuario específico. Algunos ejemplos son los resúmenes de imágenes, el procesamiento RAG contextual y la extracción del grafo de conocimiento.", - "emptyDescription": "El uso del sistema es actividad del LLM que {appName} ejecuta en segundo plano, no para un usuario específico. Algunos ejemplos son los resúmenes de imágenes, el procesamiento RAG contextual y la extracción del grafo de conocimiento.", + "description": "{start} – {end} · El uso del sistema es actividad del LLM que {appName} ejecuta en segundo plano, no para un usuario específico. Algunos ejemplos son los resúmenes de imágenes y el procesamiento RAG contextual.", + "emptyDescription": "El uso del sistema es actividad del LLM que {appName} ejecuta en segundo plano, no para un usuario específico. Algunos ejemplos son los resúmenes de imágenes y el procesamiento RAG contextual.", "error": { "title": "No se pudo cargar el uso del sistema de este periodo." }, diff --git a/web/src/i18n/messages/fr.json b/web/src/i18n/messages/fr.json index c0604cb8f2..533dc661f3 100644 --- a/web/src/i18n/messages/fr.json +++ b/web/src/i18n/messages/fr.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "Résumé des images" }, - "kgDeepExtraction": { - "label": "Extraction du graphe de connaissances" - }, - "kgDocumentClassification": { - "label": "Classification du graphe de connaissances" - }, "other": { "label": "Autres" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · L’utilisation système correspond à l’activité LLM exécutée par {appName} en arrière-plan plutôt que pour un utilisateur spécifique. Cela inclut les résumés d’images, le traitement RAG contextuel et l’extraction du graphe de connaissances.", - "emptyDescription": "L’utilisation système correspond à l’activité LLM exécutée par {appName} en arrière-plan plutôt que pour un utilisateur spécifique. Cela inclut les résumés d’images, le traitement RAG contextuel et l’extraction du graphe de connaissances.", + "description": "{start} – {end} · L’utilisation système correspond à l’activité LLM exécutée par {appName} en arrière-plan plutôt que pour un utilisateur spécifique. Cela inclut les résumés d’images et le traitement RAG contextuel.", + "emptyDescription": "L’utilisation système correspond à l’activité LLM exécutée par {appName} en arrière-plan plutôt que pour un utilisateur spécifique. Cela inclut les résumés d’images et le traitement RAG contextuel.", "error": { "title": "Échec du chargement de l'utilisation système pour cette période." }, diff --git a/web/src/i18n/messages/ja.json b/web/src/i18n/messages/ja.json index eac3df79d9..031c3b917f 100644 --- a/web/src/i18n/messages/ja.json +++ b/web/src/i18n/messages/ja.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "画像要約" }, - "kgDeepExtraction": { - "label": "ナレッジグラフ抽出" - }, - "kgDocumentClassification": { - "label": "ナレッジグラフ分類" - }, "other": { "label": "その他" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · システム使用状況とは、特定のユーザー向けではなく、{appName} がバックグラウンドで実行する LLM 処理です。例として、画像要約、コンテキスト RAG 処理、ナレッジグラフ抽出があります。", - "emptyDescription": "システム使用状況とは、特定のユーザー向けではなく、{appName} がバックグラウンドで実行する LLM 処理です。例として、画像要約、コンテキスト RAG 処理、ナレッジグラフ抽出があります。", + "description": "{start} – {end} · システム使用状況とは、特定のユーザー向けではなく、{appName} がバックグラウンドで実行する LLM 処理です。例として、画像要約やコンテキスト RAG 処理があります。", + "emptyDescription": "システム使用状況とは、特定のユーザー向けではなく、{appName} がバックグラウンドで実行する LLM 処理です。例として、画像要約やコンテキスト RAG 処理があります。", "error": { "title": "この期間のシステム使用状況を読み込めませんでした。" }, diff --git a/web/src/i18n/messages/ko.json b/web/src/i18n/messages/ko.json index 88d68dfded..7c7f83201c 100644 --- a/web/src/i18n/messages/ko.json +++ b/web/src/i18n/messages/ko.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "이미지 요약" }, - "kgDeepExtraction": { - "label": "지식 그래프 추출" - }, - "kgDocumentClassification": { - "label": "지식 그래프 분류" - }, "other": { "label": "기타" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · 시스템 사용량은 {appName}가 특정 사용자를 위해서가 아니라 백그라운드에서 실행하는 LLM 작업입니다. 이미지 요약, 컨텍스트 RAG 처리, 지식 그래프 추출 등이 포함됩니다.", - "emptyDescription": "시스템 사용량은 {appName}가 특정 사용자를 위해서가 아니라 백그라운드에서 실행하는 LLM 작업입니다. 이미지 요약, 컨텍스트 RAG 처리, 지식 그래프 추출 등이 포함됩니다.", + "description": "{start} – {end} · 시스템 사용량은 {appName}가 특정 사용자를 위해서가 아니라 백그라운드에서 실행하는 LLM 작업입니다. 이미지 요약, 컨텍스트 RAG 처리 등이 포함됩니다.", + "emptyDescription": "시스템 사용량은 {appName}가 특정 사용자를 위해서가 아니라 백그라운드에서 실행하는 LLM 작업입니다. 이미지 요약, 컨텍스트 RAG 처리 등이 포함됩니다.", "error": { "title": "이 기간의 시스템 사용량을 불러오지 못했습니다." }, diff --git a/web/src/i18n/messages/pt.json b/web/src/i18n/messages/pt.json index 97755d786b..5362eece07 100644 --- a/web/src/i18n/messages/pt.json +++ b/web/src/i18n/messages/pt.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "Resumo de imagens" }, - "kgDeepExtraction": { - "label": "Extração do grafo de conhecimento" - }, - "kgDocumentClassification": { - "label": "Classificação do grafo de conhecimento" - }, "other": { "label": "Outros" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · O uso do sistema é a atividade de LLM que o {appName} executa em segundo plano, e não para um usuário específico. Alguns exemplos são resumos de imagens, processamento RAG contextual e extração do grafo de conhecimento.", - "emptyDescription": "O uso do sistema é a atividade de LLM que o {appName} executa em segundo plano, e não para um usuário específico. Alguns exemplos são resumos de imagens, processamento RAG contextual e extração do grafo de conhecimento.", + "description": "{start} – {end} · O uso do sistema é a atividade de LLM que o {appName} executa em segundo plano, e não para um usuário específico. Alguns exemplos são resumos de imagens e processamento RAG contextual.", + "emptyDescription": "O uso do sistema é a atividade de LLM que o {appName} executa em segundo plano, e não para um usuário específico. Alguns exemplos são resumos de imagens e processamento RAG contextual.", "error": { "title": "Falha ao carregar o uso do sistema neste período." }, diff --git a/web/src/i18n/messages/zh.json b/web/src/i18n/messages/zh.json index bc3d6c7f59..bb6058d5e3 100644 --- a/web/src/i18n/messages/zh.json +++ b/web/src/i18n/messages/zh.json @@ -6114,12 +6114,6 @@ "imageSummarization": { "label": "图像摘要" }, - "kgDeepExtraction": { - "label": "知识图谱提取" - }, - "kgDocumentClassification": { - "label": "知识图谱分类" - }, "other": { "label": "其他" }, @@ -6136,8 +6130,8 @@ } }, "panel": { - "description": "{start} – {end} · 系统用量是 {appName} 在后台运行而非为特定用户运行的 LLM 活动。例如图像摘要、上下文 RAG 处理和知识图谱提取。", - "emptyDescription": "系统用量是 {appName} 在后台运行而非为特定用户运行的 LLM 活动。例如图像摘要、上下文 RAG 处理和知识图谱提取。", + "description": "{start} – {end} · 系统用量是 {appName} 在后台运行而非为特定用户运行的 LLM 活动。例如图像摘要和上下文 RAG 处理。", + "emptyDescription": "系统用量是 {appName} 在后台运行而非为特定用户运行的 LLM 活动。例如图像摘要和上下文 RAG 处理。", "error": { "title": "无法加载此期间的系统用量。" }, diff --git a/web/src/lib/admin-routes.ts b/web/src/lib/admin-routes.ts index 053caede28..ae148c9ebd 100644 --- a/web/src/lib/admin-routes.ts +++ b/web/src/lib/admin-routes.ts @@ -48,7 +48,6 @@ export interface FeatureFlags { customAnalyticsEnabled: boolean; hasSubscription: boolean; hooksEnabled: boolean; - opensearchEnabled: boolean; queryHistoryEnabled: boolean; craftAvailable: boolean; } diff --git a/web/src/lib/settings/types.ts b/web/src/lib/settings/types.ts index da31896d82..679caea81e 100644 --- a/web/src/lib/settings/types.ts +++ b/web/src/lib/settings/types.ts @@ -89,9 +89,6 @@ export interface Settings { seat_count?: number | null; used_seats?: number | null; - // OpenSearch migration - opensearch_indexing_enabled?: boolean; - // Vector DB availability flag - false when DISABLE_VECTOR_DB is set. // When false, connectors, RAG search, document sets, and related features // are unavailable. diff --git a/web/src/sections/sidebar/AdminSidebar.tsx b/web/src/sections/sidebar/AdminSidebar.tsx index af2a795d81..a2f1cd9798 100644 --- a/web/src/sections/sidebar/AdminSidebar.tsx +++ b/web/src/sections/sidebar/AdminSidebar.tsx @@ -75,7 +75,6 @@ export default function AdminSidebar() { customAnalyticsEnabled, hasSubscription: hasSubscriptionOrLicense, hooksEnabled: settings?.hooks_enabled ?? false, - opensearchEnabled: settings?.opensearch_indexing_enabled ?? false, queryHistoryEnabled: settings?.query_history_type !== "disabled" && !settings?.hide_query_history_from_admin_panel, diff --git a/web/src/views/admin/SystemUsagePanel.tsx b/web/src/views/admin/SystemUsagePanel.tsx index 60562c1cfc..051f290830 100644 --- a/web/src/views/admin/SystemUsagePanel.tsx +++ b/web/src/views/admin/SystemUsagePanel.tsx @@ -26,8 +26,6 @@ const OTHER_CATEGORY = "other"; const IMAGE_SUMMARIZATION_FLOW = "image_summarization"; const CONTEXTUAL_RAG_DOC_SUMMARY_FLOW = "contextual_rag_doc_summary"; const CONTEXTUAL_RAG_CHUNK_CONTEXT_FLOW = "contextual_rag_chunk_context"; -const KG_DOCUMENT_CLASSIFICATION_FLOW = "kg_document_classification"; -const KG_DEEP_EXTRACTION_FLOW = "kg_deep_extraction"; interface SystemUsageRow extends UsageExportTotals { category: string; @@ -56,10 +54,6 @@ function categoryLabel(category: string, t: SystemUsageTranslate): string { return t("categories.contextualRagDocumentSummary.label"); case CONTEXTUAL_RAG_CHUNK_CONTEXT_FLOW: return t("categories.contextualRagChunkContext.label"); - case KG_DOCUMENT_CLASSIFICATION_FLOW: - return t("categories.kgDocumentClassification.label"); - case KG_DEEP_EXTRACTION_FLOW: - return t("categories.kgDeepExtraction.label"); case UNATTRIBUTED_CATEGORY: return t("categories.unattributed.label"); case OTHER_CATEGORY: diff --git a/web/tests/e2e/chat/chat_message_rendering.spec.ts b/web/tests/e2e/chat/chat_message_rendering.spec.ts index fdaca441dc..8ea12cbed0 100644 --- a/web/tests/e2e/chat/chat_message_rendering.spec.ts +++ b/web/tests/e2e/chat/chat_message_rendering.spec.ts @@ -67,7 +67,7 @@ The typical indexing latency depends on your configuration: The chat system uses a **Retrieval-Augmented Generation (RAG)** pipeline: 1. User query is analyzed and expanded -2. Relevant documents are retrieved from the vector database (Vespa) +2. Relevant documents are retrieved from the vector database (OpenSearch) 3. Documents are ranked and filtered by relevance and permissions 4. The LLM generates a response grounded in the retrieved documents 5. Citations are attached to specific claims in the response